Neural ODE Processes

Abstract: Neural Ordinary Differential Equations (NODEs) use a neural network to model the instantaneous rate of change in the state of a system. However, despite their apparent suitability for dynamics-governed time-series, NODEs present a few disadvantages. First, they are unable to adapt to incoming data-points, a fundamental requirement for real-time applications imposed by the natural direction of time. Second, time-series are often composed of a sparse set of measurements that could be explained by many possible underlying dynamics. NODEs do not capture this uncertainty. In contrast, Neural Processes (NPs) are a new class of stochastic processes providing uncertainty estimation and fast data-adaptation, but lack an explicit treatment of the flow of time. To address these problems, we introduce Neural ODE Processes (NDPs), a new class of stochastic processes determined by a distribution over Neural ODEs. By maintaining an adaptive data-dependent distribution over the underlying ODE, we show that our model can successfully capture the dynamics of low-dimensional systems from just a few data-points. At the same time, we demonstrate that NDPs scale up to challenging high-dimensional time-series with unknown latent dynamics such as rotating MNIST digits. 

Reviews and author responses:

We are pleased that we have been able to provide satisfactory answers to your questions and would like to thank you again for the helpful feedback! 

The proposed NDP has two main advantages: 1- it has the capability to adapt the incoming data points in time-series (unlike NODE) without retraining, 2- it can provide a measure of uncertainty for the underlying dynamics of the time-series. NDP partitions the global latent context $z$ to a latent position $l$ and sub-context $z^\prime$. Then it lets $l$ follow an ODE, called latent ODE. This part is actually the innovation of the paper where by defining a latent ODE, the authors take advantages of ODEs to find the underlying hidden dynamics of the time-series. This assumption helps find better dynamics when the generating processes of time-series meet some ODEs. Then the authors define a stochastic process very like the idea from Neural Processes (NP) paper, that is, by defining a latent context $z$ (which here is a concatenation of $l$ and sub-context $z^\prime$) with a prior p(z) and integrating a Gaussian distribution of a function of $z$ (decoder $g(l,t,z^\prime)$ which is a neural network) over $z$. 

Overall, I liked the idea of the paper and how the authors integrate two important concepts, i.e. NODE and NP, into a single framework, which could be useful in many real-world time-series with complex underlying dynamics. However, I have some questions regarding some points in the paper:

1- The paper says that $z$ is split into two parts: $l$ and $z^\prime$, where $z^\prime$ is kept unchanged over time and only $l$ follows an ODE. I wonder why is this the case? How many dimensions should $l$ have? How does the dimension of $l$ affect the results? Why not let the whole $z$ follow an ODE? There are no explanations and clarifications for these in the paper.

2- There is no mention of how $z^\prime$ should be learned. In general, there is no mention on how to train the NDPs. It is unclear in the paper what loss function should be optimized and how the latents should be learned. If it is by variational methods, how the posteriors of $z^\prime$ and $l$ should be learned? I believe the authors should augment on these in the paper, otherwise it is very hard to know how the NDPs should be trained. 

3- What is the dimension of $l$ used for rotating MNIST experiments? Why NDP is able to extrapolate well when there is variable angular velocity and angular shift (Fig. 5) and fails to extrapolate when there is  constant angular velocity (Fig. 4)? It seems the second is an easier task and I wonder why NDP has a poor performance? Does it imply that NDP can only work well in a specific conditions?

4- typo: page 3: the the decoder --> the decoder 

########## Edit ##########

The authors have addressed all my questions. Thanks.

We are glad our reply has addressed your concerns. Thank you for updating your score!

We are happy you have found our reply satisfactory. Thank you for updating your score! 

This paper proposes a new algorithm that can adapt incoming data-points by applying Neural Ordinary Differential Equations (NODEs) to Neural Processes (NPs). It combines two algorithms properly and showed better performance than NPs through ODEs in the encoding, even with a smaller number of parameters. 

Strengths:

