Latent Neural ODEs with Sparse Bayesian Multiple Shooting
Authors: Valerii Iakovlev
Cagatay Yildiz
Markus Heinonen
Harri Lähdesmäki
Venue: ICLR 2023 poster (2023)

Abstract: Training dynamic models, such as neural ODEs, on long trajectories is a hard problem that requires using various tricks, such as trajectory splitting, to make model training work in practice. These methods are often heuristics with poor theoretical justifications, and require iterative manual tuning. We propose a principled multiple shooting technique for neural ODEs that splits the trajectories into manageable short segments, which are optimized in parallel, while ensuring probabilistic control on continuity over consecutive segments. We derive variational inference for our shooting-based latent neural ODE models and propose amortized encodings of irregularly sampled trajectories with a transformer-based recognition network with temporal attention and relative positional encoding. We demonstrate efficient and stable training, and state-of-the-art performance on multiple large-scale benchmark datasets.

Published as a conference paper at ICLR 2023
LATENT NEURAL ODES WITH SPARSE BAYESIAN MULTIPLE SHOOTING
Valerii lakovlev, Cagatay Yildiz! Markus Heinonen; Harri Lähdesmäki
ABSTRACT
Training dynamic models, such neural ODEs, on long trajectories is a hard problem that requires using various tricks, such as trajectory splitting, to make model training work in practice. These methods are often heuristics with poor theoretical justifications, and require iterative manual tuning. We propose a prin- cipled multiple shooting technique for neural ODEs that splits the trajectories into manageable short segments, which are optimised in parallel, while ensuring prob- abilistic control on continuity over consecutive segments. We derive variational inference for our shooting-basee latent neural ODE models and propose amortized encodings of irregularly sampled trajectories with a transformer-based recognition network with temporal attention and relative positional encoding. We demonstrate efficient and stable training, and state-of-the-ae performance on multiple large- scale benchmark datasets.
INTRODUCTION
Dynamical systems, from biological cells to weather, evolve according to their underlying mech- anisms, often described by differential equations. In data-driven system identification we aim to learn the rules governing dynamical system by observing the system for time interval [0, 7), and fitting a model of the underlying dynamics to the observations by gradient descent. Such optimi- sation suffers from the curse of length: complexity of the loss function grows with the length of the observed trajectory (Ribeiro et al., 2020). For even moderate the loss landscape can become highly complex and gradient descent fails to produce a good fit (Metz et al., 2021). To alleviate this problem previous works resort to cumbersome heuristics, such as iterative training and trajectory splitting (Yildiz et al., 2019; Kochkov et al., 2021; HAN et al., 2022; Lienen & Günnemann, 2022).
The optimal control literature has a long history of multiple shooting methods, where the trajectory fitting is split into piecewise segments that are easy optimise, with constraints ensure continuity across the segments (van Domselaar & Hemker, 1975; Bock & Plitt, 1984; Baake et al., 1992). Multiple-shooting based models have simpler loss landscapes, and are practical fit by gradient descent (Voss et al., 2004; Heiden et al., 2022; Turan & Jäschke, 2022; Hegde et al., 2022).
Inspired by this line of work, we develop shooting-based latent neural ODE model (Chen al., 2018; Rubanova et al., 2019; Yildiz et al., 2019; Massaroli et al., 2020). Our multiple shooting formulation generalizes standard approaches by sparsifying the shooting variables in a probabilistic setting to account for irregularly sampled time grids and redundant shooting variables. We further- more introduce an attention-basee (Vaswani al., 2017) encoder architecture for latent neural ODEs that is compatible with our sparse shooting formulation and can handle noisy and partially observed nigh-dimensional data. Consequently, our model produces sate-of-the-an results, naturally handles the problem with long observation intervals, and is stable and quick to train. Our contributions are:
We introduce latent neural ODE model with quick and stable training on long trajectories.
We derive sparse Bayesian multiple shooting a Bayesian version of multiple shooting with efficient utilization of shooting variables and continuity-inducing prior.
We introduce transformer-based encoder with novel time-aware attention and relative positional encodings, which efficiently handles data observed at arbitrary time points.
Aalto University, Finland. Conresponding author: valeri i. wwk.vleveaasts............. fi.
University of Tübingen, Germany, Code: butps:/gituub.co//yakkoee///

---

