Composability in Julia: Implementing Deep Equilibrium Models via Neural ODEs

https://julialang.org/blog/2021/10/DEQ/

117 48 5 years ago

rpmuller• 5 years ago

One of the best things about Julia is that people like Chris Rackauckas are developing great packages for it.

lytefm• 5 years ago

Definetely. I've just been working with Stan and Pyro + Python so far for modelling, but post like this encourage me to finally pick up Julia and get seroius with neural ODEs.

adgjlsfhk14• 5 years ago

This is a really cool idea, especially because it is a type of NN that can take more time for harder inputs. This makes it relatively unique since most types of NN have O(1) runtime, which is often nice, but puts limits on the types of problems they can solve.

jstx12• 5 years ago

What are some good references on neural ODEs that don't come from the Julia community? I'm looking for theory and applications - when are they good and who is using them for what?

I'm asking for sources outside of Julia because I find the coupling of algorithm types to tools kind of strange and the whole SciML trend is kind of opaque to me. (Are people applying ML as a solution to newer problems? Are they using new approaches to solve ML problems? How legit is the whole thing? I just don't know.)

UncleOxidant2• 5 years ago

There's some info from the Pyhon/Pytorch camp: https://towardsdatascience.com/neural-odes-with-pytorch-ligh...

I suspect neural ODE work was done in Julia earlier because it was easier given some language features and libraries. But there does seem to be some work on neural ODEs in Python/Pytorch.

adgjlsfhk14• 5 years ago

What do you mean by "the coupling of algorithm types to tools" (not judging, just curious).

jstx12• 5 years ago

More bluntly my question is if SciML is that good, why aren’t more people doing it yet? Why is it limited to a small group of Julia developers and packages?

(There are good possible explanations - it could be very new, have only niche applications, Julia is somehow uniquely suited for it etc. I don’t know)

ssivark• 5 years ago

The original Neural ODEs paper is quite readable, and by now there are loads of blog posts and even a few talks on the subject.

The basic idea is inspired by the “adjoint method” for ODE solving (so you don’t have to hold in memory all the intermediate layer outputs — which is otherwise necessary to compute the backpropagated gradient signal).

ChrisRackauckas4• 5 years ago

Yeah, though with the method described in that paper you do have to be very careful since it has exponential error growth with the Lipchitz constants of the ODE. See https://aip.scitation.org/doi/10.1063/5.0060697 for details. But that is generally the case in numerical analysis: there's always a simple way to do things, and then there's the way that prevents error growths. Both have different pros and cons.

orbifold5• 5 years ago

Neural ODEs are essentially a rebranding of adjoint sensitivity analysis, which has been around in various forms in established solver suites, such as Sundials, PETSC, etc. The machine learning community got a hold of it, cited one book and otherwise happily reinvented everything.

IlyaOrson• 5 years ago

That actual content of that work was good but very misleading with an excess of backpropaganda and a poor literature review. The training procedure makes sense as continuous time backprop but it is mostly a special case of adjoint sensitivity analysis. The use of a NN as defining an ODE system seems fair to be named Neural ODE, imho its a good name, although again it was not completely novel as the writing style through the paper makes it look.

ChrisRackauckas4• 5 years ago

Yes, that's mostly right (instead of PETSc put FATODE since PETSc TS Adjoint was only published in 2019, but it's based heavily on FATODE's techniques https://epubs.siam.org/doi/10.1137/130912335?mobileUi=0). And that's why the Julia tools were so ready for it: we already had adjoint sensitivity analysis (implemented for parameter estimation in systems pharmacology), so neural ODEs were a freebee. Similar to DEQs using the adjoint of a nonlinear solve: that means a DEQ is just a neural network inside of a Julia nonlinear solver and you call the gradients and you'll hit the right overloads that were originally implemented for parameter estimation of elliptic PDEs. There are definitely some aspects of the ML community that have muddied the waters so-to-speak.