1) They properly combined NODEs and NPs to fast-adapt few data points from underlying ODE over ODE distributions.

2) They showed their algorithm outperforms NPs through ODE encoding with fewer parameters.

3) They analyzed several variations like Second-Order Neural ODE Process or Latent-only version. 


Weaknesses:

1) Task details are not clearly described. I checked the appendix also, but they just mentioned: "with varying gradients and amplitudes and shifts...". 

2) Lack of comparison with previous works: For instance, one of the advantages of this work is good interpolation and extrapolation. Convolutional Conditional NP (Conv-CNP, Jonathan Gordon et al., 2019) also outperformed other NPs methods for extrapolation, but they didn't compare Conv-CNP as one of the baselines. For the rotated MNIST experiment, Sequential Neural Processes (SNPs, Singh et al., 2018) isn't compared.


The correctness of their claim and Clarity:

This paper is well written and almost correct, but the details about the experimental setting look missed.


Additional feedback:

Thank you for submitting it. I enjoyed reading it. I think that it is a well-written paper and deserved sharing in our community. However, detailed information (e.g., task details) is not clearly described, and some comparison results are missed. By updating those things, it will be more concrete. For the rotated MNIST experiment, evaluating the version applying NODEs to SNPs could be interesting also.

Minor things are

On page 5, 
"....Additional details for every task considered can be found in C." -> "Additional details for every task considered can be found in Appendix C."
Secondly, as is seen in A, NDPs train faster in a faster wall clock time than other variants. -> Secondly, as is seen in A, NDPs train faster in a wall clock time than other variants.

On page 7,
we show the mean squared errors (MSE) for the 4th rotated MNIST digit in Table 2. -> what is the meaning of the 4th rotated MNIST?

#####EDIT#####
I agree that the author's disagreement with my second comment and thank you for the update. I change my rate to 7.

Dear Reviewers, 

We have uploaded a new version of our manuscript and made our code available with the supplementary. We have incorporated into this new version all the feedback that we have received. We believe that your comments have visibly improved the quality of the manuscript. A complete changelog can be found below.

Changes to the manuscript:
- Added a new subsection called *Learning and Inference*, in which we include additional details about the training procedure and give our ELBO loss (R1, R2).
- Added a full derivation of the ELBO in Appendix A (R1, R2).
- Included pseudocode for the training procedure in Appendix B (R1, R2).
- We have included a graphical model description of NDPs in Section 3.1, which we hope will further clarify the generation and inference procedures of the model (R1, R2).
- Added further details about the tasks and the architectures that were used in the experimental section (R3).
- Added full task details for the 1D regression tasks in Appendix G.1 (R3).
- Added full task details for the Lotka-Volterra task in Appendix G.2 (R3).
- Added a detailed description of the architecture used in the low and high-dimensional experiments in Appendix F (R2, R3).
- Reformatted Table 1 to make the data easier to read. 
- Added the motivation behind the split of $z = [l, z’]$ in Section 3.1 (R1).
- Added an ablation study for the size of the latent ODE state in Appendix E (R1).
- Added final MSE on the LV task for NPs and NDPs in the main text.
- Moved the stochastic process proofs to the supplementary (R2).
- Made the discussion section slightly more concise (R2).
- Added a time unit to the training time comparisons (Table 3) and a description of the system used for our experiments in Appendix D (R2).
- Many small fixes, including an overall ‘tightening up’ of our notation and language.


This work presents a new method that combines neural ODEs, which uses neural networks to flexibly describe a non-linear dynamical system, and neural processes, which uses neural networks to parameterize a class of functions. By combining these two approaches, neural ODE processes can now adapt to incoming data-points as the latent representation parameterizes a distribution over ODEs due to the NP component of the model. The authors use different variants of the model (first order ODE, second order ODE, linear latent readout) to infer dynamics from data in the experiments section. I find this work to be an interesting and important extension to the existing literature on neural processes. 

