Augmented neural ODEs

[ChrisRackauckas]

https://arxiv.org/abs/1904.01681

[avik-pal]

Regarding adding a layer for this work, I was wondering if we need to have a new layer AugmentedNeuralODE. Or should we simply modify the NeuralODE type to store a parameter augment_dim, and dispatch based on that to use the current version if it is 0.

[ChrisRackauckas]

Yeah, that's a good question. A keyword argument for the number of augmentation dimensions is probably all that's needed, and is easy to generalize to all of the other NDE layers

[avik-pal]

Another simpler alternative might be to define a layer like

struct AugmentedNDELayer{DE<:NeuralDELayer} <: NeuralDELayer
    nde::DE
    adim::Int
end

This should allow the augmentation to work with all the existing layers, without modifying them at all.

[ChrisRackauckas]

But I think you still might need to modify the equations per-DE? Give it a try and see if it can work.

[avik-pal]

I see. I tried the Concentric Circles experiment from the paper (using NeuralODE) and it seems to work as expected.

Neural ODE

Augmented Neural ODE

I will try out the other DE layers, and open a PR soon.