Balanced Neural ODEs: nonlinear model order reduction and Koopman operator approximations

## Abstract

Abstract Variational Autoencoders (VAEs) are a powerful framework for learning compact latent representations, while Neural ODEs excel in learning transient system dynamics.
This work combines the strengths of both to create fast surrogate models with adjustable complexity.
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) 1 1 1 Our code is available under: https://github.com/juliusaka/balanced-neural-odes .
We demonstrate the effectiveness of this method on academic test cases and apply it to a real-world example of a thermal power plant.

## 1 Introduction

Motivation:
Originally, this work is driven by challenges arising when using simulation models for optimization of complex technical systems.
With increasing fidelity, the computational complexity rises.
Simulation and optimization in long-term scenarios - e.g. yearly energy system simulations - can become infeasible.
Therefore, attention to faster surrogate and reduced-order models is growing. There exists special interest in linear surrogate models, as these allow the usage of a rich set of tools from control theory and efficient optimization algorithms.
This work addresses a combination of (Dynamical) Variational Autoencoders (VAEs) with Neural ODEs, exploiting the VAE’s information bottleneck for learning reduced-order dynamical systems.
Neural ODEs for time-series modeling:
Neural ODEs(Chen et al., ) have gained significant attention in dynamic systems modeling.
Extensions include neural controlled differential (Neural CDEs) equations (Kidger et al., ) and neural stochastic differential equations (Neural SDEs) (Liu et al., ).
They excel in different areas,
but first and foremost in time-series forecasting
e.g. in geology (Shen et al., ),
medicine (Hess et al., ; Shi et al., ),
or as surrogate models in engineering (Legaard et al., ).
To continuously integrate input signals, Neural CDEs (Kidger et al., ) have been used in several publications (Seedat et al., ; Hess et al., ; Wi et al., ).
Controlled differential equations remain uncommon in dynamic system simulation, where state-space models are usually employed for the integration of inputs(Zamarreño & Vega, ), that can also be used with Neural ODEs (Rahman et al., ; Legaard et al., ).
Neural ODEs can also be extended to include time-invariant modulators, i.e. parameterized dynamics (Lee & Parish, ; Auzina et al., ).
A different application is hybrid modeling, that combine Neural ODEs with first-principle equations in different configurations (Thummerer & Mikelsons, ).
Dynamical Variational Autoencoders (DVAEs)
are the extension of VAE (see [2.2](#S2.SS2)) (Kingma & Welling, ) to sequential data (Girin et al., ).
DVAEs usually include a separation of dynamic and fixed variables (Fraccaro et al., ; Sadok et al., ; Berman et al., ), and can further be distinguished by causal, anti- or noncausal inclusion of temporal dependencies (Girin et al., ).
Deep Kalman Filters as a suitable subclass of causal DVAEs for system simulation employ a prediction-update logic for the dynamic latent variable and result in hierarchical priors (Krishnan et al., ; Fraccaro et al., ).
To achieve measurable state reduction, our method requires a non-hierarchical prior.
As a consequence, dynamics are part of the inference model, which additionally contains a temporal bottleneck, while the information bottleneck is 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., ; Shi et al., ; Dahale et al., ) combine VAE with Neural ODEs, but only assign the probability distribution to the initial latent state, and therefore fail in learning compact latent representations with input signals ([4.1](#S4.SS1)).
For DVAEs, only the contribution of (Wi et al., ) proposed (independently from us) to continuously describe the variational parameters $\mu,\sigma$, 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 et al., ).
This mostly manifests in a reduction of the system’s state dimension.
For linear systems, a variety of mathematical methods exists.
Truncated Balanced Realization ((Moore, )) is a 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, ).
Koopman theory offers to overcome this by using observables $\bm{g}(\bm{x})$ of states that can be advanced in time with a linear operator, i.e. $\bm{g}(\bm{x}(t_{k+1}))=\bm{\mathcal{K}}\cdot\bm{g}(\bm{x}(t_{k}))$, but with the drawback of this observable space being possibly infinite-dimensional (Koopman, ).
Modern Koopman theory therefore deals with finding finite-dimensional approximations (Brunton et al., ).
A popular approach for high-dimensional systems is Dynamic Mode Decomposition (DMD) (Schmid, ; Rowley et al., ).
Several extensions exist, e.g. to inputs (DMDc, Proctor et al. ()), or to low-dimensional systems (Williams et al., ), where extended DMD lifts the dimensionality by e.g. radial basis functions (Korda & Mezić, ).
Autoencoders (not VAEs) have also been applied to learn non-linear coordinate transformations to linear submanifolds (Lusch et al., ; Mardt et al., ).
A different approach, modeling the dynamics with linear and non-linear components has been recently proposed by (Menier et al., ).
A common obstacle to learning linear latent sub-manifolds is the choice of their dimensionality (Lusch et al., ; Menier et al., ; Mardt et al., ).
We avoid this choice by weakly describing an information bottleneck with a balancing parameter $\beta$.
The recently proposed Koopman VAE is a sequential DVAE with hierarchical , linear prior, and does not focus on state reduction (Naiman et al., ).

## 2 Background

### 2.1 Neural ODEs in system simulation

In dynamic systems modeling, the state $\bm{x}$ 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., ).
The states are variables that appear as derivatives and must be determined by solving the ordinary differential equation (ODE) governing the dynamical system (Cellier & Kofman, ).
While this ODE is typically formulated using first principles, in the case of a Neural ODE, the right-hand side of the equation is modeled by a neural network $f_{\phi_{\text{NODE}}}$ (Chen et al., )

$$ $\displaystyle\bm{x}^{\prime}(t)=f_{\phi_{\text{NODE}}}\left(\bm{x}(t)\right),\bm{x}(t_{0})=\bm{x}_{0};$ $\displaystyle\bm{x}(t)=\bm{x}_{0}+\int_{t_{0}}^{t}f_{\phi_{\text{NODE}}}\left(\bm{x}(t)\right)dt.$ $$