My primary qualms with the manuscript is that I found it difficult to glean some of the details about the model(s) and, in particular, the details of the inference procedure. I assume many of the same details in the original NP paper apply here, but it is not clear to what extent, and exactly how. Many of these important inference and model details seem to be missing from both the main text and the supplemental material. 

In particular, when you first discuss the aggregator some key details are missing. You mention that the mapping from r to z could be a NN but you are not clear on when/if the encoder is a neural network in the actual experiments. Also, is it the case that data from multiple contexts are trained in parallel? It is important to specify all of the details for the encoder for each of the experimental sections. The decoder and the other pieces of the model are clear.

Moreover, how exactly is this trained? SGD? Adam? Is it end to end?  I assume you are optimizing an ELBO and the inference methods are akin to a VAE (or the original NP paper), but it is not explicitly said anywhere. Stepping through or at least explaining the primary details of training the model and the training objective will be useful. 

Finally, it is unclear how long this inference takes or what kind of computing resources are needed. Though there are some comparisons of training different versions of the model in the appendix, there is no sense of how long an 'epoch' is. Because there was no code that I could see with the submission, this is doubly difficult to glean. 

I think the proofs of secion 3.2 could be moved to an appendix. Additionally, a lot of space is devoted to the discussion and the conclusion; I would rather see more clarity provided to the implementation of NDPs and their differences at every stage of the model across the experiments.

I am excited about the work and it does seem to be a useful extension of existing methods, and I think there are details that need to be clarified in order for this to be publishable. 

Minor details:
Bottom of page three "the the"



%%%%% EDIT %%%%%%

I am satisfied with the author's response and given the proposed changes will raise my score to a 7.



We would like to thank the reviewers for their thorough reviews, helpful comments, and actionable suggestions for improving our manuscript. We are of course delighted that the reviewers view our work favourably, and that they see our core innovation as interesting, important, applicable and worthy of sharing with the community. We are grateful that the reviewers have focused on concrete suggestions for improving the manuscript, particularly in the presentation of technical details, which we will be able to act upon within the discussion period and will surely improve the quality of the work. Furthermore, we appreciate the insightful questions that will help improve the clarity of our discussion.

We have responded to each review directly to address the specific comments they make. We will share an updated version of the manuscript, additional supplementary materials, and our codebase within the discussion period.

Again, we would like to express our gratitude for these constructive reviews and welcome further comments or requests.


We are grateful for the effort the reviewer has put into their feedback and we are convinced it will significantly improve the quality of our manuscript. We are pleased that the reviewer remarked on the impact our method can have in real-world applications, though we agree that the points raised by the reviewer deserve further clarification. We will release a new version of the manuscript during the discussion period to integrate this feedback as well as that of the other reviewers. We provide below a detailed response for each of the points that were raised.

**Clarification about what’s going on with L and z’**

As in NPs, $z$ captures the global uncertainty, i.e. uncertainty over functions. The purpose of the split is to factorize the global uncertainty in the dynamics into an uncertainty in the initial position (‘how things start’, given by $L(0)$) and an uncertainty in the ODE derivative function (‘how things change’, conditioned by $z’$). This inductive bias is intended to help the model adapt well to tasks where either the initial conditions or the way the system evolves is fixed. As a concrete example, consider the motion of pendulums: for a single pendulum of fixed length, the variation in trajectories is confined to the initial conditions ($L(0)$). If instead we aim to model the motion of pendulums of different length, then there is variation both in the initial conditions ($L(0)$) and the way the system evolves (determined by $z’$).

**How many dimensions should $l$ have? How does the dimension of $l$ affect the results?**

In general, the greater the dimensionality of $l$, the greater the range of dynamics that can be learned. This is the motivation behind Augmented Neural ODEs (Dupont et al, 2019), which append extra dimensions to the ODE. Extra dimensions were also shown to allow a Neural ODE to learn higher-order behaviour (Norcliffe et al, 2020). On the other hand, increasing the dimension permits overfitting. For the MNIST experiments, we found l = 40 to perform best in a (limited) hyperparameter search.

