Vincent Souveton

Arnaud Guillin

Jens Jasche

Guilhem Lavaux

Manon Michel

Fixed-kinetic Neural Hamiltonian Flows for enhanced interpretability and reduced complexity

Normalizing Flows (NF) are Generative models which are particularly robust and allow for exact sampling of the learned distribution. They however require the design of an invertible mapping, whose Jacobian determinant has to be computable. Recently introduced, Neural Hamiltonian Flows (NHF) are based on Hamiltonian dynamics-based Flows, which are continuous, volume-preserving and invertible and thus make for natural candidates for robust NF architectures. In particular, their similarity to classical Mechanics could lead to easier interpretability of the learned mapping. However, despite being Physics-inspired architectures, the originally introduced NHF architecture still poses a challenge to interpretability. For this reason, in this work, we introduce a fixed kinetic energy version of the NHF model. Inspired by physics, our approach improves interpretability and requires less parameters than previously proposed architectures. We then study the robustness of the NHF architectures to the choice of hyperparameters. We analyze the impact of the number of leapfrog steps, the integration time and the number of neurons per hidden layer, as well as the choice of prior distribution, on sampling a multimodal 2D mixture. The NHF architecture is robust to these choices, especially the fixedkinetic energy model. Finally, we adapt NHF to the context of Bayesian inference and illustrate our method on sampling the posterior distribution of two cosmological parameters knowing type Ia supernovae observations.

Introduction