The neural network $f_{\phi_{\text{NODE}}}(\bm{x}(t))$ describes the dynamics of the system.
It maps to each state $\bm{x}$ a gradient $\bm{x}^{\prime}=\frac{d\bm{x}}{dt}$ that describes the direction in which the state evolves.
In other words, a vector field as in the example in [Figure A.1](#A1.F1) is learned.
For the analysis of dynamical systems, additional classes of variables are critically important: 1) time-varying external inputs $\bm{u}(t)$ excite the system and induce a response, e.g. a motor torque. 2) Time-invariant (physical) parameters $\bm{p}$ change the dynamic properties of the system, e.g. a mass. 3) Time-varying outputs $\bm{y}(t)$, e.g. a power, can be calculated as a function of the other variables $\bm{x}(t),\bm{p},\bm{u}(t)$.
By adding an output-function neural network $g_{\phi_{\text{Out}}}$ and the respective inputs, we get a state-space model (example in [Figure A.1](#A1.F1)).

$$ $\displaystyle\bm{x}^{\prime}(t)$ $\displaystyle=f_{\phi_{\text{NODE}}}(\bm{x}(t),\bm{u}(t),\bm{p}),$ $\displaystyle\bm{x}(t_{0})=\bm{x}_{0},$ (1) $\displaystyle\bm{y}(t)$ $\displaystyle=g_{\phi_{\text{Out}}}(\bm{x}(t),\bm{u}(t),\bm{p}).$ (2) $$

We refer to this model ([1](#S2.E1)-[2](#S2.E2)) in the following as a State Space Neural ODE and train it by minimizing the MSE of predictions ([25](#A1.E25)).

### 2.2 Variational Autoencoders for Dimensional Reduction

Variational Autoencoders (VAE) are frequently used to describe high-dimensional data $\bm{x}^{\mathcal{D}}$(^2^22$\bm{x}^{\mathcal{D}}$ is general data, $\bm{x}$ is the state.) by using a smaller number of latent variables.
VAEs assume a prior distribution on the latent variables and express the data probability as marginalization over the latent variable $\bm{z}$ (Kingma & Welling, , 1.7)

$$ $p(\bm{x}^{\mathcal{D}})=\int p(\bm{x}^{\mathcal{D}}|\bm{z})p(\bm{z})d\bm{z}.$ $$

For the prior, we assume a multivariate Gaussian with zero mean and diagonal covariance one, i.e. $p(\bm{z})=\mathcal{N}(\bm{z};\mathbf{0},\bm{I})$.
To train the model $p(\bm{x}^{\mathcal{D}}|\bm{z})$, the intractability of the learning problem is resolved by approximating the true posterior probability $p(\bm{z}|\bm{x}^{\mathcal{D}})$ with an approximate model $q(\bm{z}|\bm{x}^{\mathcal{D}})$, called the encoder (Prince, , 2.2).
The model $p(\bm{x}^{\mathcal{D}}|\bm{z})$ is then the decoder.
During training, the log-likelihood of the data $\log p(\bm{x}^{\mathcal{D}})$ shall be maximized. Due to the intractability, Kingma & Welling (, 2.2) derive a lower bound to the data likelihood, which is called the evidence lower bound (ELBO, derivation in [A.3.1](#A1.SS3.SSS1)).
We use the notation (Prince, , 17.4.2)

$$ $\mathcal{G}=\mathbb{E}_{q(\bm{z}|\bm{x}^{\mathcal{D}})}\log\left[\frac{p(\bm{x}^{\mathcal{D}}|\bm{z})}{q(\bm{z}|\bm{x}^{\mathcal{D}})}p(\bm{z})\right]=\underbrace{\mathbb{E}_{q(\bm{z}|\bm{x}^{\mathcal{D}})}\log\left[p(\bm{x}^{\mathcal{D}}|\bm{z})\right]}_{\text{reconstruction quality}}-\beta\underbrace{D_{\text{KL}}\left[q(\bm{z}|\bm{x}^{\mathcal{D}})||p(\bm{z})\right]}_{\text{information bottleneck}},$ (3) $$

with the term $\beta$ being added in $\beta$-VAE (Higgins et al., ).
The functions $q(\bm{z}|\bm{x}^{\mathcal{D}}),p(\bm{x}^{\mathcal{D}}|\bm{z})$ are implemented with neural networks ([A.2](#A1.F2)).
Typically, the encoder network parameterizes a multivariate Gaussian distribution $f_{\phi_{\text{en}}}(\bm{x}^{\mathcal{D}})=\bm{\mu}(\bm{z}),\bm{\sigma}(\bm{z})$, such that $q(\bm{z}|\bm{x}^{\mathcal{D}})=\mathcal{N}(\bm{z};\bm{\mu}(\bm{z}),\bm{\sigma}(\bm{z}))$.
The decoder predicts the data $\hat{\bm{x}}^{\mathcal{D}}=f_{\phi_{\text{de}}}(\bm{z})$, with $\bm{z}$ following the conditional distribution of the encoder. For training, Kingma & Welling () proposed to draw (one) Monte-Carlo sample of the posterior distribution by reparameterizing with an external noise variable, i.e. $\bm{z}=\bm{\mu}(\bm{z})+\bm{\epsilon}\bm{\sigma}(\bm{z});\bm{\epsilon}\sim\bm{\mathcal{N}}(\bm{z};\mathbf{0},\mathbf{\mathbb{I}})$, sustaining backpropagation ability.
Burgess et al. () develop an understanding of $\beta$-VAE as a model with information bottleneck properties.
Training a VAE maximizes the similarity between the data $\bm{x}^{\mathcal{D}}$ and the reconstruction $\hat{\bm{x}}^{\mathcal{D}}$, while the least amount of information is transmitted by the intermediate latent variable $\bm{z}$.
The amount of information transmitted is measured by the Kullback-Leibler Divergence $D_{\text{KL}}$.
The balance between reconstruction and the information bottleneck is adjusted by $\beta$.
VAEs excel in learning disentangled representations Higgins et al. (); Mathieu et al. (). The diagonal approximation in the encoder and the stochastic noise enforce local orthogonality (Rolinek et al., ).
The question arises if we can measure the dimensional reduction.
The KL divergence forms with the chosen prior ([1(a)](#S2.F1.sf1)) a valley that ascends from $\mu=0,\sigma=1$ ($D_{\text{KL}}$=0) with smaller $\sigma$ and greater $\mu$.
If a latent channel has $D_{\text{KL}}=0$ for every input of the encoder, it transmits pure Gaussian noise, and therefore can not transmit any information.

Figure: ((a)) Heatmap of $D_{\text{KL}}\left(q(z|x)||p(z)\right)$, with $q(z|x)=\mathcal{N}(z;\mu,\sigma)$ and $p(z)=\mathcal{N}(\bm{z};0,1)$ (Equation [29](#A1.E29)).

To explore the latent space of a trained $\beta$-VAE, [1(b)](#S2.F1.sf2) reveals that $D_{\text{KL}}$ is $\gg 0$ for only a small and fixed subset of the overall offered latent dimensionality.
Ordering the dimensions for this example by their mean $D_{\text{KL}}$ reveals a significant drop of $D_{\text{KL}}$ ([Figure A.3](#A1.F3)) from the 5th to the 6th dimension, indicating the first 5 dimensions contain the most important information, which was verfied by experiment ([4(a)](#A1.F4.sf1)).
This motivates to use the mean $D_{\text{KL}}$ per latent channel as a measure for active latent dimensions, and we
choose for counting a value of $D_{\text{KL}}(z)>0.1$ ([4(b)](#A1.F4.sf2)).

## 3 Method: Balanced Neural ODEs

Combining VAEs with State Space Models leverages VAEs’ ability to generalize through an information 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, weakly defined by the $\beta$-value.
Additionally, the VAE’s sampling-based robustness complements stability enhancements for Neural ODEs with stochastic noise. We call the proposed model Balanced Neural ODE (B-NODE), as training results in a balance between state reduction and reconstruction quality.
In this work, we deal with continuous-time variables and their discrete-time samples.
For that, a subscript $i$ denotes the value of a variable at time $t_{i}$, e.g. $\bm{x}_{i}=\bm{x}(t_{i})$.
We suppose (but are not limited to) equidistant time steps $\Delta t$ between observations $t_{i}$ and $t_{i+1}$.
A sequence of observations in $[\bm{x}(t_{i}),\dots,\bm{x}(t_{j})]$ , is denoted as $\bm{x}_{i:j}$.
The initial time of a sequence is $t_{0}$, and the final time is $t_{T}$.
For simplicity, inputs are considered constant in a time interval, $\bm{u}(t)=\bm{u}_{i},t\in[t_{i},t_{i+1}]$.
In the following, we derive the structure of a VAE that learns sequential representations of the data $\bm{x}^{\mathcal{D}}=[\bm{x}_{0:T},\bm{y}_{0:T},\bm{u}_{0:T},\bm{p}]$ with latent variables $\bm{z}=[\bm{x}^{z}_{0:T},\bm{u}^{z}_{0:T},\bm{p}^{z}]$. We thereby follow the convention of Girin et al. () to first present the model equations and then the implementation with neural networks.

### 3.1 Inference Model and Generative Model

In variational inference, the data $\bm{x}^{\mathcal{D}}$ is assumed to be generated from the ”true” latent variables, thus it is necessary to infer these variables.
We therefore describe the inference model as an encoding of all variables in data space,

$$ $q(\bm{z}|\bm{x}^{\mathcal{D}})=q(\bm{p}^{z}|\bm{p})q(\bm{u}_{0:T}^{z}|\bm{u}_{0:T})q(\bm{x}_{0:T}^{z}|\bm{x}_{0:T},\bm{y}_{0:T},\bm{u}_{0:T}^{z},\bm{p}^{z}),$ ( 4 ) where the time-invariant parameter vector (available from simulations) is simply encoded as $q(\bm{p}^{z}|\bm{p})$ 3 3 3 $\bm{p},\bm{p}^{z}$ refer to physical parameters, $p(\bullet)$ to a probability density function, $\phi_{\square}$ to neural network parameters. . Latent controls $q(\bm{u}_{0:T}^{z}|\bm{u}_{0:T})$ are encoded sequentially and factorize in $q(\bm{u}_{0:T}^{z}|\bm{u}_{0:T})=\prod\nolimits_{i=0}^{T}q(\bm{u}_{i}^{z}|\bm{u}_{i}).$ (4a) The inference model for the latent states is a combination of initial state observer and simulator: $\displaystyle q(\bm{x}_{0:T}^{z}|\bm{x}_{0:T},\bm{y}_{0:T},\bm{u}_{0:T}^{z},\bm{p}^{z})$ $\displaystyle=\underbrace{q(\bm{x}_{0}^{z}|\bm{x}_{0:T},\bm{y}_{0:T},\bm{u}^{z}_{0:T},\bm{p}^{z})}_{\text{initial state observer}}\underbrace{\prod\nolimits_{i=0}^{T-1}q(\bm{x}_{i+1}^{z}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z})}_{\text{simulator}}.$ (4b) Having in mind that a surrogate model should simulate based on the inputs $\bm{x}_{\text{in, sim}}^{\mathcal{D}}=[\bm{x}_{0},\bm{u}_{0:T},\bm{p}]$ , we prescribe a temporal information bottleneck starting from an initial latent state $\bm{x}_{0}^{z}$ . With this, physical state properties are enforced: 1) All variables of the model at time $t_{k}$ can be described by a function of the set of latent states $\bm{x}_{k}^{z}$ , latent inputs $\bm{u}_{k}^{z}$ , and the latent parameters $\bm{p}^{z}$ . 2) The state can be advanced in time with an integrator and inputs $\bm{u}^{z}(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 to the state vector at time $t_{0}$ . The state vector unambiguously defines the system’s state, such that we can simply set $q(\bm{x}_{0}^{z}|\bm{x}_{0:T},\bm{y}_{0:T},\bm{u}^{z}_{0:T},\bm{p}^{z})=q(\bm{x}_{0}^{z}|\bm{x}_{0})$ (4ba) $$

for the initial latent state. The formulation of our inference model is general enough to also accommodate initial state observers (i.e. identify $\bm{x}^{z}_{0}$ from $\bm{y}_{k_{\text{obs.}}:0}$), latent control, or latent parameter observers, which is of use when real-world data is used.
Because $\bm{u}^{z}_{0:T}$ and $\bm{p}^{z}$ are inputs to the generative model and therefore known, the only data points of $\bm{x}^{\mathcal{D}}$ we are interested in approximating with a surrogate, are $\bm{y}^{\mathcal{D}}=[\bm{x}_{0:T},\bm{y}_{0:T}]$.
The generative model (decoder) is then simply a sequential decoding,

$$ $p(\bm{x}^{\mathcal{D}}|\bm{z})\approx p(\bm{y}^{\mathcal{D}}|\bm{z})=p(\bm{x}_{0:T},\bm{y}_{0:T}|\bm{x}_{0:T}^{z},\bm{u}_{0:T}^{z},\bm{p}^{z})=\prod\nolimits_{i=0}^{T}p(\bm{x}_{i},\bm{y}_{i}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z}).$ (7) $$

Finally, the joint distribution of the observed and latent variables is

$$ $p(\bm{x}^{\mathcal{D}},\bm{z})=p(\bm{x}_{0:T},\bm{y}_{0:T}|\bm{x}_{0:T}^{z},\bm{u}_{0:T}^{z},\bm{p}^{z})p(\bm{p}^{z})p(\bm{u}^{z}_{0:T})p(\bm{x}^{z}_{0:T}).$ (8) $$

To include state reduction features as discussed in [2.2](#S2.SS2), we set a (non-hierarchical ) Gaussian prior $p(\bm{z})=p(\bm{p}^{z}),p(\bm{x}^{z}_{0:T}),p(\bm{u}^{z}_{0:T})=\mathcal{N}(\bm{z};\mathbf{0},\mathbf{I})$.
For using a trained B-NODE as a generative model, we need to control the data generation process. For that, we condition the distributions $p(\bm{p}^{z}),p(\bm{x}^{z}_{0:T}),p(\bm{u}^{z}_{0:T})$ with the posteriors described in ([4](#S3.E4)).

### 3.2 Implementation with deep neural networks for Model Order Reduction

Figure: Figure 2: Scheme of the Balanced Neural ODE.

The parameter posterior density and control input posterior density ([4a](#S3.E4.1)) are modeled as

$$ $\displaystyle q(\bm{p}^{z}|\bm{p})$ $\displaystyle=\mathcal{N}(\bm{p}^{z};\bm{\mu}(\bm{p}^{z}),\bm{\sigma}(\bm{p}^{z})),$ $\displaystyle\left(\bm{\mu}(\bm{p}^{z}),\bm{\sigma}(\bm{p}^{z})\right)$ $\displaystyle=\bm{f}_{\phi_{\text{en,p}}}(\bm{p})$ (9) $\displaystyle q(\bm{u}^{z}_{i}|\bm{u}_{i})$ $\displaystyle=\mathcal{N}(\bm{u}^{z};\bm{\mu}(\bm{u}^{z}_{i}),\bm{\sigma}(\bm{u}^{z}_{i})),$ $\displaystyle\left(\bm{\mu}(\bm{u}^{z}_{i}),\bm{\sigma}(\bm{u}^{z}_{i})\right)$ $\displaystyle=\bm{f}_{\phi_{\text{en,u}}}(\bm{u}_{i}).$ (10) The latent initial state posterior density ( 4ba ) for $\bm{x}_{0}^{z}$ is likewise implemented. The inference of the full latent state posterior density sequence $\bm{x}_{1:T}^{z}$ ( 4b ) is implemented with a Neural ODE $\displaystyle q(\bm{x}^{z}_{0}|\bm{x}_{0})$ $\displaystyle=\mathcal{N}(\bm{x}^{z}_{0};\bm{\mu}(\bm{x}^{z}_{0}),\bm{\sigma}(\bm{x}^{z}_{0})),$ $\displaystyle\qquad\left(\bm{\mu}(\bm{x}^{z}_{0}),\bm{\sigma}(\bm{x}^{z}_{0})\right)$ $\displaystyle=\bm{f}_{\phi_{\text{en,}x_{0}}}(\bm{x}_{0})$ (11) $\displaystyle\begin{split}q(\bm{x}_{i+1}^{z}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z})&=\mathcal{N}(\bm{x}_{i+1}^{z};\bm{\mu}(\bm{x}^{z}_{i+1}),\bm{\sigma}(\bm{x}^{z}_{i+1}))\\ \begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{i+1})\\ \bm{\sigma}(\bm{x}^{z}_{i+1})\end{bmatrix}&=\mathrlap{\begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{i})\\ \bm{\sigma}(\bm{x}^{z}_{i})\end{bmatrix}+\int_{t_{i}}^{t_{i+1}}\bm{f}_{\phi_{\text{NODE}}}(\bm{\mu}(\bm{x}^{z}(t)),\bm{\sigma}(\bm{x}^{z}(t)),\bm{u}^{z}(t),\bm{p}^{z})dt,}\end{split}$ (12) $$

that propagates the distribution parameters $\bm{\mu}(\bm{x}^{z}(t)),\bm{\sigma}(\bm{x}^{z}(t))$ through time.
[Figure 2](#S3.F2) shows the scheme of the resulting model.
For $\bm{f}_{\phi_{\text{LNODE}}}(\bullet)$, different configurations are possible. We can distinguish constant variance and dynamic variance state reduction

$$ $\displaystyle\bm{f}_{\phi_{\text{LNODE}}}(\bullet)=\begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{k})\\ \bm{\sigma}(\bm{x}^{z}_{k})\end{bmatrix}^{\prime}$ $\displaystyle=\begin{bmatrix}\bm{f}_{\phi_{\text{LNODE,const}}}(\bm{\mu}(\bm{x}^{z}(t)),\bm{u}^{z}(t),\bm{p}^{z})\\ 0\end{bmatrix}$ (13) $\displaystyle\bm{f}_{\phi_{\text{LNODE}}}(\bullet)=\begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{k})\\ \bm{\sigma}(\bm{x}^{z}_{k})\end{bmatrix}^{\prime}$ $\displaystyle=\begin{bmatrix}\begin{array}[]{l}\bm{f}_{\phi_{\text{LNODE,dyn,}\bm{\mu}}}(\bm{\mu}(\bm{x}^{z}(t)),\bm{u}^{z}(t),\bm{p}^{z})\\ \bm{f}_{\phi_{\text{LNODE,dyn,}\bm{\sigma}}}(\bm{\mu}(\bm{x}^{z}(t)),\bm{\sigma}(\bm{x}^{z}(t)),\bm{u}^{z}(t),\bm{p}^{z})\end{array}\end{bmatrix}.$ (14) $$

For the latter, the mean is independent of the variance.
Like this, information with state properties can not be encoded in the variance part of the state vector. This solely models the degree of uncertainty.
For fast generative usage of the model, we set the external noise variable $\bm{\epsilon}=\mathbf{0}$, such that the variance part of $\bm{f}_{\phi_{\text{LNODE}}}$ is inactive.
Stability and convergence of Neural ODEs can be improved by applying stochasticity during training(Volokhova et al., ; Ghosh et al., ; Liu et al., ). The VAE framework includes adaptively learned variance variables, that are applied as stochastic noise during training (Kingma & Welling, ). To bring both together, we propose to set for the unsampled inputs of $\bm{f}_{\phi_{\text{LNODE}}}$ $\left(\bm{\mu}(\bm{x}^{z}(t)),\bm{\sigma}(\bm{x}^{z}(t))\right)$

$$ $\displaystyle\bm{\mu}^{*}(\bm{x}^{z}(t))$ $\displaystyle=\bm{\mu}(\bm{x}^{z}(t))+\epsilon\bm{\sigma}(\bm{x}^{z}(t)),$ $\displaystyle\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ (15) $\displaystyle\bm{\sigma}^{*}(\bm{x}^{z}(t))$ $\displaystyle=\bm{\sigma}(\bm{x}^{z}(t))+\alpha_{\sigma}\cdot\epsilon\bm{\sigma}(\bm{x}^{z}(t)),$ $\displaystyle\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I}).$ (16) $$

To learn the state reduction in $\bm{f}_{\phi_{\text{LNODE}}}$, sampling of the mean state $\bm{\mu}(\bm{x}^{z}(t))$ is required.
If a state $x_{i}^{z}$ is inactive, the model learns $\bm{\mu}(x_{i}^{z})=0$ and $\bm{\sigma}(x_{i}^{z})=1$ for that state.
Then the reparameterized state ([15](#S3.E15)) is distributed unconditioned, $\mu^{*}(x_{i}^{z}(t))=\mathcal{N}(0,1)$, i.e. no information transmission is possible and the model learns to ignore that dimension.
For $\alpha_{\sigma}$, a small value should be chosen.
The generative model (decoder) $p(\bm{x}_{i},\bm{y}_{i}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z})$ ([7](#S3.E7)) is represented with ([A.4.3](#A1.SS4.SSS3))

$$ $\hat{\bm{x}}_{i},\hat{\bm{y}}_{i}=\bm{f}_{\phi_{\text{de}}}(\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z}),$ (17) $$

and evaluated after sampling from the conditional distributions of the function inputs.

### 3.3 Adaptations to learn Koopman operator approximations

Learning Koopman operator approximations is as simple as setting for $\bm{f}_{\phi_{\text{LNODE}}}$ a linear model with parameters $\phi=\{A_{i,i},B_{i,i}\}$ instead of neural networks.
For constant-variance state reduction ([13](#S3.E13))

$$ $\displaystyle\begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{k})\\ \bm{\sigma}(\bm{x}^{z}_{k})\end{bmatrix}^{\prime}$ $\displaystyle=\begin{bmatrix}\begin{array}[]{ll}\bm{A}_{\mu,\mu}&\mathbf{0}\\ \mathbf{0}&\mathbf{0}\end{array}\end{bmatrix}\cdot\begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{k})\\ \bm{\sigma}(\bm{x}^{z}_{k})\end{bmatrix}+\begin{bmatrix}\bm{B}_{\mu,u}\\ \mathbf{0}\end{bmatrix}\cdot\bm{u}^{z}(t).$ (18) and for dynamic-variance state reduction ( 14 ) $\displaystyle\begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{k})\\ \bm{\sigma}(\bm{x}^{z}_{k})\end{bmatrix}^{\prime}$ $\displaystyle=\begin{bmatrix}\begin{array}[]{ll}\bm{A}_{\mu,\mu}&\mathbf{0}\\ \bm{A}_{\sigma,\mu}&\bm{A}_{\sigma,\sigma}\end{array}\end{bmatrix}\cdot\begin{bmatrix}\bm{\mu}(\bm{x}^{z}_{k})\\ \bm{\sigma}(\bm{x}^{z}_{k})\end{bmatrix}+\begin{bmatrix}\bm{B}_{\mu,u}\\ \bm{B}_{\sigma,u}\end{bmatrix}\cdot\bm{u}^{z}(t).$ (19) $$

Opposed to Autoencoder approaches for learning Koopman operators (Brunton et al., ),
the latent dimension must be only implicitly set with the $\beta$-value.

### 3.4 Training

Following the approach of Girin et al. () to explicitly represent the sampling dependencies, the data likelihood and KL divergence terms of the ELBO ([3](#S2.E3)) can be expanded to (derived in [A.4.1](#A1.SS4.SSS1))

$$ $\displaystyle\mathbb{E}_{q(\bm{z}|\bm{x}^{\mathcal{D}})}\log\left[p(\bm{x}^{\mathcal{D}}|\bm{z})\right]$ $\displaystyle=\mathbb{E}_{q(\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{u}_{0:T},\bm{p})}\left[\mathbb{E}_{q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})}\sum\nolimits_{i=0}^{T}\log p(\bm{x}_{i},\bm{y}_{i}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z})\right]$ $\displaystyle D_{\text{KL}}\left[q(\bm{z}|\bm{x}^{\mathcal{D}})||p(\bm{z})\right]$ $\displaystyle=D_{\text{KL}}\left[q(\bm{p}^{z}|\bm{p})||p(\bm{p}^{z})\right]+D_{\text{KL}}\left[q(\bm{u}_{0:T}^{z}|\bm{u}_{0:T})||p(\bm{u}_{0:T}^{z})\right]$ $\displaystyle\phantom{=}+\mathbb{E}_{q(\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{u}_{0:T},\bm{p})}\left(D_{\text{KL}}\left[q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})||p(\bm{x}_{0:T}^{z})\right]\right).$ (20) $$

As we simultaneously learn latent embedding and dynamics in latent space, the first step in training a B-NODE is finding initially stable ODE systems. Proven strategies for that are growing horizon (from small to long prediction horizons) and multiple shooting (splitting sequences in multiple, smaller ones) approaches.
Then, we observed robust training behavior for constant variance B-NODEs with moderate $\beta$ values (i.e. $\beta=0.1,0.01$), which is from our experience more stable than State Space Neural ODEs with the same data.
Too high $\beta$ leads to failing reconstruction, while too small $\beta$ turns the VAE integration effectively off.
In the dynamic variance case, the reparameterization of ODE inputs directly interferes with ODE stability, and more careful hyperparameter tuning (most important learning rate) is necessary.

## 4 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 a simple 1D heat flow model. It consists of 16 cells describing the temperature field of a stick whose ends are excited with external input signals ([A.5.2](#A1.SS5.SSS2)).
To generate datasets from physical models, we employ a sampling strategy that aims to capture all relevant state-input combinations ([A.5.1](#A1.SS5.SSS1)).

Figure: ((a)) Reconstruction of temperature field from latent states and latent inputs for one sample in test set, B-NODE with constant variance.

[3(a)](#S4.F3.sf1) shows the original temperature field generated by the heat flow model $\bm{x}_{0:T}$, its state-reduced representation in latent space $\bm{x}_{0:T}^{z}$ and the reconstruction of the temperature field with the decoder $\hat{\bm{x}}_{0:T}$ with an RMSE of \qty0.4 (reported loss-values are always normalized by respective means).
Taking the classical MOR perspective, we learn a nonlinear embedding of the full state space on a latent sub-manifold 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 an encoder and a decoder.
The dimensionality results from the balance between state reduction and reconstruction quality, and is modified with $\beta$ ([4(a)](#S4.F4.sf1)).
As in [3(a)](#S4.F3.sf1), the state reduction is constant over the whole time sequence.
Because of that, latent space dimensionalities are not important hyperparameters, as long as they are large enough.
Naively adding inputs $\bm{u}^{z}(t)$ to Latent ODE (Rubanova et al., ) as in [3(b)](#S4.F3.sf2) leads to failing state reduction.
The Latent ODE propagates a reparameterized initial latent state $\bm{x}^{z}_{0}$ and does not allow to enforce state reduction at times $t>0$.

Figure: ((a)) Comparison between prediction quality of B-NODEs and Truncated Balanced Realization (TBR).

Instead of learning the system’s response, we learn the system dynamics $(\bm{x}^{z})^{\prime}$ as a vector field depending on $\bm{x}^{z},\bm{u}^{z},\bm{p}^{z}$.
This enables generalization to input signals unseen in training, for example, step inputs ([4(b)](#S4.F4.sf2)).
As the discretized heat flow model is described by a linear system, we can apply the model order reduction technique Truncated Balanced Realization (TBR), which finds a balance of well-observable and controllable states (Schilders et al., ).
Our method consistently outperforms TBR ([4(a)](#S4.F4.sf1)), for example, for 4 states the BNODE has an RMSE of \qty0.3 vs. \qty14.2 for TBR.

### 4.2 Real-world use case: surrogate model of a thermal power plant

To test the surrogate qualities with a complex model, we employ a dataset generated by a power plant simulation model (Vojacek et al., ), of which we predict 102 states and 6 output signals, reacting on one input signal.
[Figure 5](#S4.F5) compares the RMSE and maximum error of different B-NODEs and a State Space NODE as surrogates for the power plant model.
The B-NODE with $\beta=0.05$ and constant variance exhibits the smallest reconstruction error of all models while only requiring 17 states instead of 102.
Taking a look into the latent space of the B-NODE with dynamic variance, [6(a)](#S4.F6.sf1) exhibits a reduction of variance at times with high dynamical behavior.

Figure: Figure 5: Error comparison key quantities of interest (outputs) for surrogates of power plant.

A possible explanation could be that dynamics are encoded by states, while steady-state behavior is modeled by the decoder. This phenomenon should be further exploited in future work.
However, we can perform uncertainty quantification in data space by sampling and generating multiple trajectories.
Finally, we make a performance comparison ([6(b)](#S4.F6.sf2)). Most significantly, the B-NODE facilitates a \qty75 reduction in computation time on CPU and a \qty99 reduction on GPU for batch simulation.
Simulating in dedicated simulation frameworks will increase speed additionally.

Figure: ((a)) Latent state trajectories for B-NODE with dynamic variance ($\beta=0.1$), ordered by mean $D_{\text{KL}}$, tested with step input.

## 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 et al., )

$$ $\displaystyle\begin{bmatrix}x_{1}\\ x_{2}\end{bmatrix}^{\prime}=\begin{bmatrix}ax_{1}\\ b\cdot(x_{2}-x_{1}^{2})\end{bmatrix},$ $\displaystyle\begin{bmatrix}x_{1}\\ x_{2}\\ x_{3,add.}\end{bmatrix}^{\prime}=\begin{bmatrix}a&0&0\\ 0&b&-b\\ 0&0&2a\end{bmatrix}\cdot\begin{bmatrix}x_{1}\\ x_{2}\\ x_{3,add.}\end{bmatrix}.$ (21) $$

System dynamics of linear systems are determined by the eigenvalues (here $\lambda=\{a,b,2a\}$).
If two systems have the same eigenvalues, they represent the same dynamics.

Figure: Figure 7: Eigenvalues of analytic solution and B-NODE.

The B-NODE can choose different states than the analytic solution, but projects them back into the state space of the analytic model with a linear decoder.
We trained a constant variance B-NODE as in [3.3](#S3.SS3) with linear layers with a dataset generated by sampling initial values with parameters $a=$-0.5, $b$=-1 for different values of $\beta$ and offered a latent dimensionality of $\text{dim}(\bm{x}^{z})$=$64$.
For $\beta$ of 0.01, we found 3 active dimensions ($D_{\text{KL}}(x_{i}^{z})>0.1$), just as the analytic solution has.
The approximated and analytic eigenvalues lie at almost the same position ([Figure 7](#S5.F7)) and therefore represent the same dynamics.

### 5.2 Real-World Use Case: Linearized Model of a Thermal Power Plant

Figure: Figure 8: Reconstruction sample of B-NODE ($\beta=0.01$) with highest error and 25 linear states (original 102 nonlinear; offered $\text{dim}(\bm{x}^{z})$=$1024$).

A model with a full linear path from model input to model output is desirable for model-predictive control (MPC) applications.
It can be learned by setting linear models with an additive bias for the control encoder $\bm{f}_{\phi_{\text{en,u}}}$ and the decoder $\bm{f}_{\phi_{\text{de}}}$.
Then, only the state encoder has non-linearities and facilitates a nonlinear state transformation into a linear latent space that can approximately represent the dynamics.
We trained constant variance B-NODEs with $\beta\in\{$0.1$,$0.01$,$0.001$\}$ for this use case with $\text{dim}(\bm{x}^{z})$=$\text{dim}(\bm{u}^{z})$=$1024$ to predict all 102 state variables and 6 output variables of interest for the thermal power plant simulation model.
The best model achieved a RMSE of \qty1.2 and a maximum error of \qty14.5 on the selected outputs (normalized by mean) ([Figure 8](#S5.F8)), while requiring 25 states.
Interestingly, even though we only have one control input signal $u(t)$, the latent control dimension increases to 18 inputs.
We observed that for short prediction horizons, the approximation quality is better, which fits the short-horizon model requirement of MPC.

## 6 Discussion & Future Work

We tailored our approach around the generation of surrogate models from simulations.
Using B-NODEs for real-world data is self-suggesting, but necessitates inference of the latent state $\bm{x}_{0}^{z}$ or parameters $\bm{p}^{z}$ with more sophisticated methods, e.g. Kalman filters.
For use cases that require uncertainty quantification, future work should exploit B-NODE’s capability of generating probabilistic time series.
As mentioned, small $\beta$-values can turn the VAE integration effectively off but might be desirable because of higher reconstruction accuracy.
Therefore, future work includes a more detailed analysis of numerical stability, disentangling, and the effect of different training strategies.
Finally, we plan to integrate surrogates in real-world simulations, e.g. enabling yearly energy system simulation and optimization.
Especially linear BNODE surrogates pose an exciting opportunity for enabling efficient model predictive control applications, whose requirement of short-horizon prediction are beneficial for accurate linear BNODE training.
It furthermore offers new approaches to global sensitivity analysis, and interpretable surrogate models and should be compared to existing methods like DMDc.

## 7 Conclusion

We present Balanced Neural ODEs (B-NODEs), balancing state reduction and reconstruction quality, a well-automatable method for global nonlinear model order reduction.
For that, we combine VAEs with state space models to leverage VAE’s information bottleneck to enable measurable dimensional 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 time-varying inputs.
To achieve this, we continuously propagate variational parameters through time with a Neural ODE (to our awareness, for the first time).
We tested our method on academic and real-world cases.
For modeling of a discretized heat flow, our method proved to have low errors, to generalize on unseen inputs, to outperform truncated balanced realization and to avoid state reduction failures, as Latent ODEs would have with time-varying inputs.
Additionally, it can accurately determine the correct Koopman operator of an analytically solvable example.
For the real-world example of a thermal power plant, we demonstrated a \qty94 state reduction with RMSE of \qty2.4, while taking \qtyrange7599 (one/batch simulation) less computation time.
Finally, we could also linearize the power plant by using 25 states with an RMSE of \qty1.4.

#### Reproducibility Statement

We provided the source code to reproduce the complete project on the first page.
In addition to the code necessary for implementing VAE, B-NODE and State Space NeuralODE, we also provide the files that implement the first-principle models along with FMUs (Modelica Association, ), 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 on sampling of first-principle models for dataset generation, data standardization, approximation of the data likelihood, and most important training settings in the appendix.

## Appendix 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. a pressure at around \qty1e5 compared to a temperature at around \qty50).
Also for loss functions, standardized values are used to avoid adapting weighting factors to the different orders of magnitudes.
In general, the not-standardized values are used only for graphical presentation of results.
To standardize the data of a data set, mean $\bm{\mu}$ and variance $\bm{\sigma}$ are calculated for each variable. Then we use the following transformation for each axis $x$ (Prince (), C.2.4):

$$ $x_{\text{stand.}}=\frac{x-\bm{\mu}}{\bm{\sigma}}$ (22) $$

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 are calculated and saved as untrainable parameter in the PyTorch-Model. It can be used as a normal layer within the model, and has also the ability to inverse standardize - i.e.

$$ $x=\bm{\sigma}\cdot x_{\text{stand.}}+\bm{\mu}.$ (23) $$

Opposed to batch-normalization (Prince (), 11.4), where the statistics $\bm{\mu}$ and $\bm{\sigma}$ 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, we chose to normalize by mean values:

$$ $x_{\text{norm.}}=\frac{x}{\bm{\mu}}$ (24) $$

### A.2 State Space Neural ODE

Figure: Figure A.1: 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 $\bm{x}^{\mathcal{D}}=[\bm{x}_{0:T},\bm{y}_{0:T},\bm{u}_{0:T},\bm{p}]$, the State Space Neural ODE is trained by minimizing the error of the predictions $[\hat{\bm{x}}(t),\hat{\bm{y}}(t)]$

$$ $\phi_{\text{NODE}},\phi_{\text{Out}}=\underset{\phi_{\text{NODE}},\phi_{\text{Out}}}{\text{argmin}}\text{MSE}(\hat{\bm{x}}_{0:T},\bm{x}_{0:T})+\text{MSE}(\hat{\bm{y}}_{0:T},\bm{y}_{0:T}).$ (25) $$

### A.3 Variational Autoencoders for Dimensional Reduction

#### A.3.1 ELBO derivation

We could learn a model $p(\bm{x}|\bm{z})$ if we would know the latent variable $\bm{z}$. But we do not know $\bm{z}$. To overcome that, we would require the ”true” model $p(\bm{z}|\bm{x})$ that represents the so called posterior probability(^4^44 To explain why the term posterior probability is used for $\text{Pr}(\bm{z}|\bm{x})$: We took the assumption that $\text{Pr}(\bm{z})$ is the known prior probability. In terms of Bayes’ Rule, after observing the generation process $\text{Pr}(\bm{x}|\bm{z})$ we can infer the posterior probability $\text{Pr}(\bm{z}|\bm{x})$. It’s a bit cumbersome to call this posterior probability. Even though this would be the first thing needed for the training of a VAE, but in terms of the VAE theory, the starting point is that the data generation process $\text{Pr}(\bm{x}|\bm{z})$ can be traced back to a prescribed prior probability $\text{Pr}(\bm{z})$.)
$\text{Pr}(\bm{z}|\bm{x})$. But we do not have that model either. The learning problem is intractable.
To make the learning problem tractable, a posterior model $q(\bm{z}|\bm{x})$ that approximates the true posterior probability model $p(\bm{z}|\bm{x})$ is introduced:

$$ $q(\bm{z}|\bm{x})\approx p(\bm{z}|\bm{x}).$ (26) $$

The model $q(\bm{z}|\bm{x})$ is called the encoder or recognition model (Kingma & Welling, , Chapter 2.1). This approximation is the variational approximation that gives the VAE its name. (Prince, , Chapter 17.5).
With this approximation, the latent variable model can be trained using data samples $\bm{x}$.

The following can be derived by using Jensen’s inequality (as in (Prince, , chapter 17.3.33)), but we choose a derivation by (Kingma & Welling, , chapter 2.2) that avoids this.

To obtain the latent variable model, express the log-likelihood $\log p(\bm{x})$ of the data as an expectation over the posterior probability $q(\bm{z}|\bm{x})$

$$ $\displaystyle\log p(\bm{x})$ $\displaystyle=\mathbb{E}_{q(\bm{z}|\bm{x})}\log\left[p(\bm{x})\right].$ (27) By applying the the product rule $\displaystyle\log p(\bm{x})$ $\displaystyle=\mathbb{E}_{q(\bm{z}|\bm{x})}\log\left[\frac{p(\bm{x},\bm{z})}{p(\bm{z}|\bm{x})}\right],$ expanding with $\frac{q(\bm{z}|\bm{x})}{q(\bm{z}|\bm{x})}$ $\displaystyle=\mathbb{E}_{q(\bm{z}|\bm{x})}\log\left[\frac{p(\bm{x},\bm{z})}{q(\bm{z}|\bm{x})}\frac{q(\bm{z}|\bm{x})}{p(\bm{z}|\bm{x})}\right],$ splitting the integrand of the expectation $\displaystyle=\mathbb{E}_{q(\bm{z}|\bm{x})}\log\left[\frac{p(\bm{x},\bm{z})}{q(\bm{z}|\bm{x})}\right]+\mathbb{E}_{q(\bm{z}|\bm{x})}\log\left[\frac{q(\bm{z}|\bm{x})}{p(\bm{z}|\bm{x})}\right]$ and applying the product rule for the first term, we can derive $\displaystyle\log p(\bm{x})$ $\displaystyle=\underbrace{\mathbb{E}_{q(\bm{z}|\bm{x})}\log\left[\frac{p(\bm{x}|\bm{z})}{q(\bm{z}|\bm{x})}p(\bm{z})\right]}_{=\mathcal{G}\text{(ELBO)}}+\underbrace{\mathbb{E}_{q(\bm{z}|\bm{x})}\log\left[\frac{q(\bm{z}|\bm{x})}{p(\bm{z}|\bm{x})}\right]}_{=D_{\text{KL}}\left(q(\bm{z}|\bm{x})||p(\bm{z}|\bm{x})\right)}.$ (28) $$

The KL-Divergence in the second term is always $\geq 0$. Therefore, the first term is a lower bound to the log-likelihood of the data. It is called the evidence lower bound (ELBO). (Kingma & Welling, , 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 KL-Divergence of the second term determines two distances (Kingma & Welling, , 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(\bm{x})$. This is called the tightness of the bound.

It is now possible to train a variational autoencoder by maximizing the ELBO $\mathcal{G}$. Maximizing ELBO will concurrently optimize two objectives (Kingma & Welling, , Chapter 2.2.1):

- •
It will improve the prediction quality $\log p(\bm{x})$ of the generative model $p(\bm{x}|\bm{z})$.
- •
It will minimize the error of approximating the true posterior distribution $p(\bm{z}|\bm{x})$ through $q(\bm{z}|\bm{x})$, meaning that the encoder will get better.

The KL-Divergence in ([3](#S2.E3)) for the comparison with the assumed prior $p(\bm{z})=\mathcal{N}(\bm{z};\mathbf{0},\bm{I})$ can be solved analytically to (Odaibo, )

$$ $D_{\text{KL}}\bigl{(}q_{{\bm{\phi}}_{\text{en}}}(\bm{z}|\bm{x}^{\mathcal{D}})||p(\bm{z})\bigr{)}=\sum_{j=1}^{n_{\bm{z}}}-\frac{1}{2}\Bigl{[}1+\log(\sigma_{j}^{2})-\sigma_{j}^{2}-\mu_{j}^{2}\Bigr{]}\geq 0,$ (29) $$

with $n^{z}$ being the chosen dimensionality of the latent variable.

Figure: Figure A.2: Scheme of a 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 5 temperature trajectories formed one sample (i.e. 5 sequences $x$ 337 timesteps per sequence $=$ 1685 datapoints per sample). We therefore re-implemented the model described in (Aka et al., ).

- •
parameter values: determined from building data
- •
parameter ranges: $\frac{1}{5}$ to $5$ times the default value
- •
timestep: \qty1800
- •
simulation length: \qty7
- •
number of timesteps: $337$
- •
number of samples: 2048 (12% test, 12% validate, 76% train)

Experiment Details

- •
Network Structure:
–
number of linear layers: 4
–
hidden dimension: 2048
–
latent space dimension: 512
–
activation function: ReLU
- •
Training Parameters:
–
beta: $1\text{\times}{10}^{-3}$
–
max epochs: $1\text{\times}{10}^{4}$
–
learning-rate: $1\text{\times}{10}^{-3}$
–
weight decay: $1\text{\times}{10}^{-4}$
–
batch size: $256$
–
early stopping patience: 300
–
early stopping threshold: 0.005
–
early stopping mode: relative
–
number of reparametrization passes for train and test: 1
–
automatic mixed precision: off
–
parameter to decoder: on (PELS-VAE), off (VAE)
- •
Code File: .\networks\pels_vae_linear\vae_train_test.py

#### A.3.3 Additional Figures

Figure: Figure A.3: Density distributions of $D_{\text{KL}},\bm{\mu},\bm{\sigma}$ for the first 7 latent space dimensions for the complete test set ordered by their mean KL divergence. Plot for the $\beta$-VAE with $\beta=$1\text{\times}{10}^{-3}$$.The model was trained with a data set of the office model

To verify that the first 5 dimension encode most information, we ordered the latent dimensions of a trained VAE by their average $D_{\text{KL}}$, masked all latent channels to 0 and consecutively lifted that mask while calculating for each configuration the MSE on the dataset ([4(a)](#A1.F4.sf1)).
As anticipated, the first 5 latent channels reduce the error by the most ($$0.85$\mapsto$0.09$$) while the last 507 dimensions only achieve marginal improvements ($$0.09$\mapsto$0.04$$).
Notably, the improvements are monotonously decreasing.
This motivates to use the mean $D_{\text{KL}}$ per latent channel as a measure for active latent dimensions, and we
To find a threshold value, we tested a variety of values in [4(b)](#A1.F4.sf2) and
choose for counting a value of $D_{\text{KL}}(z)>0.1$ ([4(b)](#A1.F4.sf2)).

Figure: ((a)) Experiment: MSE on dataset for increasing number of active latent dimensions by lifting a mask on the $D_{\text{KL}}$-ordered latent dimensions.

### 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 (), the ELBO is

$$ $\mathcal{G}=\underbrace{\mathbb{E}_{q(\bm{z}|\bm{x}^{\mathcal{D}})}\log\left[p(\bm{x}^{\mathcal{D}}|\bm{z})\right]}_{\text{Reconstruction Accuracy}}-\underbrace{D_{\text{KL}}\left[q(\bm{z}|\bm{x}^{\mathcal{D}})||p(\bm{z})\right]}_{\text{Information Bottleneck}}.$ (30) $$

First, we collect the governing equations of the model and adapt the way they are expressed
(([4](#S3.E4)),
([4a](#S3.E4.1)),
([4b](#S3.E4.2)),
([4ba](#S3.E4.2.1)),
([7](#S3.E7)),
([8](#S3.E8))). The approximated posterior (inference model) is

$$ $\displaystyle q(\bm{z}|\bm{x}^{\mathcal{D}})$ $\displaystyle=q(\bm{x}_{0:T}^{z},\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{x}_{0:T},\bm{u}_{0:T},\bm{p})=q(\bm{p}^{z}|\bm{p})q(\bm{u}_{0:T}^{z}|\bm{u}_{0:T})q(\bm{x}_{0:T}^{z}|\bm{x}_{0:T},\bm{u}_{0:T}^{z},\bm{p}^{z}),$ $\displaystyle=\underbrace{\prod_{i=0}^{T}q(\bm{u}_{i}^{z}|\bm{u}_{i})}_{\text{Control Encoder}}\underbrace{q(\bm{p}^{z}|\bm{p})}_{\text{Parameter Encoder}}\underbrace{q(\bm{x}_{0}^{z}|\bm{x}_{0})}_{\text{Initial State Observer}}\underbrace{\prod_{i=0}^{T-1}q(\bm{x}_{i+1}^{z}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z})}_{\text{Simulator}}$ $\displaystyle=\underbrace{q(\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{u}_{0:T},\bm{p})}_{\text{independent sampling possible}}\quad\underbrace{q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})}_{\text{dependent on previous sampling}},$ (31a) the generative model is $\displaystyle p(\bm{x}^{\mathcal{D}}|\bm{z})$ $\displaystyle=p(\bm{x}_{0:T},\bm{y}_{0:T}|\bm{x}_{0:T}^{z},\bm{u}_{0:T}^{z},\bm{p}^{z})=\prod_{i=0}^{T}p(\bm{x}_{i},\bm{y}_{i}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z}),$ (31b) and we prescribed the prior as $\displaystyle p(\bm{z})$ $\displaystyle=p(\bm{p}^{z})p(\bm{x}_{0:T}^{z})p(\bm{u}_{0:T}^{z}).$ (31c) $$

When inserting these terms in the ELBO, the learning problem is intractable, as the probability density functions are parameterized by outputs of neural networks that cannot be solved analytically. Therefore, sampling from the posterior distribution $q(\bm{z}|\bm{x}^{\mathcal{D}})$ in conjunction with the reparameterization trick is necessary, whenever latent variables are propagated from deep neural network to another. The propagation of the latent state distribution by the Neural ODE is an exception for this, this propagation of probability distributions is implemented tractably and state reduction ensured by the sampling of inputs ([15](#S3.E15)).
First, we derive the data likelihood part of the ELBO,

$$ $\displaystyle\mathbb{E}_{q(\bm{z}|\bm{x}^{\mathcal{D}})}\log\left[p(\bm{x}^{\mathcal{D}}|\bm{z})\right]$ $\displaystyle=\mathbb{E}_{q(\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{u}_{0:T},\bm{p})q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})}\log\left[p(\bm{x}_{0:T},\bm{y}_{0:T}|\bm{x}_{0:T}^{z},\bm{u}_{0:T}^{z},\bm{p}^{z})\right],$ Considering the sampling dependencies, $\displaystyle=\mathbb{E}_{q(\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{u}_{0:T},\bm{p})}\left[\mathbb{E}_{q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})}\log\left[p(\bm{x}_{0:T},\bm{y}_{0:T}|\bm{x}_{0:T}^{z},\bm{u}_{0:T}^{z},\bm{p}^{z})\right]\right],$ Using that $p(\bm{x}^{\mathcal{D}}|\bm{z})$ can be factorized, $\displaystyle=\mathbb{E}_{q(\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{u}_{0:T},\bm{p})}\left[\mathbb{E}_{q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})}\left[\sum_{i=0}^{T}\log p(\bm{x}_{i},\bm{y}_{i}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z})\right]\right].$ (32) $$

After this, we derive the KL divergence part. By using the derivation in [A.4.4](#A1.SS4.SSS4), we can see that the KL divergence splits up in separate terms for different random variables. Because of this,

$$ $\displaystyle D_{\text{KL}}\left[q(\bm{z}|\bm{x}^{\mathcal{D}})||p(\bm{z})\right]$ $\displaystyle=D_{\text{KL}}\left[q(\bm{p}^{z}|\bm{p})q(\bm{u}_{0:T}^{z}|\bm{u}_{0:T})q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})||p(\bm{p}^{z})p(\bm{u}_{0:T}^{z})p(\bm{x}_{0:T}^{z})\right].$ As $q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})$ depends on previous sampling, we cam write this as, $\displaystyle=D_{\text{KL}}\left[q(\bm{p}^{z}|\bm{p})||p(\bm{p}^{z})\right]+D_{\text{KL}}\left[q(\bm{u}_{0:T}^{z}|\bm{u}_{0:T})||p(\bm{u}_{0:T}^{z})\right]$ $\displaystyle\quad+\mathbb{E}_{q(\bm{u}_{0:T}^{z},\bm{p}^{z}|\bm{u}_{0:T},\bm{p})}\left(D_{\text{KL}}\left[q(\bm{x}_{0:T}^{z}|\bm{x}_{0},\bm{u}_{0:T}^{z},\bm{p}^{z})||p(\bm{x}_{0:T}^{z})\right]\right).$ (33) $$

Each variable in the KL divergence defined here expands likewise as a sum of KL divergence terms. Because we prescribe (parameterized) Gaussians, the KL divergence can be solved analytically.
The model is then trained as

$$ ${{\bm{\phi}}_{\text{en,p}}},{{\bm{\phi}}_{\text{en,u}}},{{\bm{\phi}}_{\text{en,x}}},{{\bm{\phi}}_{\text{LNODE}}},{{\bm{\phi}}_{\text{de}}}=\underset{{\bm{\phi}}}{\text{argmax}}\quad\mathcal{G}\bigl{(}{{\bm{\phi}}_{\text{en,p}}},{{\bm{\phi}}_{\text{en,u}}},{{\bm{\phi}}_{\text{en,x}}},{{\bm{\phi}}_{\text{LNODE}}},{{\bm{\phi}}_{\text{de}}},\bm{x}_{0:T},\bm{u}_{0:T},\bm{p}\bigr{)}$ (34) $$

Remark: Typically, as proposed by Kingma & Welling (), 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:
Opposed to Girin et al. (), we do not condition the true posterior probability density of the generative model for the latent states at time $t_{k}$ on the previous probability densities (for $\bm{x}^{z}_{i-1},\bm{u}^{z}_{i-1},\bm{p}^{z}$), but prescribe the latent space in general to follow a Gaussian distribution with mean zero and variance of one ($\mathcal{N}(\mathbf{0},\mathbf{\mathbb{I}})$) at all discrete times $t\in[0,T]$.
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.

#### A.4.2 Normalization of ELBO terms

We normalize the reconstruction term with $\frac{1}{2}$ to account for the summation of the state reconstruction and output reconstruction MSE,

$$ $\mathcal{G}_{\text{R}}\approx-\frac{1}{2}\bigl{(}\text{MSE}(\hat{\bm{x}},\bm{x})+\text{MSE}(\hat{\bm{y}},\bm{y})\bigr{)}.$ (35) $$

Let $D_{\text{KL},\square}$ be the KL divergence for $\square\in[\bm{p},\bm{u}_{k},\bm{x}_{k}],k\in[0,T]$. We calculate the Kullback-Leibler part of the ELBO as

$$ $\mathcal{L}_{\text{KL}}=\frac{1}{n_{\bm{p}}+n_{\bm{u}}+n_{\bm{x}}}\Bigl{(}D_{\text{KL},\bm{p}}+\frac{1}{T}\sum_{k=0}^{T}D_{\text{KL},\bm{u}_{k}}+\frac{1}{T}\sum_{k=0}^{T}D_{\text{KL},\bm{x}_{k}}\Bigr{)}$ (36) $$

Time-varying variables are normalized by the number of time steps $T$. Furthermore, we choose to normalize the components by their dimensions in physics space $n_{\square}$ in order to being able to use the same $\beta$ for different use cases.

#### A.4.3 Approximation of data log-likelihood

Optimizing the log-likelihood of $p(\bm{x}|\bm{z})$ can be represented by (Prince, , 5.3)

$$ $\underset{\phi_{\text{de}}}{\text{argmax}}\quad\prod_{i=0}^{T}p(\bm{x}_{i},\bm{y}_{i}|\bm{x}_{i}^{z},\bm{u}_{i}^{z},\bm{p}^{z})=\underset{\phi_{\text{de}}}{\text{argmin}}\quad\sum_{i=0}^{T}\text{MSE}([\hat{\bm{x}}_{i},\hat{\bm{y}}_{i}],[\bm{x}_{i},\bm{y}_{i}]),$ (37) $$

where $[\hat{\bm{x}}_{i},\hat{\bm{y}}_{i}]$ is the output of ([17](#S3.E17)) and $p$ 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(z),q(z)$ is defined as

$$ $\displaystyle D_{\text{KL}}\left(q(z)||p(z)\right)$ $\displaystyle=\int_{z}q(z)\log\left[\frac{p(z)}{q(z)}\right]dz$ For the product of two distributions of different random variables $z_{1},z_{2}$ this definition yields: $\displaystyle\mathclap{D_{\text{KL}}\left(q(z_{1})q(z_{2})||p(z_{1})p(z_{2})\right)}$ $\displaystyle=\int q(z_{1})q(z_{2})\log\left[\frac{p(z_{1})p(z_{2})}{q(z_{1})q(z_{2})}\right]d\bm{z}$ $\displaystyle=\int_{z_{2}}\int_{z_{1}}q(z_{1})q(z_{2})\log\left[\frac{p(z_{1})p(z_{2})}{q(z_{1})q(z_{2})}\right]dz_{1}dz_{2}$ Using logarithm laws $\displaystyle=\int_{z_{2}}\int_{z_{1}}q(z_{1})q(z_{2})\left(\log\left[\frac{p(z_{1})}{q(z_{1})}\right]+\log\left[\frac{p(z_{2})}{q(z_{2})}\right]\right)dz_{1}dz_{2}$ Splitting the integral $\displaystyle=\int_{z_{2}}\int_{z_{1}}q(z_{1})q(z_{2})\log\left[\frac{p(z_{1})}{q(z_{1})}\right]dz_{1}dz_{2}+\int_{z_{2}}\int_{z_{1}}q(z_{1})q(z_{2})\log\left[\frac{p(z_{2})}{q(z_{2})}\right]dz_{1}dz_{2}$ Moving the integration symbols and sorting the terms $\displaystyle=\int_{z_{2}}q(z_{2})\int_{z_{1}}q(z_{1})\log\left[\frac{p(z_{1})}{q(z_{1})}\right]dz_{1}dz_{2}+\int_{z_{1}}q(z_{1})\int_{z_{2}}q(z_{2})\log\left[\frac{p(z_{2})}{q(z_{2})}\right]dz_{2}dz_{1}$ Inserting the definition of KL-Divergence again $\displaystyle=\int_{z_{2}}q(z_{2})D_{\text{KL}}\left(q(z_{1})||p(z_{1})\right)dz_{2}+\int_{z_{1}}q(z_{1})D_{\text{KL}}\left(q(z_{2})||p(z_{2})\right)dz_{1}$ using that probability density functions always integrate to 1, we see that $\displaystyle=D_{\text{KL}}\left(q(z_{1})||p(z_{1})\right)+D_{\text{KL}}\left(q(z_{2})||p(z_{2})\right).$ (38) $$

Hence we follow that the KL-Divergence of two distributions on different random variables always yields a 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 is by imitating the data $[\bm{x}_{k},\bm{y}_{k}]$ at discretization points $k=0\dots T$ given $[\bm{x}_{0},\bm{p},\bm{u}_{k}]$.
Therefore, we need data of the input-output pairs:

$$ $(\bm{x}_{0},\bm{p},\bm{u}_{0:T})\mapsto(\bm{x}_{1:T},\bm{y}_{0:T})$ $$

We choose the control input to be constant in the discretization intervals:

$$ $\bm{u}(t)=\bm{u}_{k}=\text{const.}\quad\forall t\in[t_{k},t_{k+1}]$ $$

From experience, uniform random sampling proved to allow good learning for values that are constant over a simulation, which is why we sample the parameters and initial values like that:

$$ $\displaystyle\bm{x}_{0}$ $\displaystyle\sim\mathcal{U}(\bm{x}_{0,-},\bm{x}_{0,+})$ (39) $\displaystyle\bm{p}$ $\displaystyle\sim\mathcal{U}(\bm{p}_{-},\bm{p}_{+})$ (40) $$

with $\mathcal{U}$ denoting independent random sampling, and $\square_{-},\square_{+}$ 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 is avoided.
However, note that in general

$$ $\displaystyle\min{\bm{x}_{1:T}}\neq\bm{x}_{0,-}\quad\text{and}\quad\max{\bm{x}_{1:T}}\neq\bm{x}_{0,+}$ $$

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 or 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:

$$ $(\bm{x})^{\prime}=f(\bm{x},\bm{p},\bm{u})$ $$

This makes it possible to enter the training at arbitrary locations within the time series sequences. This works also in latent space but requires a transformation of the physical initial state $\bm{x}_{0}\mapsto\bm{x}_{0}^{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 a training sequence length $\tau\in\mathcal{N}$ with $\tau<T$ with $T\in\mathcal{N}$ being the sequence length of the training data and generate the samples in the following from it:

$$ $\tilde{\bm{x}}_{0},\bm{p},\tilde{\bm{x}}_{1:\tau},\tilde{\bm{u}}_{0:\tau},\tilde{\bm{y}}_{0:\tau}=\bm{x}_{j},\bm{p},\bm{x}_{j+1:j+\tau},\bm{u}_{j:j+\tau},\bm{y}_{j:j+\tau}\quad\quad\forall j\in[0,1,\dots T-\tau]$ (41) $$

Note that this increases the number of training samples by large magnitude, i.e. with $T=500$ and $\tau=50$, we get $T-\tau=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 to keep memory requirements low (at the expense of computational cost).
For the time-dependent inputs, simple random sampling does not work. With necessarily small discretization interval $\Delta t$, random sampling would lead to pure noise around a mean-value $\bar{\bm{u}}_{0:T}$, such that all training samples would cycle around the same position in state space. Therefore, in the following, we introduce a method called randomly clipped random sampling of controls with offset and cubic splines (RROCS). Pseudo code of the algorithm is depicted in [1](#alg1). The main idea is to ensure a smooth control input within the specified bounds through the use of cubic splines, while randomly assigning an offset and an amplitude such that the control space is filled meaningfully.

Figure: Algorithm 1 randomly clipped random sampling of controls with offset and cubic splines (RROCS) algorithm

An example for control inputs generated by this algorithm is given in [Figure A.5](#A1.F5)

Figure: Figure A.5: Six example outputs of the RROCS-algorithm for the a control input temperature_K_a with $\bm{u}_{-}=$273.15$,\bm{u}_{+}=$473.15$$.

If parameters or initial states are not sampled, they are set to a default value $\square_{\diamond}$.

#### A.5.2 Stratified Heat Flow Model (SHF) Model

Figure: ((a)) Base R2C1 thermal network with parameters R and C for the total resistance and capacitance. The stratified heat flow model is derived by repeated connection of this base network, where each segment has the parameters $R_{i}=R/\texttt{nSeg}$ and $C_{i}=C/\texttt{nSeg}$ with nSeg being the number of discretization layers.

The stratified heat flow models presented in [Figure A.6](#A1.F6) are introduced because they are a simplified version of discretized (energy) transport phenomena that often arise in the system simulation context, e.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:

- •
number of samples: 1024
- •
training fraction: \qty76
- •
test fraction: \qty12
- •
validation fraction: \qty12
- •
Control Input Sampling Ranges:
–
temperature_K_a: [273.15, 473.15] \unit
–
temperature_K_b: [273.15, 473.15] \unit
- •
Solver Settings:
–
start time: \qty0
–
stop time: \qty1.2
–
output interval: \qty0.002
–
tolerance: $1\text{\times}{10}^{-6}$
–
sequence length: 501
–
the first \qty0.2 are not included in the datasets.
- •
16 states
- •
48 outputs

Training:

- •
Hyperparameters of all networks (state encoder, control encoder, parameter encoder, Latent ODE function, decoder):
–
4 layers
–
hidden dimension of 128
–
Activation: ELU
–
latent space dimension for latent parameter, state and control space: 16
- •
Training:
–
$\beta$: different values
–
Adam Optimizer
–
learning rate: $1\text{\times}{10}^{-3}$
–
weight decay: $1\text{\times}{10}^{-5}$
–
main training
*
Batches per Epoch: 12
*
Phase 1:
·
Solver: rk4
·
training sequence length: 10
*
Phase 2:
·
Solver: rk4
·
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 n batches: 100
·
abort training sequence length increase after n stable epochs: 10

#### A.5.3 Power plant model

The model is a thermal power plant model of the ClaRa-Library (Vojacek et al., ). Specifically, the SteamCycle_01 model from the examples package of that library (^5^55[https://github.com/xrg-simulation/ClaRa-official](https://github.com/xrg-simulation/ClaRa-official) ()). We refer to this model in the following as the steam cycle model.

Figure: Figure A.7: Diagram view of the steam cycle model. (Vojacek et al., )
Refer to caption: /html/2410.10174/assets/figures/BNODE/SteamCycle_01.png

It is a model of a \qty580\mega hard coal power plant with a single reheater, one high-pressure preheater and four low-pressure preheaters. To find initialization values for the states, a static-cycle for the same power plant is included and provides these values for the start of the simulation. The model has several internal PI-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 we considered for this application.
The system has 133 states, of that some were inactive, such that surrogate models are built for 102 states. It furthermore contains 34 nonlinear systems.
The example sets two \qty10 load reductions as input signals over a total simulation time of \qty5000. 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: \qty76
- •
test fraction: \qty12
- •
validation fraction: \qty12
- •
Control Input Sampling Ranges:
–
u: [0.6, 1.1]
- •
Solver Settings:
–
start time: \qty0
–
stop time: \qty6000
–
the first \qty1000 were omitted to discard transient behavior from initialization
–
output interval: \qty5
–
tolerance: $1\text{\times}{10}^{-6}$
–
sequence length: 2500
- •
102 states
- •
7 outputs

Training:

- •
Hyperparameters of all networks (state encoder, control encoder, parameter encoder, Latent ODE function, decoder OR Neural ODE function, output network):
–
4 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 $<$ 0 (B-NODE Koopman)
- •
Training:
–
$\beta$: different values
–
Adam Optimizer
–
learning rate: $1\text{\times}{10}^{-4}$ (B-NODE, B-NODE Koopman), $1\text{\times}{10}^{-5}$ (NODE)
–
weight decay: $1\text{\times}{10}^{-5}$, $1\text{\times}{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 n 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 model(^6^66ClaRa official: [https://github.com/xrg-simulation/ClaRa-official](https://github.com/xrg-simulation/ClaRa-official)) (example 1), the model was exported as a co-simulation FMU with CVODE solver with Dymola (^7^77Dymola: [https://www.3ds.com/products/catia/dymola](https://www.3ds.com/products/catia/dymola)) and simulated with FMPy(^8^88FMPy: [https://github.com/CATIA-Systems/FMPy](https://github.com/CATIA-Systems/FMPy))

- •
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:
–
computation time: 720 seconds
–
hardware specifications:
*
16-core processor with up to 5.7 GHz
*
128 GB RAM
–
technical setup:
*
parallel processing using Dask(^9^99Dask: [https://www.dask.org/](https://www.dask.org/))
*
simulations executed via FMPy with per-time-step communication
- •
simulation with surrogate models on GPU:
–
balanced neural ODE (B-NODE): 7 seconds
–
state space neural ODE (SS-NODE): 15.5 seconds
–
hardware specifications:
*
Nvidia RTX4080 with 16 GB of GPU memory
–
technical setup:
*
all samples are solved concurrently in a single batch, provided that sufficient GPU memory is available.

#### A.5.5 Koopman analytic example

We implemented ([21](#S5.E21)) as Modelica model.

Dataset Generation:

- •
number of samples: 1024
- •
training fraction: \qty76
- •
test fraction: \qty12
- •
validation fraction: \qty12
- •
initial state sampling ranges:
–
x1: [-50, 50]
–
x2: [-50, 50]
- •
solver settings:
–
start time: \qty0
–
stop time: \qty10
–
output interval: \qty0.1
–
tolerance: $1\text{\times}{10}^{-5}$
–
sequence length: 99
- •
2 states

Training:

- •
latent state dimension: 64
- •
Training:
–
$\beta$: different values
–
Adam Optimizer
–
learning rate: $1\text{\times}{10}^{-3}$
–
weight decay: $1\text{\times}{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
·
increase training sequence length over n batches: 1200
·
abort training sequence length increase after n stable epochs: 10
*
Phase 3:
·
Solver: rk4
·
training sequence length: 80
·
increase training sequence length over n batches: 1200
·
abort training sequence length increase after n stable epochs: 10

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 2
questions: 1. Are there other related works that combine model order reduction and latent space evolution, other than latent ODE?
2. How does the balanced neuralODE compare to Dynamic Mode Decomposition method (for linear operator learning)?
soundness: 4
strengths: 1. Method is fully described, intuitive to understand, many derivations are fully given
2. Balanced neural ODE outperforms latent neural ODE and is able to learn latent dynamics.
2. Numerical experiments demonstrate the breadth/diversity of application of the method.
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.
weaknesses: 1. Method combines two existing approaches (however, the combination is novel)
2. Presentation feels brisk at times (might be alleviated with an extra page)
3. The model involves many components with architecture choices, and there is not enough ablation on

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: **Technical questions**:
1. 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. 
2. 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?

**Suggestions for presentation**:
1. 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.
2. 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.
soundness: 3
strengths: - The paper shows promising results with using B-NODE on certain tasks.
- The paper provides certain handy illustrations and detailed figures explaining their approach.
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). 

The 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.

As 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.
weaknesses: **Technical Weaknesses**:
1. *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]).

**Presentation Weaknesses**: 

1. *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.

2. *Vague exposition.* Below I list certain examples where the exposition is vague.

Example 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?

Example 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. 

Example 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).

Example 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.

Example 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.

Overall, 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.


```
[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.
[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).
```

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - Please provide more comprehansive understanding of β-value?  A discussion of the theoretical implications of different β values
- It seems that β-value is small does it mean that information bottleneck become negligible?
- if you set β=0 what will happen? please discuss the theoretical implications of setting β=0

If the authors can clarify my concern, I could change my opinion.
soundness: 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.
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.
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.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
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?

- Could author further explain how they arrive equation (4ba)? You assume that the $x_0^z$ only related to initial state $x_0$?

-  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.
soundness: 3
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. 

- The inclusion of real-world test cases adds practical value to the study and demonstrates its potential applicability.
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.
weaknesses: ## Presentation 
- 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.
- 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).
- Please ensure there is proper spacing between different paragraphs and sections. 

## Experiments 
- 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.

## Minor issues 
some examples:
- Lines 100 the integral should be $\int_{t_0}^{t}f_{\phi NODE}(\mathbf{x}(\tau))d\tau$
- Line 225, is -> are
- Line 312, Observed -> observe 

I'd encourage authors redo proofreading.

Meta-review: The authors introduce a new method called Balanced Neural ODEs (B-NODE) which combine $\beta$-VAE and neural ODEs for modelling dynamical systems. Two primary applications are considered: state reduction and Koopman operator approximation. Reviewers appreciated the impressive benchmarks, showing B-NODE outperforming latent neural ODEs. While the presentation of the work was found to be quite dense, leading to some vague exposition, reviewers and authors were in agreement that the increased page limit and provided feedback should improve this aspect. This is good work, and should be in a state suitable for publication, provided that reviewer feedback is incorporated to improve the presentation.

Decision: Accept (Poster)