We will include an ablation study to show the effects of $l$ on performance for the Sine task.

**Clarification of the learning procedure. How should $z’$ be learned?**

We train end-to-end using an ELBO loss, similar to NPs. Specifically, we use the following lower bound for the log-likelihood of the target set $t_{m+1:n}$, $y_{m+1:n}$ given a context set $t_{1:m}$, $y_{1:m}$:

$\log p(y_{m+1:n} | t_{1:n}, y_{1:m}) \geq  E_{q(z|t_{1:n}, y_{1:n})}\Bigg[\sum_{i=m+1}^{n} \log p(y_i | z, t_i) + \log \frac{q(z|t_{1:m}, y_{1:m})}{q(z|t_{1:n}, y_{1:n})} \Bigg] $

As usual, $Z$ is the latent variable from the amortised variational inference procedure, and the approximate posterior over $Z$ is given by the encoder. As in NPs, during training, we sample context and target sets of different sizes such that the model can become sensitive to the size of the context. The size of these sets is drawn from a uniform distribution over $\{1, …, N\}$, where $N$ is the maximum context size.

We will add these details to the main text as well as pseudocode and a full ELBO derivation in the supplementary, and make the code available to reviewers together with the new manuscript.

**Why NDP is able to extrapolate well when there is variable angular velocity and angular shift (Fig. 5) and fails to extrapolate when there is constant angular velocity (Fig. 4)**

This is an interesting question. First, we note that NPs are not able to learn periodic functions that extrapolate indefinitely (a widely known result that has been formalised and explored recently by Ziyin et al. in Neural Networks Fail to Learn Periodic Functions and How to Fix It, NeurIPS 2020.) Neural ODEs however, are not constrained by this finding and are able to learn periodic functions (e.g. $(\dot{x},\dot{v}) = (v,-x)$ produces periodic motion in $x$).

For the original task (Rotating MNIST, Fig.4) the variation is only typographic, as the angular shift and velocity are fixed, so the dynamics are fixed. The NP model produces reasonable outputs over the first 16 frames because it’s learned a good approximation of sine over the interval $(0,2\pi)$ but does not extrapolate. As there is no variation in the dynamics, the NDP is not learning to model the function as being periodic and its extrapolation closely matches that of the NP i.e. stops rotating and becomes noisy after 16 frames. So we conclude that the NDP model has collapsed to something like the NP as a result of the task not having any variation in the dynamics.

In the case of the task we introduce (Fig.5) there is also variation over the dynamics. In this case the NP performs much worse, with the outputs being indistinct blurs. However, as the NDP is able to learn periodic functions, it is able to produce legible reconstructions, and, as the model is now being tasked with actually learning a distribution over periodic dynamics, the periodic extrapolation quality improves.


We would first like to thank the reviewer for the detailed and useful feedback they have provided. We were glad to see that the reviewer has appreciated the usefulness of our method as well as the experimental validation of its advantages. We completely agree that more information should be provided about the tasks. However, we respectfully disagree regarding the missing comparisons and our reasoning is provided below. We will integrate these changes in a revised version of the manuscript, which we will release within the discussion period. We provide below a detailed response for the individual points that were raised

**More task details needed**

We agree that further task and dataset details should be included. We will include both the used architectures and the dataset details. Additionally, our code will be included in our supplementary material. More details can be found below: 

*One-dimensional regression experiments*

Each task is based on a random function described by a set of parameters that are uniformly sampled. A trajectory example is formed by sampling from the parameter distributions and then sampling from that function at evenly spaced timestamps, $t$, over a fixed range to produce 100 data points $(t,y)$ per function. 

