9 ms·
Trade-Offs in Automatic Differentiation: TensorFlow, PyTorch, Jax, and Julia
- cl3misch 5y ago> fun fact, the Jax folks at Google Brain did have a Python source code transform AD at one point but it was scrapped essentially because of these difficulties I assume you mean autograd? https://github.com/HIPS/autograd https://github.com/HIPS/autograd
- ChrisRackauckas 5y agoNo, 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 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.
- brilee 5y agoFYI - Tangent evolved into TF2's AutoGraph.
- ChrisRackauckas 5y agoIndeed, and that makes a lot of sense. The qualm of tangent is that you get the source code translation but without the additional optimizations that the technique can provide. It was then natural to just target TensorFlow/XLA to do a similar thing but get the performance of TensorFlow as a result. The downside is that it loses the one true upside of tangent which was that, by generating Python code, it could in theory be easier for a Python programming to debug. But this was probably the right sacrifice to make for most people.
- albertzeyer 5y agoThis misses some discussion on tf.function which does a Python AST level transformation to a static TF computation graph, including dynamic control flow like loops and conditional branches.
- chillee 5y agoTf.function is largely morally equivalent to Torchscript, which he does discuss.
- snicker7 5y agoEven though this posts thesis is “trade offs”, it doesn’t really talk about any technical advantages that the Python’s AD ecosystem (Tensorflow, PyTorch, JAX) has over Julia’s (Zygote.jl, Diffractor.jl).
- adgjlsfhk1 5y agoIt touches on the main one which is simplicity. it's much easier to write an AD system for a more static language.
- civilized 5y agoMaybe 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.
- jaggirs 5y agoOne disatvantage of the language itself is the need for compilation, which isn't that fast in my limited experience. But I would love to hear how much this affects iteration speed.
- civilized 5y agoYeah, I'd imagine that for things both Python and Julia can do AD-wise, Python may be preferable since it's interpreted and thus instant-feedback, but all the numerical heavy lifting in packages like Jax and PyTorch is done in fast C++. So you should be getting a more appealing environment for experimentation without losing out on speed.
- calaphos 5y agoThe same issue exists with Jax. XLA compilation can take up quite a bit of time, especially on larger NN models. And theres no persistent compile cache, so even if you don't change the jitted function you need to wait for compilation again as you restart the process.
- mark_l_watson 5y agoInteresting read. I was very disappointed when the Swift TensorFlow project withered away. A good general purpose programming language combined with deep learning seemed like a great idea. Apple actually provides a good dev experience with Swift and CoreML (for fun I wrote a Swift/SwiftUI/CoreML app that uses two deep learning models that is in the App Store). Wolfram Language takes an approach similar to Apple’s in providing a good number of pre trained models, but I haven’t yet discovered any automatic differentiation examples. Of the frameworks described in the article, I find Julia most interesting but I need to use Python and TensorFlow in my work.
- 2sk21 5y agoI 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.
- hnarayanan 5y agoThis is really interesting to me. Could you please share your learning pathway?
- ChrisRackauckas 5y agoI 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 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 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 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 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/9332c513ef44b682e9347822c2e457ac-Paper.pdf 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.
- spacetracks 5y ago"The second factor, and probably the more damning one, is that most ML codes don't actually use that much dynamism." I would argue that this in true precisely because it is not available in an AD system. When I tell friends and coworkers about what zygote can do they light up and start describing different use cases they have that could benefit from AD. Diff eq solving is a big one.
- KKKKkkkk1 5y agoThis 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.
- ogogmad 5y agoI don't know whether this belongs here, but... Formally, there is a generalisation of differentiation which can handle functions like ReLU (i.e. locally Lipschitz non-differentiable functions) by allowing a derivative to be set-valued. It's called the Clarke gradient. The Clarke gradient of ReLU at 0 is the closed interval [0,1]. Note that the Clarke gradient doesn't satisfy the chain rule (except in a weakened form) which might seriously mess up some assumptions about autodiff. Is this generalised derivative useful in autodiff? I imagine that this is a largely theoretical tool that's useful in analysing algorithms but useless for actually computing things.
- medo-bear 5y agoi havent heard of clarke gradient before. convex analysis has something called subgradient [0], is it different? [0] https://en.m.wikipedia.org/wiki/Subderivative https://en.m.wikipedia.org/wiki/Subderivative
- ogogmad 5y agoThe subgradient in convex analysis is a special case of the Clarke gradient. The subgradient is precisely the Clarke gradient for convex functions. Convex functions are always locally Lipschitz except in weird cases. [edit] Question: Are there numerical applications in which the subgradient is actually computed, or is it a purely analytical tool?
- fault1 5y agoit's quite interesting how at least in ML, the transformer architecture has 'won out', at least for the time being, it appears to be everywhere these days: https://threadreaderapp.com/thread/1468370605229547522.html https://threadreaderapp.com/thread/1468370605229547522.html the advantage of transformers (computationally) seems to be how little sophistication the attention mechanism needs from AD systems (and how well it appears to scale with data). it's also a very static architecture in terms of a data flow/control flow perspective. as far as I understand, this is far different from systems needing to be modeled in continuous time, especially things like SDEs. I am curious if things like delay embeddings will ever be modeled in terms of mechanisms similar to attention however.
- jowday 5y agoOutside 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).
- xiphias2 5y agoIt will be interesting when Tesla and Waymo moves to transformer architecture, but as you wrote my guess is that it's not yet in production for vision tasks.
- jowday 5y agoI’m not sure they will, at least not with the research in the state it is presently. Researchers are interested in vision transformers because they’re competitive with CNNs if you give them enough training data - they don’t drastically outperform them. Right now switching over to them would require a ton of code changes, relearning intuitions, debugging, profiling, etc. for not a ton of benefit.
- xiphias2 5y ago
- secondcoming 5y agoThis [0] video has a quite interesting proposal to add AD to compilers. [0] https://youtu.be/1QQj1mAV-eY https://youtu.be/1QQj1mAV-eY
- eigenspace 5y agoIMO 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.
- sockfish 5y agoI often wonder why so much effort is being put into shoehorning everything into a single language. Wouldn't it make much more sense to use a fully differentiable DSL for machine learning / xla, then call it from whatever host language you use? This approach has worked really well for SQL for the past couple of decades.
- niklasd 5y agoHas 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.
- handzhiev 5y agoImo ORM is mostly a sign that (for some odd reason) many developers don't want to learn / use SQL. But what actual problem is ORM solving beyond that?
- baq 5y agoBoilerplate. Writing serializers and deserializers by hand is not an efficient use of developer time. Related to ORMs, but not quite on topic - query building. Type checked queries, parts of which can be passed around business logic, are very powerful and flexible.
- ithkuil 5y agoThere are more and more libraries that let you to write SQL and bind the results into native records (objects, structs) in the host language. I find it an interesting middle ground
- viraptor 5y agoThey solve 99% of your queries in much less time while allowing to drop down to SQL when you really want/need it.
- dnautics 5y ago
- carterschonwald 5y agoPart of the challenge is that most formulations of (reverse mode )auto diff wind up requiring extra runtime data structures for the backwards computation step. There’s been some great work in this space in the past 5 years. I’ve got some stuff I worked out this fall I’m overdue to write up and share some prototypes : there is a way to do reverse mode auto diff isolated to just being an invisible compiler pass! Without any of the extra complexity in what are otherwise equivalent formulations
- The_rationalist 5y agohttps://github.com/breandan/kotlingrad#coroutines https://github.com/breandan/kotlingrad#coroutines
- carterschonwald 5y agoThe coroutines link doesn’t work. And that aside it has a different strategy than the one I’ve been playing with ;)
- bmc7505 5y agoHey! Thanks for the link. We decided to postpone our coroutine-based implementation until native support for multi-shot delimited continuations becomes more stable. The linked section may be found here [1], for posterity. [1]: https://github.com/breandan/kotlingrad/tree/c02ac55325c05a2ecee6ec26b24bbd61c138e7e2#coroutines https://github.com/breandan/kotlingrad/tree/c02ac55325c05a2e...
- The_rationalist 5y agoKotlingrad make any autodiff library looks pale in comparison..
- jimsimmons 5y agoFrom a quick skim of it’s readme I don’t see what’s so special about it. Can you provide additional context?
- vkkhare 5y agoWhat do people think of automatic differentiation support facebook was trying for Kotlin. They called it differentiable programming https://ai.facebook.com/blog/paving-the-way-for-software-20-with-kotlin/ https://ai.facebook.com/blog/paving-the-way-for-software-20-...
- amcoastal 5y agoThank you so much for this post. I was going to have to self discover a lot of this information for an upcoming project, and now its laid right before me. Merry Christmas!!
- ogogmad 5y agoThis is possibly irrelevant to practice: But is perturbation confusion a problem which affects reverse-mode autodiff as well as forward-mode autodiff? With dual-number based approaches, the usual solution is tagging, but this is known to be incorrect when higher order functions get used. For people who don't know, perturbation confusion is a problem that affects naive implementations of autodiff where derivatives of order bigger than 1 are not computed correctly.
- ChrisRackauckas 5y agoYou 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 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.