Trade-Offs in Automatic Differentiation: TensorFlow, PyTorch, Jax, and Julia

This is because continuous optimization is useless when crossing a discontinuity, which is what control flow creates. Even in a trivial situation like ReLU, where the control flow is mimicking a continuous transition, you have the "dead ReLU" problem, where you have to start training on the correct side of the discontinuity and make sure to never cross.

I remember reading an early S4TF manifesto talking about natural language processing, image processing, etc. thinking, that cannot be your audience because that audience already has AD systems which support their domain. Building something that is more general for the sake of being more general is never a good idea, that's bad engineering. It had a lot of great ideas, and indeed the dev experience seemed nice. But I would venture to guess that the Google overlords had to question what the true value of S4TF in that light. "Standard ML" cannot be your target audience if you want to work on AD extensions.

Following this thread, you can also see what how the Julia tools evolved. If you see the paper that was the synthesis for Zygote.jl, it was all about while loops and scalar operations (https://arxiv.org/abs/1810.07951). Why did that not completely change ML? Well, ML doesn't use those kinds of operations. I would say the project kind of started as a "tool looking for a problem". It did get a bit lucky that it found a problem: scientific applications need to be able to use automatic differentiation without rewriting the whole codebase to an ML library, leading to the big Julia AD manifesto of language-wide differentiable programming by directly acting on the Julia source itself rather than a language subset (http://ceur-ws.org/Vol-2587/article_8.pdf). Zygote was a good AD, but not a great AD, why? Because it could not hit this goal, mostly because of its lack of mutation handling. Yes, it does handle standard ML just fine, but does not justify its added complexity.

What has actually kept Julia AD research going is that some scientific machine learning applications, specifically physics-informed neural networks (PINNs), require very high order derivatives. For example, to solve the PDE u_t = u_xx with neural networks, you need to take the third derivative of the neural network. With Jax this can only be done with a separate language subset (https://openreview.net/pdf?id=SkxEF3FNPH), and thus a new AD for Julia to replace Zygote, known as Diffractor.jl, was devised to automatically incorporate higher order AD optimizations as part of the regular usage (https://www.youtube.com/watch?v=mQnSRfseu0c). It is these PINN SciML applications that have funded its development and is its built-in audience: it solves a problem nothing else does, even if it is potentially niche. Similarly with Enzyme, it solved the problem of how to do mutation well, which is where you can see in the paper that its applications are mostly ODE and PDE solvers (Euler, RK4, the Bruss semilinear PDE) (https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b...). Torchscript and Jax do not handle this domain well, so it has an audience, which may (or may not?) be niche.

A big part of writing this blog post was to highlight this to the Julia AD crew that I regularly work with. What will keep these projects alive is understanding the engineering trade-offs that are made and who the audience is. The complexity has a cost so it better have a benefit. If that target is lost, if any benefit is a theoretical "but you may need more features some day", then the projects will lose traction. The project needs to be two-fold: identify new architectures and applications that would benefit from expanded language support from AD and build good support for those projects. Otherwise it is just training a transformer in Julia vs training a transformer in Python, and that is not justifiable.

I am really glad to hear this. As it happens, my main post-retirement has been to learn Swift to try out CoreML. I'm really enjoying learning this so far.

Note however that Kotlin is backed by Facebook https://ai.facebook.com/blog/paving-the-way-for-software-20-... and that there is a mature library for autodiff https://github.com/breandan/kotlingrad

Outside of research transformers are rarely used for computer vision problems and CNNs remain the go to architecture. And you actually need to do some hacks to get transformers to work with computer vision at a meaningful scale (splitting images into patches and convoluting the patches to produce features to feed into the transformer).

Like other comments, CNNs and LSTMs are still in wide use today. If you dig deep enough, position encoding doesn't really capture time-based series information that well.

It touches on the main one which is simplicity. it's much easier to write an AD system for a more static language.

Maybe it has no technical advantages, unless you count being very popular and in an accessible language as a technical advantage (which it definitely could be depending on your definition of "technical").

Julia is designed for advanced numerical computing and Python isn't. The metaprogramming affordances needed for AD are much better developed in Julia than they ever will be in Python. And let's not forget the immense utility of multiple dispatch in Julia, another feature Python will probably never have. So it's not surprising that Julia is simply way more capable.

Has it worked really well? I feel ORMs are a sign it hasn't. Though I really enjoy having learned SQL and being able to interact with almost all relational databases.

You might like dex[1].

[1]: https://github.com/google-research/dex-lang

that seems one of the points of enzyme[1], which was mentioned in the article.

[1] - https://enzyme.mit.edu/

being able to do interprocedural cross language analysis seems awesome considering how much code is written in C++, but used in higher level languages.

Tf.function is largely morally equivalent to Torchscript, which he does discuss.

IMO automatic differentiation shouldn’t be baked into the compiler, because that forces a bit of a monoculture around one AD system and forces you to accept the trade offs of that particular AD system.

The Julia approach has been instead to expose an interface for compiler plugins that any AD (or other ‘nonstandard interpretation / code transformation) library can access. There’s negative trade offs to to this as well, but I really like the way it’s turned out and I think it’s given us some fantastic tools for non-AD purposes as well.

No, autograd acts similarly to PyTorch in that it builds a tape that it reverses while PyTorch just comes with more optimized kernels (and kernels that act on GPUs). The AD that I was referencing was tangent (https://github.com/google/tangent). It was an interesting project but it's hard to see who the audience is. Generating Python source code makes things harder to analyze, and you cannot JIT compile the generated code unless you could JIT compile Python. So you might as well first trace to a JIT-compliable sublanguage and do the actions there, which is precisely what Jax does. In theory tangent is a bit more general, and maybe you could mix it with Numba, but then it's hard to justify. If it's more general then it's not for the standard ML community for the same reason as the Julia tools, but then it better do better than the Julia tools in the specific niche that they are targeting. That generality means that it cannot use XLA, and thus from day 1 it wouldn't get the extra compiler optimizations that some which uses XLA does (Jax). Jax just makes much more sense for the people who were building it, it chose its niche very well.

You probably could get perturbation confusion in reverse mode, but it's not an easy trap like it is with forward mode. The problem with forward mode AD is that it's deceptively easy: you transform every number a into a pair of numbers (a,b) and you change every function from f(a) to (f(a),f'(a)), put the chain rule on it, and you're done. Whether you call it dual number arithmetic or a compiler transformation it's the same thing in the end. The issue with perturbation confusion is that you've created this "secret extra number" to store the derivative data, and so if two things are differentiating code at the same time, you need to make sure you turn (a,b) into (a,a',b,b') and that all layers of AD are always grabbing the right value out of there. Note that in the way I wrote it, if you assumed "the b term is always num[2] in the tuple", oops perturbation confusion, and so your generated code needs to be "smart" (but not lose efficiency!). Thus the fixes are proofs and tagging systems that ensure the right perturbation terms are always used in the right places.

With reverse mode AD, this is much less likely to be an issue because the AD system isn't necessarily storing and working on hidden extensions to the values, it's running a function forwards and then running a separate function backwards having remembered some values from the forward pass. If the remembered values are correct and never modified, then generating a higher order derivative is just as safe as the first. But that last little detail is thus what I think is most akin to perturbation confusion in reverse mode: reverse mode has the assumption that the objects captured in the forward pass will not be changed (or will be at least be back in the correct state) when it is trying to reverse. The easy way to break this assumption doesn't even require second derivatives. The easiest way to break it is mutation: if you walk forward by doing Ax, then the reverse pass wants to do A'v so it just keeps the pointer to A, but if A gets mutated in the meantime then using that pointer is incorrect. This is the reason why most AD systems simply disallow mutation except in very special unoptimized cases (PyTorch, Jax, Zygote, ...).

Enzyme.jl is an exception because it takes a global analysis of the program it's differentiating (with proper escape analysis etc. passes at the LLVM level) in order to know that any mutation going forward will be reversed during the reverse path, so by the time it gets back to A'*v it knows A will be the same. Higher level ADs could go cowboy YOLO style and just assume the reversed matrix is correct (and it might be a lot of the time), though that causes some pretty major concerns for correctness. The other option is to simply make a full copy of A every time you mutate an element, so have fun if you loop through your weight matrix. The Diffractor.jl near future approach is more like Haskell GHC where it just wants you to give it the non-mutating code so it can try and generate the mutating code when that would be more efficient (https://github.com/JuliaLang/julia/pull/42465).

So with forward-mode AD there was an entire literature around schemes of provable safety to perturbation confusion, and I'm surprised we haven't already started seeing papers about provable safety with respect to mutation in higher-level reverse-mode AD. I would suspect that the only reason why it hasn't started is that the people who write type-theoretic proofs tend to be the functional programming pure function folks that tell people to never mutate anyways, so the literature might instead go the direction of escape analysis proofs to optimize immutable array code to (and beyond) the performance of mutation code on commonly mutating applications. Either way it's getting there with the same purpose in mind.

From a quick skim of it’s readme I don’t see what’s so special about it. Can you provide additional context?