- Sines - $y=a\sin(t-b)$: $a$ range = (-1.0, 1.0), $b$ range = (-0.5, 0.5), $t$ range = (-$\pi$, $\pi$).
- Exponentials - $y = \frac{a}{60}\exp(t-b)$: $a$ range = (-1.0, 1.0), $b$ range = (-0.5, 0.5), $t$ range = (-1, 4)
- Straight Lines - $y=at+b$: $a$ range = (-1.0, 1.0), $b$ range = (-0.5, 0.5), $t$ range = (0, 5).
- Harmonic Oscillators - $y = a\sin(t-b)\exp(-0.5t)$: $a$ range = (-1.0, 1.0), $b$ range = (-0.5, 0.5), $t$ range = (0, 5)

For each task, we generate 490 such training sequences and use 10 for testing. To compute the standard error, each model was trained 5 times on each dataset, with a batch size of 5. 

*Lotka-Volterra*

To generate samples from the Lotka-Volterra system, we sample different starting configurations, $(u_{0}, v_{0}) = (2E, E)$, where $E$ is sampled from a uniform distribution in the range (0.25, 1.0). We then evolve the Lotka Volterra system parametrised by $(\alpha,\beta,\gamma,\delta)=(\frac{2}{3},\frac{4}{3},1,1)$. This is evolved from $t=0$ to $t=15$ and then the times are rescaled by dividing by 10. 

We train on a set of 40 trajectories using a batch size of 5. We test on 10 separate trajectories. To compute the standard error, we train the models on 5 different seeds. 

*Rotating MNIST*

For the rotating MNIST experiment, we follow the approach from ODE2VAE (Yildiz et al, 2019). We remove the digit corresponding to the 4th rotation from the dataset and use it as a test frame. Additionally, we remove four random rotation angles from each sequence in the dataset to simulate irregularly sampled data. At testing time, when we use a context set of size one (first part of Table 2), we supply only the first frame in the sequence (same as ODE2VAE). When we use a larger context set (the second part of Table 2), we supply seven randomly sampled frames. In both cases, we report the MSE for the predicted test frame over all the sequences.   


**Missing comparisons with ConvCNPs and SNPs**

We respectfully disagree that the work lacks comparisons with these models, primarily as the approaches are orthogonal to our own.

ConvCNPs aim to improve the way context sets are represented and encoded in the related family of Conditional NP models. Our method does not rely on a particular choice of encoder or representation. (In fact, we use a regular NP encoder in all of our experiments.) The methods (ConvCNP and NDP) are orthogonal, and the contributions from the ConvCNP paper could be employed in our model by simply replacing the encoder. In the interests of providing a clear exposition of the benefits of our approach, we opted to focus on improving directly upon the vanilla NP model, rather than including additional comparisons between baseline ConvNPs and ConvNDPs. However, we agree that it would be interesting to measure the potential synergistic effects of the two methods. NOTE: to clarify, the convolutional models we use for the image experiments are using regular convolutional encoders (i.e. a CNN instead of an MLP) and do NOT use a ConvCNP-like encoder. This is just an unfortunate choice of nomenclature. 

Sequential NPs (SNPs) are a different kind of model again, as, unlike our method, which models a single stochastic process (as theoretically proven in the paper), SNPs model a dynamically changing sequence of stochastic processes, each of which is not necessarily defined over time. Therefore, the SNP method naturally targets a different set of applications from ours. Still, SNPs are also orthogonal to our approach. For instance, one could consider a hybridization that captures a dynamically evolving sequence of NDPs.

We have included a discussion of these orthogonal directions in our related work section.


We appreciate the positive feedback and we are happy the reviewer has highlighted the clarity of the structure, the novelty of our method as well as its applicability to a broad range of real-world datasets. We will also make available during the discussion period an improved version of the manuscript, which will incorporate all the feedback we have received. 

We would like to thank the reviewer for the time they’ve put into this thorough and actionable review, which we believe will significantly improve our paper. We’re pleased that the reviewer appreciates the value in combining the NODE and NP frameworks, and we appreciate and agree with the various suggestions for improving the clarity of the manuscript. We will make a new version available within the discussion period that will address these comments as well as the suggestions from the other reviewers. We provide a detailed answer below for each of the points that were raised.

