Keno Fisher did a presentation in a discussion moderated by Simon Peyton Jones here: https://www.youtube.com/watch?v=mQnSRfseu0c
Also combined with enzyme which can differentiate through static paths at the llvm bitcode level: https://enzyme.mit.edu/
I wonder what this will enable at the frontier of what is computationally tractable in the combo of physics + ml.
It seems like AD is a solved problem, which is not. I can personally think two instances in which I have hand rolled my own gradients: in the context of recommender systems (dense matrices are evil) and right now in the context of real time collision detection (memory allocations are evil).
One of the domains where I hope Julia will excel is precisely as you comment on physics + ml (I check frequently if MuJoCo source code is available at last).
Anyway, I encourage everyone with any background in scientific computing or DS to have a look at Julia. The ecosystem is nowhere near Python yet, but the language is very good and the tooling is getting better. The performance Julia provides without interfacing C/C++ or Fortran is not just a convenience, it has architectural consequences. It's not about of coding faster, it's about coding "further".
Is that a compliment or an insult? When Haskell is involved I can never tell.
Let me explain by quoting these paragraphs Roger Ebert's review[1] of 24h party people:
> As the film opens, Wilson is attending the first, legendary Sex Pistols concert in Manchester, England. (...) Wilson is transfixed by the Pistols as they sing "Anarchy in the U.K." and sneer at British tradition. He tells the camera that everyone in the audience will leave the room transformed and inspired, and then the camera pans to show a total of 42 people, two or three of them half-heartedly dancing in the aisles.
Sounds like the average language designer entranced and inspired by their first time grokking Haskell.
> Wilson features the Pistols and other bands on his Manchester TV show. Because of a ban by London TV, his show becomes the only venue for punk rock. Turns out he was right about the Pistols. They let loose something that changed rock music. And they did it in the only way that Wilson could respect, by thoroughly screwing up everything they did, and ending in bankruptcy and failure, followed by Sid Vicious' spectacular murder-suicide flameout. The Sex Pistols became successful because they failed; if they had succeeded, they would have sold out, or become diluted or commercial. I saw Johnny Rotten a few years ago at Sundance, still failing, and it made me feel proud of him.
I could rephrase that last sentence at "I checked out a contalk by Simon Peyton Jones from a few years ago, still avoiding success at all costs[2], and it made me feel proud of him" and it would be absolutely true.
And no, I also did not expect to find a parallel between Haskell and punk music, but there you go.
[0] https://openculture.com/2015/06/the-sex-pistols-1976-manches...
[1] https://www.rogerebert.com/reviews/24-hour-party-people-2002
And instead of people going berserk over this, all he got was some friendly emails from people who just wanted to inform him they bumped into this problem, but it's fine, they had a workaround. Because Haskell people in those days were used to Breaking changes with a capital B.
I think that's at least a little bit punk
Which makes all the sense because Haskell is about avoiding bugs at compile time
Scala (experimenting with various styles) comes to mind.
How about proto punk? Something like ML?
These statements set the expectations really high. Yet they omit the ugly truth of increasingly slow compilation times, pure compilation caching, JIT which might affect overall performance due to type instability.
Yeah, like Haskell has a nice Quick-sort example prominently displayed at the front page, Julia is sleek for scientific computing, yet it can bite you really bad as soon as you start wrapping the developed model into a real world REST app.
Do not get me wrong, I see some of the benefits of Julia. On the other hand, I do not think this uncritical hype about "coding 'further'" does good for the language.
https://github.com/SciML/DifferentialEquations.jl/issues/786
There may well have also been improvements in the Julia compiler.
Somebody already did a comparison of DifferentialEquations.jl's standard usage against MuJoCo and DiffTaichi and showed that it outperformed them by about 5x-10x.
https://arxiv.org/abs/2012.06684 https://homes.cs.washington.edu/~thickstn/ctpg-project-page/...
That's all showing the raw iteration count to show that it algorithmically is faster, but the time per iteration is also fast for many reasons showcased in the SciMLBenchmarks routinely outperforming C and Fortran solvers (https://github.com/SciML/SciMLBenchmarks.jl). So it's excelling pretty well, and things like the automated discovery of black hole dynamics are all done using the universal differential equation framework enabled by the SciML tools (see https://arxiv.org/abs/2102.12695 for that application).
What we are missing however is that, right now these simulations are all writing raw differential equations so we do need a better set of modeling tools. That said, MuJoCo and DiffTaichi are not great physical modeling environments for building real systems, instead we would point to Simulink and Modelica as what are really useful for building real-world systems. So it would be cool if there was a modeling language in Julia which extends that universe and directly does optimal code generation for the Julia solvers... and that's what ModelingToolkit.jl is (https://github.com/SciML/ModelingToolkit.jl). That project is still pretty new, but there's already enough to show some large-scale models outperforming Dymola on examples that require symbolic tearing and index reduction, which is far more than what physical simulation environments used for non-scientific purposes (MuJoCo and DiffTaichi) are able to do. See the workshop for details (https://www.youtube.com/watch?v=HEVOgSLBzWA). And that's just the top level details, there's a whole Julia Computing product called JuliaSim (https://juliacomputing.com/products/juliasim/) which is then being built on these pieces to do things like automatically generate ML-accelerated components and add model building GUIs.
That said, MuJoCo and DiffTaichi have much better visualizations and animations than MTK. Our focus so far has been on the core routines, making them fast, scalable, stable, and extensive. You'll need to wait for the near future (or build something with Makie) if you want the pretty pictures of the robot to happen automatically. That said, Julia's Makie visualization system has already been shown to be sufficiently powerful for this kind of application (https://nextjournal.com/sdanisch/taking-your-robot-for-a-wal...), so we're excited to see where that will go in the future.
We had a paper that included MuJoCo (with custom defined adjoints) within Julia's DiffEq framework to learn continuous control policies (https://arxiv.org/abs/2012.06684) which revealed a whole host of other issues, namely that gradients are not all that helpful when your optimization landscape is insanely non-linear. But that's a different problem.
Open to chatting about; Julia is definitely the tool for this kind of work.
Did you try multiple shooting? https://diffeqflux.sciml.ai/dev/examples/multiple_shooting/
I want to restructure the docs because I have come to think that any non-trivial fitting problem requires multiple shooting in order to be stable. Otherwise the loss becomes dominated by "loss values calculated beyond where the simulation has already fallen off". So I'd like to dig into some of the non-fitting examples and see if this or some other tricks are the right answer.
Now some uninteresting story: the reason to have a look at MuJoCo is personal: I studied Mechanics I, II and Analytical Mechanics at college and thought non-continuos mechanics a solved problem. When doing my masters thesis about helicopter simulation I needed at some point to model ground contact and I thought "easy, just add a damped spring". I feel that thesis was nice for an undergrad overall but that part as you can imagine was a complete and utter piece of crap. And now at my forties I want to have vengeance.
basically the big takeaway that i walked away with was that they defined a sort of pseudoclosure in the guts of the language before the compiler which relaxes some of the behaviors of a closure (ie what gets captured) which then enables the creation of larger optimizable regions for the compiler in the context of doing autodiff.
the demo by matt bauman where they do autodiff on a function that prompts for and has the user type in the name of another julia mathematical function was really impressive!
Yes! That's essentially correct.
Reading the front page it seems to focus on efficient higher order derivatives, is that it?
Essentially, what Zygote does is insert itself into the compiler pipeline at the lowered code stage to perform the AD transform. That is it operates on the form of Julia code before we perform any type analysis, devirtualization or optimizations. Essentially, you can think of Zygote as a lisp-style macro applied automatically/dynamically to every function being called starting at a particular entry point. This works pretty well, but has a couple limitations.
The first is that it has absolutely no semantic information available (since it operates on non-inferred code), so it can't use that for optimizations or things like data layout planning, which are important optimizations for a production-grade AD system. Essentially, it's not allowed to know that the `+` symbol it sees is actually the `Base.+` function that does addition, so must make the most pessimistic assumptions. This issue doesn't actually show up that much in machine learning use cases, but it's a big issue when you need to AD scalar code (which happens frequently in various differentiable programming contexts).
The next issue with running at this stage is that a lot of existing julia code is written with some mental model of the capabilities of type inference and the optimizer. For example, destructuring code like `a, b = f(x)`, people don't think about at all, but semantically that allocates several tuples and then indexes into them to take them apart. By running the AD transform, you basically double (at least), the complexity of every operation, so in a number of cases patterns that used to completely optimize away are now terribly slow, because they are no longer optimized (and then AD transformed on top of that).
A related issue is that because you cannot interleave optimization with the AD transform, if you want to perform nested (i.e. higher order) differentiation, you're gonna get exponential code generation, which you then have to hope the optimizer will cut down again for you, which it often can't all the way, but even if it could, generating exponential code in the first place is bad, because you're gonna wait an exponential amount of time for the compiler to be done with it (and super-exponential in practice, because the compiler is not linear in the size of the input problem), so Zygote is basically unusable for anything beyond second order (and even at second order is a struggle).
Lastly, there is also a technical issue, which is that code that operates on lowered IR isn't technically allowed to build any closures, but the AD transform has to do that in order to put the code for the backwards pass somewhere. Zygote gets around this by taking advantage of the fact that the AD transform is pure, so it basically runs everything twice (once to generate the forward pass, once to generate the backwards pass when execution gets to it) and it knows that things match up because the input code is the same. That mostly works, but this dependency isn't visible to the runtime system, so you can get into issues where code is updated in between the forwards and the backwards pass, which breaks Zygote in all sorts of ways.
Anyway, Diffractor is designed to fix all of this by "simply" moving the AD transform stage from lowered IR to post-inference (where the optimizer sits). The issue is that Julia the language, currently doesn't really allow semantic transformations post inference (after all optimizers are supposed to make things faster, but not change the outcome of things) and in particular, running optimizers is always optional and the language may choose not to do it. So to get there, we need to do a couple of things:
1. We need to have some sort of wedge into the semantics of the language that allows for optimization-time changes that are semantic. For this, I added `OpaqueClosures`, which are essentially like regular closures, except that they do not have semantics as to their capture lists or the code that they run. Now, this may be confusing to people from some functional languages where all closures have such semantics (you can see SPJ ask this question in the linked discussion we had - didn't think about it ahead of time, because I'm so used to our semantics, so my answer was a bit muddled), but essentially in Julia, the contents of closures is semantically visible and optimizations over closure boundaries are mostly prohibited, because ordinary closures do not close over the world age (i.e. if you have `f = x->sin(x)` and then say `sin(x) = 9999`, then the next execution of `f` will return 9999). So opaque closure basically change this and say "nothing in the system is allowed to look at the code they contain, or the capture list, and also we capture the world age". More importantly though, we now have a datastructure (the capture list of an opaque closure) that is allowed to be changed by the optimizer, so we drive basically drive a truck through that.
2. We need the actual mechanism to move the transform to inference time. This isn't super hard, but we haven't quite finished this yet, and it's something I'm hoping to get to over the next couple of months. It's somewhat intertwined with making the compiler in general more accessible to packages outside the core language. There's about 8-10 different packages that want to do compiler-y things on Julia IR and we really need to figure out a good solution to make that generally possible. But as always, designing core language APIs is a bit tricky, because you're gonna be stuck with them for a while.
There's also some other nice bits and pieces in Diffractor. The big one is that I did a fair bit of theoretical exploration at the top of this project to really understand higher order AD. When I first took differential geometry (10 years ago now), I used to joke with my classmates that I had no idea what a second derivative was, because all the textbooks basically just say "look, the tangent bundle of a smooth manifold is a smooth manifold" and never actually really go into higher order derivatives at all. Anyway, I really sat down to work all that out and in the course of that I came across a way to do higher order derivatives more efficiently (under suitable assumptions of what the compiler does) than just nesting the first order transform. In retrospect, it seems pretty obvious, but I can honestly say that I didn't think about that until I worked out the theory. In the course of it, I also managed to make very precise the notion that reverse-mode AD and pullbacks of cotangent vectors are the same thing. This connection was pretty well known in the oral tradition, but I never really saw a convincing writeup of it (and my result extends to higher orders of course).
Other than that, it directly uses ChainRules.jl as its AD rule system, where Zygote used to keep its own rules (Zygote and ChainRules developed concurrently, and Zygote was later adapted to look in both places, but its a bit of a mess) and it also has a unified forward mode that (in theory, once it's robust enough), subsumes both ForwardDiff.jl and TaylorSeries.jl.
In theory, this should all make for a pretty good AD system, but there's a fair bit of work still to be done to make it robust (though all the pieces except the julia-level mechanism for actually moving the AD pass are done and tested) and as always, time is very limited :).
How broad will Diffractor's support of the language be? E.g. will it support mutation / exceptions / recursion?
As for exceptions and recursion, shouldn't be a problem, just needs to be implemented.
I am sure they take a huge amount of work but my appetite for them has never been satisfied!
I've always thought that AD needs something akin to a "compile" step.
- https://internals.rust-lang.org/t/automatic-differentiation-...