Improving Generalization with Physical Equations

Antoine Wehenkel∗
University of Liège
Hsiang Hsu
Harvard

Guillermo SapiroApple

Jens Behrmann
Apple
Gilles Louppe
University of LiègeJörn-Henrik Jacobsen
Apple

Abstract
Hybrid modelling reduces the misspecification of expert physical models with amachine learning (ML) component learned from data. Similarly to many ML algorithms, hybrid model performance guarantees are limited to the training distribution.
To address this limitation, here we introduce a hybrid data augmentation strategy,
termed expert augmentation. Based on a probabilistic formalization of hybrid
modelling, we demonstrate that expert augmentation improves generalization. We
validate the practical benefits of expert augmentation on a set of simulated andreal-world systems described by classical mechanics.

1

Introduction

Recently, hybrid learning demonstrated success in complementing partial physical models with amachine learning component, e.g., [1, 2, 3, 4, 5, 6]. These works have shown that hybrid models
are more faithful than their stand-alone physical counterparts. Moreover, training hybrid models is
more sample-efficient and generalizes better than purely data-driven models. Indeed, the underlying
physical models are often valid beyond the training data, and therefore, the corresponding hybrid
models shall generalize to some unseen scenarios.
In this work, we first observe that current hybrid learning algorithms are sub-optimal in the amortized
inference setting – when we aim to build hybrid models that are valid for various test configurations.
Contrary to the common belief that hybrid learning achieves better generalization than black box ML
models, we argue that hybrid learning algorithms do not yet meet their promise regarding robustness
in amortized settings. Motivated by this observation, we introduce expert augmentations to extend
the range of validity of hybrid models. Our experiments demonstrate that the proposed augmented
hybrid models (AHMs) achieve generalization superior to standard hybrid learning algorithms.

2

Hybrid learning

We formalize hybrid learning with the probabilistic model depicted in Figure 1a. The expert model,
often derived from first-principles physics, is a conditional density p(ye |x, ze ) that describes the
distribution of the expert response Ye to an input x together with the physical parameters ze . Theinteraction model is a conditional distribution p(y|x, ye , za ) that aims to correct the expert model.
Our goal is to create a robust predictive model p(y|x, (xo , yo )) of the random variable Y , given the
input x and independent observations (xo , yo ) of the same system. Given an accurate estimation ofza and ze , we can predict the distribution of Y for any input x to the same system. In particular, the∗Work done as an intern at Apple.
Corresponding author: awehenkel@apple.com

Machine Learning and the Physical Sciences workshop, NeurIPS 2022.

p(ze |x, y) ≈ qψ (ze |x, y)|{z}Expert model
+

Ze

Ye

X

Za

+

Ze

Z̃e

Ω

Ω̃

Ze

∗

Za

Ω

Za

