3 ms·
I remember reading an early S4TF manifesto talking about natural language processing, image processing, etc. thinking, that cannot be your audience because that
by ChrisRackauckas 5y ago
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 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.
- p1esk 5y agoMy impression from your comment is that you don't care that much about "standard" ML users. As a "standard" ML user (pytorch/jax), and a potential Julia user in the future, this is not what I like to hear.
- borodi 5y agoThe idea, I imagine, is to differentiate what the julia ML stack offers over what is already on python, if it offers the same thing, but without the funding from facebook or google, why bother switching? It has to offer something more.
- ChrisRackauckas 5y agoIf the purpose of the AD was to do something simply for standard ML workflows, the current AD tools are not the right design for that. They are too complex and solve a harder problem than they would need to. A better approach would be to use the abstract interpretation afforded by Symbolics.jl, mixed with the array symbolics, and use MetaTheory.jl to define simplification rules similar to XLA. You'd essentially get a slightly expanded Jax/TensorFlow with graph simplification rules that could be adjusted/improved directly from the host language. There's a prototype (ReversePropogation.jl), but it needs array support to actually be useful for this application. Or if you don't need the whole stack to be modifiable from Julia, a better interpreter on XLA.jl would get you there. A sufficiently decent lone coder could get those up and optimized fairly quickly for standard ML applications, if that's the goal. Diffractor.jl has a much loftier goal: optimized differentiable programming of any code from any package in the Julia ecosystem. Because it's building typed IR, it will need a full set of Julia-based analysis tools (escape analysis, loop-invariant code motion, etc.) to approach the amount of optimization XLA can do when XLA optimizations are applicable. While such passes are being developed (for example, this is the PR for putting immutable array optimizations into the language so that Diffractor-friendly immutable can generate the optimized mutable form: https://github.com/JuliaLang/julia/pull/42465 https://github.com/JuliaLang/julia/pull/42465), it's at least a few years away before it's doing something like reliably combining multiple matrix-vector products into a matrix-matrix BLAS3 call. That would put it on even footing today to compete against PyTorch in a kernel vs kernel optimization battle, but not against TensorFlow code in cases where XLA optimizations are doing something more. "Just wait 3 years and it will be really cool" is not a good way to start building a robust community, instead those interested in it need to ask how to demonstrate the improvements afforded by the added generality today. That's why what's making the project tick right now is the ARPA-E projects for physics-informed neural networks, the DJ4Earth project to do direct differentiation of the CLIMA climate model without changing any of the model code (https://dj4earth.github.io/ https://dj4earth.github.io/), etc. Those kinds of projects are what is keeping a lot of the dev team open to be full time on these AD and compiler optimization projects. But if successful, it will also give an AD that is great for standard ML.