Hamiltonian Graph Networks with ODE Integrators

[ChrisRackauckas]

https://arxiv.org/pdf/1909.12790.pdf

Need a good name for it. Tracking it here, and we should make sure to have a good implementation of it that's easy to plug and play for "non-diffeq people".

[yuehhua]

Count on me! :smile: Graph network is introduced by Relational inductive biases, deep learning, and graph networks by Google. Graph network is composed of meta layer, which processes data in three layers in edge level, vertex level and global level, respectively. I will implement Meta layer in GeometricFlux, then integrate with ODE integrators.

[yuehhua]

OK, I think I have time to work on it. For meta layer, I will have Meta which is abstract type with three main abstract method update_edge, update_vertex and update_global. If DiffEqFlux is loaded, ContinuousMeta is loaded, which is a subtype of Meta but still abstract. I hope any DEs supported in DiffEqFlux could be used under this framework. Those update functions for graph are interfaces for users to implement their own neural network and plug in, so we have to put those update functions into a DE problem to make them work. Time is also an important parameters to put in, but I still don't have idea how to passing these arguments.

[yuehhua]

I prototyped the Meta layer in a good abstraction and tried to integrate with DiffEqFlux. Then I realized that I don't have to do that! Because GeometricFlux is an extension of Flux layer, it is already integrated by DiffEqFlux and it is fully composable! All users need to do is replace the regular deep learning layers into GeometricFlux's layers! Awesome Julia! All I have to do is to make a example for it and maybe I will need some data.

[yuehhua]

So far, Meta layer abstraction and message passing scheme is done in GeometricFlux. However, GeometricFlux utilize Flux with Zygote, so DiffEqFlux should also have Zygote support. Related issues: #52 and JuliaDiffEq/DiffEqSensitivity.jl#71.

[Zymrael]

https://arxiv.org/abs/1911.07532 "Graph Neural Ordinary Differential Equations" (GDE)

Sharing here our recent work on a generalized version of GNNs with a continuous depth domain. We have an implementation that relies on torchdiffeq and dgl (geometric deep learning library on top of PyTorch) https://github.com/Zymrael/gde ; however these models would benefit from a greater ODE solver selection (in particular stiff solvers) so we would like to have an implementation in DiffEqFlux for the experiments.

[Zymrael]

We're still in the process of familiarizing ourselves with Julia and Flux; we're also more than willing to help with the implementation or the model details if that would be helpful.

[ChrisRackauckas]

I think the main blocker right now is just that we need to update to Flux 0.10 with Zygote.jl. Once that's done then GeometricFlux should just work. We might need @MikeInnes 's help getting this completed though, hopefully in the next week or so.

[Zymrael]

That's exciting. With DiffEqFlux + GDEs together we would have a way to seamlessly incorporate both relational biases as well as partial knowledge of the underlying governing equations driving the system. Looking forward to tackling some interesting problems with this combination.

[ChrisRackauckas]

@yuehhua our release now uses Zygote. Let's give GDEs a go!

[Zymrael]

@yuehhua @ChrisRackauckas For what it's worth, we remain available to share successes/failures of GDEs and insights we collected in the last 3 months of playing with differential equation on graphs (all on PyTorch + torchdiffeq). This includes code if needed (mostly unreleased atm). We experimented on the standard node classification benchmarks (Cora, Pubmed, Citeseer) which I see are already available in GeometricFlux, as well as traffic forecasting and simple observation of multi-particle dynamical systems.

[ChrisRackauckas]

What's a good example problem to do with this?