z
}|
{
p(x, y|ze ) = Ep(za )p(ye |x,ze ,) [p(x)p(y|x, ye , za )]Y≈ Ep(za )p(ye |x,ze ) [p(x)pθ (y|x, ye , za )](a) A hybrid probabilistic model relates (b) Visualization of the distribution shifts considered in this work. The train support (in blue) Ωthe output Y to the input X provided
the latent variables Ze and Za . The expert model defines the conditional density p(ye |x, ze ), where Ye is an approximation of Y . Hybrid learning aims
at learning the conditional distributionp(y|x, ye , za ).

of (x, y) results from (za , ze ) ∈ Za × Ze . The test supports (in red) are denoted with a tilde
symbols as Z̃e for ze and Ω̃ for (x, y). The augmented support (in green) includes both train and+test scenarios and corresponds to (za , ze ) ∈ Za × Ze . The outer violet domain depicts a situationin which the domain of za is also shifted. Hybrid modelling algorithms alone may learn a mapping
+

+

+

+

pθ : Ze → Ω but augmentation is necessary to learn the inverse mapping qψ : Ω → Ze .

Figure 1
Bayes optimal hybrid predictor pB is
pB (y|x, (xo , yo )) = Ep(za ,ze |(xo ,yo )) [p(y|x, za , ze )] .

(1)

In the amortized setting, we aim to learn a model of both the predictive model p(y|x, za , ze ) andof the posterior over the parameters p(za , ze |(xo , yo )). We will see that existing hybrid learningalgorithms neglect the importance of building a robust encoder p(za , ze |(xo , yo )) to make predictions
in out-of-distribution (OOD) settings.
2.1

Hybrid generative modelling

We consider deterministic expert models fe : X × Ze → Ye , for which pθ (ye |x, ze ) is a Dirac
distribution. Given a dataset D = {(x(i) , y (i) )}N
i=1 of N IID samples, we aim to learn the interaction
model pθ (y|x, ye , za ) that fits the data well and minimally corrects the expert model. Two recent
approaches, the APHYNITY algorithm [1] and the Hybrid-VAE [2, HVAE], compete in this setting.
APHYNITY. Yin et al. [1] augment an expert ordinary differential equation (ODE) with an additive
correction. They assume the hybrid models can fit the data perfectly. This assumption is equivalent totaking a Dirac distribution for pθ (y|x, ye , za ). They propose solving the following problem
dytmin ||Fa || s.t. ∀(x, y) ∈ D, ∀t,= Fe (yt ) + Fa (yt ) with y0 := x,(2)ze ,Fadt
where || · || is a norm on the function space, Fa : Yt × Za → Yt is a learned function, Fe : Yt × Ze →Yt is the expert model, and D is a training dataset, which contains initial states x := y0 and k-longsequences y ∈ Y := (Yt )k . We amortise APHYNITY as suggested in the supplementary material of[1]. The encoder network gψ (·, ·) : X × Y → Za × Ze corresponds to a Dirac distribution located atgψ . The interaction model is the solution of the augmented ODE in (2).
Hybrid-VAE (HVAE). In contrast to APHYNITY, the HVAE is not limited to additive interactions,nor to ODEs. The HVAE is a variational auto-encoder (VAE) in which the decoder specificallyfollows Figure 1a. The encoder gψ (x, y) predicts a posterior distribution over za and ze . The modelis trained with the classical Evidence Lower Bound on the likelihood (ELBO) and three regularizers
RP P C , RDA,1 , and RDA,2 that encourage the hybrid model to stay close to the expert model. We
refer the reader to Takeishi and Kalousis [2] for more details. The predictor takes the form describedby (1) where p(za , ze |(xo , yo )) is approximated by the encoder qψ (za , ze |x, y) and p(y|x, za , ze ) by
the learned hybrid generative model.

3

Robust hybrid learning

We consider the important sub-class of distribution shifts for which the marginal train distributionp(ze ) and test distribution p̃(ze ) may be different but the marginals of za and x are constant. The
2

encoder of classical hybrid models fails in this context; it is unable to predict accurately ze . In
consequence, the hybrid model derived from (1) does not generalize. We address this failure byintroducing expert augmentation to fine tune the encoder. Our goal is to learn a predictive model thatis exact on both the train and shifted test domains. A predictive model p̂(a|b) is E-exact, or exact on
the sample space E, if p̂(a|b) = p(a|b)∀(a, b) ∈ E for a given ground-truth p(a|b). A model is robust
to a test scenario if its exactness on the training domain implies exactness on the test domain.
++

We now define an augmented distribution p(ze ) over the expert variables whose support Ze includes
the joint support Ze ∪ Z̃e between the train and test distribution of the physical parameters. As+

depicted in Figure 1b, we denote the corresponding support over the observation space X × Y as Ω, Ω,and Ω̃, respectively. In this context, and with A1, we demonstrate that even under perfect learning,classical hybrid learning algorithms do not produce an Ω̃-exact predictor while our augmentation
strategy does.
+

Assumption 1 (A1): Hybrid modelling learns an interaction model pθ (y|ye , x, za ) that is Ω-exact.
The hypothesis A1 is consistent with the recent literature on hybrid modelling [1, 2], which assumes
that pθ (y|ye , x, za ) should not be overly complex. For example, extrapolation to unseen ye should+

hold for additive defects. That said, the exactness of the interaction model pθ on Ω is insufficient+

to prove that the predictive model pθ,ψ is Ω-exact. Indeed, the encoder qψ is only trained on the
training data and cannot rely on a strong inductive bias in contrast to pθ . Thus, even if the encoder+

is exact on the training distribution, the corresponding predictive model is not Ω-exact. While thedecoder’s performance are not limited to the training scenarios thanks to the broader validity of theexpert model, the encoder does not generalize to unseen settings as it is purely data-driven.
3.1

Expert augmentation

We increase the encoder’s generalization domain by training it on additional synthetic configurations
generated by the hybrid decoder. Once trained, the hybrid model is composed of an encoder qψ and+an interaction model pθ that are respectively Ω- and Ω-exact. We create an augmented training set++

with support over Ω by sampling physical parameters ze ∈ Ze . Then, we finetune the encoder qψ on++Ω. Under perfect training, the predictive model becomes robust; pθ,ψ (y|x, (xo , yo )) is Ω-exact.
+

After training a hybrid model, we have to decide on a realistic distribution p(ze ) in order to create++

the dataset D by sampling from the hybrid model. The augmented training set D contains groundtruth values for the expert’s variables ze , in contrast to the original dataset. We freeze the interaction++

model, as we assume it is Ω-exact. We only fine-tune the encoder qψ on D with a combination of
the loss function ℓ of the original hybrid learning algorithm and supervision on the latent variableobjective. We then learn a decoder that solvesψ ⋆ = arg min E + [ℓ(x, y; θ, ψ) − log qψ (ze |x, y)] .
ψ

4DExperimental results

We assess the benefits of expert augmentation on two synthetic problems described by an ODE
dyt= Fe (yt ; ze ) + Fa (yt ; za ).
dt

(3)
t /∆t

1
The initial state y0 corresponds to X, and the sequence of states y1:t1 := [yi∆t ]i=1
to Y .

T
The damped pendulum: the system’s state at time t is yt = θt θ̇t , where θt is the angle ofthe pendulum at time t and θ̇t its angular speed. The evolution of the state over time is described by

T

T
(3), where ze := ω, za = α, Fe := θ̇ −ω02 sin θ , and Fa := 0 −αθ̇ . The train domain is
generated from Ze := [1.5, 3.1] and Za := [0, 0.6] and at testing Z̃e := [0.5, 1.5]. The 2D reaction3diffusion: this is a 2D FitzHugh-Nagumo model on a 32 × 32 grid. The system’s state at timeTt is a 2 × 32 × 32 tensor yt = [ut vt ] . The evolution of the state over time is described byT
T
Fe := [a∆ut b∆vt ] and Fa := [Ru (ut , vt ; k) Rv (ut , vt )] , where ze := {a, b}, za = {k}, ∆is the Laplace operator, and the reaction terms are Ru (u, v; k) = u−u3 −k −v and Rv (u, v) = u−v.
The train domain is generated from Ze := [1e−3, 2e−3] × [3e−3, 7e−3] and Za := [3e−3, 5e−3]
and at testing Z̃e := [2e−3, 4e−3] × [1e−3, 1e−1]. The double pendulum: we use the real-world
dataset of a double pendulum introduced by Asseman et al. [7] and consider a friction-less system toderive an expert model; see, e.g., [8] for the details. This idealized model misses non-negligible effects
such as errors in extracting the pendulum’s positions, frictions, or vibrations. This experiment allows
us to validate the effectiveness of expert augmentation when the exact nature of the mispecificationis unknown. We use low-energy configurations for the training set and high-energy for the test to
enforce a distribution shift. The validation set contains samples whose energy is between the train
and test examples. In our experiment, the initial positions of the pendulum are given. The encoder
network must predict the initial angular speed given a partial observation of the pendulum position
for a few consecutive measurements.
DatasetPendulum
Val.
Test

RLC
Val. Test

APH.
APH.+

6±2
6±2

2±0
2±0

66±9
10±4

27±2
3±1

HVAE. 3±1 117±10 2±0 32±10
HVAE+ 2±1 11±2 2±0 2±0

Figure 2: The average log-MSEs over ten runs. We compare HVAE (in red) and Table 1: The relative errors on ze (in %) over ten
APHYNITY (in green), in light colours, to their expert augmented versions HVAE+ runs. The + denotes augmented versions. The accuracy
and APHYNITY+, in darker colours, and to a simple baseline (in grey). On the test of APHYNITY and HVAE collapses on the OOD test set
sets, the proposed AHMs outperform the original models. Augmentation conserves while the augmented versions’ performance is stable.
good performance on the validation set.
MetricDatasetTrain

Log-MSE
Valid

Test

Train

Rel. Err.
Valid

Test

Exp.
−4.4±0−3.9±.4
2.8±.4324±3186±1 222±3
APH. −6.7±1.2 −5.2±.8 −3.4±.4 211±90 106±41 147±61
APH.+ −6.1±.6 −5.3±.5 −4.4±.3 157±27 71±10 72±10

Table 2: The average log-MSEs and average relative errors on the initial angular speeds of the double pendulum over three runs for three
predictive models: the expert model only (Exp.), APHYNITY (APH.) and APHYNITY followed by expert augmentation (APH.+). Expert
augmentation outperforms other models except for predicting the state evolution on the training set.

Our synthetic experiments show the effect of expert augmentation on APHYNITY and HVAE. All
models explicitly assume that the interaction model follows the structure of (3) where Fa is a small
neural network. We select the best models from validation performance.
Figure 2 and Table 1 demonstrate that HVAE and APHYNITY are not robust to OOD test scenarios incomparison to the corresponding proposed AHMs. We only compare performance in OOD settings
in Figure 2 as hybrid models have already demonstrated greater performance than non-hybrid models
in the literature [1, 2]. Both algorithms strongly benefit from expert augmentation. Table 1 showsthat the encoder does not predict the physical parameters perfectly. This failure indicates that theencoder is not Ω-exact, and neither should the decoder. However, this departure from an ideal setting
does not preclude the effectiveness of the proposed expert augmentation. Expert augmentation also
improves the parameter estimation in the OOD setting. The augmented encoder accurately estimates
the physical parameters in the IID and OOD settings. Finally, Table 2 show that the proposed expert
augmentations is also effective in practical settings. Overall, AHMs outperform classical hybrid
learning both in the predictive accuracy and in inferring the expert variables.

5

Conclusions

We have presented a simple augmentation strategy, termed expert augmentation, to improve thegeneralization capabilities of hybrid learning models to scenarios where the expert model is valid.

4

Broader Impact Statement
This work simplifies the combinations of domain knowledge within machine learning models. As
with most machine learning algorithms, the real-world impact may be positive or negative. However,improving robustness and interpretability of machine learning models may eventually help prevent
harms caused by the failure and unpredictability of purely data-driven models under real-world
distribution shifts.

References[1] Yuan Yin, Vincent Le Guen, Jérémie Dona, Emmanuel de Bézenac, Ibrahim Ayed, Nicolas
Thome, and Patrick Gallinari. Augmenting physical models with deep networks for complex
dynamics forecasting. Journal of Statistical Mechanics: Theory and Experiment, 2021(12):124012, 2021.
[2] Naoya Takeishi and Alexandros Kalousis. Physics-integrated variational autoencoders for robust
and interpretable generative modeling. Advances in Neural Information Processing Systems, 34,
2021.
[3] Zhaozhi Qian, William R Zame, Mihaela van der Schaar, Lucas M Fleuren, and Paul Elbers.
Integrating expert odes into neural odes: Pharmacology and disease progression. arXiv preprint
arXiv:2106.02875, 2021.
[4] Viraj Mehta, Ian Char, Willie Neiswanger, Youngseog Chung, Andrew Oakleigh Nelson, Mark D
Boyer, Egemen Kolemen, and Jeff Schneider. Neural dynamical systems: Balancing structureand flexibility in physical prediction. arXiv preprint arXiv:2006.12682, 2020.
[5] Chon Lok Lei and Gary R Mirams. Neural network differential equations for ion channel
modelling. Frontiers in Physiology, page 1166, 2021.
[6] Markus Reichstein, Gustau Camps-Valls, Bjorn Stevens, Martin Jung, Joachim Denzler, Nuno
Carvalhais, et al. Deep learning and process understanding for data-driven earth system science.
Nature, 566(7743):195–204, 2019.
[7] Alexis Asseman, Tomasz Kornuta, and Ahmet Ozcan. Learning beyond simulated physics.
In Modeling and Decision-making in the Spatiotemporal Domain Workshop, 2018. URL
https://openreview.net/pdf?id=HylajWsRF7.
[8] Tomasz Stachowiak and Toshio Okada. A numerical analysis of chaos in the double pendulum.
Chaos, Solitons & Fractals, 29(2):417–422, 2006.

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [No] Due to the 4-page limit. However,
we will do it in the corresponding long version.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See thebroader impact statement.
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [No] We will do it forthe corresponding long version of this paper.
5(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [No] Due to space limit and strong discouragement of appendices,it was difficult to provide the necessary details to reproduce our experiment exactly.
However, the corresponding long version will describe the experiments thorougly inthe appendices.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [No] Final results were generated on a
32GB MacBook Pro from 2020.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes](b) Did you mention the license of the assets? [No](c) Did you include any new assets either in the supplemental material or as a URL? [N/A](d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [No] The dataset is publicly available and we cite properly the authors.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]

6