Published as a conference paper at ICLR 2023
1
t
S1
21
2
t2
y (81)
222
t3
t3
94ho (82)
t4
Encoder
Decoder
14
$2
07.
275
t5
t5
Latent Space
Figure 2: Method overview with two blocks (see Section 3.1). The encoder maps the input sequence
91:5 observed at arbitrary time points f155 to two distributions Pv, ($1), 96 (82) from which we sample shooting variables 81 82. Then, 81, 82 are used to compute two sub-trajectories that define the latent trajectory #1:5 from which the decoder reconstructs the input sequence.
2 PROBLEM SETTING AND BACKGROUND
Data. We observe a dynamical system at arbitrary consecutive time points f1.N = (f1...., tx), which generates observed trajectory y1:N (y11.... yN)) where y: yt,) E RD. Our goal to model the observations and forecast the future states. For brevity we present our methodology for single trajectory, but extension to many trajectories is straightforward.
Latent Neural ODE models. L-NODE models (Chen et al., 2018; Rubanova et al., 2019) relate the observations y1:N to latent trajectory U1:N (æ1, .... xNN), ). where 2; == (t,,) E Rd, and learn dynamics in the latent space.
N-30
Loss
An NODE model is defined as:
12006,
Reration
2; ODEsolve(xx
= 2,..., N, (1) SQ
y: [x; p(y (90 (a:)),
i= 1, N. (2)
Variable æ1 is the initial state at time t1. Dynam- ies function feo. is the time derivative of =(t), and ODEsolve (x t1, ti feom) ) is defined the solution of the following initial value problem at time t;:
Train loss of L-NODE model using iterative training heuris- tic. We start training on a short trajec- tory (N = 10), and double its length every 3000 iterations. The training fails for the longest trajectory. Bot- tom: 1-D projection of the loss land- scape around the parameters which the optimizer converged for a given trajectory length. Complexity of the loss grows dramatically with N.
Figure 1: Top:
dxx(t) dl
Sa (1, x((), ax11) = 21, L (1, (3)
Decoder gaw: maps the latent state æ; to the parameters of 199 (x. Dynamics and decoder functions neu- ral networks with parameters Bayn and f In typical appli- cations, data is high-dimensional whereas the dynamics are modeled in low-dimensional latent space, i.e., d < D.
L-NODE models are commonly trained by minimizing loss function, e.g., evidence lower bound (ELBO), via gradient descent (Chen et al., 2018; Yildiz et al., 2019). In gradient-based optimization complexity of the loss landscape plays a crucial role in the success the optimization. However, has been empirically shown that the loss landscape of L-NODE-like models (i.e., models that com- pute latent trajectory I1:N from initial state x1) is strongly affected by the length of the simulation interval [t,, tw] (Voss et al., 2004; Metz et al., 2021; Heiden al., 2022). Furthermore, Ribeiro et al. (2020) show that the loss complexity in terms of Lipschitz constant can grow exponentially with the length of (1,, 4x]. Figure 1 shows example of this phenomenon (details in Appendix A).
2

---

Published as a conference paper at ICLR 2023
3 METHODS
In Section 3.1, we present our latent neural ODE formulation that addresses the curse of length by sparse multiple shooting. In Section 3.2 we describe the generative model, inference, and forecasting procedures. In Section 3.3 we describe our time-aware, attention-baser encoder architecture that complements our sparse multiple shooting framework.
3.1 LATENT NEURAL ODES WITH SPARSE MULTIPLE SHOOTING
Multiple shooting. A simple and effective method for solving opti- misation problems with long simulation intervals is split these in- tervals into short, non-overlapping sub-intervals that are optimised in parallel. This the main idea of a technique called multiple shooting (Hemker, 1974; Bock & Plitt, 1984). To apply multiple shooting to an L--ODD model we introduce new parameters, called shooting variables, $1:N ($1...., SN-1) with 8; IRd that correspond to time points 11 and redefine the L-NODE model as
22
U4
111
32
$3
04
S1
S2
(4) Figure 3: Top: Trajec- tory over [t, fa], 2; is (6) computed from æ1. Bot- tom: [t1, f4] is split into three sub-intervals, computed from 8:-1.
C; (i ti foon):
(5)
y x ~ plvlloe.. (z:)).
The initial state is represented by the first shooting variable 81, and the latent state x; is computed from the previous shooting variable 8;-1- This gives short simulation intervals 4.). which greatly reduces complexity of the loss land- scape. Continuity of the entire piecewise trajectory is enforced via constraints on the distances between o; and s; (see Figure 3), which discuss in Section 3.2. Multiple shooting leads to new optimisation problem over Oayn: Odec, and 81:N11.
is
Sparse multiple shooting. Multiple shooting assigns shooting vari- able to every time point (see Figure 3). For irregular or densely sampled time grids this approach might result in redundant shooting variables and excessively short and uninformative sub-intervals due to high concentra- tion of time points in some regions of the time grid.
82
fy
We propose to fix these problems by sparsifying the shooting variables. Instead of assigning shooting variable to every time point, we divide the time grid into B non-overlapping blocks and assign single shooting variable to each block. For block b (1. B}. we deñine an index set In containing indices of consecutive time points associated with that block such that U,TI, = (2....,N' N). We do not include the first time point 4 in any of the blocks. With every block b we associate observa- {5,6).
block
block
Figure 4: An exam-
ple of sparse multiple shooting with B 2,
I₁ = (2,3, 4] and Ig =
tions {3, heI,. time points {t, hte, and a shooting variable sy placed at the first time point before the block. The temporal position of s denoted by tpj. Latent states {ii are computed from 86 as
3; ODEsolve(ss, (10) 44 feay.), i E Iz.
As shown in Figure 4, this approach reduces the number of shooting variables and grants finer control over the length of each sub-interval to ensure that it is both sufficiently long to contain enough dynamics information and sufficiently short to keep the loss landscape not too complex.
As illustrated in Figure 4, an ODE solution (Eq. 7) does not necessarily match the corresponding shooting variable. Standard multiple shooting formulations enforce continuity of the entire trajectory via a hard constraint or penalty term (Voss et al., 2004; Jordana et al., 2021; Turan & Jäschke, 2022). Instead, we propose to utilize Bayesian inference and naturally encode continuity as prior which leads to sparse Bayesian multiple shooting which discuss in the next section.
3

---

Published as a conference paper at ICLR 2023
3.2 MODEL, INFERENCE, AND FORECASTING
Model. Our model is a latent neural ODE with sparse multiple shooting (Section 3.1). To infer the parameters 81:B, Oayn. and Oäyn P(81:B, Oayn: (dec) p($1:Bl (ayn) ) (Oay p((aee))
Bayesian inference with the following prior:
(8)
where pl p(@ayn), (aaa)).(p@dee.) p(ddec) are Gaussians, and the continuity inducing prior p( is defined as
B
B
((8 Oayn) - p(8)) Il p( 18-- Osyn) - p(81) II N ( 1p- (b), Sona) of
b=2
where p(81) is diagonal Gaussian, A is the Gaussian distribution, [ is identity matrix, and parameter of controls the strength of the prior. The continuity prior forces the shooting variable sy and the final state of the previous block b- which is obtained using the dynamics model, be
close (e.g., $2 and x(((2)) 24 in Fig. 4), thus promoting continuity of the entire trajectory.
With the priors above, we get the following generative model Bayn: Boc ~ p((aya))((aac) $1.B Baya ~ p(s1 (yyy))
(10)
31 81:
(11)
x; = ODEsolve(st: 1pj tu feam);
be (I,....B B), , i Ib.
(12)
y: y ~ p(y: Bes (rr,), Since II.N are deterministic functions of 81.B and Odyn. we have the following joint distribution (see Appendix B for more details)
1, ..... N.
(13)
p( (1:N, $1:B, Oayn: Odec) = p(y1:N 81:B, Oayn: dee (((1.B
(14)
Inference. We use variational inference (Blei et al., 2017) to approximate the true posterior p(Odyn: Odees $1.0 (1:N) by approximate posterior
g(Oayns Odces 8I.B) =:
Qapagn (or (Oaec)
II Yor (86) b=1
(15)
with variational parameters &yayn thace and LL..B = (vi..... BB). Note that contrary to standard VAEs, which use point estimates of Bayn and faec. extent the variational inference to these param- eters to adequately handle the uncertainty. To avoid direct optimization over the local variational parameters 1/I.B. we use amortized variational inference (Kingma Welling, 2013) and learn an encoder ha with parameters Oene which maps observations yI.N to 1/1.n (see Section 3.3), We denote the amortized shooting distributions lv., (8ol91-1) (y1.N O. where Th, he.. (yI.N), simply as q(s)) or 94p, (st) for brevity. We assume Quhaya Iehas and 9p,, to be diagonal Gaussians.
With a fully factorised (($1B) we can sample the shooting variables $1.B independently which allows to compute the latent states I.NN in parallel by simulating the dynamics only over short sub- intervals. If the posterior ((8).B) followed the structure of the prior P (0yy) we would not be able to utilize these benefits of multiple shooting since to sample 81.1 we would need simulate the whole trajectory $1.B starting at $1.
In variational inference we minimize the Kullback-Leible divergence between the variational ap- proximation and the true posterior,
KL [I[yyOo Osec, 81.B) 81:B))lp(0ss)) Osees 81.B (eIN) )] which is equivalent maximizing the ELBO which for our model is defined
(16)
B
C = E (0 log p(v1 |1. Bacec]
E Sul0 ) log p(y. y (s. Oayn: dsee)
(17)
b 1 EET
data likelibood
(ii) data likelihood
KL(g(s))lp((s))
KL [L 0sya))
(18)
2
(127) initial state prior
(1v) cootinuity prior
KL [L[ [g((taa)))) | p (laym)
KL [KL[q((an) (0(cc)) (apPacellppppuu (pp(Oee)))
(19)
(v) dynamics prior
(vi) deccnder prior
4

---

Published as a conference paper at ICLR 2023
y1:N
@1:N
214
heomp
a1:NN
01
02
12
11
Attention
(a)
1X
(b)
b1:B
PLN
b1
by
hyead
b1:B
2/1 1:B
Figure 5: (a) Encoder structure. (b) Encoder with two blocks (i.c., B 2) operating on input sequence V1. with shooting variables 81, 82 located at (1, lg-
Appendix B contains detailed derivation of the ELBO, and fully specifies the model and the approx- imate posterior. While terms (iii), (v) and (vi) have closed form, computation of terms (i), (ii) and (iv) involves approximations: Monte Carlo sampling for the expectations, and numerical integration for the solution of the initial value problems. Appendix details the computation of ELBO.
Forecasting. Given initial observations Vim of test trajectory time points timm we predict the future observation yn at a time point t > ty as the expected value of the approximate posterior predictive distribution
P B1IN) ~o / plyiils, si Odyn d )w; (87) (Odyn) (yran) (ddce dldn
(20)
where wi = he... (y..)) The expectation estimated Monte Carlo integration (Appendix C). Note that inferring st, instead of sf could lead to more accurate predictions, but in this work of to simplify implementation of the method.
3.3 ENCODER
We want design an encoder capable of operating on irregular time grids, handling noisy and partially observed data, and parallelizing the computation of the local variational parameters dh1.B- Transformer (Vaswani al., 2017) satisfies most of these requirements, but not directly applica- ble to our setup. We design transformer- relative positional encodings. These modifications provide useful inductive biases and allow the en- coder to effectively operate on input sequences with temporal component. The encoder computes 1h with (see Figure (a-b)):
encoder with time-aware attention and continuous
V1.B he... (3I.N) hread hreag (heomp (11.N) ))).
(21)
where
1. heomp RDD + B Dim compresses observations VIEN E DxN into a low-dimensional sequence a1:N E IR Dom xN where Dew << D.
2. hage D - Pam aggregates information across aiN into b E R xB where b, is located the temporal position of (Figure 5 (b)).
3. brend Dr - RR. reads the parameters t/i. E PPx from bi.n.
Transformations hemmp and Pread are any suitable differentiable functions. Transformation hage is transformer encoder (Vaswani et al., 2017) which sequence-to-sequenee mapping represented (E)
by stack of L layers (Figure 5 (a)). Each layer é E {1, {I,...,L) contains a component called attention sub-layer which maps an input sequence Ci:N := (ax)
( R XN to an (L) output sequence ( := (3(²) ..... a(2) Di. x except for the last layer which maps a to ((5)) to match the number of shooting variables. For the first layer, (1) = a1., and for the last layer, b (F( where FF(.) is feed-forward network with residual connection. In the
5

---

Published as a conference paper at ICLR 2023
following, drop the index l for notational simplicity since each layer has the same structure. The attention sub-layer for the standard, scaled dot-product self-attention (assuming a single attention head) is defined using the dot dot-product (C))), ) softmax (C2y) and weighted average (B,) (Vaswani et al., 2017):
Cijy (Woor, WkaS) VDiww
exp (CDP)
Ci
k=l exp (CR) B Σ (Wva)),,
(22)
where Wo. WK. Wy E Diix Doe are learnable fayer-specifi parameter matrices, and C C VxN the attention matrix. This standard formulation of self-attention works poorly irregularly sam- pled trajectories (see Section 4). Next, we discuss modifications that we introduce to make it appli- cable on irregularly sampled data.
Temporal attention Dot product attention has no notion of time hence can attend to arbitrary elements of the input sequence. To make G dependent mostly on those input elements that are close to t, we augment the dot dot-product attention with temporal attention CJA and redefine the attention matrix as
exp(C (
ps Do
nzl
Ity (1)
P
exp (C C
C = (c)
CJ =
(23)
6,
E exp (Co + CTA) N and 8,
R> are constants. Since exp (C + CTA exp (C exp (C ) the main purpose of temporal attention is to reduce the amount of attention from 3, to aj as the tempo-
where E (0, 1|.
(b)
Pas
w
ral distance |44 - 14] grows. Parameter 8, defines the distance beyond which exp(Ch) i is scaled by at least e, while p defines the shape the scaling curve. Figure (a) demonstrates shapes of the scaling curves for various values of p.
Figure 6: (a) Tem- poral attention. (b) Relative position encoding.
Relative positional encodings To make 3r independent of its absolute tem- poral position t; replace the standard global positional encodings with relative positional encodings which we define as
(D 8r
P.j @ hardtanh
and redefine
Σ Cij(Wva + P.y), (24)
j=1
where w E IRd is a vector trainable parameters, ② is point-wise multiplication, and 8, is the same as for temporal attention. This formulation is synergistic with temporal attention as it ensures that B; has useful positional information about Or; only if 14 Lil < 8, which further forces 3; to depend on input elements close 4; (see Figure 6 (b)). In this work we share across attention sub-layers. For further details about the encoder, see Appendix E. In Appendix F we investigate the effects of and 6,. In Appendix we compare our transformer-bis aggregation function with ODE-RNN of Rubanova et al. (2019).
Note that our encoder can process input sequences varying lengths. Also, as discussed in Section 3.2, at test time we set B = 1 so that the encoder outputs only the first parameter vector tb since we are only interested in the initial state 8 from which we predict the test trajectory.
4 EXPERIMENTS
To demonstrate properties and capabilities of our method we use three datasets: PENDULUM, RMNIST, and BOUNCING BALLS, which consist of nigh-dimensional (D 1024) observations of physical systems evolving over time (Figure 7) and are often used in literature on modeling of dynamical systems. We generate these datasets on regular and irregular time grids. Unless otherwise stated, use the versions with irregular time grids. See Appendix D for more details.
We train our model for 300000 iterations with Adam optimizer (Kingma Ba, 2015) and learning rate exponentially decreasing from 3 1 to 10-5, To simulate the dynamics we use an ODE solver from torchdif ffeq package (Chen et al., 2018) (dopri5 with rtol = atol 10 We second-orde dynamics and set the latent space dimension d to 32. See Appendix E for detailed description of seaniin/validattion/eee
setup and model architecture. Error bars are standard errors evaluated with five random seeds. Training is done single NVIDIA Tesla V100 GPU.

---

Published as a conference paper at ICLR 2023
3 3
3
3
3
m
3
3
3
3
mm
3
Figure 7: Top row: PENDULUM dataset consisting of images pendulum moving under the influence of gravity. Middle row: RMNIST dataset consisting of images of rotating digits 3. Bottom row: BOUNCING BALLS dataset consisting of images of three balls bouncing in box.
4.1 REGULAR AND IRREGULAR TIME GRIDS
Here we compare performance of our model on regular and irregular time grids. As Figure 8 shows, for all datasets our model performs very similarly on both types of the time grids, demonstrating its strong and robust performance on irregularly sampled data. Next, to investigate how design choices our encoder affect the results on irregular time grids, we do an ablation study where we remove temporal attention (TA) and relative positional encodings (RPE). Note that when we remove RPE we add standard Linuoidal-cosinn positional encodings as in Vaswani et al. (2017). The results are shown in Table 1. We that removing temporal attention, RPE, both tends noticeably increase test errors, indicating the effectiveness of our modifications.
0.030
0.025
Irregular
Test MSF
0.020
Regular
0.015
0.010
Model
Pendulum
RMNIST
Bouncing
0.005
-RPE
-TA
0.036 0.007
0.068 0.000
0.079 = 0.001
0.000
+RPE
-IA
0.013
0.010
0.062 = 0.002
0.013 = 0.013
Pendulum RMNIST Bounciing
-RPE
+TA
0.009
0.001
0.017 0.002
0.024 = 0.002
Balls
+RPE
+1A
0.004 = 0.001
0.015 0.002
0.024 = 0.001
Figure 8: Test errors for our model on regular and irregular time grids.
Table 1: Test MSEs for different ablations.
4.2 BLOCK SIZE
Our model operates on sub-trajectories whose lengths are controlled by the block sizes, i.c., the number of observations in each block (Section 3.1). Here we set the size of all blocks to given value and demonstrate how affects the performance of our model. Figure shows test errors and training times for various block sizes. We see that the optimal block size is much smaller than the length of the observed trajectory (51 our case), and that in some cases the model benefits from increasing the block size, but only up to some point after which the performance starts to drop. We also see how the ability to parallelize computations across block improves training times.
Bouncing
Pendulum
IMNIST
Balls
Training time (hours)
0.07
90
0.06
75
0.05
60
0.04
45
0.03
0.02
30
0.01
15
0.00
2
5
25
Block Size
Block Size
Figure 9:
Test errors and training times for different block sizes.

---

Published as a conference paper at ICLR 2023
4.3 CONTINUITY CONSTRAINT
Our model divides training sequences into blocks and uses the con- tinuity prior (Equation 9) enforce continuity of the latent trajec- tories across the blocks. Here we investigate how the strength of the
0.14
Pendulum
0.12
prior (in terms of σe)) affects the model's performance. In Figure 10 we show results for different values of de- We see that stronger con-
IMNIST
0.08
Bouncing
Balls
0.06
tinuity prior tends to improve the results. For BOUNCING BALLS
0.04
with 2 10-5 the model failed to learn meaningful latent dy- namics, perhaps due to excessively strong continuity prior. For new datasets the continuity prior well as other hyperparameters can be set e.g. by cross-vies validation. In appendix we also show how the value of de affects the gap between the blocks.
0.00
2e-5 2e-4
2e-3
20-2
20-1
Figure 10: Test enrors vs. de.
4.4 CONSTRAINING THE APPROXIMATE POSTERIOR
We found that constraining variance of the approximate poste- riors Qu, (8:) be least minn > 0 (in each direction) might
0.010
noticeably improve performance of our model. In Figure 11
0.035
Tiin 0
0.030
we compare the results for Tmin and Tmin 0.02. As can
0.025
Tmiis 0.02
be seen, this simple constraint greatly improves the model's performance on more complex datasets. This constraint could
0.0220
0.015
be viewed as an instance of noise injection, technique used improve stability of model predictions (Laskey et al., 2017: Sanchez-Gonzalez et al., 2020; Pfaff al., 2021). Previous works inject noise into the input data, but we found that inject- ing noise directly in the latent space produces better results. Details are in Appendix E.4.3.
0.005
0.000
Pendulum RMNIST Bouncing
Balls
Figure 11: Errors for constrained and unconstrained approximate posteriors.
4.5 COMMON HEURISTICS
As discussed previously, models that compute II.N directly from æ without multiple shooting (so called single shoot- ing models) require various heuristics to train them in prac- tice. Here we compare two commonly used heuristics with our multi-block model. First, we train our model with a single block (equivalent to single shooting) and use it as the base- line (SS). Then, we augment SS with the two heuristics and train it on short sub-trajectories (SS+sub) and on progressively increasing trajectory lengths (SS+progr). Finally, we train our sparse multiple shooting model (Ours) which is identical to SS, but has multiple blocks and continuity prior. See Appendix G for details. The results in Figure 12. The baseline single shooting model (SS) tends to fail during training, with only a few runs converging. Hence, SS produces poor predictions average. Training single shooting model on short sub-trajectorie tends to make results even worse in our case. With relatively easy training, SS+sub produces unstable test predictions that quickly blow up. In our case SS+progr was the most effective heuristic, with stable training and reasonable test predictions (with a few getting bit unstable towards the end). Compared to our model, none of the heuristics was able to match the performance of our sparse multiple shooting model.
SS+progr
SS+suh
Ours
0.25
Test MSE
0.15
0.10
0.05
0.00
Pendulum
RMNIST
Bouncing
Balls
Figure 12: Errors for different
heuristics.
4.6 COMPARISON TO OTHER MODELS
We compare our model to recent models from the literature: ODE2VAE (Yildiz et al., 2019) and NODEP (Norcliffe et al., 2021). Both models learn continuous-time deterministic dynamics in the latent space and use an encoder to map observations to the latent initial state. For comparison we datasets regular time grids since ODE2VAE's encoder works only regular time grids. All models are trained and tested on full trajectories and use the first 8 observations to infer the latent initial state. We the default parameters and code provided in the ODE2VAE and NODEP papers.
8

---

Published as a conference paper at ICLR 2023
0.10
NODEP
3
0.09
OIDEVVEE
Test MSB
0.07
0.06
3
3
0.04
0.03
M
0.01
M
0.00
Pendichum
RMNIST
Bouncing
Ballss
Test errors for different models and datasets. Right: For each dataset, we plot data and predictions for NODEP, ODE2VAE and model (top bottom). Each sub-plot shows data as the first row, and prediction as the second row. We show prediction with the median test error See Appendix H.4 for more predictions.
Figure 13: Left:
All models are trained for the same amount of time. See Appendix H for more details. Figure 13 shows the results. We see that NODEP produces reasonable predictions only for the PENDULUM dataset. ODE2VAE performs slightly better and manages to learn both PENDULUM and RMNIST data quite well, but fails on the most complex BOUNCING BALLS dataset (note that ODE2VAE uses the iterative training heuristic). Our model performs well all three datasets. Also, Appendix H.5 for a demonstration of the effect of the training trajectory length on NODEP and ODE2VAE.
5 RELATED WORK
The problem with training on long trajectories is not new and multiple shooting (MS) was pro- posed as a solution long time ago (van Domselaar & Hemker, 1975; Baake al., 1992; Voss al., 2004). Recent works have tried to adapt MS to modern neural-network-basss data regimes. Jordana et al. (2021) and Beintema al. (2021) directly apply MS in latent space in fully deterministic setting, but use discrete- time dynamics without amortization or with encoders applicable only to regular time grids, and also both use ad-hoc loss terms to enforce continuity (see Appendix H.6 for comparison against our method). Hegde et al. (2022) proposed probabilistic formulation of MS for Gaussian process based dynamics, but do not use amortization and learn dynamics directly in the data space. While not directly related to this work, recently Massaroli al. (2021) proposed to use MS to derive a parallel-in-time ODE solver with the focus on efficient parallelization of the forward pass, but they do not explicitly consider the long trajectory problem.
models and large
Different forms relative positional encodings (RPE) and distance-basse attention were introduced in previous works, but usually for discrete and regular grids. Shaw et al. (2018) and Raffel et al. (2020) use discrete learnable RPES which they add to keys, values or attention scores. Both works clipping, i.e., learn RPEs only for & closest points, which is some function. Press et al. (2022) use discrete distance- based attention which decreases linearly with the distance. Zhao al. (2021) use continuous learnable RPEs which represented as an MLP which maps difference between spatial positions of two points to the corresponding RPEs which are then added to values and attention scores without clipping.
similar using hardtanh
Variants of attention-basee models for irregular time series were introduced in Shukla & Marlin (2021) and Zhang et (2020), but they are based on global positional encodings and do not con- strain the size and shape of the attention windows.
CONCLUSION
In this work we developed a method that merges classical multiple shooting with principled prob- abilistic modeling and efficient amortized variational inference thus making the classical technique efficiently applicable in the modern large-data and large-model regimes. Our method allows to learn large-scale continuous-time dynamical systems from long observations quickly and efficiently, and, due to its probabilistic formulation, enables principled handling of noisy and partially observed data.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY STATEMENT
Datasets and data generation processes are described in Appendix D. Model, hyperparameters, chitectures, training, validation and testing procedures, and computation algorithms are detailed in Appendices B, C, E. Source code accompanying this work will be made publicly available after review.
ACKNOWLEDGMENTS
This work was supported by NVIDIA AI Technology Center Finland.
10

---

Published as a conference paper at ICLR 2023
REFERENCES
Ellen Baake, Michael Baake, HG Bock, and KM Briggs. Fitting ordinary differential equations to chaotic data. Physical Review A, 45(8):5524, 1992.
Gerben I. Beintema, Roland Toth, and Maarten Schoukens. Non-linear state-space model identifi- cation from video data using deep encoders. IPACCPPppesOOnLine, 54(7):697-701, 2021. ISSN 2405-8963, doi: 10. if001/...acoo..221000..... ifacol. 2021 08....... 442.
wttp:///.o...//...00//// URL /ttp////////ddooo.....oo//./.../
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statis-
ticians. Journal of the American Statistical Association, 112(518):59-887,, Apr 2017. ISSN 1537-274X. doi: 10011000011114992211.22220
1000//000/
01621459. 2017 1285773.
Hans Georg Bock and J. Plitt. A multiple shooting algorithm for direct solution of optimal control problems. IFAC Proceedings Volumes, 17:1603-1660 1608, 1984.
Francesco Paolo Casale, Adrian V Dalca, Luca Saglietti, Jennifer Listgarten, and Nicolo Fusi. Gaus- sian process prior variational autoencoders. 32nd Conference on Neural Information Processing Systems, 2018.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary dif- ferential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran As- sociates, Inc., 2018. URL htwps:ttpp.ooomm..p. http:////ooeeeeeeeeeeeooe meurrps. wwwpappr///a..e/eeeee. httppappr22188 /2018/f111 pdf.
Zhe Gan, Chunyuan Li, Ricardo Henao, David Edwin Carlson, and Lawrence Carin. Deep temporal sigmoid belief networks for sequence modeling. In NIPS, 2015.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Yee Whye Teh and Mike Titterington (cds.), Proceedings of the Thirteenth Interna- tional Conference on Artificial Intelligence and Statistics, volume Proceedings of Machine Learning Research, pp. 249-256 Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR. URL https:tt/////..... http://pcocceeeeeeeeeeee mlr www.:ppess/ee............ wttpg//gggooooooooooooo Da html.
Jung-Su Ha, Young-Jin Park, Hyeok-Joo Chae, Scon-Seo Park, and Han-Lim Choi. Adaptive path- integral autoencoder: representation learning and planning for dynamical systems. Journal of Statistical Mechanics: Theory and Experiment, 2019(12):1240M) 124008, Dec 2019, ISSN 1742-5468. doi: 10. 10.1086//74225546/555555 URL http: dx doi. org /tlgg/10.00 1088 5468/
ab3455.
XU HAN, Han Gao, Tobias Pfaff, Jian-Xun Wang, and Liping Liu. Predicting physics in mesh- reduced space with temporal attention. In International Conference Learning Representattonss 2022. URL https:: https://
Pashupati Hegde, Çağatay Yıldız, Hami Lähdesmäki, Samuel Kaski, and Markus Heinonen. Vari- ational multiple shooting for bayesian ODEs with gaussian processes. The 38th Conference Uncertainty in Artificial Intelligence, 2022. URL https: www.revveeeeeeeeeeoeoeeeee.ee.ee net forum?
(d=r2NuhTUoces
Eric Heiden, Chris Denniston, David Millard, Fabio Ramos, and Gauray S. Sukhatme. Probabilistic inference of simulation parameters via parallel differentiable simulation. In ICRA, 2022.
Pieter W. Hemker. Nonlinear parameter estimation in initial value problems. 1974.
Jun-Ting Hsich, Bingbin Liu, De-An Huang, Li Fei-Fei, and Juan Carlos Niebles. Leamning to decompose and disentangle representations for video prediction. In NewrIPS, 2018.
Armand Jordana, Justin Carpentier, and Ludovic Righetti. Learning dynamical systems from noisy sensor measurements using multiple shooting. ArXiv, abs/2106. 1712, 2021.
11

---

Published as a conference paper at ICLR 2023
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Con- ference on Learning Representations, 2017. URL https:/ http///ppeeeeeeeeeeeeeeeeeeeeeep http//ewwwee/.orrnmee.e
https://openrvview....
(d=-yyqHH.5xxg
Diederik P. Kingma and Jimmy Ba. abs/1412.6980,, 2015.
Adam: A method for stochastic optimization. CoRR,
Diederik Kingma and Max Welling. Auto-encoding variational bayes, 2013.
Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer. Machine coeminng-ccceeratee computational fluid dynamics. Proceedings of the National Academy of Sciences of the United States of America, 118, 2021.
Michael Laskey, Jonathan N. Lee, Roy Fox, Anca D. Dragan, and Ken Goldberg. Dart: Noise injection for robust imitation learning. In CoRL, 2017.
Marten Lienen and Stephan Günnemannn. Learning the dynamics of physical systems from sparse observations with finite element networks. In International Conference on Learning Representa- tions (ICLR), 2022.
William Lotter, G. Kreiman, and David D. Cox. Unsupervised learning of visual structure using predictive generative networks. ArXiv, abs/1511.063380 2015.
Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Dissect- ing neural odes. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 3952-3963, Curran Asso- ciates, Inc., 2020. URL httpp::/pproceeteee neurip cctccc///a/a htc/pappr/2220// 198856227555554442e22e222
pdf.
Stefano Massaroli, Michael Poli, Sho Sonoda, Taiji Suzuki. Jinkyoo Park, Atsushi Ya- mashita, and Hajime Asama. Differentiable multiple shooting layers. In M. Ran- Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances
zato, A. Beygelzimer, in Neural Information Processing Systems, volume 34, pp. 16532-16544, Curran Asso- ciates, Inc., 2021. URL https:::ppp//oooo https::/prroceedingg www.oiiiiinnnnnnnnnneeeeom oc catss.cccppppeeeee cc/ppaper/222//1L... pdf.
Luke Metz, C. Daniel Freeman, Samuel S. Schoenholz, and Tal Kachman. Gradients are not all you need. ArXiv, abs/21 11.211100580, 2021.
Alexander Norcliffe, Cristian Bodnar, Ben Day, Jacob Moss, and Pietro Liò. Neural {ode} processes. In International Conference on Learning Representations, 2021. URL https:
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito,o Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp- 8024-8035, Curran Associates, Inc., 2019. URL http: pett::/ppapee neurips cc cc/paper/ 9015 1:15---ytrcc---o an imperative style high- performance deep- paaringg-ibbaan pdf.
Tobias Pfaff, Meire Fortunato, Alvaro Lanchez-Gonzalez and Peter W. Battaglia. Learning mesh- based simulation with graph networks. In International Conference on Learning Representations, 2021.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL https: http///pp/eeeeeeeeeeeeeeeeeeeeee net forumm &
12

---

Published as a conference paper at ICLR 2023
Colin Raffel, Noam M. Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. ArXiv, abs/1910.10683, 2020.
Antônio H. Ribeiro, Koen Tiels, Jack Umenberger, Thomas Schön, and Luis A. Aguirre. On the smoothness of nonlinear system identification. Automatica, 121:109158, Nov 2020. ISSN 0005- 1098. doi: 10. www.l//aa.ttmmmccc..a..... 2020.11918...... URL http: /xx doi. ozg. 1016/J. AUTOMATICA, 2020. 109158.
Yulia Rubanova, Ricky T. Q Chen, and David K Duvenaud. Latent ordinary differential equations for irreullrly-samplled time series. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, vol- ume 32. Curran Associates, Inc., 2019. URL https: http::::/pcccceednngeee neurips. ca/ 7be paper/201 9/ file//eee//// Totaaaaas55aattaaaaaaaaaa4441111 To
Alvaro Gonzalz, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In International Confer- ence on Machine Learning. 2020.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. In NAACL, 2018.
Satya Narayan Shukla and Benjamin Marlin. Multi-time attention networks for irregularly sampled time series. In International Conference on Learning Representations, 2021. URL https: / //openreview. net/ fott/foounm wtt.//ooddmmddommmm
llya Sutskever, Geoffrey E. Hinton, and Graham W. Taylor. The recurrent temporal restricted boltz- mann machine. In NIPS, 2008.
Evren Mert Turan and Johannes Jäschke. Multiple shooting for training neural differential equations on time series. IEEE Control Systems Letters, 6: 1899921900, 2022.
van Domselaar and Piet W Hemker. Nonlinear parameter estimation in initial value problems. Stichting Mathematisch Centrum, Numerieke Wiskunde, (NW 18/75), 1975.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, Fergus. nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL https:/ http:s////oooeeedeeeoo 3f5ee243 (iceee24347deee (ateeeee@ddoo311caa55aaaaa Paper. pdf.
Vishwanathan, and R. Gar-
wwwniiiisss.. inerrreeeednnns....r.pp.p.p.. http:r/ppp..pppppe///////. ccc/paper/211//f1ee
Henning U Voss, Jens Timmer, and Jürgen Kurths. Nonlinear dynamical system identification from uncertain and indirect measurements. International Journal of Bifurcation and Chaos, 14(06): 1903-1933, 2004.
Çagatay Yildiz, Markus Heinonen, and Harr Lähdesmäki. Ode2vae: Deep generative second order odes with bayesian neural networks. In NeurIPS, 2019.
Lin Zehui, Pengfei Liu, Luyao Huang, Junkun Chen, Xipeng Qiu, and Xuanjing Huang. Dropatten- tion: A regularization method for fully-connected self-attention networks, 2019.
Qiang Zhang, Aldo Lipani, Omer Kimap, and Emine Yilmaz. Self-attentive Hawkes process. In
Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference Ma- chine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 11183-11193. PMLR, 13-18 Jul 2020. URL https: pttpeeeieceeee...ooooep//o mml. prwss setessvvvlo https//vv11///hha html.
Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip H. S. Torr, and Vladlen Koltun. Point transformer. 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 16239 16248,, 2021.
13

---

Published as a conference paper at ICLR 2023
A DEPENDENCE OF LOSS LANDSCAPE ON THE OBSERVATION INTERVAL
Here we demonstrate how complexity of the loss landscape grows with the length of the training trajectory.
For simplicity, we train a neural ODE model which is similar to the L-NODE model in Equations 1-2, but with S. being the identity function. The dynamics function is represented by an MLP with two hidden layers of size 16 and hyperbolic tangent nonlinearities.
The training data consists of single 2-dimensional trajectory observed over time interval of [0, 20] seconds (see Figure 14). The trajectory is generated by solving the following ODE
d'x(t) 9.81 sin (ar(1)
(25)
dt
with the initial position being 90 degrees (relative to the vertical) and the initial velocity being zero. The training data is generated by saving the solution of the ODE every 0.1 seconds.
We train the model with MSE loss using Adam (Kingma & Ba, 2015) optimizer and dopriS adaptive solver from the torchdi ff feq package (Chen et al., 2018). We start training the first 10 points of the trajectory and double that length every 3000 iterations (hence the spikes the loss plot in Figure 15). At the end of each 3000 iterations cycle (right before doubling the training trajectory length) we plot the loss landscape around the parameter value to which the optimizer converged. Let 9 be the point which the optimizer converged during the given eycle. We denote the corresponding loss value by marker in Figure 15. Then, we plot the loss landscape around @ by evaluating the loss parameter values c,, where E 1-4, 6]. For the given observation time interval, the trajectory of length 10 is easy to fit, hence is considered to be "short".
Pendulum Data
Ang- position
Ang. velocity
-4
0.0
2.5
5.0
7.5
10.0 Time
12.5
15.0
17.5
20.0
Figure 14: Pendulum data.
10a
N=10
N=20
Nodo
N=
No160
Train Loss
10
10-1
10-6
2000
4000 6000
8000 Iteration
0000 12000
Figure 15: Top: Training loss of NODE model. We start with short training trajectory (N 10) and double its length at iterations denoted by the markers. Note that training fails for long enough trajectory. Bottom: One-dimensional projection of the loss landscape around the parameter values to which the optimizer converged for given trajectory length. Note that complexity of the loss landscape grows with the trajectory length.
14

---

Published as a conference paper at ICLR 2023
B MODEL, APPROXIMATE POSTERIOR, AND ELBO
Here we provide details about our model, approximate posterior and derivation of the ELBO.
Joint distribution
The joint distribution is
plyiNN 81, Oayn, Osec) p(y):N 81:13 Odyn, 0 (81.) Odyn (0a0y)))) ) Odec)
(26)
with
p(Oayn)
-N(0dnn M0nm Bay !), p(Odec) -V M1 B 1);
(27)
B
p(s 11.11100am = p(()) II
( 1 Oayn)
(28)
b=2
B =N(ss 1 II (OLsovvvvvlv 1/611 116), frnm) o21).
(29)
6=2
B II
p(y1:N (81BB Odyn, Oucc) = p((/1 $1, d
p( (ys hez, 186, Oayn, Osec)
(30)
b=1
p(y1 $1, Bose)
p(y 1s6, Oaya Bose)
(31)
6
= =W (81), (8
II b= VET,
I
N(y: I0om (ODFsolve(s,, 11), 41 fana)) oz T)
(32)
=N (21), 131) II II
(1 b=1 1EZ,
(33)
where A is the Gaussian distribution, / dgd is identity matrix, and oद is the observation noise variance that is shared across data dimensions.
Approximate posterior
The family of approximate posteriors is defined
B
((Osyn Osce, S1:B) = ((0syn)g( (00co) (dac)
II (8a)
(34)
B
= -( Yours diag 72 (fute diag( 2 Teme ))
II
(ssb diag(T;) r2)),
b=
(35)
where diag(T)) (T)) is matrix with vector the main diagonal.
15

---

Published as a conference paper at ICLR 2023
ELBO
The ELBO can be written as
/ ((yny Odces 81.13) In P(yI.N, 81.B, Odyn: Odce)
d0 d0
(36)
g(0dyn- Odec $1:0)
/ O'oyyn daces $1.B) In p(y1.N 81.B. Oayn. dee )p( (111 (o
d8 d8 dyn ds 1::B
q(81.B (0ay ) (dcc))
(37)
(Oayas Odces $1.B) In (8..B
Odya, Ssce do00
(38)
I (Oay, Odec; 81:8) In q( (881.))
d0 d0c
(39)
p($1:13 (aya)
f g(Oayn Odcc: $1.B) In q( (dya)) d0 do
(40)
p(Odyn)
I g(odoc. Osce: $1.B) In (Bsec) p(dce))
d6 dsc
(41)
C1 L2 - L3 Cs.
(42)
Let's look at each term C, separately.
/ g(Odyn Osco $1.B) In p( $1.8, Oayn, Oac
C1
(43)
B II b=1
f g(Osyn: Oscor $1.B) In p(y s1, O
pl pllylhee 18., Ooyn loo) doondlaodds.. (44)
f g(Oayn: Osec 81.B) in p(y1 d00
(45)
/ ((Oaym B $1.B) In Il p( daya @ace) do (46)
/ @(Oayn: Odec, 81.13 In p(yn (8,, Osec d00
(47)
Σ / g(Oayn: Oacc s1.n) In p( (s,, Oayn: Odec )d0 dec 1:1
(48)
(Oite 81) plri (81, Oace )d0 d8
(49)
B / g(Oayn: O 8n) In Iu (ep(((tiice Fiez, 18,, Oäyn, dde
(50)
Eg(onn..r) 01) [1np(y1/s2, awee
Σ
(Inp((y: (ee.eeeeeee.............. Ogyms Osec)l
(51)
12
Eg(Oar)))) 1mpp(y/ss 0scc)) +
Pert [In p(ysi sb,, Oxyn, Ocec)]
(52)
b=1
16

---

Published as a conference paper at ICLR 2023
C2 / Q(dynn does $1.B) In
(
d0 1:
(53)
p(81.1
(yyn)
/ (Odyyn Odees $1:B) In
((51)
((8b)
d0 d0cydd0dodd
(54)
((81)
bm p(8y 18 Odyn)
g(81)
a(Odyns Odeos 81.13) In
d0
d0synelsec
(55)
[p($1)
/ ((Oaams Odec; 81:11) In
((8b)
d0 @0yndOlln d00daadd11110
(56)
p( 86
186 0dyn)
b=
/ ((dynn Odces $1.B) In
(1$1)
d8 dyn dec
(57)
[p(s1)
B
/ g(Odyn: Odce, $1.B) In
((s))
dl d0
(58)
p(ss) 86-161111111 Odyn)
b-2
/ ((11)1 In
(((11) p(s1)
d8i
(59)
B
((86)
g g(Oayn 86) In
dbayndso-1da
(60)
p(s 86-1180-1100 Odya)
b=2
/ (((1) In
(((11) [(())
dai
(61)
B
((Odyn 96-1) (ʃ ((sn) In
((bb)
deb
d0dyndsb-1
(62)
((88)) 8b Oayn)
/ ((s))
[((11] p(())
dsi
(63)
B
q (Odayn S
JKL (g(5s))ll((ss)ss. Odyn)) dsn
(64)
b=2
KL (c(s)))(((s) +
Ey Ey(0 [KL (n
(65)
where KL is Kullback-Leibler divergence.
L3 = (L((((Oy))II(((am)))
L1 = (L(((Pcc))l((Ooc)))
(66)
Computing ELBO All expectations are approximated using Monte Carlo integration with one sample, that is
Ep(a)[(()) ~ f((), where ( is sampled from p(2).
(67)
The KL. terms contain only Gaussian distributions, so can be computed in closed form.
17

---

Published as a conference paper at ICLR 2023
C
COMPUTATION ALGORITHMS
C.1 ELBO
To find the approximate posterior which minimizes the Kullback-Leibler divergence
(Lyya Bdec $1:B) $1:B)llp((dpn) Odes $1:B (11.NN)))
(68)
we maximize the evidence lower bound (ELBO) which for our model is defined as
C (0(0un..s) 01) log p 81, Odcc)
E Bygtara as)) [log ply vile 8b. Oäyn: Odec
(69)
b= REIn
(i) data likelihood
(i) data likelihood
[L(g(s))ll(ss))
By(Bsa
KL (Lo Odyn))
(70)
(il initial state prior
(iv) continuity prior
KL (a( (Daaa) (ly(da)))
KL [LL[g((baaa (q(00ss)llp((asss [pp(Occc)]
(71)
(v) dynamics prior
(vi) docoder prior
The ELBO is computed using the following algorithm:
1. Sample Buya= fec from guhay (Baya); Perrn (9dec).
2. Sample 81.B from Qup (81), ... lub (8B) with 1/1.B hew (V1:N).
3. Compute 21:N from $1:B as in Equations 11-12.
4. Compute ELBO L (KL terms are computed in closed form, for expectations we use Monte Carlo integration with sample).
Sampling is done using reparametrization to allow unbiased gradients w.r.t. the model parameters,
We observed that under some hyper-parameter configurations the continuity-promotingg term (iv) might cause the shooting variables to collapse to single point hence preventing the learning of meaningful dynamics. Downscaling this term helps to avoid the collapse. However, in our experi- ments we did not use any scaling.
C.2 FORECASTING
Given initial observations vi-N, at time points ti:N, we predict the future observations y +1:N, at time points N the expected value of the (approximate) posterior predictive distribution
+1NN+ MiN MIN) R
feppppuii plyNN s7. Osyn: Psco 10p (s (ap) (0dyn Iria (9
(72)
where के has- ). The expected value is estimated via Monte Carlo integration, the algorithm for predicting +N,NN is
1. Sample Odya- Osee from Jehu. (Dayn), Ferres (0sec)-
2. Sample oj from Gep (si) with 2/o how (yT:Nn).
3. Calculate latent states æ; ODEsolve
6i Soon) iE (Ni +,..,,
4. Sample J from p(y [go.. (x;)), E (N1 + 1, 1,.,,,aa,,
5. Repeat steps 1-4 times and average the predicted trajectories VN,+1:N (we n 10).
18

---

Published as a conference paper at ICLR 2023
D DATASETS
Figure 16: Examples of trajectories from the PENDULUM dataset.
3
3
3
3
3
3
3
3 3
3
3
3
3
f
3
3
3
3
3
3
3
3
C
3
3
3
3
Figure 17: Examples of trajectories from the RMNIST dataset.
Figure 18: Examples of trajectories from the BOUNCING BALLS dataset.
Here we provide details about the datasets used in this work and about the data generation pro- cedures. The datasets we selected are commonly used in literature concerned with modeling of temporal processes (Karl et al., 2017: Ha et al., 2019; Casale et al., 2018; Yildiz et al., 2019; Nor- cliffe al., 2021; Sutskever al., 2008; Lotter et al., 2015; Hsieh et al., 2018; Gan et al., 2015). To the best of our knowledge, previous works consider these datasets only on regular time grids (i.e., the temporal distance between consecutive observations is constant). Since in this work are mostly interested in processes observed at irregular time intervals, we generate these datasets on both regular and irregular time grids. The datasets and data generation scripts can be downloaded at httss//github..om/yaeeee 1/msvi.
D.1 PENDULUM
This dataset consist of images of pendulum moving under the influence of gravity. Each trajectory is generated by sampling the initial angle æ and angular velocity è of the pendulum and simulating its dynamics over period time. The algorithm for simulating one trajectory is
1. Sample x ~ Uniform(0, 2T] (in rads) and t ~ Uniform[=7/2,/// (in rads/second),
2. Generate time grid (tI....,ty tx). Regular time grids are generated by placing the time points
at equal distances along the time interval (x] with the first time point placed at t and the last time point placed at (x. Irregular time grids are generated by sampling N points from the time interval [41 4x] uniformly at random with the first time point placed at 4 the last time point placed at t.. and also ensuring that the minimum distance between time points is larger than -y-tu (i.e.. quarter of the time step of regular time grid).
4(N-1)
19

---

Published as a conference paper at ICLR 2023
02
-02
00
0.
10
15
20
25
30
022
00
-0.2
13
25
0.
0.5
1.
1.5
20
25
30
(a) Regular time grid.
02
-02
00
0.
10
20
25 30
(b) Irregular time grids.
Figure 19: Examples of regular and irregular time grids for PENDULUM dataset. At test time, observations before the red lines are used to compute the latent initial state.
3. Solve the ODE d'=(t)
-9.81 sin (zxt)) with initial state æ, Ø at time points (f1:-.., tN) dt
4. Create sequence of observations (y1...., yN) with = observe(< (xx))), where xfx²) is the solution of the ODE above at time point & and observe(-) is mapping from the pendulum angle to the corresponding observation.
The training'validatiionnest sets contain 400/50/50 trajectories. Regular time grids are identical across all trajectories. Irregular time grids are unique for each trajectory. The only constraint we place the time grids is that they contain N time points (for efficient implementation and meaning- ful comparison). We set t = O, tN = 3, and 51. Each observation y; is (024-dimensional vector (flat 32 32 image).
D.2 RMNIST
This dataset consist of images of rotating digits sampled from the MNIST dataset. Each trajectory generated by sampling a digit 3 from the MNIST dataset uniformly at random without replace- ment, then sampling the initial angle and angular velocity and simulating the frictionless rotation of the digit. The algorithm for simulating one trajectory is
1. Sample digit 3 from the MNIST dataset uniformly at random without replacement.
2. Sample x Uniform(0, (0, 2x] (in rads) and ē Uniform[T, (T, 2m]] (in rads/second).
3. Generate time grid (t1.....ty) tx). Regular time grids generated by placing the time points at equal distances along the time interval [41, AN] with the first time point placed at I and the last time point placed Ix. Irregular time grids generated by sampling N points from the time interval [t,, tx] uniformly at random with the first time point placed at t1, the last time point placed at t.. and also ensuring that the minimum distance between time points is larger than Ey -f
4(N-1) (i.e., a quarter of the time step of regular time grid).
4. Solve the ODE dan(t)
ē with initial state at time points (t,,..,, tN)-
5. Create sequence of observations (y1...., yN) with y: = observe( (eevee(((())) (t:)), where xtxa) is the solution of the ODE above at time point t; and observe(-) is a mapping from the digit angle to the corresponding observation.
The vaining/validatioonnest sets contain 4000/500/500 trajectories. Regular time grids are identical across all trajectories. Irregular time grids are unique for each trajectory. The only constraint we place the time grids is that they contain N time points (for efficient implementation and meaning- ful comparison). We set t = 0, ty = 2, and 51. Each observation y; is dimenssiii vector (flat 32 32 image).
20

---

Published as a conference paper at ICLR 2023
D.3 BOUNCING BALLS
This dataset consist of images of three balls bouncing in frictionless box. Each trajectory is gener- ated by sampling the initial positions and velocities the three balls and simulating the frictionless collision dynamics. The algorithm for simulating one trajectory is
1. Sample initial positions of the three balls uniformly at random such that the balls do not overlap and do not extend outside the boundaries of the box.
2. Sample initial velocities of the three balls E R as v v' Ter where e' is sampled from the standard normal distribution.
3. Generate time grid (t1...., tt). Regular time grids are generated by placing the time points at equal distances along the time interval [4,,4x] with the first time point placed at and the last time point placed at fy. Irregular time grids are generated by sampling N points from the time interval [L,, Ex] uniformly at random with the first time point placed at t the last time point placed at iN. and also ensuring that the minimum distance between time points is larger than _x-t (i.e., a quarter of the time step of regular time grid).
1]
4. Solve the ODE representing the frictionless collision dynamics at time points (f1,...,tx) (N) (see the data generating script for details).
5. Create sequence of observations (y),..,, yv) with y, = observe(0(()). where ((4:) )) is the solution of the ODE above time point t, and observe(...) is mapping from positions of the balls to the corresponding observation.
The training/validationn/ee sets contain 10000/1000/ (000 trajectories. Regular time grids are identi- cal across all trajectories. Irregular time grids are unique for each trajectory. The only constraint we place the time grids is that they contain N time points (for efficient implementation and meaning- ful comparison). We set t = O, ty = 20, and N = 51. Each observation y; is 1024 dimennsoon vector (flat 32 32 image).
E SETUP
E.1 TRAINING, VALIDATION, TESTING
E.1.1
DATA PREPROCESSING
We normalize the observations by the maximum absolute value in the training set.
1.
TRAINING
We train our model for 300000 iterations using Adam optimizer (Kingma & Ba, 2015) with learning rate exponentially decreasing from 3e-4 1e-5. To simulate the model's 's dynamics we use differ- entiable ODE solvers from torchdi E Efeq package (Chen et al., 2018). In particular, we use the dopri5 solver with rtol atol 10 without the adjoint method. For PENDULUM, RMNIST, and BOUNCING BAL.SS datasets the batch size is set to 16, 16, and 64, respectively. while the block size set to 1, 1, and 5, respectively. For some datasets we use data augmentation: PENDULUMM horizontal flip, BOUNCING BALLS- vertical and horizontal flips. For each dataset set 6, to 15% of the corresponding observation interval [f],
1.3 VALIDATION
We use validation set to track performance the model during training and save the parameters that produce the best validation performance. As performance measure we use the mean squared error at predicting the full validation trajectories given some number of initial observations. We use all observations within the interval [t1, + deal latent initial state. As during training, we set & to 15% of the observation interval [t,, ty]. The predictions made described in Section 3.2 but with single sample from the posterior.
initial observations from which we infer the
21

---

Published as a conference paper at ICLR 2023
E.1.4 TESTING
Predictions for the test trajectories made as described in Section 3.2. Similarly to validation, use all observations within the interval [&,, & + deaa] as initial observations from which we predict the latent initial state. We set des 15% of the observation interval [t].
E.2 PRIORS
As discussed in Appendix B, use the following priors: p(Oagn) =
(Oayn M4 2 I), p(Osec) = (Osea M0 deac f);
(73)
p($1:B (dya) (V(s o II [ M(sss/DDssovresss-r// 1]: [(b) fe..)) (74)
6m2
We set /AB... = P0xx = 0, Ob.e = 0a = 1, μr = 0, ou = 1, and o. = va
& where & denotes the
required average distance between and 2, and d is the latent space dimension. In this work we use 32. The parameter & is dataset specific, for PENDULUM and RMNIST we set g 10-4 for BOUNCING BALLS we set { = 10-3
E.3 VARIATIONAL PARAMETERS
As discussed in Appendix B, we use the following family of approximate posteriors:
(00yn Oec: $1:B) = (00y/l Yea, diag(TT (O Mo diag(T ))
II N Nsshhr diag(T?))
(75)
While Yb and are provided by the encoder, other variational parameters are directly optimized. We initialize Yea. and You using default Xavier (Glorot & Bengio, 2010) initialization of the dy- namics function and decoder (see PyTorch 1.12 (Paszke et al., 2019) documentation for details). We initialize Ta and Ti as vectors with each element equal to 10-4
E.4 MODEL ARCHITECTURE
E.4.1 DYNAMICS FUNCTION
Many physical systems, including the ones we consider this work, are naturally modeled using second order dynamics. We structure the latent space and dynamics function so that we include this useful inductive bias into our model. In particular, we follow Yildiz et al. (2019) and split the latent space into two parts representing "position" and "velocity". That is, represent the latent state a(t e gd as a concatenation two components:
(p(() a., (1)
)
xxt)
(76)
where (1) E Rd// the position component and (1) E R2/2 is the velocity component.
Then, we represent the dynamics function Sea. (()) as
xx(() feam (t, xt)) (1r (1, (() (()
(77)
where fi (()) : I Rd Rd/2 is the dynamics function modeling the instantaneous of change of the velocity component.
In all our experiments remove the dependence of f on time and represent it as multi-layer perceptron whose architecture depends on the dataset:
PPNNUUUM: input size d, output size d/2, two hidden layers with size 256 and ReLU nonlinearities.
22

---

Published as a conference paper at ICLR 2023
imput size d, output size d/2, two hidden layers with size 512 and ReLU non- linearities.
RMNIST:
BoUNCING BALLS: input size d, output size d/2, three hidden layers with size 1024 and ReLU nonlinearities.
In this work we use 32.
E.4.2 DECODER
The decoder 96a maps the latent state I; to parameters of p(ys:gemm (æt)). As we discussed in Appendix B, we set p p( (a:)) Nyy.lgem.. (z,), 1f-)), so the decoder outputs the mean of a Gaussian distribution. We treat σy as hyperparameter and set trying to learn resulted in overfitting. Following Yildiz et al. (2019), our encoder utilizes only the "position" part æ. of the latent state æ; since this part is assumed to contain all the information required to reconstruct the observations (see Appendix E.4.1).
10-3 In our experiments,
We represent a0 as composition of convolutional neural network (CNN) with a sigmoid function to keep the mean in the interval (0, 1). In particular, 98s. has the following architecture: linear layer, four transposed convolution layers (2x2 kernel, stride 2) with batch norm and ReLU nonlincarities, convolutional layer (5x5 kernel, padding 2), sigmoid function. The four transposed convolution layers have 8n, 4n, 2n and n channels, respectively. The convolution layer has n channels. For datasets PENDULUM, RMNIST, and BOUNCING BALLS n to 16, and 32, respectively.
E.4.3 ENCODER
Encoder maps observations y1, yN to parameters tb,,..., th of the ap- proximate posterior (Equation 75). In particular, it returns the means "i..... YB and standard deviations Ti. TB of the normal distributions (sslm,,
y1:N
(s YB. diag(T (T)). Using second order dynamics natu- rally suggests splitting the parameters into two groups. The first group contains pa- rameters for the "position" part of the latent space, while the second group contains parameters for the "velocity" part. So, we split the means and standard deviations into position and velocity parts
heamap
a1:N
Haggs
Hauss
UY
1:8
)
2 %6
TPP
.
b be (1, B),
(78)
where the position and velocity parts occupy half of the latent space each (have dimension d/2). Then, we simply make each u/; contain the means and standard deviations as:
b1:B
Bread
1/.BB
(79) Figure 20: Encoder for 2nd order dynamics. In Section 3.3 we described the structure our encoder. For the ease of exposition we omitted overly general descriptions and presented simple understand overall architecture (Figure 5 (a)). However, in practice we use a slightly more general setup which we show in Figure 20. As can be seen, simply
B
BB
aggregation function bigg and Agge to aggregate information for the position and velocity components separately. Then, we concatenate by and by to get b Other components remain exactly the same as described in Section 3.3.
Now, we describe the sub-components of the encoder:
convolutional neural network (CNN). In particular, heomp has the fol- lowing architecture: three convolution layers (5x5 kernel, stride 2, padding 2) with batch norm and ReLU nonlinearities, convolution layer (2x2 kernel, stride 2) with batch norm and ReLU non- lincarities, linear layer. The four convolution layers have n, 2n, 4n and 8n channels, respectively. 8, 16, and 32, respectively.
heamp is represented
For datasets PENDULUM, RMNIST, and BOUNCING BAL.LS
23

---

Published as a conference paper at ICLR 2023
Ragg and agg transformer encoders with our temporal dot product attention and relative po- sitional encodings (Section 3.3). The number of layers (i.e., L in Figure 5) is for hagg and 8 for Bagg We set Diow = 128, e = 10 110-2,2 p = oo (i.e., use masking), and finally we set 6, to 15% of the training time interval [+1. f.]. For both aggregation functions we use only temporal attention at the first layer since we found that it slightly improves the performance. In Appendix F we investigate the effects that p and 6, have on the model's performance.
hread is mapping from b, to v,,, Recall that we define b,
6
b; = 6
(80)
so hread is defined as
Linear(b.)
exp (Linear(b()) Linear(b) (b.)
T
hread (b.) =
= 2/.
(81)
exp (Linear(b)))
where Linear() is a linear layer (different for each line).
Constraining variance of the approximate posteriors As we showed is Section 4.4, forcing the variance of the approximate posteriors lv. (8;) to be at least min 0 in each direction might greatly improve the model's performance. In practice, we implement this constraint by simply adding Tmin rf.. We do add Tmin T. found that it tends make long-term predictions less accurate.
Structured Attention Dropout We found that dropping the attention between random elements of the input and output sequences improves performance of our model on regular time grids and for block sizes larger than one. In particular, at each attention layer we set an element of the unnormal- ized attention matrix CDP C to -o with some probability (0.1 in this work). This ensures that the corresponding element of C. is zero. This is similar to DropAttention of Zehui et al. (2019), however in case we do not drop arbitrary elements, but leave the diagonal of Ciy and one of the first off diagonal elements unchanged. This is done to ensure that the output element has access to least the i' ttt element of the input sequence and one of its immediate neighbors.
F PROPERTIES OF THE ENCODER
Our encoder has parameters p and 6, which control the shape and size, respectively, of the temporal attention windows (see Section 3.3). Here we investigate how these parameters affect our model's performance. At test time we assume to have access observations within some initial time interval [f1. f + frea] Figure 21 (left) shows that there seems to be no conclusive effect from the shape of the attention window. On the other hand, have noticeable effect on all three datasets. We see that the curves have the U-shape with the best performance being at 5, 6 We also see that too wide attention windows (i.e., large 6,) tend to increase the error.
Figure 21 (right) shows, parameter d, seems to
0.00
0.05
Purdulum
Pendulum
MSE.
0.44
IMNIST
Bounding
0.04
0.05
Bouncing
0.0
Balla
Balls
0.01
25
Figure 21: Test errors for different values of and 6,
24

---

Published as a conference paper at ICLR 2023
G COMMON HEURISTICS
Here we provide details about our heuristics comparison setup in Section 4.5.
G.1 SETUP
In all cases (SS, SS+sub, SS+progr, Ours), training, testing and model setups are described in Ap- pendix E. The only difference between the single shooting version of our model (SS) and the multiple shooting version (Ours) is the number of blocks. For SS we use single block, while for Ours we use multiple blocks (see Appendix E).
G.2 HEURISTICS
Training on sub-trajectories. Here, instead of training on full trajectories, at each training iter- ation we randomly select short sub-trajectory from each full trajectory and train on these sub- trajectories. For PENDULUMMMMNSTT/BOUXCC
BALLS datasets used sub-trajectories of length 2/2/6. These sub-trajectory lengths were selected such that they are identical to the sub- trajectories used the multiple shooting version our model (Ours).
Increasing training trajectory length. Here, instead of starting training full trajectories, start training on small number of initial observations, and then gradually increase the training trajectory length. particular, for PENDULUM and RMNIST datasets we start training on first observations, and then double that length every 10k iterations until we reach the full length. For BOUNCING BALLS dataset start training on first observations, and then double that length every 10k iterations until we reach the full length.
H COMPARISON TO OTHER MODELS
Here provide details about model comparison setup in Section 4.6 and show predictions from different models.
H.1 NODEP
NODEP is similar to our model the sense that it also uses the soaeee-siulateedde approach, where takes the deterministic latent dynamics, and then maps the latent trajectory to the observation space via decoder. The encoder works by concatenating the initial observations and their temporal positions, mapping each pair to representation space and averaging the individual representations to compute the aggregated representation from which the initial latent state is obtained. This encoder allows NODEP to operate on irregular time grids, but, due to its simplicity (it is roughly equivalent to single attention layer), might be unable to accurately estimate the latent initial state.
number of initial observations, maps them latent initial state, simulates
NODEP reported results on variant of RMNSIT dataset, so we use their setup directly with our RMNIST and PENDULUM datasets. For our BOUNCING BALLS dataset we used 32 filters for the encoder and decoder (close to our model), and the same dynamics function as for our model.
We train NODEP using random subsets of the first observations to infer the latent initial state. We found this approach to generalize better than training strictly the first 8 observations. For validation and testing we always use the first observations.
H.2 ODE2VAE
ODE2VAE is similar to our model in the sense that also uses the encode-simullae-dccooee ap- proach, where it takes some number of initial observations, maps them to a latent initial state, sim- ulates the deterministic second-order latent dynamics, and then maps the latent trajectory to the observation space via decoder. The encoder computes the latent initial state by stacking the initial observations and passing them thought CNN. This encoder flexible, but restricted to regular time grids and constant number of initial observations.
25

---

Published as a conference paper at ICLR 2023
ODE2VAE reported results variants of RMNIST and Bouncing Balls datasets, setup directly with our RMNIST and BOUNCIGN BALLS datasets. For our PENDULUM we use ODE2VAE with the same setup as for RMNIST. We tried to iecrasse the sizes of the ODE2VAE components, but resulted in extremely long training times.
their
For training, validation and testing the first 8 observations infer the latent initial state.
H.3
OUR MODEL
Our model followed the same setup described in Appendix E.
H.4 MORE PREDICTIONS
In the model comparison experiment (Section 4.6) we showed only the median test predictions. Here, we plot test predictions corresponding different percentiles. Figures 22, 23, and 24 show predictions of NODEP, ODE2VAE, and our model.
100'th MSC=0.046:
100'th MSE=0.029:
100'th MSE=0.011
75tth 0.330.033:
75thh MSE-0.024:
75 MSE=0.00m
50'th MSE70.027:
50tt 0.E-00017:
50tth MSE-0.003.
25tth MSE-0.022.
25thh MSE-0.011:
25thh 0.000.001
LAVAIAZ
(a) NODEP
(b) ODE2VAE
Ours
Figure 22: Predictions on PENDULUM dataset. Shown are test predictions corresponding to different percentiles test MSE. The first snapshot t, the last one is t-.. The distance between snapshots is five time points. First row is ground truth, second row is the prediction.
pere. MSF=00106.
pert.
pare.
C CEOWEENN M
EC
DEE
3
M
W
DE W
75tth perc. MSE-0 066:
75tth pere. 0SE-
75thh perc. MSE=0.020:
8 RO
QU
3 m
& 3
E 3
3
G
0
3
50'th perc.
058:
50 pero 0.031
perc. 0.01-
M
S
3
3
W
M
E
W
3
25tth perc. MSE-0.050:
25tth pere. 0.040
25tth perc. 0.09
i
W
3
A
W
1
W
& 3
M
S
3
3
W
3
3
M
3
3
W
3
3
3
3
O'th
perc. MSEE0.025.
O'th perc. MSE=0.004:
O'th pert. MSE=0.004:
3
3
3 3 3
m
3
E
3
F
3
3
(a) NODEP
(b) ODE2VAE
(c) Ours
Figure 23: Predictions on RMNIST dataset. Shown test predictions corresponding to different percentiles wrt test MSE. The first snapshot at t1, the last one is at ts-. The distance between snapshots is five time points. First row is ground truth, second row is the prediction.
26

---

Published as a conference paper at ICLR 2023
10 pere. MSR1=0.081:
100'th pert. MSE=0.1
part.
75'th pere. MSE-0.07
75 porc. MSE-0.0788
75tth pere. MSE-0.031
50'th perc. MSE- 069:
50 pare. MSE-0.071:
50 perc. MISE-0.072
25tth perc. MSE-0.067:
25tth pero. 0.063
25 pero. MSE 009:
2
0'th perc. MSE=0.060:
0'th pere. 0.000
O'th pert. MSE=0.00 0.001:
(a) NODEP
(b) ODE2VAE
(c) Ours
Figure 24: Predictions on BOUNCING BALLS dataset. Shown are test predictions corresponding to different percentiles wrt test MSE. The first snapshot is at t1, the last one is at f51. The distance between snapshots is five time points. First row is ground truth, second row is the prediction.
H.5 TRAINING WITH DIFFERENT SUB-TRAJECTORY LENGTHS
We train our model on full trajectories. Other models are trained sub-trajectories of length N. Note that this experiment we remove the iterative training heuristic from ODE2VAE to study the sub-trajectory length effects directly. All models are tested on full trajectories and the first 8 observations to infer the latent initial state. Figure 25 shows results for different values of N. We that our model outperforms NODEP and ODE2VAE in all cases. We also that both NODEP and ODE2VAE perform poorly when trained on short sub-trajectories; in figures below we show that for N 10 both models perform well the first N time points, but fail to generalize far beyond the training time intervals, which is in contrast to our model which shows excellent generalization. Increasing the sub-trajectory length tends to provide some improvement, but only up to a certain point, where the training starts to fail; in figures below we show how NODEP and ODE2VAE fail for large N.
Figures 26, 27, and 28 show predictions of NODEP and ODE2VAE trained on sub-trajectories of different lengths.
Overall, we see that NODEP and ODE2VAE tend to perform well when trained and tested on short trajectories, but do not generalize beyond the training time interval very well. Simply training these models on longer sequences does not necessarily help as the optimization problem becomes harder and training might fail. Our model provides a principled solution to this dilemma by splitting long trajectories into short blocks and utilizing the continuity prior enforce consistency of the solution across the blocks thus ensuring easy and fast training with stable predictions over long time intervals.
H.6 COMPARISON AGAINST ANOTHER BULTIPLE-SHOOTING BASED METHOD
We compare the performance of our method against Jordana et al. (2021) which use a deterministic discrete-time latent dynamics model and apply multiple shooting directly in the latent space without amortization. After training the model, the optimized shooting variables are used to train a discrete- time RNN-based recognition network to map observations to the corresponding shooting variables. The recognition network is then used at test time map initial observations to the latent initial state.
We use the official implementation from Jordana et al. (2021). For PENDU- BALLS datasets
the penalty constant of 1e3/1e3/1e4, learning rate of 1e-3/1e-3/3e-4 batch size 16/16/64, number of training epochs of 600/600/3000. In all cases the number of shooting variables is set to
LUW/RMNIST/BOUNCII
In all cases, architecture of the dynamics function and decoder is the same as for our model. The encoder of lordana et al. (2021) first maps the images low-dimensional vectors using CNN (we
27

---

Published as a conference paper at ICLR 2023
NODEP
ODE2VAE
Our
0.07
00220 0.06
0.05
Tol
100
003
102
0.00
0.16
0.14
3200
0.12
0.10
0.08
0.06
w
c wi m
0.00
0.16
0.14
0.12
(bounette
0.10
0.06
0.02
0.00
10
20
30
40
50
traj. longth
Test errors for different models and datasets. Right: For each dataset, we plot ground truth and predictions for NODEP, ODE2VAE and our model (top to bottom). Each sub-plot shows the ground truth as the first row, and the prediction as the second row. We plot test prediction with the median test error (for each model and dataset select the value of N which gives the best predictions).
Figure 25: Left:
N=10:
N=10:
N=20:
N=20:
N=30:
N=30:
N=40:
N=40:
5=50:
N=50:
(a) NODEP
(b) ODE2VAE
Figure 26: Predictions of NODEP and ODE2VAE on PENDULUM dataset when trained of sub- trajectories of length N. Shown are test predictions with the median test emror. The first snapshot is I,, the last one is at 151- The distance between snapshots five time points. First row is ground truth, second row is the prediction.
28

---

Published as a conference paper at ICLR 2023
N=10:
1=
J
3
3
3
3
3
27
3
2
3
3
3
3
w
3
3
3
3
3
(
3
C
3
8
2=
N=
3
33
3
3
3
3
3
3 3
3
3
3
3
w
3
3
3
3
3
3
3
3
3
3
3
3
C
Q
N=30:
N=30:
3
3
3
8
3
W
3
3
3
3
3
6
3
4
3
3
3
3
3
3
3
N=40:
N=40:
3
8
3
3
M
2
2
3 3
3
3
3
3
M
3
3
3
3
3
3
1
N=50:
N=50:
3
3
3
3
3
m
3 3
3
3
3
3
(a) NODEP
(b) ODE2VAE
Figure 27: Predictions of NODEP and ODE2VAE RMNIST dataset when trained of sub- trajectories of length N. Shown are test predictions with the median test error. The first snapshot is 1 the last is at tzg- The distance between snapshots is five time points. First row is ground truth, second row is the prediction.
N=10:
N=10:
N=20:
N=20:
N=30:
N=30:
N=40:
N=40:
N=50:
N=50:
NODEP
(b) ODE2VAE
Figure 28: Predictions of NODEP and ODE2VAE BOUNCING BALLS dataset when trained of sub- tro---ejeetoiissseeee of length N. Shown are test predictions with the median test error. The first snapshot is 4, the last one is at fs1 The distance between snapshots is five time points. First row ground truth, second row is the prediction.
29

---

Published as a conference paper at ICLR 2023
Dataset
Test MSE (Ours)
Test MSB (Jordana al. (2021))
Pendulum (reg.)
0.004
0.005
RMNIST (reg.)
0.016
0.020
Bouncing Balls (reg.)
0,023
0,081
Pendulum (irreg.)
0.004
0,029
RMNIST (irreg.)
0.015
0.072
Bouncing Balls (irreg.)
0.024
0.096
Table 2: Comparison results.
used the same architecture as for our model), and then applies an LSTM (we used the latent state of dimension 1024) map these vectors shooting variables. Note that the encoder trained after the model. The latent space dimension is the same as for our model. At test time we use the first observations infer the latent initial state.
We applied the method of Jordana et al. (2021) on our datasets with regular and irregular time grids and report the results in Table 2. We found that Jordana et al. (2021) performs quite similarly to our method on regularly sampled PENDULUM and RMNIST datasets, but fails to produce stable long-term predictions on the BOUNCIGN BALLS dataset. Also, due being a discrete- time method, Jordana al. (2021) fails irregularly sampled versions of the datasets.
I STRENGTH OF THE CONTINUITY PRIOR VS GAP BETWEEN BLOCKS
We investigate how the strength of the continuity prior (as measured by o..) affects the gap between consecutive blocks of latent trajectory. compute the mean squared gap between the end of a current block and the beginning the next block (i.e.. between the latent state æ at a time |i and the shooting variable 81;)). We report the results in Table 3. We see that stronger continuity prior (i.c., smaller oa) tends to result in smaller gap between the blocks and, consequently, in better continuity of the whole trajectory. We also see that better continuity tends to result in smaller prediction errors.
train our model with different values of a, and
Pendulum
RMNIST
Bouncing Balls
Test MSE
Avg. gap
Test MSE
Avg. gap
Test MSE
Avg. gap
2e-1
0.189
1.3223
0.104
6.2465
0.0805
0.0929
2e-2
0.028
0.0326
0.062
0.5094
0.0724
0.0849
2e-3
0.012
0.0017
0.027
0.0101
0.0475
0.0121
2c-4
0.002
0.0004
0.017
0.0009
0.0243
0.0012
2e-5
0.004
0.0004
0.015
0.0004
0.0825
0,0002
Table 3: Dependence of test MSE and inter-block continuity de
USING ODE-RNN AS AGGREGATION FUNCTION
Here test the effect of replacing our transformer-hasee aggregation function hagg by ODE-RNN (Rubanova et al., 2019). For each dataset, we set ODE-RNN's hyperparameters such that the number of parameters similar to that of our transformer-basse hagg- We report the results Table 4. We see that on the PENDULUM dataset ODE-RNN works par with our method, while on other datasets it has higher test error. The training time for ODE-RNN tends to be much larger than for our method highlighting the effectiveness of parallelization provided by the Transformer architecture.
30

---

Published as a conference paper at ICLR 2023
Dataset
Test MSE (Ours)
Test MSE (ODERNN)
Training time (Ours)
Training time (ODERNN)
Pendulum RMNIST
0.004
0.007
hours
68 hours
0.015
0.027
hours
98 hours
Bouncing Balls
0.024
0.036
hours
133 hours*
Table 4: Test MSE and training times for transformer-based and RNN-based aggregation functions. *Trained with block size of due long training times.
31


Reviews:
[{"review_id": "ayokwWBnwjv", "reviewer": "Reviewer_9Vai", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Neural ODEs (NODE) are very powerful to model the dynamical systems. In particular, latent NODE models are great candidate when we deal with very high-dimensional data, as they learn the dynamics in a much lower dimensional space. This paper (along with some other papers they have cited) claims that when the length of trajectories (time span) starts to grow, the loss landscape of the NODE models become very complicated leading to poor estimates and generalization. To solve this problem, they employ a technique from optimal control called \"multiple shooting\" (MS), which works has follows. The authors chop the whole trajectory to multiple blocks such that each block contains fewer observation. This allows learning ODEs for smaller periods of time and also parallelization which makes the training much faster. There is a risk that smaller blocks might be disconnected. In order to tackle this issue, they add a continuity prior in a Bayesian framework which also learns a sparse blocks. Overall, the results show that this method is effective in both improving the test errors and decreasing the training time.", "summary_of_the_review": "The paper is well written and clear. I am mostly convinced that this method is helpful in learning the dynamics for the long trajectories and appreciate the innovative solutions like MS and using Bayesian framework to enable continuity prior and uncertainty of the dynamics. However, I have some questions:\n\n- I wonder where the complexity of long trajectories come from? Is it because of error accumulation on ODE solvers for long times? If this is the case, can we shrink the time span by preserving the ordering of the observation over time? For example, if the original time span is [0,50] second, can we scale it to [0,5] and then solve the ODEs, or will we still see the same problem?\n\n- The new loss function proposed by the authors in (17-19) has six different terms. Usually, in VAE models, the optimal performance is achieved by some non-uniform weights of these terms. I wonder if the authors have played with weighting these terms? Have the authors found that optimizing the loss as is in (17-19) gives the best results? I feel this complicated loss should be tricky to train without knowing the optimal weights and curious to see how that is the case with the authors' experience.\n\n- Why have the authors chosen to use transformer instead of RNNs in the encoder? The original work of latent ODE uses RNN and also consider the case of irregular time-series. I wonder if the authors have tried RNN as well and not found it satisfying? I would recommend adding it as another comparison so that we can see how much adding transformers help in encoders.\n\n\n", "strength_and_weaknesses": "Strength:\n\n- Improved generalization and test error. \n- Improved training time.\n\n- Loss function has a lot of terms which can make finding the optimal wighting of terms more difficult.\n- Lack of evaluation based on RNN encoders.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "KP0huWZi9i", "reviewer": "Reviewer_XZSe", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper focuses on the problem of improving neural ODEs for learning from long trajectories. It shows that the loss landscape of latent neural ODEs is adversely affected by the length of the time interval and  complexity of loss can grow dramatically with increase in the length of the trajectory. The paper proposes bayesian multiple shooting technique to address this issue. Multiple shooting technique splits the trajectory into multiple short segments which are optimized in parallel while controlling the continuity between multiple segments. Standard multiple shooting approaches enforces continuity of the entire trajectory by applying a hard constraint or penalty term while training. The paper proposes to utilize Bayesian inference and naturally encoded continuity as a prior. The papers also proposes temporal attention and relative positional encodings to build a time-aware attention based encoder to deal with irregularly sampled input data.  Experiments on multiple datasets show that the proposed approach is able to achieve better performance and that multiple shooting clearly helps improve the performance and reduce training times via parallel computations.", "summary_of_the_review": "My recommendation for this paper is weak accept. Although the paper proposes a novel solution and shows clear motivation behind the proposed approach, rigorous experimentation is missing. I would be happy to increase my score if my concerns are addressed.", "strength_and_weaknesses": "### Strengths\n1. The paper presents a novel approach involving Beyesian inference for a quick and stable training of latent ODE models on long trajectories. It also introduces a transformer-based encoder with novel time-attention and relative positional embeddings for irregulalry sampled data.\n2. The motivation of the paper is clear and it focuses on solving an important problem. The experimental results indeed show that multiple shooting is an interesting and effective way to deal with long trajectories while reducing training times.\n3. The paper shows ablations to show the effectiveness of the proposed time-aware attention and relative positional encodings. \n\n### Weaknesses\n1. Different feature in irregularly sampled sequences can have observations at different times. It is not clear how does the proposed encoder takes care of this. The paper is also missing an important related work [1] in this area which proposes an attention based encoder for irregularly sampled time series.\n2. Another concern with the paper is the lack of rigorous experimentation to show the usefulness of the proposed method for irregularly sampled time series datasets. The paper does not include any experiments on real-world irregularly sampled time series datasets (e.g. PhysioNet, MIMIC-III). I would expect to see at least the comparison on the PhysioNet experiments form Latent ODE paper. \n3. The paper defines short trajectory with N = 10 in Figure 1 but the Latent ODE is running in continuous time. N = 10 can constitute of much longer trajectory based on the time difference between the observations or a much shorter if the observations are close to each other in time. \n4. The paper is missing important comparisons with the Latent ODE work with ODERNN encoder. \n5. How does the block size impact the MSE and training time on irregularly sampled time series? Empirical study with PhysioNet or MIMIC-III/IV would be very useful here. \n6. The paper is also missing comparisons with Hedge et al(2022) or Jordana et al(2022) to show the effectiveness of their proposed bayesian multiple shooting approach.\n\n[1] Shukla, Satya Narayan, and Benjamin M. Marlin. \"Multi-time attention networks for irregularly sampled time series.\" arXiv preprint arXiv:2101.10318 (2021).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "D48w_deIcb", "reviewer": "Reviewer_Lhfw", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes an amortized inference approach combined with a temporal attention mechanism for parallel training of dynamical models given long irregularly sampled trajectories. Experimental results on three synthetic datasets demonstrate improved mean squared error compared to baselines.", "summary_of_the_review": "The technical novelty is limited. Additionally, experimental evaluations are missing several baselines and rely on clean synthetic datasets. Also, the paper proposes several training heuristics with little justification. Hence the source of performance gain is unclear.", "strength_and_weaknesses": "**Strengths**\n- The paper tackles an important challenge of training models given long and irregularly sampled datasets\n- The paper is well-written and easy to follow\n- The proposed approach seems easy to implement\n\n**Weakness** \n\n*Technical Limitations*\n-  The novelty is limited. The proposed approach seems to be a straightforward application of amortized inference and attention\n- The discussion of the proposed temporal attention needs to be expanded. There are several approaches for temporal encoding. Did the paper consider alternative formulations to Eqn. 23? \n-  How is the block size chosen given observational data?\n- The performance gains seem marginal when compared to other models such as ODE2VAE and NODEP\n-  Eqn 17: What is the assumed data likelihood? It's unclear why the data likelihood depends on $s_1$ only.\n- Overall the paper proposes several heuristics for training without proper justification, *e.g.*, the paper claims constraining posterior variance is crucial without providing specific details of the actual approach. Hence, it's unclear where the performance benefits come from.\n   \n*Experimental Evaluation*\n- The paper needs to discuss and compare with other latent-ODEs approaches for irregularly sampled observations, *e.g.*, [1, 2] and temporal attention mechanisms [3]\n- The paper evaluates relatively clean synthetic datasets. The evaluation should also consider challenging real-world datasets such as Physionet\n\n**References**\n- [1]  De Brouwer et al., \"Gru-ode-bayes: Continuous modeling of sporadically-observed time series\",  NeurIPS, 2019\n- [2] Rubanova et al.,  \"Latent odes for irregularly-sampled time series\", NeurIPS, 2019\n- [3] Zhang et al., \"Self-attentive Hawkes process\", ICML, 2020", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "1ms12d6CnY", "reviewer": "Reviewer_AJZS", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work tackles the issue of modelling time-series data tractably and efficiently. Modelling long, irregular time sequences can be very challenging to optimise (in terms of the underlying landscape of the loss function), and inefficient (simulation of the time series in serial can be very slow). This work adapts the well-known technique of multiple-shooting, but reframes with a Bayesian lens. This allows modelling a time series as a series of blocks that can be simulated in parallel, which resolves both issues of efficiency and optimisation.\n\nA variational framework is used. This allows factorising the shooting variables so they can be sampled independently (hence in parallel).\n\nFor the encoder, the authors develop a novel attention mechanism and positional encoding.", "summary_of_the_review": "Overall, the paper is very solid. The method is novel and interesting, and thorough ablations and comparisons are done to a high-standard with plenty of information for reproducibility. I have a few questions (listed in Strengths and Weaknesses), but overall I’m happy to recommend this paper for acceptance.\n", "strength_and_weaknesses": "Overall, the paper is very strong.\n\nThe model is clearly explained, both on a theoretical level, and how to make it work in practice. There are thorough ablations and comparisons to other methods.\n\nI only have one question/concern:\n\nHow continuous are the models in practice? Does better continuity correlate with better performance, or is the model learning to make it discontinuous and then learning that?\nI know there is an experiment showing this for different strengths of the continuity prior, but I’m wondering if you have an experiment that measures continuity itself directly, rather than a parameter that is supposed to encourage it.\nHow correlated is the continuity prior with actual continuity (e.g. MSE between shooting points and trajectory points)?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster