Balanced Neural ODEs: nonlinear model order reduction and Koopman operator approximations
Authors: Julius Aka
Johannes Brunnemann
Jörg Eiden
Arne Speerforck
Lars Mikelsons
Venue: ICLR 2025 Poster (2025)
Keywords: Neural ODE
VAE
Dynamical Variational Autoencoder
state space models
Koopman theory
surrogate models
B-NODE
Balanced Neural ODE
Primary area: learning on time series and dynamical systems
TL;DR: We present Balance Neural ODEs, that learn dynamics with control inputs as continous, compact latent representation of adjustable complexity.

Abstract: Variational Autoencoders (VAEs) are a powerful framework for learning latent representations of reduced dimensionality, while Neural ODEs excel in learning transient system dynamics. This work combines the strengths of both to generate fast surrogate models with adjustable complexity reacting on time-varying inputs signals. By leveraging the VAE’s dimensionality reduction using a non-hierarchical prior, our method adaptively assigns stochastic noise, naturally complementing known NeuralODE training enhancements and enabling probabilistic
time series modeling. We show that standard Latent ODEs struggle with dimensionality reduction in systems with time-varying inputs. Our approach mitigates this by continuously propagating variational parameters through time, establishing fixed information channels in latent space. This results in a flexible and robust method that can learn different system complexities, e.g. deep neural networks or
linear matrices. Hereby, it enables efficient approximation of the Koopman operator without the need for predefining its dimensionality. As our method balances dimensionality reduction and reconstruction accuracy, we call it Balanced Neural ODE (B-NODE). We demonstrate the effectiveness of this methods on several academic and real-world test cases, e.g. a power plant or MuJoCo data.

Published as a conference paper at ICLR 2025
BALANCED NEURAL ODES: NONLINEAR MODEL OR- DER REDUCTION AND KOOPMAN OPERATOR APPROXI- MATIONS
Julius Aka
Johannes Brunnemann Jörg Eiden
Chair of Mechatronics University of Augsburg, Germany Jul www....uuuiiiaee....... -eee
XRG Simulation GmbH
Hamburg, Germany
(brunnemann, e eiden si
Arne Speerforek
Lars Mikelsons
Institute of Engineering Thermodynamics Hamburg University of Technology, Germany speerforck@tuin wh.etthhhcco de
Chair of Mechatronics
University of Augsburg. Germany
lars
de
ABSTRACT
Variational Autoencoders (VAEs) are powerful framework for learning latent representations of reduced dimensionality, while Neural ODEs excel in learning transient system dynamics. This work combines the strengths of both gen- erate fast surrogate models with adjustable complexity reacting on time-varying inputs signals. By leveraging the VAE's dimensionality reduction using hierarchical prior, our method adaptively assigns stochastic noise, naturally com- plementing known NeuralODE training enhancements and enabling probabilistic time series modeling. We show that standard Latent ODEs struggle with dimen- sionality reduction systems with time-varying inputs. Our approach mitigates this by continuously propagating variational parameters through time, establish- ing fixed information channels in latent space. This results in flexible and robust method that can learn different system complexities, e.g. deep neural networks linear matrices. Hereby, it enables efficient approximation of the Koopman oper- ator without the need for predefining its dimensionality. As method balances dimensionality reduction and reconstruction accuracy, we call Balanced Neural ODE (B-NODES - We demonstrate the effectiveness of this methods on several academic and real-world test cases, e.g. power plant or MuJoCo data.
INTRODUCTION
Motivation: Originally, this work is driven by challenges arising when using simulation models for optimization of complex technical systems. Increasing the fidelity of models generally results in a significant increase in computational expense. Consequently, simulation tasks such as real-time simulation for control purposes, long-term scenario analyses, optimization processes requiring hundreds or thousands simulations often become computationally prohibitive. For instance, in the automotive industry, detailed models of subsystems such air conditioning, batteries, and en- gines must operate real time for applications like predictive thermal control. Similarly, in the building sector, long-term simulations of detailed heat pump models are essential for design opti- mization and controller tuning. Therefore, attention to automated generation of faster surrogate and reduced-order models is growing. Reducing the model order directly lowers the computational pense of the solver to evaluate the underlying ODE. There exists special interest in linear surrogate models, as these allow the usage of rich of tools from control theory and efficient optimization algorithms.
This work describes combination of (Dynamical) Variational Autoencoders (VAEs) with Neural ODEs, exploiting the VAE's information bottleneck for learning reduced-order dynamical systems.
Our code is available under: ht https:://githbb.ccmm/// usaka
aaaaaaaaa/aa /ba 1

---

Published as a conference paper at ICLR 2025
Neural ODEs for time-series modeling: Neural ODEs (Chen al. 2018) have gained significant attention in dynamic systems modeling. Extensions include neural controlled differential (Neural CDEs) equations (Kidger et al., 2020) and neural stochastic differential equations (Neural SDEs)
(Liu et al. 2020). They excel in different areas, but first and foremost in time-series forecasting
e.g. in geology (Shen al.] 2023), , medicine (Hess al. 2023 Shi et al.. 2022), as surrogate models in engineering (Legaard et al. (022). To continuously integrate input signals, Neural CDEs
(Kidger al., 2020) have been used in several publications (Seedat al. 2022 Hess et al. 2023
Wi | Controlled differential equations remain uncommon in dynamic system simulation,
where state-space models are usually employed for the integration of inputs
& Vegal
1998). that can also be used with Neural ODEs (Rahman et al. 2022 Legaard et al., (2022). Neural ODEs can also be extended to include time-invariant modulators, i.e. parameterized dynamics (Lee)
& Parish, 2021 Auzina et al., (023). Convergence of Neural ODE
training can be improved by
applying stochastic noise (Volokhova et al. 2019; Ghosh et al. 2020 Liu ct al. 2020).
Dynamical Variational Autoencoders (DVAEs) are the extension of VAE (see 2.2 (Kingma &
Welling, 2014) to sequential data (Girin et al. 2020). DVAES usually include a separation of dy-
namic and fixed variables (Fraccaro et al.) 2017 Sadok al. 2024] Berman et al.) 2024), and can further be distinguished by causal, anti- or noncausal inclusion of temporal dependencies (Girin al. 2020). Deep Kalman Filters a suitable subclass of causal DVAEs for system simulation employ a prediction ceddiction-uddtte logic for the dynamic latent variable and result in hierarchical priors (Krishnan et al., 2015] Fraccaro et al., 2017). To achieve measurable state reduction, our method re- quires a non-hierarchical prior. As consequence, dynamics are part of the inference model, which additionally contains a temporal bottleneck, while the information bottleneck represented by the latent space. To our knowledge, in DVAEs, probabilistic dynamic latent representations have never been modeled with Neural ODEs. On the other hand, LatentODEs (Rubanova et al. 2011 (Shi et al.] 2022 Dahale et al. 2023) combine VAE with Neural ODEs, but only assign the probability distribu- tion to the initial latent state, and therefore fail in learning compact latent representations with input signals (4. For DVAEs, only the contribution of (Wi et al., (024) proposed (independently from us) to continuously describe the variational parameters μı a. but with a Neural CDE.
Model Order Reduction (MOR) and Koopman Theory: MOR deals with the challenge of generating models of reduced computational complexity from highly detailed simulation models (Schilders al. 2008). This mostly manifests in a reduction of the system's is state dimension. For linear systems, a variety of mathematical methods exists. Truncated Balanced Realization (Moore) 1981 is prominent example and finds a state-space transformation that orders states by balancing its controllability and observability and truncates non-important. These techniques can also be used for nonlinear systems, but only in sufficiently small regions around a linearization point (Verriest) 2008). Koopman theory offers to overcome this by using observables
g(x) of states that can be ad- vanced time with a linear operator, i.e. g(æ((1+11) = =K g( )), but with the drawback of this observable space being possibly sefinite-dimessione (Koopman 1931). Modern Koopman theory therefore deals with finding finite-dimensional approximations (Brunton et al.) (02T). A popular approach for gige-dimensional systems is Dynamic Mode Decomposition (DMD) (Schmid 2010 Rowley et al. (009). Several extensions exist, e.g. to inputs (DMDe, |Proctor et al. (2016)) or to low-dimensional systems (Williams et al. 2015), where extended DMD lifts the dimensionality by e.g. radial basis functions (eDMDc, Korda & Mezić (2018). ). Autoencoders (not VAEs) have also been applied to learn non-linear coordinate transformations to linear submanifolds (Lusch et al., 2018 [aardt et al. 2018). A different approach. modeling the dynamics with linear and non-linear components has been recently proposed by (Menier et al., 2023). common obstacle learning linear latent sub-manifolds is the choice of their dimensionality (Lusch et al. 2018, Menier et al., 2023 Mardt et al. 2018). We avoid this choice by weakly describing an information bottleneck with a balancing parameter B.
2 BACKGROUND
2.1 NEURAL ODES IN SYSTEM SIMULATION
In dynamic systems modeling, the state æ is a vector of variables that is sufficient to reconstruct all other variables of a system at time t. Important for the remainder of this work, the choice of state variables is not unique (Schilders et al. 22088). The states are variables that appear as derivatives and must be determined by solving the ordinary differential equation (ODE) governing the dynamical
2

---

Published as a conference paper at ICLR 2025
system (Cellier & Kofman 2006). While this ODE is typically formulated using first principles, in the case of a Neural ODE, the right-hand side of the equation modeled by a neural network føsonn (Chen et al. 2018),
x' (1) = SeNoee (ar(1) (o()) = mo;
+(6=
3
+
S (ar( dt.
to
The neural network Sawon. (ar(1)) describes the dynamics of the system. It maps to each state æ a dæ gradient x' = at that describes the direction in which the state evolves. In other words, a vector field as in the example in Figure A. is learned.
For the analysis of dynamical systems, additional classes of variables are critically important: 1) time-varying external inputs (t) excite the system and induce a response, e.g. a motor torque. 2) Time-invariant (physical) parameters change the dynamic properties of the system, e.g. mass. 3) Time-varying outputs y()) e.g. a power, can be calculated as a function of the other variables a<)), P, u()). By adding an crpput-fnntiion neural network Peou and the respective inputs, we get state-space model (example in Figure A.I..
x' (t) fenoee (zx(t), u(t),)) yt) = 9oom (z)) u(t). p).
x(4)) Bu.
(1)
(2)
We refer to this model 2 in the following as State Space Neural ODE (SS-NODE) and train by minimizing the MSE of predictions 23
2.2 VARIATIONAL AUTOENCODERS FOR DIMENSIONAL REDUCTION
Variational Autoencoders (VAE) are frequently used to describe high- digh-dimmesional data x = by using a smaller number latent variables. VAEs assume a prior distribution on the latent variables and express the data probability as marginalizaaion over the latent variable z (Kingma & Welling 2019 1.7)
p(x²) - p(< p(<<)>>)((()))))))) (z)p(2)dz.
For the prior, multivariate Gaussian with zero mean and diagonal covariance one is typically as- sumed, i.e. p(z) = NN 0, T). To train the model p=x:1=), the intractability of the learning problem is resolved by approximating the true posterior probability p(zla:) with an approximate model q(z) (() called the encoder (Prince 2023 2.2). The model p(x²12) [2) is then the decoder.
During training, the log-likelihood of the data log p(xx2) shall be maximized. Due to the intractabil- ity, Kingma& & Welling (2019, 2.2) derive lower bound to the data likelihood, which is called the
evidence lower bound (ELBO, derivation in I. We use the notation (Prince 2023 17.4.2)
p(x l2)
g = B) log
)) p(z)
= =E Ey(e|u log [p(x B DKL [g(x)=))[l<<>)) (3)
reconstrucionn quality
information bottleneck
with the term 3 being added in 3- (Higgins et al. 2017). The functions ((=|x=), p(x [z)
are implemented with neural networks (A2)) Typically, ll,, the encoder network parameterizes a multi- variate Gaussian distribution fa..
(2; μız), or (z)). The with following the conditional distribution of the encoder. For training, Kingma & Welling 2014) proposed to draw (one) Monte-Carlo sample of the poste- rior distribution by reparameterizing with an external noise variable, i.e. pu(z) cα(();
(x P)
pa (2) 2 o (z). such that q(
decoder predicts the data _D ou
\22:0, 0, I), sustaining backpropagation ability.
Burgess et al. (2018) present 8-VAE as a model with information bottleneck properties. Training a VAE maximizes the similarity between the data x and the reconstruction a.,,, while the least amount of information is transmitted by the intermediate latent variable z, which can be measured by the Kullback-Leibler Divergence DKL- The balance toward higher reconstruction accuracy can be adjusted by lowering B. The data locality property of VAE describes that points close in data space are also encoded close in latent space (Burgess et al.] 2018). Due to this, they excel in learn- ing disentangled representations Higgins et al. (2017); Mathicu et al. (2019), and actually exhibit local orthogonality of the latent variables (Rolinek et al. 2019).
2xp is general data, is the state.
3

---

Published as a conference paper at ICLR 2025
To identify latent dimensions carrying high information amount, KL divergence can be used. For the Gaussian prior with diagonal covariance, KL divergence can be analytically calculated (Odaibo 2019). The overall DRt. is sum of contributions from every latent channel
nz DR
nz Σ
DKL. (1o (2 )lp(2)
DKL
1 log log( o - 2 0. (4)
It forms valley Figure Ia) that ascends from ju(z;) (2,) (DKL... 0). If latent channel has D for every input the encoder, transmits pure Gaussian noise, and therefore can not transmit any information. Therefore, only latent channels with DKL.,,; 0 can transmit information. Because of this, for dimensional reduction, a fixed prior is required. Hierarchical priors would learn characteristics of the data and therefore contradict using them as fixed standard comparision.
KL divergence
divergence
500
Latent dimension
8
475
450
4
425
400
10
16
10-3
0
20
40
60
80
100
((A)
sample
(e)leatmp of DKL DKa (g(a (m) [p(e;) with g(2:/1))
(b) Heatmap Dxa different latent dimensions
VAE
((Erp()) and p(=) [2=;0,1
for given dataset: only 4 dimensions (in total 5) show values " O.
Figure I: Insights in the information transmission through the latent space. Figure 1b is the result of a 3-VAE (3 = 1 10-3) with latent space of fz = 512 trained on dataset of thermal network model, generated by random sampling of 14 simulation model parameters (details in A.3.2).
To explore the latent space of trained 3- B-E, Figure 1b reveals that DRL., is > O for only small and fixed subset of the overall offered latent dimensionality. Ordering the dimensions for this example by their mean DRI..;; reveals significant drop of DKL.2, from greater 1 to approximately between the 5th and the 6th dimension (Figure A.3). indicating the first 5 dimensions contain the most important information, which was verified by experiment (Figure A.4a). This motivates to use the mean DXL,;, per latent channel as a measure for active latent dimensions, and we choose for counting value of DKL,2, > 0.1 Fiiuure A.4b).
3 METHOD: BALANCED NEURAL ODES
Combining VAEs with State Space Models leverages VAEs' ability to generalize through infor- mation bottleneck with a numerical pre-determined latent space that promotes data locality and latent orthogonality. The information bottleneck enables dimensional reduction without à priori defined latent dimensions, with the balance toward higher reconstruction accuracy adjusted by lowering the 3-value. Additionally, the VAE's saneling-based robustness complements stability enhancements for Neural ODEs with stochastic noise. We call the proposed model Balanced Neural ODE (B- NODE), as training results balance between state reduction and reconstruction quality. In this work, we deal with continuous-time variables and their discrete-tima samples. For that, subscript denotes the value of a variable at time f. e.g. C; = srt.)). We suppose (but are not limited to) equidistant time steps At between observations &; and f;+1- A sequence of observations in (zt(t)...... x (t, ) is denoted as Tiyz. The initial time
sequence is to. and the final time is lr. For simplicity, inputs are considered constant in time interval, u(t) u,, E [4i 1.+1].
In the following, we derive the structure of a VAE that learns sequential representations of the data [ up.r- p²I. We thereby follow the vention of Girin et al. (2020) to first present the model equations and then the implementation with neural networks.
x" (o:rr Yo:r. Mg:T p) with latent variables
4

---

Published as a conference paper at ICLR 2025
3.1 INFERENCE MODEL AND GENERATIVE MODEL
In variational inference, the data x is assumed to be generated from the "true" latent variables, thus it is necessary infer these variables. We therefore describe the inference model as encoding of all variables in data space,
(<2|=>) q(p" () (or yorr uh:r p²):
where the time-invariant parameter vector (available from simulations) is simply encoded as q(p Ip It is natural to assume that the control inputs are independent from and p, and thus can be encoded sequentially as
q( ((urrl (o::) II q u.).
(5a)
The inference model for the latent states is a combination of initial state observer and simulator:
q( Forr, your: V p²) (
Forr, Vo.r, N6:T p²)
(( |a ui p²).
(5b)
initial state observer
simulator
Having in mind that a surrogate model should simulate based on the inputs x sim [zo. uo:T p), we prescribe a temporal information bottleneck starting from an initial latent state zú. With this, physical state properties
enforced: 1) All variables of the model time & can be described by a function of the set of latent states oi latent inputs uf, and the latent parameters " 2) The state can be advanced in time with an integrator and inputs u (t) that arrive later in time. For the initial state, we can argue that if we want to build a surrogate model of a physical simulation model, we have access the state vector time to. dynamic systems the initial state is in general independent from the control input and the parameters, such that we have
q To.T: your, W p²) (00))=o)
(Sba)
for the initial latent state, although integration of an initial state observer (identify æn from y could also be accommodated.
Because o and p² are inputs to the generative model and therefore known, the only data points of D we are interested approximating with surrogate, are yd [Zo.r, yo:r]. The generative model (decoder) is then simply a sequential decoding,
p(æ P (z) ~ p(y [ = pl (o.r, Vo.r Kū.T Vo:T, p²) 0 pla,, Valee uzl
(5)
Finally, the joint distribution of the observed and latent variables is
p(x = p(Zo:r Yo:r p "o p (pp(p))((urr)))))) )p(arr))p(aa.r)
(6)
To include state reduction features as discussed in 2.2 we set mnon-hieaarchica ) Gaussian prior p(z) p(p"). )(pr))(()))))) p.
Az 0, I). For using trained B-NODE a generative model, we need to control the data generation process. For that, we condition the distributions p(p*), p(rg.r))))))))) p(u).)) with the posteriors described
3.2 IMPLEMENTATION WITH DEEP NEURAL NETWORKS FOR MODEL ORDER REDUCTION
The parameter posterior density and control input posterior density
Sa
modeled
q(p" (p) N/pp;: ((pp).((p"))
(u(p").o(p") fyor (p) (u(u)), (u²)) f focm (u.).
(7)
q(u u.))) (N("; (u
(8)
The latent initial state posterior density (5ba) for wj is likewise implemented. The inference of the full latent state posterior density sequence wi. (5b) implemented with Neural ODE
q (r)) Na (xx6), (xj)),
(x()) fehearg (20)
q | u p²)
H x ++
(10)
Sexme (1(x (1), (²)
J
PP p² refer physical parameters, p()) to probability density function, du neural network parameters.
5

---

Published as a conference paper at ICLR 2025
Data Space
Latent Space
Reconstructed Data Space
P
four
u(p") a
D0
p
uo:r
plu
Tonn
oi
"i
f
(L)
pa (ta))
your
Toma
or (ta))
(m))
(xf)
(x²))
to
tr
Figure 2: Scheme of the Balanced Neural ODE.
that propagates the distribution parameters 1a(#* (1)), (a*(())) through time. Figure shows the scheme of the resulting model.
For feruue (o), different configurations are possible. We can distinguish constant variance and dynamic variance state reduction
[((x)) [x)))
f fernone. (m(x*(1)) (t)): u (t).
fórsoue
(11)
(((z) o (aff)
M (u(x)(1)) (()) u (1),p²)
(12)
fonscommnn (11(2 (1)) (ax (I), u (0)p')
For the latter, the mean is independent of the variance. Like this, information with state properties not be encoded in the variance part of the state vector.
Stability and convergence of Neural ODEs can be improved by "smoothing" the ODE's vector field with application of stochasticity during training (Volokhova al.) 2019; Ghosh et al. 2020: Liu iet al., (020). The VAE framework includes adaptively learned variance variables, that are applied as stochastic noise during training (Kingma & Welling, 2019). To bring both together, we propose to set for the unsampled inputs of feLNcu (1((x
(t)) during training
μ¹ (xx (1)) = u( (1)) + ea (xx (1)),
<~~ ()
(13)
o' (1)) or (m= (t)) + 0, ca (x (t)),
c ~ I.,
(14)
To learn the state reduction in for somm sampling of the mean state (t)) is required. If state x inactive, the model learns م (x=²) [3) is distributed unconditioned, /a (x; = (O, 1), i.e. no information transmission is possible and the model learns to ignore that dimension. For Cur, small value should be chosen.
and o (x:) 1 for that state. Then the reparameterized state
The generative model (decoder) p( p( y a 2x p²) (5) is represented with AA.4.3
Z. 91 foa (x. u.
(15)
and evaluated after sampling from the conditional distributions of the function inputs. For accurate generative usage of B-NODES, we set all noise variables e O.
3.3 ADAPTATIONS TO LEARN KOOPMAN OPERATOR APPROXIMATIONS
Learning Koopman operator approximations is simple setting for f linear model with parameters o {Ai,, Br,}) instead of neural networks. For constant-variance state reduction
[ o x
A p.p
0
M [ o
B B
u² (t).
(16)
O
0
(
and for dynamic-variance state reduction (12)
A p.
O
j ao o a
B Bp.u B
2 (t).
(17)
a (x
A
A
Opposed to Autoencoder approaches for learning Koopman operators (Brunton et al. latent dimension must only be implicitly set with the 3-value (see 2.2
2021 the

---

Published as a conference paper at ICLR 2025
3.4 TRAINING
Following the approach of Girin et al. (2020) to explicitly represent the sampling dependencies, the data likelihood and KI. divergence terms of the ELBO (3) can expanded to (derived in (A.4.1)
(a(ppP log [p(x |=)] = E (u -rppp" (uo :p)
E By (oo"o Σ
p(a, Yo [ u p")
DKL [[=x)])|p[(z]] [[p(2)]] DKL (9(p (p)pp))pp((p)) + DKL (q(w):r (ooo) (pp()))l((((rr))))
+ +Ruquu (uuu: (p) (DeL (1(f)rra " p (1P(()))))
(18)
As simultaneously learn latent embedding and dynamics in latent space, the first step in training a B-NODE finding initially stable ODE systems. Proven strategies for that are growing hori- zon (from small to long prediction horizons) and multiple shooting (splitting sequences in multi- plc. smaller ones) approaches. Then, we observe robust training behavior for constant variance B-NODEs with moderate values (i.e. 3 0.1, 0.01), which is from our experience more stable than State Space Neural ODEs with the same data.
APPLICATIONS: B-NODES FOR STATE REDUCTION
4.1 ACADEMIC USE CASE: DISCRETIZED HEAT FLOW
First, we illustrate state reduction, generalization, and performance of B-NODEs with simple 1D heat flow model. It consists of 16 cells describing the temperature field of stick whose ends are excited with external input signals (A.5.2). To generate datasets from physical models, we employ sampling strategy that aims to capture all relevant state-input combinations (A.5.I).
Figure 3a shows the original temperature field generated by the heat flow model To.r, its state-
RMSB
77.6
Dauie(aijerr)) (p(=))
mat Sinumdlyy
stale dim. C
dim.
DEL
0
0.5
0.5
0
0.5
200
0.0
0.5
1.0
[s]
[s]
[s]
sample set
[s]
(a) Reconstruction of lemperature field from latent states and latent inputs for one sample test set, B. B-NODE with constant variance.
(b) Liven though D for several dimensions, Latent
ODE with inputs ua does achieve state reduction.
Figure 3: Comparison of state reduction of B-NODE and vanilla Latent ODE Rubanova et al, 2019 for discretized heat flow example.
reduced representation in latent space æj.r and the reconstruction of the temperature field with the decoder žorr with a RMSE of 0.4 % (loss-values normalized by respective means). Taking the MOR perspective, we learn nonlinear embedding of the full state space with reduced dimensionality, requiring only 4 instead of 16 states, sufficient to globally describe the full state space. This embedding is described bi-directional with encoder and decoder. The dimensionality results from the balance between state reduction and reconstruction quality, and is modified with 3 (Figure 4a). As Figure 3al, the state reduction constant over the whole time sequence. This also holds when transtorming the weakly applied information bottleneck (by KL divergence) to a hard one (by masking latent channels), see FFiuure A.10 Because of that, latent space dimensionalities are not important hyperparameters, given they are large enough.
latent sub-manifold
Naively adding inputs 2 (t) to Latent ODE (Rubanova et al. 2019) as in Figure 3b leads to failing state reduction. The Latent ODE propagates reparameterized initial latent state zū and does not allow to enforce state reduction at times > 0.
Instead of learning the system's 's response, we learn the system dynamics ( vector field de- pending on x², u² p². This enables generalization to input signals unseen in training, for example, step inputs (Figure 4b).
7

---

Published as a conference paper at ICLR 2025
++ B-NODE, coust.
<
Latent ODE
R35SE-0,4R
m0X -30 16 Dad onr cnmrieed
uift: tempentuur
A B-NODE, dyn.
Y
TBR
100
20
RMSE (1b)
Bi partitsonl heaCCpucioor
10P
8lt pertison BJ CapnCipel
12
Number of states
()) Comparison between prediction quality of R-NODEs and Thin-
(b) Generalization: model prediction for sampled step inputs; sample
cated Balanced Realization (TBR).
with highest RMSE.
Figure 4: Generalization and comparison of methods for discretized heat flow example.
As the discretized heat flow model is described by a linear system, we can apply the model order re- duction technique Truncated Balanced Realization (TBR), which finds a balance of well-observable and controllable states (Schilders et al., 2008). Our method consistently outperforms TBR (Fig- ure 4a for example, for states the B-NODE has an RMSE of 0.3 % vs. 14.2 % for TBR.
4.2
REAL-WORLD USE CASE: SURROGATE MODEL OF A THERMAL POWER PLANT
To test the surrogate qualities with complex model, we employ a dataset generated by a power plant simulation model (Vojacek et al. 2023), of which we predict 102 states and 6 out- put signals, reacting on one input signal. Fig- ure 5 compares RMSE and maximum error of different B-NODEs and a SS-NODE as surro- gates. The B-NODE with 8 0.05 and con- stant variance exhibits the smallest reconstruc- tion error of all models while only requiring 17 states instead of 102.
+ B-NODE, const.
NODE
A B-NODE, dyn.
7.5
5.0
RMSE
2.5
6:00
m1 69
0.0
0
10
15
20
25
100
106
number of states
Figure 5: Error comparison key quantities of in- terest (outputs) for surrogates of power plant.
Taking look into the latent space of the B- NODE with dynamic variance, |Figure 6a ex- hibits a reduction of variance times with high dynamical behavior. A possible explanation could be that dynamics are encoded by states, while steady-state behavior is modeled by the decoder. However, we can perform uncertainty quantifica- tion in data space by sampling and generating multiple trajectories.
Finally, we make a performance comparison (Figure 6b). Most significantll, the B-NODE facilitates a 32x speed-up in computation time on CPU and 103x speed-up on GPU for batch simulation. It also simulates several factors faster than SS-NODE. Both models were evaluated using an adaptive steps size ODE solver. As consequence of model order reduction and training with stochastic noise B-NODE requires 10 times less ODE function evaluations for integration than SS-NODE.
14)
p(x + (t)
Scenario
Sim.
Time
ODE eval.
0.0
-2.5
Baseline with EMPy
[FMyy
2023)
2.5
0.0
CPU, 1 core
16s
-25
CPU, 16 core
1021
720s
2-
Surrogate Methods with torchdiffeq Chen
2018
00
B-NODE (CPU)
1
0.5s, 32 x
1012
1000
2006
3000
4000
5000
6000
SS-NODE (CPU)
1
3.7 4x
10033
Time [s]
N-NODE (GPU)
1024
7.0 103x
4316
Lalent state traeectorie for B. NODE with dynamic varionce (3 = 0. ordered by mean DEL, tested
NO-N (GPU)
1024
68.0 10x
654
with stop input.
Computation comparison. Technical details
A.5.4
Figure 6: Application results for the power plant simulation model.
8

---

Published as a conference paper at ICLR 2025
4.3 BENCHMARK
provides a benchmark of the proposed methodology for different datasets (details for Mu- JoCo and Waterhammer in An..5.7 A.5.8). As shown, B-NODE only marginally increases the predic- tion error compared to other methods but vastly reduces the required state dimension for all datasets. Not controlling the latent state embedding as for B-NODE with 3 0.0 and Latent Neural ODE leads to higher errors or failing prediction if inputs are incorporated.
Table
Table 1: Benchmark results (RMSE and state dimension) for nonlinear surrogates. A.5.9
SHF
Power Plant RMSE
MuJoCo
Water Hammer
RMSE
dim.
dim.
RMSE
dim.
RMSE
dim.
SS-NODE
0.01
16
0.06
102
0.43
27
0.02
400
B-NODE 3 0.0
0.01
16
18.2
128
0.51
128
0.05
512
B-NODE 3 = 0.01
0.02
5
0.02
11
0.28
25
0.03
8
B-NODE 3 = 0.1
0.04
4
0.04
6
0.50
11
0.05
4
Latent NODE B 0.1
0.05
16
0.14
128
0.43
128
0.06
512
5 APPLICATIONS: B-NODES FOR KOOPMAN OPERATOR APPROXIMATION
5.1
ACADEMIC USE CASE
An example for a nonlinear system that can be linearized by adding a state is (Brunton al. 2021)
21
a
0
T1
T2
T1
b
(I2
22
0
0
0
8
2a
-
L3,
T2
(19)
System dynamics of linear systems are determined by the cigenvalues (here À = b, b,,2).. If two systems have the same eigenvalues, they represent the same dynamics. The B- NODE choose different states than the analytic solution, but projects them back into the state space of the analytic model with linear decoder. We trained constant variance B-NODE as in 3.3 with linear layers with a dataset generated by sam- pling initial values with parameters =00.5, b=-1 for different values of 3 and offered latent dimensionality of dim =664 664. For B of 0.01, we found active dimensions (DxL (xx) > 0.1), just as the analytic solution has. The approximated and analytic eigenvalues lie at almost the same position (Figure 7p and therefore represent the same dynamics.
B-NODE Analyti
Imf A]
NRMSE
I.IS min-max domo
0.. 30, L.)
lestact
0
Re(A]
Figure 7: Eigenvalues of analytic solution and B-NODE.
5.2 REAL-WORLD USE CASE: LINEARIZED MODEL OF A THERMAL POWER PLANT
A model with full linear path from model input model output is desirable for model-predictive control (MPC) applications. It be learned by setting linear mod- els with an additive bias for the control encoder f and the decoder fo.. Then, only the state encoder has non-linearities and facilitates a nonlinear state transfor- mation into linear latent space that can approximately represent the dynamics. We trained constant variance B- NODES with 3 E (0.1, 0.001 for this use case with "dimm="=)m =1024 to predict all 102 state variables and output variables of interest for the thermal power plant simulation model. The best model achieved a RMSE of 1.2 % and a maxi- mum error of 14.5 % on the selected outputs (normalized
600
old or womten d/tk
pec
1PP
simpleomeericc shaft
Wil qusdrpicl
Wibh qunduplel
381
Figure 8: Reconstruction sample of B- NODE (3 0.01) with highest and 25 linear states (original 102 non- linear; offered dim -i022.))1024).

---

Published as a conference paper at ICLR 2025
by mean) (Figure 8), while requiring 25 states. Interestingly, though we only have control input signal u(t), the latent control dimension increases to 18 inputs. We observed that for short pre- diction horizons, the approximation quality is better, which fits the short-horizon model requirement of MPC.
5.3 BENCHMARK
Table 2: Benchmark results (RMSE and state dimen- sion) of linear surrogates A.5.IO).
A benchmark with other popular linear sur- rogate methods is given in Table 2] For the Power Plant, B-NODE achieves the same accuracy than eDMDDe, completely avoid- ing observable function tuning. The Small Grid model [ of computational fluid dynamics (CFD) can be learned with higher accuracy than with DMDc.
Power Plant RMSE
Small Grid
dim.
RMSE
dim.
B-NODE 0.01
0.20
25
0.48
example for MOR
B-NODE 0.1
0.20
12
0.52
17
eDMDe
0.20
45
DMDc
0.60
2
DISCUSSION
Table 3: Comparison of B-NODE and other surrogate methods.
Control inputs
Nonlinear Model- order reduction
Koopman operator approximation
Low-resolution data
Uncertainty
Quantification
B-NODE
SS-NODE
DMDc
eDMDe
LatentODE
Balanced Neural ODEs are model order reduction method and therefore result a compromise be- tween reconstruction accuracy and dimensional reduction to lower the computational burden. As showed in our experiments, in comparison to other state-of-the-ar methods, B-NODEs can signifi- cantly lower the model order, while only marginally increasing the prediction error. Computational speed can be increased by a factor of 8 compared to SS-NODE, as the ODE solver requires signifi- cantly less ODE function evaluations. As B-NODE is GPU compatible, a speed-up of 100 compared to physical simulation models is possible for batch simulations. The B-NODE framework is versa- tile enough accommodate nonlinear model order reduction as well as linear Koopman operator approximation, and additionally allows for uncertainty quantification. Compared to MOR methods, B-NODEs provide a well-automatable method as extensive hyperparameter tuning (e.g. bottleneck dimension or observable functions) can be avoided.
Thus, BNODE offers robust and versatile approach for surrogating detailed simulation models, making it applicable across diverse domains such as energy and automotive. Its flexibility enables efficient handling of a wide range of simulation tasks, including optimization, control, and other computationally demanding scenarios.
7 CONCLUSION
We present Balanced Neural ODEs (B-NODEs), balancing state reduction and reconstruction qual- ity, a well-automatable method for global nonlinear model order reduction. For that, combine VAEs with state space models to leverage VAE's information bottleneck to enable measurable di- mensional reduction without à priori defined latent dimensions. Additionally, the VAE's sampling- based robustness complements in literature reported stability enhancements for Neural ODEs with stochastic noise. We are motivated by the need for fast surrogate models with tire-varying inputs. To achieve this, we continuously propagate variational parameters through time with Neural ODE (to our awareness, for the first time). We tested our method on academic and real-world cases and show promising results in terms of reconstruction quality, generalization, model order reduction, and reduction of simulation time. We also showed that the B-NODE framework is well suited to learn Koopman operator approximations linearize simulation models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTI
A large part of this work was carried out as Julius Aka's master thesis at Hamburg University of Technology within the company XRG Simulation GmbH, substantially supported by University of Augsburg. J.A. would like thank XRG Simulation GmbH for supporting this project and the pos- sibility to make its content public. J.A. would also like to thank Yi Zhang and Tobias Thummerer from University of Augsburg for their support for this research.
The continuation oft this research was partially funded by the ITEA 4 project OpenScaling (Open standards for SCALable virtual engineerING and operation) Nº22013 under grant number 011S23062H. See https: http:://cpenccaaa ing. org/ for more information.
REPRODUCIBILIT STATEMENT
We provide the source code ba gellancd-neuuall. www.rnuuurlleeeee
In addition to the code necessary for implementing VAE, B- NODE and SS-NODE, we also provide the files that implement the first-principle models along with FMUs (Modelica Association, 2022), containing simulators of those. We also provide the configuration files, that were used for the data generation process and training of all models.
Apart from this, we provided detailed explanation of B-NODEs in this paper, and included more details sampling of first-principle models for dataset generation, data standardization, approxi- mation the data likelihood, and most important training settings in the appendix.
REFERENCES
Julius Aka, Johannes Brunnemann, Svenne Freund, and Ame Speerforck. Efficient global multi parameter calibration for complex system models using machine-learning surrogates. In Pro- ceedings of the 15th International Modelica Conference 2023, pp. 107-120, 2023. ISBN 978-91- 8075-505-4. doi 1..3384//c//
Modelica Association. Modelica standard library 4.0.0, 2023. URL https:///ithhbb com/ aodellccc/ooeeloo@Saaaaaa..__a
1lze Amanda Auzina, Çagatay Yıldız, Sara Magliacane, Matthias Bethge, and Efstratios Gavves. Modulated neural odes, 2023. URL https::/openrevii net wwwp:/eeefffffmlofmmm
0p9z20fXbC
Nimrod Berman, Ilan Naiman, Idan Arbiv, Gal Fadlon, and Omri Azencot. Sequential Disentan- glement by Extracting Statie Information From A Single Sequence Element, June 2024. URL https::/arxiv.orgg
wttp://r//bb//......e..// 1813
Steven L. Brunton, Marko Budišić, Eurika Kaiser, and J. Nathan Kutz. Modern koopman theory for dynamical systems, 2021. URL httttttts///axiv/e org/ abtp:/aab//112.2...
Christopher Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 9-vae, 2018. WRlphttts://axiv
orq htg/abs//88 1804 03599
François E Cellier and Ernesto Kofman. Continuous system simulation. Springer Science & Busi- ness Media, 2006. 10.00.0007700
Ricky T. Q. Chen. Eorchdiffeg
torchdiffeq. 2018 URL httta:://ggtaub..coo/ttttttt
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary dif- ferential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Biandr and R. Garnett (eds.). Advances in Neural Information Processing Systems, volume 31. Cur- ran Associates, Inc., 2018. URL htt.http:://prooeeee httpss::/proceediing.. nwwr wwtpppprrpcpppppppppppppp.ppp..p fi File paper/2018/f111 69386 6f (5
Shweta Dahale, Sai Munikoti. Balasubramaniam Natarajan, and Rui Yang. Latent neural ode for integrating multi-timescale measurements in smart distribution grids. In 2023 IEEE Power
11

---

Published as a conference paper at ICLR 2025
& Energy Society Innovative Smart Grid Technologies Conference (ISGT), pp. 1-5, 2023.
doi 100
FMPy. Fmpy, a free python library to simulate functional mock-up units (fiuus, 2023. URL https://githin httpa////ichuuub.c/aacaaaaa Sooo//CATASSSstemmssss SSy-ttssttteeteee
Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther. A Disentangled Recognition and Nonlinear Dynamics Model for Unsupervised Learning. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Arnab Ghosh, Harkirat Singh Behl, Emilien Dupont, Philip Н. S. Torr, and Vinay P. Namboodiri. STEER Simple temporal regularization for neural odes. CoRR, abs/2006. 2020. URL lattps://proceedinggg uet
neurips catasscccppppeee cc/ppperr/202// haa79990csc5aa22004eallleele2222
Laurent Girin, Simon Leglaive, Xiaoyu Bie, Julien Diard, Thomas Hueber, and Xavier Alameda- Pineda. Dynamical variational autoencoders: A comprehensive review. CoRR, abs/2008. 12595, 2020. URL/ttps://axxv..rr/////
12595
Konstantin Hess, Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Bayesian neu- ral controlled differential equations for treatment effect estimation, 2023. URL https://
openreview ner/forum openrreview.net//oorumt www.oruumnndwww00aaaaaa
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. heta-VAE: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations, 2017. URL
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equa- tions for irregular time series. In H. Larochelle, M. Ranzato, Hadsell, M.F. Balcan, and Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 6696-6707. Cur- ran Associates, Inc., 2020. /ttp///pse//ppoeeeeee....ooe.oe
files/
paper/2020/Eiil 4a587 /AA.A///51100000111111600772e2e0
Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2014. URL https:
7/arxiv. wwt./org//aaa abs 1312.66 6114
Diederik P. Kingma and Max Welling. An introduction to variational autoencoders. Foun-
dations and Trends@ in Machine Learning, 1244):307-392, 2019. ISSN 1935-8237.
dbi100156112200000066
B. O. Koopman. Hamiltonian Systems and Transformation in Hilbert Space. Proceedings of the National Academy of Sciences, 17(5):315-118, May 1931. doi 10.103//assss 17.5.315
Milan Korda and Igor Mezić. Linear predictors for nonlinear dynamical systems: Koopman op- crator meets model predictive control. Automatica, 93:149-160, July 2018. ISSN 0005-1098.
Rahul G. Krishnan, Uri Shalit, and David Sontag, Deep Kalman Filters, November 2015. URL https::/azzt
Kookjin Lee and Eric J. Parish. Parameterized neural ordinary differential equations: applications to computational physics problems. Proceedings of the Royal Sociery A: Mathematical, Physical and Engineering Sciences, 477(02233::0210166) 2021.
Christian Møldrup Legaard, Thomas Schranz, Gerald Schweiger, Ján Drgoña, Basak Falay, Cláudio Gomes, Alexandros Iosifidis, Mahdi Abkar, and Peter Gorm Larsen. Constructing neural network- based models for simulating dynamical systems, 2022. URL httpps://dl acm. org/do..... doi/
10. 11 15.1145/3555 44553556557
Oliver Lenord. Proper hybrid model generation of fluid systems: From research to- wards accesible solutions, 2024. URL htttps://tts http://www../aaattttt/eeeeemnnr/ P PHyMoS Resu tumumma ry_ The rmoS i pdd
12

---

Published as a conference paper at ICLR 2025
Xuanqing Liu, Tesi Xiao, Si Si. Qin Cao, Sanjiv Kumar, and Cho-Jui Hsieh. How does noise help robustness? explanation and exploration under the neural sde framework. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
Bethany Lusch, J. Nathan Kutz, and Steven L. Brunton. Deep learning for universal linear em- beddings of nonlinear dynamics. Nature Communications, 9(1):4950, November 2018. ISSN
2041-1723. doi IO.
Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. VAMPnets for deep learning of molecular kinetics. Nature Communications, 9(1):5, January 2018. ISSN 2041-1723. 10.100.033//44
017-02388-1
Emile Mathieu, Tom Rainforth, N Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autoencoders. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings tfe 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 4402 4412. PMLR, 09-15 Jun 2019. URL https: https://proceediingg
mlr. wwwpppresssss. htt////athhhal.eeee html
Emmanuel Menier, Sebastian Kaltenbach, Mouadh Yagoubi, Marc Schoenauer, and Petros Koumoutsakos. Interpretable learning of effective dynamics for multiscale systems, September 2023. URL /ttpp::/axx.v.ooggg///oo///
05812
Modelica Association. Functional mock-up interface for model exchange and co-simulation, 2022. att.aiiesttaacarrreeeeeeseeeeees oad/
URL http::://git.mb..omm///o v2. 0. htt.///--Spee.iicaaioooooo---ee. 2..
B. Moore. Principal component analysis in linear systems: Controllability, observability, and model reduction. IEEE Transactions on Automatic Control, 26(1):17-32, February 1981. ISSN 1558- 2523. TO. 109/TAC. 1981.1 1102688
Stephen Odaibo. Tutorial: Deriving the standard variational autoencoder (vae) loss function, 2019.
Shaowu Pan, Furika Kaiser, Brian M. de Silva, J. Nathan Kutz, and Steven L., Brunton. Pykoopman: A python package for data-driven approximation of the koopman operator. Journal of Open Source Software, 994))55811, 2024. 10.1 URL htttps::/doi..zgg/...ot
21105/joss 05881
Simon J.D. Prince. Understanding Deep Learning. MIT Press, 2023. ISBN 9780262048644. URI.
http://udlbook. com
Joshua L. Proctor, Steven L. Brunton, and J. Nathan Kutz. Dynamic Mode Decomposition
with Control. SIAM Journal on Applied Dynamical Systems, 15(1):142-161 January 2016.
10.1170117/5
Aowabin Rahman, Ján Drgoña, Aaron Tuor, and Jan Strube. Neural ordinary differential equations for nonlinear system identification. 2022 American Control Conference (ACC), pp. 3979-3984, 2022. doi 11..3399///CC55... CCC1 2022.8862289988888
Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pea direc- tions (by accident). In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 6 2019. doi 100.10099CVPR2229900022000
Clarence W. Rowley, Igor Mezić, Shervin Bagheri, Philipp Schlatter, and Dan S. Henningson. Spec- tral analysis of nonlinear flows. Journal of Fluid Mechanics, 641:115-127 1227 December 2009. ISSN 1469-7645, 0022- 120. 110100000070002211000999
Yulia Rubanova, Ricky T. Q. Chen, and David K Duvenaud. Latent ordinary differ- ential equations for sareuee time series. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- -uc, E. Fox, and R. Garnett (eds.), Advances Neu- ral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL httpa//rrrppppppppppaaaaaaaa///a 2245445555eooF774aacccc33444444 httpee/eesee//aee////////L1L.1
13

---

Published as a conference paper at ICLR 2025
Samir Sadok, Simon Leglaive, Laurent Girin, Xavier Alameda-Pineda and Renaud Séguier. A multimodal dynamical variational autoencoder for audiovisual speech
representation learning. Neural Networks, 172:106120, 2024. ISSN 0893-6080.
106120 URL https:: www.scciencedee rect.
com/ fttp:///sseooooeerrreeeee///////
Wilhelmus HA Schilders, Henk A Van der Vorst, and Joost Rommes. Model order reduction; theory,
research aspects and applications, volume 13. Springer, 2008. d6:110.1007/978-3-5400--
Peter J. Schmid. Dynamic mode decomposition of numerical and experimental data, Journal of Fluid Mechanics, 656:5-28, August 2010.
ISSN 1469-7645, 0022-1120.
IO. 1017/50022112010001111
Nabeel Seedat, Fergus Imrie, Alexis Bellot, Zhaozhi Qian, and Mihaela der Schaar. Continuous- Time Modeling of Counterfactual Outcomes Using Neural Controlled Differential Equations,
June 2022. URL bttps::/proceeding: mml.. wwe.s.peesseee wwwss/vvesse..eeee seedat 22b. html
Chaopeng Shen, Alison P Appling, Pierre Gentine, Toshiyuki Bandai, Hoshin Gupta, Alexandre Tartakovsky, Marco Baity-Jesi, Fabrizio Fenicia, Daniel Kifer, Li Li, et al. Differentiable mod- elling to unify machine learning and physical models for geosciences. Nature Reviews Earth & Enviromment, 4(8):552-567, 2023. 10. 110001088//3311--020005555
Ruian Shi, Haoran Zhang, and Quaid Morris, PAN-cODE: COVID-19 forecasting using conditional latent ODEs. Journal of the American Medical Informatics Association 201e:aatioon2911222229999-99 2022.
ISSN 1527-974X. doi 10. 10.1093/aamia/ocacccooooo 160
ClaRa Development Team. Clara library. URL https: /tRLLt/t/////ttu..coooooo/ www.:/coomiggiiiii
ClaRa-official
Mark Towers, Ariel Kwiatkowski, Jordan Terry John U. Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, Rodrigo Perez-Vicente, An- drea Pierré, Sander Schulhoff, Jun Jet Tai, Hannah Tan, and Omar G. Younis. Gymnasium: A standard interface for reinforcement learning environments, 2024. URL https://arxiv: org/ abs/ 2407.1703
E. I. Verriest. Time Variant Balancing and Nonlinear Balanced Realizations. In Wilhelmus H. A. Schilders, Henk A. van der Vorst, and Joost Rommes (eds.), Model Order Reduction: Theory, Research Aspects and Applications, pp. 213-250. Springer, Berlin, Heidelberg, 2008 ISBN 978- 3-540-78841-6,
Ales Vojacek, Johannes Brunnemann, Tim Hanke, Thomas Marx-Schubach, and Jörg Eiden. Status of the clara library: Detailed transient simulation of complex energy systems. In Proceedings of the 15th International Modelica Conference 2023, pp. 617-626, 2023. ISBN doi 10.344443//ccp204
Alexandra Volokhova, Viktor Oganesyan, and Dmitry Vetrov. Stochasticity in neural ODEs: An empirical study. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2019. URL
Hyowon Wi, Yehjin Shin, and Noseong Park. Continuous-time Autoencoders for Regular and Ir- regular Time Series Imputation. In Proceedings of the 17th ACM International Conference on Weh Search and Data Mining, WSDM '24, pp. 826-835, New York, NY, USA, March 2024. Association for Computing Machinery. ISBN 9798400703713 13. 1655
Matthew O. Williams, Ioannis G. Kevrekidis, and Clarence W. Rowley. A Data-Driven Approxima- tion of the Koopman Operator: Extending Dynamic Mode Decomposition. Journal of Nonlinear Science. 25(6):1307-134 December 2015. ISSN 1432-1467 1467.
Omar G. Younis, Rodrigo Perez- Vicente, John U. Balis, Will Dudley, Alex Davey, and Jordan K Terry. Minari, September 2024. TRLLattss://dat wtt.:/i..rgoo1.. http://1o.521//eecccccee 13767625777755
Jesús M. Zamarreño and Pastora Vega. State space neural network. properties and application.
Neural Networks, 1196:1099-1112, 1998. ISSN 0893-6080. doi/https:/dooiorr//1.. 1016/800999.
6000494
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 DATA STANDARDIZATION
In this work, all values that are fed to the neural network models are standardized. This is especially of importance, as physical quantities might have very different order of quantities (e.g. pressure at around 1 x 10 Pa compared to a temperature at around 50% CC). Also for loss functions, standardized values used to avoid adapting weighting factors to the different orders of magnitudes. In general, the not-standardize values are used only for graphical presentation of results.
To standardize the data of data set, mean μ and variance we use the following transformation for each axis x. (Prince (023), C.2.4):
calculated for each variable. Then
I 1
staand
(20)
To ensure the same standardization for every data sample, the mean and variance values have to be saved. Because of that, we introduce a Standardization Activation Layer. If not previously initialized, with the first of batch of data, mean and variance calculated and saved untrainable parameter in the PyTorch-Model. It can be used as normal layer within the model, and has also the ability to inverse standardize i.c.
sstand. + μ..
(21)
Opposed to batch-normalization (Prince 2023), 11.4), where the statistics μ and are calculated for each batch independently, the statistics are constant.
Physical model parameters are standardized along each feature axis. All other time-varying values are standardized along the feature axis and the time axis.
For presentation and calculation of summary statistics in this publication. chose normalize by mean values:
Zporm.
(22)
A.2
STATE SPACE NEURAL ODE
I(9)
x' = f(x)
d!
-Up-- R-1
(1,(1)
(mi(()
25
dLc
+
0
L
Ue($)
0
[UR]
R-
0
-20
y/ =
UL
Uc
R-
U
Uk(+)
-0.5
0.0
0.5
[[A]
(1)) = (Ua(())); p (R, C)
Figure A.I: The series resonant circuit with an voltage source as external input as an example of a dynamical system. We show its vector field (for no input) and its state space model.
A.2.1 TRAINING
Given a dataset D = To: Vo.T, uo:T, p). the State Space Neural ODE trained by minimizing the error of the predictions (t))
ΦNODE; δoout argmin MSE(#o. To.r) + MSE (o:r, Morr).
(23)
A.3 VARIATIONAL AUTOENCODERS FOR DIMENSIONAL REDUCTION
A.3.1 ELBO DERIVATION
We could learn a model p(æ)2) 22 if we would know the latent variable z. But we do not know z. To overcome that, we would require the "true" model p(z(a) (a)) that represents the so called poste- rior probability Pr(2(x). But we do not have that model either. The learning problem is intractable.
'To explain why the term posterior probability is used for Pr(z m): the known prior probability. In terms of Bayes' Rule, after observing the generation process Pr(z|2) we
We took the assumption that Pr(z) is
15

---

Published as a conference paper at ICLR 2025
To make the learning problem tractable, a posterior model g(2)22) that approximates the true posterior probability model p(z[z) is introduced:
g(x[2) R p(z|z).
(24)
The model g(=[=) is called the encoder recognition model (Kingma & Welling 2019
Chapter
2.1). This approximation is the variational approximation that gives the VAE its name.
Prince
2023 Chapter 17.5).
With this approximation, the latent variable model can be trained using data samples æ.
The following can be derived by using Jensen's incquality (as in (Prince 2023 chapter 17.3.33)), but we choose a derivation by (Kingma & Welling 2019 chapter 2.2) that avoids this.
To obtain the latent variable model, express the log- eigllikelihood log p (x) of the data expectation over the posterior probability g=x/2>)
log p(x) Bx(=(m) log [p(x)).
(25)
By applying the the product rule
p( x) [((2(()
log p(æ) Bg(((()) log
expanding with a(=(=)
p(x, 2) ((=(x)) <(((() (x) p(z|z))
(=(m)) log
splitting the integrand of the expectation
p(x, 2)
7 ((2 [p(z|x)) xx)
(=( log
+ (a()) log
{(( xx)
and applying the product rule for the first term, we can derive
p(x]22
[(2 (r) p(z (x)
p(x) Ex((((m) log
((2 xx) p(z)
E (((mm))) log
(26)
(9BBBO
The KL-Divergence in the second term is always 2 0. Therefore, the first term is a lower bound to the log-likelihood of the data. is called the evidence lower bound (ELBO). (Kingma & Welling 2019 Chapter 2.2) With the chosen variational approximation, it is possible to compute all parts of the ELBO. Before showing that, we note that the Li-eeeieeeecee of the second term determines two distances (Kingma & Welling 2019 Chapter 2.2):
The KL-Divergence between the approximate posterior and the true posterior,
* The gap between the ELBO and the marginal likelihood log p(x). This is called the tight- ness of the bound.
It is now possible to train variational autoencoder by maximizing the ELBO g. Maximizing ELBO will concurrently optimize two objectives (Kingma & Welling 2019 Chapter 2.2.1):
It will improve the prediction quality log p(x) of the generative model p(x]22) It will minimize the error of approximating the true posterior distribution p(z(ar) through
q( meaning that the encoder will get better.
The KL-Divergence in 3 for the comparison with the assumed prior p(z) = -N 0, I) can be solved analytically to (Odaibo 2019)
ngg
DKL. (gom. x (x [p(=)))
Σ
2 +++1oge 1og(())))
m3 2 0,
(27)
infer the posterior probability Pr(z x)). It's a bit cumbersome to call this posterior probability. Even though this would he the first thing needed for the training of VAF, but in terms of the VAF theory, the starting point that the data generation process Pr(æ)] z) he traced hack preseribed prior probahility Pr(z).
16

---

Published as a conference paper at ICLR 2025
with n² being the chosen dimensionality of the latent variable.
+ 0 010,1)
Decoder
Encoder
Pen
Ode
Figure A.2: Scheme of Variational Autoencoder
A.3.2 EXPERIMENT DETAILS
Dataset
The dataset was generated by random sampling of parameters of a physical simulation model with 14 parameters. The VAE model itself has no concept of sequences, therefore the sequences for temperature trajectories formed one sample (i.e. 5 sequences 337 timesteps per sequence = 1685 datapoints per sample). We therefore re-implemented the model described in (Aka al. 2023
parameter values: determined from building data
parameter ranges: } to 5 times the default value
timestep: 1800
simulation length: 7d
number of timesteps: 337
number of samples: 2048 (12% test, 12% validate, 76% train)
Experiment Details
Network Structure:
number linear layers: 4
hidden dimension: 2048
latent space dimension: 512
activation function: ReLU
Training Parameters:
beta: x 10
max epochs: 104
learning-rate: x 10-3
weight decay: x 10-1
batch size: 256
carly stopping patience: 300
early stopping threshold: 0.005
early stopping mode: relative
number reparametrization passes for train and test: I
automatic mixed precision: off
parameter decoder: (PELS-VAE), off (VAE)
Code File: \networks wtw.e///wwwweee...eeeeeeeeeeeeee
17

---

Published as a conference paper at ICLR 2025
A.3.3 ADDITIONAL FIGURES
distribution density
distribution density
KL distribution density
dim())=449
diml 435
dimi()) [= 107
dhm(=)=4110
ditn(z)=47)
dim 193
0.5
100
KL
Density distributions of DKL, ,4, a for the first 7 latent space dimensions for the com- plete test set ordered by their mean KL divergence. Plot for the 3-VAE with B 1 10-3 The model was trained with a data set of the office model
Figure A.3:
To verify that the first 5 dimension encode most information, we ordered the latent dimensions of a trained VAE by their average DKL. masked all latent channels and consecutively lifted that mask while calculating for each configuration the MSE on the dataset (Figure A.4a). As anticipated, the first 5 latent channels reduce the error by the most (0.85 + 0.09) while the last 507 dimensions only achieve marginal improvements (0.09 + 0.04). Notably, the improvements are monotonously decreasing. This motivates to use the mean DKL per latent channel as a measure for active latent dimensions, and we To find threshold value, tested variety of values in Figure A.4b and choose for counting value of DRI. (x) > 0.1 (Figure A.4b).
0.35
threshold
20
MSE
0.35 @1
2.2
0.09
03
0.06
0.06 e 125
Q/512
10
#5
10a
10at
10
10
10P
active latent dimensions
KL Uhreshold
(a) Experiment MSE on dataset for increasing number of active latent dimensions by lifting mask on -xt-ordered latent dimensions.
(b) Number of latent dimensions that have KL divergence greater certain threshold, averaged over one dataset.
Figure A.4: Insights in the information transmission through the latent space. Results 3- B-V
(3 1 x 10-3) with latent space of nz 512 trained dataset of thermal network model, generated by random sampling of 14 simulation model parameters (details in A.3.2
18

---

Published as a conference paper at ICLR 2025
A.4 BALANCED NEURAL ODES
A.4.1 B-NODE ELBO DERIVATION
In the following we expand the ELBO to our use case. As provided by Kingma & Welling 2014))
the ELBO
g (= log [p(x*))>)
DKI. (((2()"))ll(())
(28)
Reconstruction Accurncy
Information Bottleneck
First, collect the governing equations of the model and adapt the way they are expressed (5 Sah. 5b). (bb). (5), (6)) The approximated posterior (inference model) is
(<2(x)) q( Mo.r. (xorr Uo-r p) g(p" ())(((
((p")))(((n / (uo: (a Forr (.:
II g((()u;) <(p" pp)
( 2r
II
qlafft 22 up,p"
Parameter Encoder Initial Observer
Control Encoder
Simulator
=g(un:r lu
g ( Po- uo:T
(29a)
independent sampling possible
dependent en previous sampling
the generative model is
T pla ua uZ,p"),
p(æ" [2) = p(ZooTT your Fi::: "o) = II
(29b)
1=0
and we prescribed the prior as
p(z) "epp(((ppppppp )p(rrr))((rrr)))
(29c)
When inserting these terms in the ELBO, the learning problem is intractable, the probability density functions are parameterized by outputs of neural networks that cannot be solved analytically. Therefore, sampling from the posterior distribution (=2[=²) in conjunction with the reparameter- ization trick is necessary, whenever latent variables are propagated from deep neural network to another. The propagation of latent state distribution by the Neural ODE an exception for this, this propagation of probability distributions is implemented tractably and state reduction ensured by the sampling of inputs
First, we derive the data likelihood part of the ELBO,
Bg(=|eP) log [p(2*1 [x)] = g Imo p))((p "p log [p(xo:r, yorr kerr wi:r: p²)],
Considering the sampling dependencies,
Bygug. - pp u pp)
Byta 1e log [p(20.T, yorr 2ō.r uar p²)]
Using that p(= 2)) can be factorized,
T
Eggug y- (Mo.T
log p(x, y1 p²) Σ
(30)
After this, we derive the KL divergence part. By using the derivation in A.4.4, we can see that the KL divergence splits in separate terms for different random variables. Because of this,
DRL. [q(=l=P))lpd [[]())] DRL. [q(p" pp) (o. ( (zo, uĀ. p²) )ptp::")))"""""..rreree (p(u).))p((
As q q( Po, uor: p²) depends on previous sampling, we cam write this as,
DKLL [q(p" (p)lp((p))) )] + DKL [g(wū:r/u / )p((ur)))
+ (y(" (uorrpp) (DeL. (g (oo u p=)p==((()))))))
(31)
19

---

Published as a conference paper at ICLR 2025
Each variable in the KL divergence defined here expands likewise sum of KL divergence terms. Because we prescribe (parameterized) Gaussians, the KL divergence can be solved analytically. The model is then trained
Deapp: Penu, Penx, LLNODE, dede argmax
g ( Pen,s Peax: LNODEE Øde To.T, Vo.T, Pp) (32)
Typically, as proposed by Kingma & Welling 2014), we sample each distribution only one time during training and train with mini-batches. Although we elaborated here on the sampling dependencies, it is not necessary to obey to this when always sampling once.
Remark:
do not condition the true posterior probability density of the generative model for the latent states at time tę on the previous probability densities (for 1 u but prescribe the latent space in general to follow Gaussian distribution with mean zero and variance of one (N(0, I)) at all discrete times E [0, 7) This means that we do not condition the generative model in equation 4.3 in their review on anything. and means as well that all conditional probabilities evolving an advance in time are collected in the inference model.
Remark: Opposed to Girin et al. 2020),
A.4.2 NORMALIZATION OF ELBO TERMS
We normalize the reconstruction term with 2 to account for the summation of the state reconstruction and output reconstruction MSE,
1 GR ~ [MSE( x) MSE(Ú. ()).
(33)
Let DKL, be the KL divergence for O [p, uk Fz), [0, 7). We calculate the Kullback-Leibhe part the ELBO as
T + T
T
CKL
pp + Nw Thet (DKL.P
Σ DKL
Σ
DKL.
(34)
T
k=0
k=0
Time- varying variables are normalized by the number of time steps Furthermore, we choose to normalize the components by their dimensions in physics space no in order to being able to use the same 3 for different use cases.
A.4.3 APPROXIMATION OF DATA LOG-L.KEIJJHOO
Optimizing the log- likelhoodd of p( ((x 2) can be represented by (Prince 2023 5.3)
argmax
II
p(a, y u ui p²)
argmin
MSE [ 9s), [m. vs)).
(35)
1=0
tue
1=0
where [ Di] is the output of 155 and is assumed to be a Gaussian with diagonal, fixed covariance.
A.4.4 KL-DIVERGENCE OF MULTIPLE VARIABLES
The KL-Divergence between two distribution p(=), q(z) defined as
[p2)]
DEL (q(z))lp(=)) = ((=)) log
dz
[q(2)
For the product two distributions of different random variables 21, 22 this definition yields:
DEL (g(2))g (22)))((22))p((())))()
/ ((2)))(22) log
p(z))o(22) dz
(
-f ((21)g(22) log
p(21 (22))
dz
( (2))
20

---

Published as a conference paper at ICLR 2025
Using logarithm laws
J ((21)((22) log
p
p(22) g(22)
+ log
dz d22
g(21)
Splitting the integral
J/ ((21) ((21))(22) log
[p(1)) dz1d22 + S ((<1)((22) log
[p(2))
d2 d21d22
[9(2))
[g(2))
Moving the integration symbols and sorting the terms
p(21)
d2 d2p d22 + / ((21) / ((22) log
(
( 21
((<1) log
dz2d2
((11)
(((2)
Inserting the definition of KL-Divergence again
((22) Dxr. (a(z) (((()))l(==)))d dz2 + / Dra. (1(2))llp(22) dzi
using that probability density functions always integrate 1, we see that
DKi. (α(=))lp(=ı)) DRi. (((=))][((z2)).
(36)
Hence we follow that the KI.-Divergence of two distributions on different random variables always yields sum of KL-Divergence terms for the single variable types.
A.5 EXPERIMENTS
A.5.1 DATASET GENERATION FROM PHYSICAL SIMULATION MODELS
The surrogate models shall replace the physical models based on first principles. The way we choose to train these surrogates by imitating the data [a's; ya at discretization points k 0... given (wo. p, un]. Therefore, we need data of the input-output pairs:
(zo, p, uo:T) + (I1.T, yo:r)
We choose the control input to be constant in the discretization intervals:
(l) Ut = const. A [4,,44
From experience, uniform random sampling proved allow good learning for values that stant over simulation, which is why we sample the parameters and initial values like that:
Fu -=(, Lo.+)
(37)
p~ p~u(p_,Pt)
(38)
with 24 denoting independent random sampling, and 04 denoting the user-specified upper and lower values for sampling. These values should be chosen sufficiently larger than the area of interest, such that extrapolation avoided. However, note that in general
min 11. * Z0.- and max #1. * 20.+
as the states can not be bounded priori.
Sampling the initial state can be seen as a way to increase the expressiveness of the training data to special regions of states space, which are otherwise rarely seen unseen during training. When generating training data, it should be ensured that required initial states for the surrogate purpose are included and well present.
However, variation of initial states for time-stepper methods in typical model regions can always be included with the following. Recall that time-stepper methods build up vector fields:
(x)' f(x p,
This makes it possible to enter the training at arbitrary locations within the time series sequences. This works also in latent space but requires transformation of the physical initial state 2o + z to latent space. We can always adapt the initial states during training to typical state vectors within the specific physical model use case, by choosing training sequence length E A with < T with
21

---

Published as a conference paper at ICLR 2025
T N being the sequence length of the training data and generate the samples in the following from it:
Co. p, T1.7, uo your = Jú P. #1+1:j+ts Ujz+t Ujs+t
Aj E 10.
(39)
Note that this increases the number of training samples by large magnitude, i.e. with 7 500 and 450 new samples per sequence. Therefore in a software implementation, the generated samples should not be saved on file, but returned on the fly from the original time series sequences keep memory requirements low (at the expense of computational cost).
т 50, we get T
For the dime-dependent inputs, simple random sampling does not work. With necessarily small discretization interval Åt, random sampling would lead to pure noise around mean-value uorr, such that all training samples would cycle around the same position in state space. Therefore, in the following, we introduce a method called rundomly clipped rundom sampling of controls with offset and cubic splines (RROCS). Pseudo code of the algorithm is depicted in I The main idea is to ensure a smooth control input within the specified bounds through the use of cubic splines, while randomly assigning offset and amplitude such that the control space filled meaningfully.
Algorithm randomly clipped random sampling controls with offset and cubic splines (RROCS) algorithm
for i in nampless do for j in nu do
sampled
random sequence with min. and max. frequency
Usampled
--
sequence with random offset and amplitude(u-.. u+)
Uspline
Cubic (laine(ssampeed, Usampled to:r)
Ensure smoothness
Uporm f Normalize Uspline to be in [0,1]
b, A f- random(u u_, u+)
b. Push controls to different regions Ensure constraints on input
b. A - clip such that within bounds(b, A,u ut)
U.J.0:T + Unorm A end for end for
An example for control inputs generated by this algorithm is given in Figure A.5
400
300
0.0
8.2
0.4
0.6
0.8
1.0
Time [s]
Figure A.5: Six example outputs of the RROCS-algorithm for the control input temperature K.a with u 273.15, u+ 473.15.
parameters or initial states are not sampled, they are set to default value Do.
A.5.2 STRATIFIED HEAT FLOW MODEL. (SHF) MODEL.
The stratified heat flow models presented in Figure A.6 are introduced because they are a simplified version of discretized (energy) transport phenomena that often arise in the system simulation context, c.g. fluid transport in pipes, discretized heat transfer in heat exchangers, discretized modeling of passenger cabins or battery packs. These kind of discretizations are often necessary for exact modeling of the different use-case, but often result in models that are computationally too expensive for tasks like optimization, calibration or controller tuning. We anticipate that depending on the specific implementation of the discretized segments correlations between some states can be found. For example, for the stratified heat flow model controlled at both ends, the segments in the middle of the model show probably very similar temperature.
Dataset Generation:
22

---

Published as a conference paper at ICLR 2025
heatcapacitor
port
IheistResistorr
ResisalResistoor
port b
R=R2
R=R2
(a) Base R2CI thermal network with parameters and for the total resistance and capacitance. The stratified heat flow model derived by repeated connection of this network, where each segment the parameters R. R/aSeg and C/nßeg with nãeg being the number of discretization layers.
temperature_ K_a
preseribed Tempo...
presorbedTempe...
temperature K_b
staaffoo.....
K
(b) Stratified heat flow model with parameters 2 for the total resistance and capacitance, and control inputs sempereture R. and cmmperature.. K.b. The beat flow discretized into näeg 16 segments.
Figure A.6: Stratfied heat flow models created with Modelica (Association 2023
number of samples: 1024
training fraction: 76 %
test fraction: 12 %
validation fraction: 12
Control Input Sampling Ranges:
temperature K.a: [273.15, 473.15) K
temperature K b: [273.15, 473.15) K
Solver Settings:
start time: 0s
stop time: 1.2s
output interval: 0.002s
tolerance; 1 x 10-6
sequence length: 501
the first 0.2 not included in the datasets.
16 states
48 outputs
Training:
Hyperparameters of all networks (state encoder, control encoder, parameter encoder, Latent ODE function, decoder):
layers
hidden dimension of 128
Activation: ELU
latent space dimension for latent parameter, state and control space: 16
Training:
3: different values
Adam Optimizer
learning rate: 10-3
weight decay: x 10
main training
Batches per Epoch: 12
* Phase I:
Solver: rk4
training sequence length: 10
* Phase 2:
Solver: rk4
23

---

Published as a conference paper at ICLR 2025
training sequence length: 250
increase training sequence length over n batches: 100
abort training sequence length increase after n stable epochs: 10
* Phase 3:
Solver: rk4
training sequence length: 400
increase training sequence length over batches: 100
abort training sequence length increase after n stable epochs: 10
A.5.3 POWER PLANT MODEL
The model is thermal power plant model of the ClaRa-Library (Vojacek et al. 2023). Specifically, the SteamCycle. the following as the steam cycle model. is model of a 580 MW hard coal power plant with single reheater, one high-pressurs preheater
model from the examples package that library We refer this model in
Modic
cysl....
Σ
Figure A.7: Diagram view the steam cycle model. Vojacek et al. 2023
and four low-pressure preheaters. To find initialization values for the states, static-cycle for the same power plant included and provides these values for the start of the simulation. The model has several internal Pl-controllers that control the pressure-levels of the preheaters and condensers. The heating power and feedwater pump power is controlled by an input signal, the relative power. This is the only signal considered for this application.
The system has 133 states, of that some were inactive, such that surrogate models are built for 102 states. furthermore contains 34 nonlinear systems.
The example sets two 10s load reductions input signals over total simulation time of 5000s. To make the model compatible with this work, an input connector was added and the initialization was modified to take the initial value from the input connector signal.
Dataset Generation:
number of samples: 1024
training fraction: 76 %
test fraction: 12 %
validation fraction: 12 %
Control Input Sampling Ranges:
u: [0.6, 1.1]
Team
24

---

Published as a conference paper at ICLR 2025
Solver Settings:
start time: Os
stop time: 6000s
the first s were omitted to discard transient behavior from initialization
output interval: 5s
tolerance: 1 x 10-6
sequence length: 2500
102 states (Prediction results only shown in ??)
outputs (RMSE this work for output prediction error)
Training:
Hyperparameters of all networks (state encoder, control encoder, parameter encoder, Latent ODE function, decoder OR Neural ODE function, output network):
layers
hidden dimension of 128
Activation: ELU
latent space dimension for latent parameter, state and control space: 128
For Koopman approximation: latent control, latent state dimension: 1024
Initialization: random (B-NODE, NODE), eigenvalues (B-NODE Koopman)
Training:
3: different values
Adam Optimizer
learning rate:
10-4 (B-NODE, B-NODE Koopman), 10-5 (NODE) weight decay: 1 x 10-5, 1 x 10-8 (B-NODE Koopman, NODE)
clip grad norm: 2.0
main training
# Batches per Epoch: 12
* Phase 1:
Solver: euler (rk4 B-NODE Koopman) training sequence length: 5 (10 NODE) break after loss of: 0.2 (0.05 NODE)
Phase 2:
Solver: rk4
training sequence length: 80
increase training sequence length over n batches: 6000 (1200 NODE)
abort training sequence length increase after n stable epochs: 1000
* Phase 3:
Solver: rk4
training sequence length: 250
increase training sequence length over batches: 1200
abort training sequence length increase after n stable epochs: 1000
A.5.4 DETAILS FOR COMPARISON OF COMPUTATION TIME FOR POWER PLANT MODEL.
For generation of data from the power plant mode (example 1), the model was exported as co- simulation FMU with CVODE solver with Dymola and simulated with FMP
solving one system:
computation time: 16 seconds
hardware specifications:
* 8-core processor with up to 4.9 GHz
technical setup:
solved using FMPy with per-time-step communication
training data generation of 1024 samples:
official: htc.hhhkh..k..oa.........eoeeeen
*PMPy: https://ggthhbb.co/////// http:///gggihabooooooooo////////
25

---

Published as a conference paper at ICLR 2025
computation time: 720 seconds
hardware specifications:
16-core processor with up to 5.7 GHz
128 GB RAM
technical setup:
* parallel processing using Dask
simulations executed via FMPy with per-time-step communication
simulation with surrogate models GPU:
model trained with dopri5 solver in the last period
for adaptive step size control, the sample with the highest error determines step size control.
hardware specifications:
* Nvidia RTX4080 with 16 GB of GPU memory technical setup:
all samples solved concurrently in single batch, provided that sufficient GPU memory is available.
A.5.5 KOOPMAN ANALYTIC EXAMPLE
We implemented
I9 as Modelica model.
Dataset Generation:
number of samples: 1024
training fraction: 76 %
test fraction: 12 %
validation fraction: 12 %
initial state sampling ranges:
x1: [-50, 50]
x2: [-50, 50]
solver settings:
start time: Os
stop time: 10s
output interval: 0.18
tolerance: 1 x 10-5
sequence length: 99
2 states
Training:
latent state dimension: 64
Training:
3: different values
Adam Optimizer
learning rate: x 10-3
weight decay: 1 x 10-5
clip grad norm: 1.0
main training
Batches per Epoch: 12
* Phase 1:
Solver: euler
training sequence length: 5
break after loss of: 0.2
Phase 2:
Solver: rk4
training sequence length: 30
Dask: httph:t////// htt.pp:/wwww.ask..oo
26

---

Published as a conference paper at ICLR 2025
increase training sequence length over batches: 1200
abort training sequence length increase after n stable epochs:
Phase 3:
Solver: rk4
training sequence length:
increase training sequence length over batches: 1200
abort training sequence length increase after n stable epochs:
A.5.6 SMALL GRID MODEL
The small grid model describes simplified 3 dimensional CFD simulation of air flow cubic grid consisting of 3x3x3 grid cells. The cube in the center is modeled solid mass. Two grid cells at the sides are connected sources of prescribed temperature and mass flow, and a third cell connected outflow port, such that directed flow pattern evolves in the grid, cycling around the cube in either one or another direction (Figure A.8)
SG3
B23
Wavd
Hart
(a) First flow pattem.
(b) Second flow pattern.
Figure A.8: Visualization of simulation results of the small grid model.
The model is implemented using the Modelica HumanComfort Library 10 sampling to generate a datset.
We apply RROCS-
Dataset Generation:
number of samples: 1024
training fraction: 76 %
test fraction: 12 %
validation fraction: 12 %
Control Input Sampling Ranges:
T_in: [273, 310] K
mn.flowin:: [0.0, 1.0] kgs-1 T-inl: [273, 310] K
miflow.ini: [0.0, 1.0] kgs
Solver Settings:
start time: 0s
stop time: 1100s
output interval: 0.25s
tolerance: 1 10-4
sequence length: 4001
the first 100 s are not included in the datasets.
100 states
27

---

Published as a conference paper at ICLR 2025
A FMU of the model or the generated dataset is available request.
A.5.7 WATER HAMMER MODEL
The Waterhammer model consists of pipe filled with water, whose ends are connected to pressure and a mass flow source (Lenord] (024). Through the pressure source, pressure change can be applied, such that pressure wave is generated that is reflected at the ends of the pipe. Evaporation effects of water are neglected. To model the spatial distribution of mass flow and pressure, a high discretization of the pipe's 's volume is necessary. We choose discretization of 100 control volumes over pipe length of 10m. The Model is implemented using the ClaRa Library (Vojacek et al. 2023
h,
Q2
Figure A.9: Diagram view of the water hammer model.
Dataset Generation:
number of samples: 1024
training fraction: 76 %
test fraction: 12 %
validation fraction: 12 %
Control Input Sampling Ranges:
sampling of pressure ramp with parameters:
* p-start: [5e5,20e5] Pa
p.increase: [1e4,80e5] Pa
* t time_for fommrforrramp: [0.0001,0.2]
constant mass flow rate between 0.001 kgs 1 0.1 kgs"
Solver Settings:
start time: 0s
stop time: 0.5
output interval: 1 x 10-3
tolerance: 1 10-6
sequence length: 351
the first 0.15 not included in the datasets.
400 states
2 control imputs
A FMU of the model or the generated dataset is available on request.
A.5.8 MuJoCo DATASET
The MuloCo dataset provided through Minari Younis et al. 2024 and is available as download 13
The data is obtained from simulation of the Gymnasium (Towers et al. 2024) Ant Mode [4] The Ant model describes a 3D robot that consists of torso with four legs where each leg consists of two parts, and the applied torques at the legs can be controlled.
lttpss:/imiisttaddara.
28

---

Published as a conference paper at ICLR 2025
The dataset consists of 2026 samples of different length, where the majority has length of 1000. To make it compatible with our code, we select all sequences length 1000 and have in total 1986 samples for training. The dataset has observation space with 105 variables, which we assume to be the state space, and an action space of variables, which in our wording the control input space. From the state space, only the first 27 variables proper system states, while the remaining 78 elements are external forces, which is why we exclude those from our training dataset.
A.5.9 BENCHMARK NONLINEAR SURROGATE MODELS
The benchmark shows the RMSE normalized by respective variances of prediction of the surrogates for the dataset's s states. For the Power Plant dataset, we predict the RMSE on output variables. The trained B-NODE models are constant variance B-NODEs with nonlinear layers. For the Latent ODE, we added to the ODE function of (Rubanova et al. 2019) an input vector, i.e.
x (1) forNODE (z) 2 (())
(40)
Dataset details for MuJoCo and Water Hammer can be found in A.5.8
A.5.7
A. BENCHMARK LINEAR SURROGATE MODELS
The benchmark shows the RMSE normalized by respective variances of prediction the surrogates for the dataset's states. As DMD methods implemented for states prediction, predict also states for the Power Plant dataset. The trained B-NODE models are constant variance B-NODEs with linear layers. Dataset details for MuloCo and Small Grid can be found in A.5.8 and A.5.6 The DMD methods are implemented using the a-Kooppman-package (Pan et al., 2024). For DMDe, we performed the algorithm for all ranks possible. For eDMDc, we employed radial basis functions as observables and implemented a genetic algorithm to search for best basis functions.
A.6 ADDITIONAL. EXEPERIMENT RESULTS
A.6.1 MASKING OF IDENTIFIED LATENT CHANNELS
To evaluate if B-NODEs are able to learn dimensionality reduced representation, that also holds when transforming the weakly defined information bottleneck (by KL divergence) into a hard one (by manually setting latent channels with DKL,,, < 0.1 zero), we performed training in which we set after sufficient training the mask. Specifically, we trained a nonlinear and linear B-NODE for SHF model. After setting the mask, ,90 is set to and the B-NODE is trained as NODE latent state space. As shown in Figure A010 the error does not increase after masking the latent channels, it decreases additionally.
29

---

Published as a conference paper at ICLR 2025
109
B-NODE
mask
10-1
RMSE
10--2
10-3
100
200
300
400
500
600
700
109
lincar B-NODE
mask on
10"1
RMSE
10-2
10-3
0
100
200
300
400
500
600
700
800
steps
Figure A.10: Manually masking the identified latent channels to O.
A.6.2 HYPERPARAMETER SENSITIVITY OF STRATIFIED HEAT FLOW MODEL (SHF)
Table 4: Results of hyperparameter sensitivity analysis for the stratified heat flow (SHF) model.
Hyperpemeeter
Com. Test
Test
Train
States Test
Outputs Text
Baseline
0.0163
n. lineor Jayne 4+
0.0171
0.0167
0.0168
0.0012
0. 0.0323
2
n lines layers
-
0.0425
0.0431
0.0429
0.0021
0.084
2
lincar JaddJiiddeeiiii 128
32
0.0180
0.0164
0.0169
0.0011
0.0317
2
hidhen Jim:
0.0164
liearr liaear_hidlnn hiddeedddddddi
0.0114
0.0210
activation ELU LeskyRel.U
0.0187
0.0178
0.0179
0.0026
0.0330
2
activation: ELU
ReI.
0.0462
0.0439
0.0432
0.0025
0.0853
2
ELU
0.9918
activation: ELU CELU
0.0186
0.0158
0.0153
0,0010
0 0.0036
2
sctivation ELU + Tanh
0.0174
0.0170
0.0168
0.0011
0. 0328
2
hilben.dim.cotuuuum 128
0.0163
0.0150
0.0149
0.0009
0.0290
2
biiden_lim.wuput.mmm
0.0167
0.0147
0.0162
0.0008
0.0284
2
control_inyeeell Thue Fales
0.3972
0.3803
0.3569
0.1476
0.6131
D
Jimemindsslim:: 12
0.0162
0.0149
0.0155
0,0009
0.0289
2
br.contross.im: 12 24.
0.0185
0.0156
0.0168
0.0010
0.0302
2
basaates.imm: 12 + 24
sujenntnsslim:: 12-+
0.0165
0.0161
0.0161
0.0010
0.0311
2
Ialsaatesdimm:
suber_atol.nwemile 0.0001 1e-05,
0.0163
0.0163
0.0179
0.0011
0.0315
2
solver soleer-tol_-vemdd: 0.001 - 0 00001
beta start owriidd 0.1 - 1.0
0.1407
0.1377
0.1470
0.0163
0.2591
tem.gaar.overidd: 0.1 - 001
0.0030
0.0035
0.0036
0.0002
0.0074
2
beta override: 0.001
0.0032
12
7
30


Reviews:
[{"review_id": "cL78uafcuY", "reviewer": "Reviewer_kk7L", "summary": "A method for model order reduction using $\\beta$-VAEs and state space NeuralODEs called **balanced neural ODE** is described and tested on a few examples. The hyperparameter $\\beta$ controls the number of active latent space dimensions, i.e. the reduced dimension of balanced neural ODE. Two main methods of evolution are proposed for the latent VAE parameters: constant-variance and dynamic-variance. Koopman operator approximation is shown to be a special case of balanced neural ODE.", "questions": "1. Are there other related works that combine model order reduction and latent space evolution, other than latent ODE?\n2. How does the balanced neuralODE compare to Dynamic Mode Decomposition method (for linear operator learning)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "1. Method is fully described, intuitive to understand, many derivations are fully given\n2. Balanced neural ODE outperforms latent neural ODE and is able to learn latent dynamics.\n2. Numerical experiments demonstrate the breadth/diversity of application of the method.", "weaknesses": "1. Method combines two existing approaches (however, the combination is novel)\n2. Presentation feels brisk at times (might be alleviated with an extra page)\n3. The model involves many components with architecture choices, and there is not enough ablation on", "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": "", "extra_scores": {}}, {"review_id": "OSaYpG2CIs", "reviewer": "Reviewer_w89j", "summary": "This paper proposes a novel method, called B-NODE, for time-series modeling of dynamical systems (e.g. power plants). B-NODE is claimed to be fast (in terms of runtime) and accurate (in terms of reconstruction error on the original data). The main idea behind B-NODE is a combination of $\\beta$-VAE (for dimensionality reduction) with a Neural ODE (for dynamical systems' modeling). \n\nThe faster runtime of B-NODE results from dimensionality reduction via the VAE. Previous approaches (like standard latent ODEs) fail to provide such dimensionality reduction. B-NODE also claims to offer the user with a knob to balance the trade-off between dimensionality reduction and reconstruction accuracy.\n\nAs for the paper structure, section 3 describes the methodology. Sections 4 and 5 demonstrate the advantages of B-NODE on certain synthetic and real-world tasks.", "questions": "**Technical questions**:\n1. How exactly does B-NODE \"balance\" between dimensionality reduction and reconstruction accuracy? Is it simply just varying the $\\beta$ in $\\beta$-VAE? If yes, it would be nice to explicitly mention this. \n2. In equation 4, why is it $q(u_{0:T}^z|u_{0:T})$ instead of $q(u_{0:T}^z|u_{0:T}, p^z, p)$? Since the latter is what one would get from the definition of $x^\\mathcal{D}$ and $z$ in line 188. Is it because the external input $u$ is independent of the time-invariant parameters $p$? Is this a standard assumption?\n\n**Suggestions for presentation**:\n1. In the first paragraph explaining motivation, it would be nice to present more examples so that the reader can conceptualize the problem. Currently, only the words \"yearly energy system simulations\" are mentioned as examples.\n2. In Figure 4(a), the reader needs to zoom in a lot to read the values of $\\beta$ around the markers on the plot. Please make the plots easier to read.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper shows promising results with using B-NODE on certain tasks.\n- The paper provides certain handy illustrations and detailed figures explaining their approach.", "weaknesses": "**Technical Weaknesses**:\n1. *Limited evaluation.* Since this is a paper that proposes a novel method to solve a problem, the authors should aim to provide a comprehensive suite of experiments showing the advantages of the proposed methodology over existing baselines. The only real-world use case is presented in Section 4.2, which is just one single task. In Figure 5 (in the same section), B-NODE is compared against only one baseline, which is NODE. As a reference, two closely related papers in this field are [1] and [2]. Both provide much more detailed evaluation (Table 3 in [1] and Tables 2,3 in [2]).\n\n**Presentation Weaknesses**: \n\n1. *Unclear on the main advantage of B-NODE.* Is your main contribution speed of the proposed methodology, or accuracy, or both? It would be good to explicitly state it. The paper abstract mentions the word \"fast\", but the only runtime comparison I could find was Figure 6(b), which is only discussed briefly in the text in lines 388-392. In contrast, Figures 4,5,7,8 all seem to focus on accuracy via the RMSE.\n\n2. *Vague exposition.* Below I list certain examples where the exposition is vague.\n\nExample 1: The \"B\" in B-NODE stands for \"Balanced\", but what exactly is the knob that lets one choose the balance of the trade-off between reconstruction accuracy and latent dimension size?\n\nExample 2: In lines 150-155, how is the $D_{KL}$ measured per channel? The overall $D_{KL}[q(z|x^\\mathcal{D})||p(z)]$ is well-defined, but to measure for certain channel $i$, do you measure $D_{KL}[q(z_i|x^\\mathcal{D})||p(z_i)]$? Are the marginals on a \"channel\" $i$ well-defined? It would have been nice to see this in proper notation. \n\nExample 3: It is unclear what the following statement means: \"The VAEs sampling-based robustness complements stability enhancements for Neural ODEs\". This statement is mentioned twice, in line 180 (Methodology section) and line 476 (Conclusion section).\n\nExample 4: Line 280 says \"state $x_i^z$ is inactive\", but no notion of active/inactive-ness is defined? Likely by inactive, the authors mean that the $D_{KL}$ of the $i^{th}$ channel is below the threshold, but this is sloppy exposition.\n\nExample 5: Line 60 (and a few other places) mention that B-NODE *requires* a non-hierarchical prior, but the paper does not explain why this is the case.\n\nOverall, I found the paper difficult to understand. Sentences that follow one another do not necessarily build the same idea in a linear/sequential fashion. This makes it hard to extract what the authors are trying to convey.\n\n\n```\n[1] Naiman, I., Erichson, N. B., Ren, P., Mahoney, M. W., & Azencot, O. (2023). Generative modeling of regular and irregular time series data via koopman VAEs. arXiv preprint arXiv:2310.02619.\n[2] Wi, H., Shin, Y., & Park, N. (2024, March). Continuous-time Autoencoders for Regular and Irregular Time Series Imputation. In Proceedings of the 17th ACM International Conference on Web Search and Data Mining (pp. 826-835).\n```", "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": "", "extra_scores": {}}, {"review_id": "hutpzs89DU", "reviewer": "Reviewer_1CNC", "summary": "This work presents a novel approach combining Variational Autoencoders (VAEs) and Neural ODEs to address model order reduction (MOR) for dynamic systems. The proposed model, named Balanced Neural ODE (B-NODE), utilizes VAEs for latent dimensionality reduction and Neural ODEs for capturing transient dynamics, resulting in a compact, efficient surrogate model that can dynamically adapt to the complexity of input data.", "questions": "- Please provide more comprehansive understanding of β-value?  A discussion of the theoretical implications of different β values\n- It seems that β-value is small does it mean that information bottleneck become negligible?\n- if you set β=0 what will happen? please discuss the theoretical implications of setting β=0\n\nIf the authors can clarify my concern, I could change my opinion.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper effectively combines Variational Autoencoders (VAEs) and Neural ODEs, leveraging their complementary strengths for dimensionality reduction and transient dynamics modeling. This approach could significantly impact surrogate modeling in dynamic systems.", "weaknesses": "Although the paper acknowledges areas for future work, a more detailed discussion on scenarios where B-NODE might underperform would be helpful. Specifically, it would be beneficial to outline any cases where the model might struggle with particular types of inputs or dynamic systems.", "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": "", "extra_scores": {}}, {"review_id": "GOpa9U7J9H", "reviewer": "Reviewer_LoSy", "summary": "This paper proposes a novel method called Balanced Neural ODE (B-NODE), which combines Variational Autoencoders (VAEs) with Neural ODEs to reduce the complexity of dynamical systems while maintaining accuracy. This approach extends Latent ODEs to handle time-varying inputs by continuously propagating latent variables. It also aims to approximate Koopman operators by leveraging the weight factor of KL divergence during VAE training. The proposed method is evaluated on both academic and real-world scenarios.", "questions": "- In line 159, the statement \"Combining VAEs with State Space Models leverages VAEs' ability to generalize through an information bottleneck with a numerically pre-determined latent space that promotes data locality and latent orthogonality\" is unclear. Specifically, what is meant by \"data locality\" and \"latent orthogonality\"? Could the authors provide further explanation, relevant citations, or supportive numerical results to clarify these concepts?\n\n- Could author further explain how they arrive equation (4ba)? You assume that the $x_0^z$ only related to initial state $x_0$?\n\n-  Lines 297-299: Could the authors further justify why they do not need to pre-determine the dimensionality of the latent space? The weighting factor might contribute to sparsity in certain contexts, but it is unclear how it directly adapts the dimensionality. Additional explanation on this point would be helpful.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The problem addressed is both interesting and crucial for applying Neural ODEs to dynamical systems with time-varying inputs. The combination of VAE and Neural ODEs to achieve dimensionality reduction alongside dynamic modeling provides an approach for managing complex systems. \n\n- The inclusion of real-world test cases adds practical value to the study and demonstrates its potential applicability.", "weaknesses": "## Presentation \n- The layout of this paper is somewhat disorganized, which makes it very hard to follow. The presentation is unclear without a main thread, with complex concepts that are not explained adequately and unnecessary.\n- The concepts and some formulations appear to be justified solely by the authors themselves, lacking sufficient external validation or references to prior work (see points in Questions).\n- Please ensure there is proper spacing between different paragraphs and sections. \n\n## Experiments \n- There are no baseline comparisons with Koopman-based models, Neural ODEs, or VAE-based forecasting approaches (which are highly relevant to this work), which limits the ability to assess how much improvement the proposed model actually provides.\n\n## Minor issues \nsome examples:\n- Lines 100 the integral should be $\\int_{t_0}^{t}f_{\\phi NODE}(\\mathbf{x}(\\tau))d\\tau$\n- Line 225, is -> are\n- Line 312, Observed -> observe \n\nI'd encourage authors redo proofreading.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)