Generative models are widely used for sampling highdimensional probability distributions with applications from molecular biology [START_REF] Lopez | Enhancing scientific discoveries in molecular biology with deep generative models[END_REF] to cosmology [START_REF] Rodriguez | Fast cosmic web simulations with generative adversarial networks[END_REF] or medical science [START_REF] Frazer | Disease variant prediction with deep generative models of evolutionary data[END_REF]. Traditional architectures like Generative Adversarial Networks (Goodfellow et al., 2014) have shown impressive results in image generation but since their adversarial loss seeks a saddle point rather than a local minimum, GANs are notoriously hard to train and may suffer from mode-collapse (Lin et al., 2018;[START_REF] Arjovsky | Towards principled methods for training generative adversarial networks[END_REF][START_REF] Berard | A closer look at the optimization landscapes of generative adversarial networks[END_REF]. More robust techniques like Normalizing Flows (NF) were developed [START_REF] Tabak | Density estimation by dual ascent of the log-likelihood[END_REF][START_REF] Dinh | Nice: Non-linear independent components estimation[END_REF][START_REF] Rezende | Variational inference with normalizing flows[END_REF]. NF consist in training a neural network to map a simple prior distribution onto the desired target through a chain of invertible transformations. They come with interesting characteristics, such as stability and correctness, see for example [START_REF] Papamakarios | Normalizing flows for probabilistic modeling and inference[END_REF]. The main limitation comes from the design of an invertible function for the mapping. In particular, computing the Jacobian determinant in the change of variable formula may be costly. Furthermore, while explainability is now under a growing concern within the community [START_REF] Gilpin | Explaining explanations: An overview of interpretability of machine learning[END_REF], and in particular regarding applications in natural science, the transformation learned by NF models is commonly hard to interpret.

Exploiting the Newtonian evolution in classical Mechanics, Neural Hamiltonian Flows (NHF) [START_REF] Toth | Hamiltonian generative networks[END_REF] are NF models that use Hamiltonian transformations. The Hamiltonian of a system is composed of a kinetic and potential energy terms, which sets its dynamical evolution, which is reversible and has a Jacobian determinant equal to one. They come with performance similar to the ones obtained with Real-NVPs in sampling 2D distributions [START_REF] Toth | Hamiltonian generative networks[END_REF]. Being Physics-driven models, they are expected to enhance interpretability. However, the NHF architecture is made of four neural networks black-boxes that render difficult the interpretation of the learned dynamics. Furthermore, even if they have been numerically shown to transfer multimodality from the target distribution to the potential energy in some cases [START_REF] Toth | Hamiltonian generative networks[END_REF], this property is not guaranteed. However, the Hamiltonian formalism could offer a powerful formalism to design more easily interpreted architectures. Thus, some recent works have proposed to exploit the Hamiltonian properties to design interpretable flows that are invariant under symmetrical transformations (Jimenez [START_REF] Jimenez Rezende | Equivariant Hamiltonian Flows[END_REF].

In this work, we focus on the transfer of the negative logarithm of the target distribution into the learned potential, which should be the case in any physical systems. This leads us to propose a fixed-kinetic version of NHF where the kinetic-energy term in the Hamiltonian is set to its true form in classical mechanics, where momenta follow a Gaussian distribution. We also discuss the impact of the hyperparameters on such transfer and study the overall robustness of NHF towards the choice of the numerical integration scheme and of the prior distributions. In more details, this work provides the following four main contributions:

• We introduce a fixed-kinetic version of NHF that, thanks to the Hamiltonian evolution, enhances interpretability of the model and does so at a cheaper computational price.

• We analyze the effect of multiple parameters of the architecture on sampling a 2-dimensional multimodal distribution and show that NHF is robust to the choice of hyper-parameters, especially the fixed-kinetic model.

• We show that the choice of prior has an influence on the learned dynamics. The fixed-kinetic model also allows better robustness to the choice of prior distribution.

• Finally, aside from Generative modeling, flow-based models have been shown to be suited as inference algorithms [START_REF] Rezende | Variational inference with normalizing flows[END_REF][START_REF] Winkler | Learning likelihoods with conditional normalizing flows[END_REF]. We test a framework for Bayesian inference using NHF and present numerical experiments for inferring cosmological parameters from astronomical observations. The methodology we propose is inspired by Boltzmann generators [START_REF] Noé | Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning[END_REF].

The manuscript is organized as follows: in Section 2, we present and review the related works. In Section 3, we describe the theoretical framework and the practical implementation of NHF. In Section 4, we discuss the choice of models for the kinetic energy and introduce a fixed-kinetic version for enhanced interpretability and reduced complexity. Section 5 discusses the choices one can make to maximize expressivity given a fixed computational budget, with tests on a 2D Gaussian mixture and analysis of the impact of the Leapfrog-hyperparameters and model complexity. We also show how the choice of base distribution affects the learned energies and thus the interpretability of the model. Finally, in Section 6, we adapt NHF for Bayesian inference and illustrate our method on a toy model from cosmology.

Related works

Generative models. These methods make the assumption that data can be represented by an underlying probability distribution. The aim of generative models is to learn such distribution in order to produce original samples similar to the training dataset. Various architectures have been presented such as Generative Adversarial Networks (Goodfellow et al., 2014) or diffusion networks [START_REF] Sohl-Dickstein | Deep unsupervised learning using nonequilibrium thermodynamics[END_REF]. In this paper, we will focus on Normalizing Flows techniques [START_REF] Tabak | Density estimation by dual ascent of the log-likelihood[END_REF][START_REF] Dinh | Nice: Non-linear independent components estimation[END_REF][START_REF] Rezende | Variational inference with normalizing flows[END_REF] as a way of smoothly transforming a simple prior distribution into the target posterior.

Learning Hamiltonians. Learning Hamiltonian potentials, i.e. physical conserved quantities, is a first step towards a better understanding of the physical processes that have governed the data generation. Multiple architectures have been proposed, such as Hamiltonian Neural Networks [START_REF] Greydanus | Hamiltonian Neural Networks[END_REF] or Hamiltonian Generative Networks [START_REF] Toth | Hamiltonian generative networks[END_REF]. These methods parameterize the Hamiltonian of the system with neural networks and come with useful properties such as exact reversibility and smoothness. They have inspired applications from domain translation [START_REF] Menier | Continuous Methods : Hamiltonian Domain Translation[END_REF] to fault-detection in industry [START_REF] Shen | Machine fault classification using hamiltonian neural networks[END_REF]. It is worth mentioning that they can be combined with Markov-Chain Monte Carlo (MCMC) methods, for instance as proposals in the Hamiltonian Monte Carlo (HMC) algorithm [START_REF] Duane | Hybrid monte carlo[END_REF][START_REF] Dhulipala | Bayesian inference with latent hamiltonian neural networks[END_REF]. In this work, we learn artificial Hamiltonians for sampling and our goal is to extract the negative logarithm of the target distribution into the potential.

Neural Hamiltonian Flows. NHF is a Normalizing Flow architecture derived as a modification from Hamiltonian Generative Networks. As generative models, they have been used to sample from 2D distributions [START_REF] Toth | Hamiltonian generative networks[END_REF]. The Hamiltonian formalism is also a natural framework for exploiting symmetries inside data and build interpretable flows [START_REF] Jimenez Rezende | Equivariant Hamiltonian Flows[END_REF]. In this paper, we discuss its robustness with respect to hyperparameters and choice of prior distribution. Also, we propose an alternative version of NHF to enhance interpretability while reducing the complexity of the model.

Inference with NHF. Inferring the posterior distribution of parameters conditioned on data is an important problem in science. Traditional MCMC methods [START_REF] Robert | Monte Carlo Statistical Methods[END_REF] are very popular because they come with guarantees in terms of convergence and many progress have been made regarding their tuning [START_REF] Homan | The no-u-turn sampler: Adaptively setting path lengths in hamiltonian monte carlo[END_REF][START_REF] Carpenter | A probabilistic programming language[END_REF]. NF architectures have also been proposed in this framework [START_REF] Rezende | Variational inference with normalizing flows[END_REF][START_REF] Winkler | Learning likelihoods with conditional normalizing flows[END_REF]. Here, we adapt NHF to sampling Bayesian posterior distributions by transforming the prior distribution into the posterior with no access to samples from the target.

Explainable AI. XAI deals with the problem of understanding the decisions made by an Artificial Intelligence (Samek & Müller, 2019). Indeed, complex architectures made of multiple (deep) neural network are often easier to train than to understand. Some solutions involve surrogate techniques [START_REF] Ribeiro | Explaining the predictions of any classifier[END_REF], local perturbations [START_REF] Ancona | Gradient-Based Attribution Methods[END_REF] or meta-explanations [START_REF] Lapuschkin | Unmasking clever hans predictors and assessing what machines really learn[END_REF]. Including physical prior knowledge into neural networks may be another solution to understand the model [START_REF] Raissi | Physicsinformed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations[END_REF][START_REF] Toth | Hamiltonian generative networks[END_REF]. In this work, we build on that idea and try to make the model as explainable as possible by fixing the kinetic energy of the model and thus enforcing a classical Mechanics knowledge into the architecture.

Normalizing Flows with Hamiltonian transformations

Normalizing Flows

Normalizing flows are generative models that map a complex target distribution π onto a known prior distribution π 0 from which it is easy to sample [START_REF] Papamakarios | Normalizing flows for probabilistic modeling and inference[END_REF]. This mapping is a series of smooth invertible transformations. Once the model is trained, one can reverse the learned dynamics to generate samples from the target distribution starting from the prior.

If X = T (Z), where T is a C 1 -diffeomorphism and Z ∼ π 0 , then the density fol- lowed by X reads m(x) = π 0 T -1 (x) × |det J T -1 (x)| , where det J f -1 (x)
is the Jacobian determinant of the inverse transformation of f evaluated in x. For more expressivity, it is possible to chain several invertible transformations x = T L • ... • T 1 (z). In this case, the change of variable formula is now:

m(x) = π 0 T -1 1 • ... • T -1 L (x) × L k=1 det J T -1 k (x)
. The model parameters to optimize are denoted Θ. The goal is to minimize the Kullback-Leibler divergence between the target distribution π and the model distribution m with respect to Θ, i.e. minimizing:

L(Θ) = Eπ [log π(X) -log m(X; Θ)] = -Eπ log π T -1 1 • ... • T -1 T (X; Θ) + T k=1 det J T -1 k (X) + C.
As we seek to find the minimum of this quantity, C can be discarded. Then, one can use samples x 1 , ..., x m from the target distribution for estimating the above loss via Monte Carlo. The function to minimize becomes:

L(Θ) = - 1 m m k=1 log π T -1 1 • ... • T -1 L (x k ; Θ) + T k=1 det J T -1 k (x k ) . (1)
At this point, transformations are to some extent arbitrary. Any smooth invertible function would be suited but the main computational cost comes from the Jacobian determinants.

The first goal is then to reduce this computational cost, and the second one to enhance interpretability, by a proper choice of the transformation, and of its induced inverse.

Neural Hamiltonian Flows

To alleviate these issues, Neural Hamiltonian Flows (NHF, [START_REF] Toth | Hamiltonian generative networks[END_REF] is a NF technique that uses a series of Hamiltonian transformations as normalizing flows. In classical Mechanics, a system is fully described by its coordinates (q, p) in phase-space. From that description, it is possible to define a scalar quantity called a Hamiltonian (Landau & Lifshitz, 1982). It can be seen as the total energy of the system and, in this paper, we make the assumption that it is written as the sum of a potential energy V , solely depending on the generalized positions q, and a kinetic energy K, solely depending on the momenta p. The system evolves in phase-space following Hamilton's equations that read:

dq dt = ∂H ∂p , dp dt = - ∂H ∂q . (2)
Hamiltonian transformations present at least two main advantages that make them suited for normalizing flows:

• they are invertible by construction and inversion is easy by using a classical numerical integrator, i.e. just reversing the speed;

• their Jacobian determinant is equal to 1, removing the necessity to compute such determinant for each transformation.

Numerically, the continuous solution can be approached by a symplectic, invertible and stable integrator as a Leapfrog:

     p n+ 1 2 = p n -∇V (q n ) × dt 2 , q n+1 = q n + ∇K(p n+ 1 2 ) × dt, p n+1 = p n+ 1 2 -∇V (q n+1 ) × dt 2 .
(3)

NHF is trained on a dataset consisting in realizations from the target distribution. In order to simulate a Hamiltonian dynamics, one must extend the position space in which live the samples into the phase space, by adding artificial momenta: this is the role of the Encoder. The dynamics is integrated in phase-space with the Leapfrog integrator [START_REF] Toth | Hamiltonian generative networks[END_REF]. More precisely, during training, NHF takes batches of q T from the training dataset as inputs. For each q T , one p T is drawn from a Gaussian distribution whose mean µ(q T ) and standard deviation σ(q T ) depend on the q T . The resulting point in phase-space is then evolving through a series of L Leapfrog steps with integration timestep dt.

The outputs consist in the final position q 0 and momenta p 0 , as well as the initial mean µ(q T ), standard deviation σ(q T ) and p T that are used in the loss computation. Once trained, one can easily define a sampling function that transforms q 0 , p 0 into q T , by changing the sign of integration timestep and moving the system through the learned dynamics. The architecture is illustrated in Figure 1. Following the previous notations, let us call f (.|q T ) the density of a normal distribution N (µ(q T ), σ(q T ) 2 ), and T -1 the backward transformation of phase-space performed by NHF i.e.

T -1 (q T , p T ) = (q 0 , p 0 ). Also, we denote Π 0 the joint distribution of q 0 , p 0 . By adding artificial momenta p T [START_REF] Toth | Hamiltonian generative networks[END_REF], the distribution modeled by our NHF is

m(q T ) = M (q T , p T )dp T = Π 0 (T -1 (q T , p T ))dp T .
This integral is intractable so instead one can maximize the following ELBO:

L(q T ) = E f log π(T -1 (q T , p T )) -log f (p T |q T ) . ( 4
)
This quantity is approximated via Monte Carlo integration.

Once it has learned the transformation, one can reverse the sign of timesteps and use the same potentials to transform the base distribution into the target distribution.

The first part of the architecture consists in adding artificial momenta to simulate a Hamiltonian dynamics. This is done by the Encoder. Here, µ and σ are approximated by two neural networks. As for the Hamiltonian transformations, they are made by chaining Leapfrog steps. To do so, one must design the potential energy V and the kinetic energy K of the system. In previous work [START_REF] Toth | Hamiltonian generative networks[END_REF], authors have proposed to parameterize each potential by a neural network. We will discuss this choice in the following section.

For now, let us highlight the fact that the choice of neural network for the energies or within the Encoder is without constraint because the properties needed for normalizing flows are ensured by the use of a Leapfrog scheme. This is an illustration of the model flexibility.

Designing the kinetic energy for NHF

MLP-kinetic NHF. If the kinetic energy is chosen to be a MLP [START_REF] Toth | Hamiltonian generative networks[END_REF], then the model contains two black-boxes that are not easy to interpret a priori, namely the kinetic and potential energies K and V . In particular, when sampling a multimodal distribution from a unimodal prior, one may have some troubles understanding where the transfer of multimodality occurs in the process.

Fixed-kinetic NHF. By fixing the kinetic energy inside NHF, we expect to gain interpretability on the learned flow by forcing the latter to obey some Physics principles. Ideally, one would like to find a way of enforcing these energies to be classical from a classical Physics perspective, i.e. making learned kinetic to be a quadratic form and learned potential to be the negative logarithm of the target distribution (or an approximation). This is an important aspect because we would like to keep track of the transformation dynamics. This way, we can play with the familiar classical mechanics framework and interpret the model more easily. By doing so, both interpretability and maniability (less parameters) are gained through the Fixed-kinetic NHF. In this model, K is no longer a MLP but a fixed function. In practice, one can impose the kinetic energy to be written as:

K(p) = 1 2 p T Mp, (5)
with M a mass matrix. We impose M to be a positive symmetric definite matrix, as it is the case in classical mechanics. For instance, in 2D, it comes down to optimizing three real parameters, rather than a complete network. Starting from (q 0 , p 0 ) drawn from a unimodal prior distribution and imposing a quadratic kinetic energy significantly reduces the number of possibilities for the potential energy if one wants to recover a multimodal q T . Let us examine what happens for a single Leapfrog transformation. We have ∇K(p) = Mp. Using the first step of the Leapfrog algorithm (3), the final expression for (q 1 , p 1 ) reads:

q 1 = q 0 + M p 0 -∇V (q 0 ) dt 2 dt p 1 = p 0 -∇V (q 0 ) dt 2 -∇V q 0 + M p 0 -∇V (q 0 ) dt 2 dt dt 2 .
So, if V is quadratic, then the overall transformation is linear and multimodality cannot be obtained from the original Gaussian measure. Thus, multimodality can be achieved only through the potential V in order to sample from the correct target.

It is possible to create many different versions of NHF by fixing its kinetic energy, which thus leads to analysis of the learned dynamics in light of a different Physics framework. One could, for instance, use a relativistic kinetic energy instead of a classical one. We now numerically show how the classical choice for kinetic energy yields an interpretable potential V and such in a robust manner.

Robustness and optimization of NHF

We present the results of numerical experiments for analyzing the robustness of NHF. The experiments consist in sampling a 2D Gaussian mixture with 9 modes (see figure 2). Such toy example, similarly studied in [START_REF] Toth | Hamiltonian generative networks[END_REF], is interesting in the sense that it will allow us to discuss various aspects, from memory usage to interpretability. Also, traditional generative models like GAN may suffer from mode-collapse problems even in simple multimodal 2D settings [START_REF] Eghbal-Zadeh | Mixture density generative adversarial networks[END_REF]. It should be noted that we The energies are MLPs with size (D, H, H, 1), D being the dimension of data, except for the fixed-kinetic version of NHF where K is a fixed quadratic form. Models are trained on a 5,000 points dataset with minibatches of size 512. Weights and biases are optimized with Adam [START_REF] Kingma | A method for stochastic optimization[END_REF], setting the learning rate to 5 × 10 -4 .

Impact of Leapfrog-hyperparameters and model complexity

We first discuss the effect of Leapfrog-hyperparameters L (number of Leapfrog steps) and T = L × dt (integration time) on the optimization, but also the impact of the model complexity. The latter is governed by the total number of neurons in the model, this number being an increasing function of H, the number of neurons per hidden layer in each MLP of the model. If the model is complex enough, we expect it to learn how to adjust to the number of Leapfrog steps and choice of integration time. If not, we expect the model to perform better by increasing the number of computations, i.e. increasing L. As L increases, we do not expect to see any difference as the number of computations becomes sufficient to solve the problem.

We tested both fixed-kinetic and MLP-kinetic NHF with various choices of L, T and H. We use a soft-uniform prior ∝ s(x + 3)s(-x + 3), where s is the sigmoid function. The corresponding loss decays are illustrated in Figure 3 and additional details can be found in Appendix A.

Firstly, the fixed-kinetic NHF model is shown to be more robust than the MLP-kinetic one to the choices of L and T , at fixed H, as discrepancy in the loss decay more clearly appears only with H = 8. Then, regarding the tuning of the Leapfrog scheme, at fixed-integration time, models with L = 1 always reach higher final value of the loss, this effect being less visible with the fixed-kinetic model. Increasing the number of leapfrog steps leads to better final performance of the model even if this effect disappears once the number of Leapfrog steps gets sufficient and no further expressivity can be achieved. Finally, as for the effect of integration time T , it barely appears for the fixedkinetic model, showing that the latter efficiently adjusts to this parameter. As for the MLP-kinetic model, the effect of the integration time is clearer, but mostly at H = 8, where performance improves for T = 1, 10 compared to T = 0.1. Overall, as the number of parameters in the model is increasing (H = 32 or 128), the impact of the choice of the integration time is very limited. when using the fixed-kinetic variant.

Impact of the prior distribution on the learned dynamics

Now we illustrate the impact of the prior choice on the transfer of characteristics of the target distribution on the potential V , especially regarding the multimodality nature. All models were trained for 15,000 epochs using H = 128, T = 1 and L = 10, with a 5,000 points training dataset.

In terms of sampling, all considered schemes recover the nine correct modes from the target distribution, as illustrated in Figure4. We now consider the learned potential V , as illustrated in Figure5. As the Hamiltonian evolution only involves its derivative, we also include a shifted version in Figure6. When choosing a relatively flat soft-uniform prior distribution that covers the target region, multimodality transfers to the potential energy for both fixed-kinetic and MLP-kinetic NHF. The potential exhibits indeed local extrema centered at the modes of the target, which can either be minima or maxima for the MLP-kinetic NHF but are minima for the fixed-kinetic one. Indeed, with a MLP-kinetic model, the orientation of the learned energies may change from one numerical experiment to another, as we do not enforce the positiveness of the output of V and K making this variant of NHF indifferent to the sign of the learned energies. Finally, similar results were obtained using a Gaussian base N (0, 2.5 2 I 2 ) with variance large enough to cover the support of the target distribution, which stresses the impact of the spatial expansion rather than the nature of the prior distribution.

On the other hand, when choosing a "peaked" prior distribution N (0, I 2 ), we remark that, for the MLP-kinetic NHF, the momenta p T generated by the Encoder inherit from the multimodality of the target distribution, with the same number of modes. As a consequence, learned energies are different from the classical Physics ones. Learned potential energy is different from one model to another but in the case of fixed-kinetic NHF, multimodality is always transferred to the potential energy, showing the robustness of the model to the choice of prior distribution.

Thus, using a fixed-kinetic model allows to more robustly transfer important properties of the target distribution into the learned potential. For both models though, choosing a wide prior that covers the support of the target distribution is improving such transfer. Finally, when the learned potential is not multimodal, it is an indication that multimodality has been transferred instead to the artificial momenta p T generated by the Encoder.

Adapting NHF for Bayesian inference

Methodology, derivation of the new loss function

We present how NHF can be used to perform Bayesian inference. It consists in using Hamiltonian flows to transform the prior distribution, in the sense of Bayes' theorem, π 0 of some vector of parameters q into the target posterior distribution π(q|d) of these parameters, knowing some data d and likelihood ℓ. The main difference with the above described NHF lies in the loss inspired from the KL phase in Boltzmann Generators [START_REF] Noé | Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning[END_REF], as well as in the learning procedure. During training, this NHF takes batches of q 0 from the prior distribution as inputs. For each q 0 , one p 0 is drawn from a Gaussian distribution whose mean and standard deviation depend on the q 0 . The resulting point in phase-space is then evolved through a series of L Leapfrog steps with integration time dt. The outputs consist in the final positions q T and momenta p T , as well as the initial mean µ(q 0 ), standard deviation σ(q 0 ) and p 0 . All these outputs, as well as the data d, are used in the loss computation. Once trained, it is able to transform the prior into the desired posterior distribution of the parameters. Thus, both training and sampling are now made following the forward-direction flow from the prior to the posterior.

In order to compute the loss, one needs to have access to the likelihood distribution ℓ of the model, which encapsulates the covariance matrix of the data as well as the underlying physical mapping between vectors of parameters and the corresponding data. Then, in the framework of Hamiltonian dynamics, the full system is made of both positions (the parameters of interest) and artificial momenta. Let us call q 0 , p 0 the initial position and momentum, respectively, and q T , p T the corresponding final position and momentum, respectively, obtained after L Leapfrog transformations T dt 1 , ..., T dt L with timestep dt, i.e: (q T , p T ) = T dt L • ... • T dt 1 (q 0 , p 0 ) := T (q 0 , p 0 ). Also, we introduce the notations for the projections along the final positions and momenta, i.e. q T := T q (q 0 , p 0 ) and p T := T p (q 0 , p 0 ). By the change of variables formula and the fact that the Jacobian determinant of each Hamiltonian transformation is one, the model joint distribution M is written as:

M (q T , p T ) = 1 × Π 0 (T -dt 1 • ... • T -dt L (q T , p T )) = Π 0 (q 0 , p 0 ) = π 0 (q 0 ) × f (p 0 |q 0 ),
where Π 0 is the joint prior distribution, π 0 the prior distribution of the parameters of interest and f the Gaussian distribution of the Encoder. We fix the target density of the final momenta g(p) (e.g. Gaussian). We now seek to minimize the KL-divergence between the model joint distribution and the desired target joint distribution conditioned on data Π(q, p|d) = π(q|d)g(p). We write the latter as the product of a density depending on q and an other density depending on p so that the two are independent. Using Bayes' theorem, we have (see Appendix B for more details):

DKL(M (qT , pT ) || π(qT |d)g(pT )) = π0(q0, p0) log π0(q0) + log f (p0|q0) -log π0(Tq(q0, p0))
log ℓ(d|Tq(q0, p0)) -log g(Tp(q0, p0) dq0dp0 + cst.

(6)

Application to cosmology

We apply the herein-above architecture to cosmological analysis. One typical case is the determination of the cosmic expansion, and more generally of the cosmological param-eters, from the observation of brightness and recession velocity of Type Ia supernovae (e.g. [START_REF] Riess | Observational Evidence from Supernovae for an Accelerating Universe and a Cosmological Constant[END_REF][START_REF] Betoule | Improved cosmological constraints from a joint analysis of the SDSS-II and SNLS supernova samples[END_REF]. We use it as a toy example, however, while the model used so far has been simple, it may be expanded in very complicated direction for which sampling from the probability distribution becomes very complex. New observatory are presently being built which is expected to deliver tens of thousands of new supernovae Ia over the next decade (LSST Science Collaboration, 2009). We now present the essential feature of the model that we need before applying our method to the toy example.

According to the Λ-CDM model, the relation between the distance and the brightness of Type Ia supernovae is of great interest because it depends on two cosmological parameters: the matter density parameter Ω m and the adimensional Hubble parameter h. To be more specific, database of type Ia supernovae report the distance modulus µ. This quantity is defined as the difference between the apparent and the absolute magnitude of an astronomical object, and is directly related to luminosity distance [START_REF] Weinberg | Gravitation and Cosmology: Principles and Applications of the General Theory of Relativity[END_REF] and thus a function of the redshift z, Ω m and h:

µ(z, Ω m , h) = 5 log 10 D * L (z, Ω m ) h10pc where D * L (z, Ω m ) = c(1+z) H0 z 0 ds √ 1-Ωm+Ωm(1+s) 3
, and H 0 = 100 km s -1 Mpc -1 , c being the speed of light. In practice, we avoid computing the integral in D * L by using an approximation [START_REF] Pen | Analytical fit to the luminosity distance for flat cosmologies with a cosmological constant[END_REF] which is only valid for a flat Universe:

D * L (z, Ω m ) = c(1 + z) H 0 η(1, Ω m ) -η 1 1 + z , Ω m , with η(a, Ω m ) = 2 1 + s 3 1 a -0.1540 s a 3 +0.4304 s 2 a 2 + 0.19097 s 3 a + 0.066941s 4 .
The formal definition of these quantities imposes constraints on the possible values of the parameters, that can only be comprised between zero and one. We avoid the problem by outputting a sigmoid of q T .

We aim to sample from the posterior distribution π(Ω m , h|data) quantifying the probability that we are living in a universe whose mean density and expansion is equal to Ω m and h given D observations data = {z i , µ i } 1≤i≤D of type Ia supernovae, and the covariance matrix C of the observed distance moduli. The likelihood of the problem is supposed to be Gaussian, i.e. the observed data and the simulated output from parameters are different up to a Gaussian noise. The final momenta distribution g is set to a Normal distribution. The cumulative plots in figure 8 compare the performance of a fixed-and MLP-kinetic NHF with a HMC. They are both biased which is expected since they minimize the KL-divergence between the model distribution and the target, but less than with an ELBO. We leave for future work a possible correction using importance sampling methods at the end of training.

Conclusion

In this work, we analyzed how well Hamiltonian Normalizing Flows are suited to solving some issues in Generative modeling. The main advantage of these methods is twofold. First, the volume-preservation in phase-space avoids the costly computation of Jacobian determinants. Then they allow for flexibility in the type of neural networks that is used since reversibility is ensured by the symplectic integrator.

It is this flexibility that allows us to propose a NHF variant based on classical kinetic energy. By exploring 2D sampling problem, we showed that both variants are robust to the choice of hyperparameters and prior distribution. In particular, we illustrate how the explicit classical design of the kinetic energy is a way to increase such robustness and make the model more easily interpretable. Finally, we explained how to adapt NHF to the context of Bayesian inference to obtain a sampler of the posterior distribution. Further work could then address how the offsets could be corrected by importance sampling techniques on a trained model. More generally, a next step would be testing such architectures on more complex problems such as image generation to attest their behavior in a high-dimensional setting.

Figure 1 .

 1 Figure 1. NHF architecture

Figure 2 .

 2 Figure 2. Density estimation of the 2D multimodal target distribution that is used in Section 5 with its marginals. It consists in 9 equally-weighted Gaussians with same covariance matrix 0.5 2 I2.

Figure 4 .

 4 Figure 4. Density estimation of the samples produced by the six models previously defined. First row: MLP-kinetic NHF; second row: Fixed-kinetic model. From left to right: wide soft-uniform prior; peaked Gaussian prior; wide Gaussian prior. Models were trained for 15, 000 epochs on a 5,000-point dataset with minibatches of size 512, H = 128, L = 10 and T = 1.

Figure 5 .

 5 Figure 5. Potential energies learned by the six models previously defined. First row: MLP-kinetic NHF; second row: Fixed-kinetic model. From left to right: wide soft-uniform prior; peaked Gaussian prior; wide Gaussian prior. Models were trained for 15, 000 epochs on a 5,000-point dataset with minibatches of size 512, H = 128, L = 10 and T = 1.

Figure 6 .

 6 Figure 6. Shifted potential energies learned by the six models previously defined. First row: MLP-kinetic NHF; second row: Fixedkinetic model. From left to right: wide soft-uniform prior; peaked Gaussian prior; wide Gaussian prior. Models were trained for 15, 000 epochs on a 5,000-point dataset with minibatches of size 512, H = 128, L = 10 and T = 1.

Figure 7 .

 7 Figure 7. Density estimation of the artificial momenta drawn by the Encoder for the six models previously defined. First row: MLPkinetic NHF; second row: Fixed-kinetic model. From left to right: wide soft-uniform prior; peaked Gaussian prior; wide Gaussian prior. Models were trained for 15, 000 epochs on a 5,000-point dataset with minibatches of size 512, H = 128, L = 10 and T = 1.

Figure 8 .

 8 Figure8. Cumulative plots of means and standard deviations of both Ωm and h produced by trained NHF models and a HMC on a 20,000-sample dataset, compared to the ground truth. Soft-uniform prior, 30,000 training epochs, g chosen to be a N (0, I2)

  Training loss as a function of epochs for models with different H (number of neurons per hidden layer in each neural network of the model), L (number of Leapfrog steps) and T (integration time). First row: MLP-kinetic NHF; second row: fixed-kinetic model. From left to right: H = 128, 8. Models were trained for 15, 000 epochs on a 5,000-point dataset with minibatches of size 512.

	5.0
		T=0.1, L=1	T=1, L=10
	4.8	T=0.1, L=2	T=1, L=50
		T=0.1, L=10	T=10, L=1
	4.6	T=0.1, L=50	T=10, L=2
		T=1, L=1	T=10, L=10
	4.4	T=1, L=2	T=10, L=50
	4.2
	4.0
	3.8 5.0
	4.8
	4.6
	4.4
	4.2
	4.0
	3.8
	0	2500 5000 7500 10000 12500 15000	0	2500 5000 7500 10000 12500 15000
	Figure 3.

Thus, for a given training dataset, there are basically four hyperparameters that require tuning: three are usual in learning (minibatch size, learning rate and number of neurons per hidden layer, i.e. number of learning parameters of the model) and only one is specific to NHF (minimum number of Leapfrog steps), whose tuning is especially less sensitive

Acknowledgments

All the authors are grateful to the support from CNRS through the MITI-Prime 80 project "CosmoBayes". All the authors thank the Mésocentre Clermont Auvergne University (https://mesocentre.uca.fr/) and the HPC cluster at Fysikum (it.fysik.su.se/hpc/) for providing help, computing and storage resources. This work was supported by the French ANR under the grant ANR-20-CE46-0007 (SuSa project) and by the Simons Collaboration on "Learning the Universe". JJ acknowledges support by the Swedish Research Council (VR) under the project 2020-05143 -"Deciphering the Dynamics of Cosmic Structure".

We acknowledge the use of the following packages in this analysis: Numpy (Harris et al., 2020), IPython (Pérez & Granger, 2007), Matplotlib (Hunter, 2007), Pytorch (Paszke et al., 2019), SciPy (Jones et al., 2001-).

A. Additional plots about robustness and optimization of NHF

We present the results of additional experiments with H = 32 in figure 9. Removing more than 90% of parameters (passage from H = 128 to H = 32), the final values are always higher by less than 4%, for both models. The different final values of the loss function can be represented on scatter plots, see Figure 10. The latter clearly illustrate the robustness of the fixed-kinetic model, see Figure 9. It also shows that models with L = 1 perform poorer than with L = 2, 10, 50.

B. Derivation of the KL-divergence and ELBO for the inference problem

The KL-divergence suited to the inference problem is derived as follows: We can also adapt the ELBO from [START_REF] Toth | Hamiltonian generative networks[END_REF] to our inference framework:

ln Π 0 (q 0 , p 0 )ln f (p 0 |q 0 )] = E f [ln M (T (q 0 , p 0 ))ln f (p 0 |q 0 )] Then, expliciting M (q, p) = π 0 (q)ℓ(d|q)g(p): ELBO(q 0 ) = E f [ln [π 0 (T q (T (q 0 , p 0 )))ℓ(d|T q (T (q 0 , p 0 )))g(T p (T (q 0 , p 0 )))]ln f (p 0 |q 0 )] (8)