The recent one that I found funny was the Second-Order Neural ODE (https://arxiv.org/abs/2109.14158). The second order adjoint is rather old, with a canonical implementation in Sundials (https://github.com/LLNL/sundials/raw/master/doc/cvodes/cvs_g...) based on a very good analysis of how to do second order adjoints fast (https://epubs.siam.org/doi/abs/10.1137/030601582?journalCode...). But what about putting a neural network in there? In Julia it's just composing forward-over-reverse to get the optimal adjoint that matches Sundials, so there's been a tutorial on it in DiffEqFlux since 2019 (https://diffeqflux.sciml.ai/dev/examples/second_order_adjoin...) with both Newton and Newton-Krylov methods. The rest of the optimizations from the paper follow using BacksolveAdjoint and relying on dead code elimination (DCE), which is a compiler pass that wouldn't exist in Python so I guess they have to care? I assumed that was too trivial to publish given that history of prior work, but somehow that paper got a NeurIPS Spotlight with "a novel computational framework for computing higher-order derivatives of deep continuous-time models". At this point I just kind of shrug though, I think it's a symptom of conference culture not giving people enough time to review the literature thoroughly so random things seem to seep through. Take it as one reason among many to treat any ML conference paper similar to an unreviewed preprint.

All of that together, that's why the Julia universe of tools has mostly been focusing on improvements and performance vs Sundials, PETSc, etc. since those are the real challengers. You can see that we spend our time benchmarking against Sundials all of the time in our stiff ODE benchmarks and recently started outperforming it with QNDF pulling 2x-5x wins in various ways (https://benchmarks.sciml.ai/html/Bio/BCR.html, see https://sciml.ai/news/2021/05/24/QNDF/). The stiff neural ODE paper describes an 3 ways to achieve an improvement on the Sundials adjoint in terms of complexity (https://aip.scitation.org/doi/10.1063/5.0060697), and our adjoint benchmarking paper shows how putting all of these various pieces together leads to about 2-3 orders of magnitude improvement over the naive CVODES adjoint (https://arxiv.org/abs/1812.01892) (against the method, the follow up of course will be against the direct wrapping).

What's funny though is that people then get upset when we show benchmarks against some of the Python tools, like >100x improvements in solver speeds on physical and biological problems (https://gist.github.com/ChrisRackauckas/cc6ac746e2dfd285c28e...) and adjoints (https://gist.github.com/ChrisRackauckas/4a4d526c15cc4170ce37...). I don't understand why anyone would be surprised though: we've spent years "competing" against the C and Fortran codes and only recently started pulling ahead due to the combined effort of a whole community, while the Python tools were just a few people with simple methods who never benchmarked against the previous tools. If they benchmarked enough they would see that we're not an outlier claiming to be 100x faster than everyone else, instead we're in and slightly ahead of the pack but they are the outlier that is 100x behind the whole group. Personally, I would require every paper to at least have a benchmark against Sundials (and/or its methods) as a baseline which is the standard we tend to hold.

cs7022• 5 years ago

Very cool. I have only one question:

Has anyone successfully applied DEQs to larger-scale cognitive tasks or benchmarks, as opposed to MNIST, which is a tiny trivial task by today's standards?

Think ImageNet-1000, COCO, LVIS, WMT language translation, ..., SuperGLUE. There's a long list of datasets and benchmarks that regular boring fixed-depth NNs tackle with remarkable ease these days.

Has anyone anywhere applied DEQs to any of those datasets / benchmarks?

avikpal1410• 5 years ago

MDEQ work applies DEQ to some of the large scale benchmarks you mention: https://arxiv.org/pdf/2006.08656.pdf

cs7022• 5 years ago

Thank you! The results don't look that great (e.g., EfficientNet models achieve greater accuracy on ImageNet-1000 with ~5x fewer parameters), but the works looks interesting and worthwhile. I'll take a look.

Tarrosion3• 5 years ago

I've been part of and/or following the Julia community since 2015, and Julia is my favorite programming language by a wide margin. Seems like every two months there's a new blog post, usually with some of these folks as authors, that describes...something...to do with ODEs, machine learning, neural ODEs, GPUs, adjoints, scientific machine learning, ...

I have never once been able to follow one of these blog posts. Seems like universally these posts have terrible curse of knowledge [1]. To be fair, I know only a light-to-moderate amount about machine learning and very little about differential equations, but still— I'm a long-term Julia fan, professional data scientist, mathy PhD, would hope that's at least table stakes. Maybe I just need to...try harder? I wonder how many people would be excited but are totally lost by these posts.

To that point, if anybody has a recommendation for a gentle introduction to these topics (preferably Julia), I'd be most appreciative.

[1] https://en.wikipedia.org/wiki/Curse_of_knowledge

agumonkey• 5 years ago

What's funny is how most of them are revolving around a similar culture of concepts (ODE, automated differentiation, analysis) which is very unlike mainstream computing.

ps: I need to read those adjoints articles.

uoaei4• 5 years ago

but still— I'm a long-term Julia fan, professional data scientist, mathy PhD, would hope that's at least table stakes.

I don't understand this part. Are you saying you're a "professional data scientist" with a "mathy PhD" who has a "light-to-moderate amount" of machine learning knowledge? How did you get the job?

I would expect anyone with a mathy PhD to understand ODEs and PDEs, and neural ODEs are commonly understood (by those who read the papers, where I think it is an appropriate assumption that people who want to understand this stuff would do) to be effectively infinite-depth neural networks where every layer represents the same function.

diskzero2• 5 years ago

Have you been in interview loops or worked with bread and butter data scientists performing common tasks? I am curious what your view of what most data scientists do day in and day out?

eigenspace• 5 years ago

This is a pretty shitty, non-constructive response. I think neural differential equations are not easy to wrap one's head around even if you have a solid understand of deep learning and differential equations.

Sure, if they spent a lot of time wading through the literature they'd probably understand fine, but the point they were making was that the post was quite unapproachable without having delved into the specific literature on neural differential equations.

I think this is a reasonably valid complaint, and does not warrant you implying that they don't deserve their job.

Tarrosion3• 5 years ago

Are you saying you're a "professional data scientist" with a "mathy PhD" who has a "light-to-moderate amount" of machine learning knowledge? How did you get the job?

Indeed! My job involves approximately zero machine learning, at least for a narrow or stereotypical definition of machine learning. I work on optimization, domain-specific models inspired by queueing theory, various phsyically-motivated structural models, writing production code for data-heavy products, testing hypotheses in data, brainstorming how existing data can be used to solve new customer problems, writing documentation, communicating with customers, etc.

If you think this isn't data science but know of better nomenclature: please share! I've struggled to write a great job posting for this kind of work, and surely improved nomenclature would help :)