**Being more thorough about the inference procedure, training procedure, and encoder architecture**

We will significantly expand on these aspects in the main text and will include pseudo-code in the supplementary material for the inference and training procedures. We are also working on making our code available to the reviewers in the next few days.

*Encoder architecture*: the encoder is left unchanged in our method with respect to a normal NP. For the low-dimensional experiments, the encoder is an MLP that processes each $(t_i, y_i)$ pair. When $y_i$ is an image, the encoder is a convolutional network and we append $t_i$ (the time) to the convolutional features of the later layers. The embeddings $r_i$ for all $(t_i, y_i)$ pairs in the context are aggregated in a single embedding $r = mean(\{r_i\})$, which takes the mean of the individual embeddings. This $r$ is used to parametrise the normal distribution from which $z$ is sampled: $\mathcal{N}(Linear(r), Linear(r))$. 

*Training & inference procedure*: indeed, we train end-to-end using an ELBO loss similar to the ELBO loss from NPs. Specifically, we use the following lower bound for the log-likelihood of the target set $t_{m+1:n}$, $y_{m+1:n}$ given a context set $t_{1:m}$, $y_{1:m}$:

$\log p(y_{m+1:n} | t_{1:n}, y_{1:m}) \geq  E_{q(z|t_{1:n}, y_{1:n})}\Bigg[\sum_{i=m+1}^{n} \log p(y_i | z, t_i) + \log \frac{q(z|t_{1:m}, y_{1:m})}{q(z|t_{1:n}, y_{1:n})} \Bigg] $

We will add a full derivation of this ELBO loss in the supplementary and hope this will make our training procedure precise. As in NPs, during training, we sample context and target sets of different sizes such that the model can become sensitive to the size of the context. The size of these sets is drawn from a uniform distribution over $\{1, …, N\}$, where $N$ is the maximum context size. For optimization, we use RMSProp with the default PyTorch parameters. 

**More information about the computational cost**

Along with the ratios of times NDP/NP, we will also include the training times for NPs over 30 epochs for the 1D datasets in the supplementary material. These were between 20 seconds and 100 seconds. For this, we used a machine equipped with an Nvidia Titan XP GPU.

**Are different contexts trained in parallel?**

Yes. They are trained in parallel in the sense that we train over batches of sampled contexts from multiple time-series. In other words, our $t$ (time) batch has shape (batch_size, context_size, 1) and our corresponding $y$ batch has shape (batch_size, context_size, dim($y$)) . 

Because Neural ODEs do not support batching in the usual sense, we use the following computational trick. We consider a bigger ODE with a state of dimension [latent_size * batch_size], which concatenates the independent states of the ODEs corresponding to each dimension in the batch. That allows us to evolve all the independent ODE states in the batch concurrently. We integrate this extended ODE state over the union of all the time steps in the batch. This “batching” approach is common for Neural ODEs applied to (irregularly sampled) time-series (Rubanova et al., 2019)

**Proofs could be moved to an appendix. A lot of space is devoted to the discussion and the conclusion.**

We will move the proofs to the appendix as the reviewer suggested. We also agree that the discussion and conclusion can be more concise. We will use the space gained from these changes to clarify the aspects from above.


This paper proposes a new class of stochastic processes determined by a distribution over Neural ODEs. The overall structure of the paper is clear. I find the newly defined process interesting and applicable to many real data sets.

Meta-review: This work proposes a stochastic process variant that extends existing work on neural ODEs. The resulting method allows for a fast data-adaptive method that can work well fit to sparser time series settings, without retraining. The methodology is backed up empirically, and after the response period, the reviewers' concerns are sufficiently addressed and reviewers are in agreement that the contributions are clear and correct.

Decision: poster-presentations