jimsimmons• 5 years ago

I consider myself to be reasonably good at ML and work for one of the cool labs and this stuff is impenetrable to me. Stuff like this are deceptively written and a cynic might say that they are primarily to show off rather than help.

ViralBShah2• 5 years ago

The SciML tutorials are a good start: https://github.com/SciML/SciMLTutorials.jl

And also the 18.337 lecture notes (probably also has videos available): https://github.com/mitmath/18337

gugagore2• 5 years ago

. For example, when we apply convolution filters on images the network consists of repetitive blocks of convolutional layers, and one linear output layer at the very end. It's essentially f(f(f(...f(x))...)) where f is the neural network, and we call this "deep" because of the layers of composition. But what if we make this composition go to infinity?

This really does not jive with my understanding. Each layer of, e.g. VGG 16 [1] does not implement the same function. Each layer has its own weights. There are certain architectures that tie weights across layers, but not all of them.

[1] https://neurohive.io/en/popular-networks/vgg16/

taeric2• 5 years ago

I think this works if you consider subscripted fs. That is, it is always the same arity, and the output is the same. So, would be better it they said f_n, where n is the layer if they network.

(I mean this as a question, but don't see an obvious place for a question mark...)

gugagore2• 5 years ago

well, I think it has to be a fixed f for the interpretation of a (time-invariant) dynamical system x_{n+1} = f(x_n) to work.

ChrisRackauckas4• 5 years ago

Correct, not all models have equivalent weights at each step. But it is an abstraction that needs to be done to send it to infinity.

savant_penguin• 5 years ago

15000x speed up compared to what?

adgjlsfhk14• 5 years ago

the previous implantation used by Nasa in I believe simulink. See https://www.youtube.com/watch?v=tQpqsmwlfY0 for more details.

snicker7• 5 years ago

Simulink -- probably the most popular modeling toolbox.

Composability in Julia: Implementing Deep Equilibrium Models via Neural ODEs

Related Stories