Etrit Haxholli
email: etrit.haxholli@inria.fr

Marco Lorenzi
email: marco.lorenzi@inria.fr

Enhanced Distribution Modelling via Augmented Architectures For Neural ODE Flows

While the neural ODE formulation of normalizing flows such as in FFJORD enables us to calculate the determinants of free form Jacobians in O(D) time, the flexibility of the transformation underlying neural ODEs has been shown to be suboptimal. In this paper, we present AFFJORD, a neural ODE-based normalizing flow which enhances the representation power of FFJORD by defining the neural ODE through special augmented transformation dynamics which preserve the topology of the space. Furthermore, we derive the Jacobian determinant of the general augmented form by generalizing the chain rule in the continuous sense into the cable rule, which expresses the forward sensitivity of ODEs with respect to their initial conditions. The cable rule gives an explicit expression for the Jacobian of a neural ODE transformation, and provides an elegant proof of the instantaneous change of variable. Our experimental results on density estimation in synthetic and high dimensional data, such as MNIST, CIFAR-10 and CelebA (32 × 32), show that AFFJORD outperforms the baseline FFJORD through the improved flexibility of the underlying vector field.

Introduction

Normalizing flows are diffeomorphic random variable transformations providing a powerful theoretical framework for generative modeling and probability density estimation (Rezende & Mohamed, 2015).While the practical application of normalizing flows is generally challenging due to computational bottlenecks, most notably regarding the O(D 3 ) computation cost of the Jacobian determinant, different architectures have been proposed in order to scale normalizing flows to high dimensions while at the same time ensuring the flexibility and bijectivity of the transformations (Rezende & Mohamed, 2015;[START_REF] Dinh | Density estimation using real nvp[END_REF][START_REF] Kobyzev | Normalizing flows: An introduction and review of current methods[END_REF]. The common strategy consists of placing different architectural restrictions on the model, to enforce special Jacobian forms, with less computationally demanding determinants. A noteworthy approach is based on neural ODEs [START_REF] Chen | Neural ordinary differential equations[END_REF], as they enable us to calculate the determinants of free form Jacobians in O(D) time [START_REF] Grathwohl | Ffjord: Freeform continuous dynamics for scalable reversible generative models[END_REF]. More specifically, the rule for the instantaneous change of variable [START_REF] Chen | Neural ordinary differential equations[END_REF] provides an important theoretical contribution to normalizing flows, as it yields a closed form expression of the Jacobian determinant of a neural ODE transformation. In this case, calculating the Jacobian determinant simplifies to calculating the integral of the divergence of the vector field along the transformation trajectory. Such models are known as Continuous Normalizing Flows (CNFs). In [START_REF] Grathwohl | Ffjord: Freeform continuous dynamics for scalable reversible generative models[END_REF], these ideas are further explored and computational simplifications are introduced, notably the use of Hutchinson's trace estimator [START_REF] Hutchinson | A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines[END_REF]. The resulting model is named FFJORD.

In [START_REF] Dupont | Augmented neural odes[END_REF], it is shown that there exist functions which neural ODEs are not capable of representing. To tackle this issue and enhance the expressiveness of the neural ODE transformation, they propose to lift the data into a higher dimensional space, on which the neural ODE is applied, to subsequently project the output back to the original space. Such augmented neural ODEs (AN-ODEs) have been experimentally shown to lead to improved flexibility and generalisation properties than the non-augmented counterpart. Additional background material is provided in Appendix A. Inspired by [START_REF] Dupont | Augmented neural odes[END_REF], in this paper we develop a theoretical framework to increase the flexibility of ODE flows, such as FFJORD, through dimension augmentation. To this end, we derive an explicit formula for Jacobians corresponding to neural ODE transformations. This formula represents the continuous generalization of the chain rule, which we name the cable rule, that ultimately allows the derivation of the Jacobian expression and determinant for the composition of the operations defining the ANODE flow: augmentation, neural ODE transformation, and projection. To enable computational feasibility and ensure that the ANODE transformation is diffeomorphic, we allow the augmented dimensions to parameterize the vector field acting on the original dimensions (but not vice-versa). We name the resulting model augmented FFJORD (AFFJORD). In AFFJORD, the evolution of time is high dimensional, and is learnt via the vector field defined by the augmented component, contrasting the linear time evolution of FFJORD. This setup coincides with the framework introduced by (Zhang et al., 2019), but in the context of normalizing flows. Our experiments on 2D data of toy distributions and image datasets such as MNIST [START_REF] Lecun | The mnist database of handwritten digits[END_REF], CIFAR-10 (Krizhevsky et al., 2010) and CelebA [START_REF] Liu | Deep learning face attributes in the wild[END_REF], show that the proposed ANODE flow defined by AFFJORD outperforms FFJORD in terms of density estimation, thus highlighting the improved flexibility and representation properties of the underlying vector field.

Proposed Framework

In what follows, we first give the generalisation of the chain rule in the continuous sense which we refer to as the cable rule, which is analogous to forward sensitivity. Then, we derive the expression of the loss function for general augmented neural ODE flows, and finally introduce our model as a special case of such augmented neural ODEs.

The Cable Rule. If we define a chain of transformations z i = g i (z i-1 ) for i ∈ {1, ..., n}, (1) due to the chain rule we have:

dz n dz 0 = ∂z n ∂z n-1 dz n-1 dz 0 = ∂z n ∂z n-1 ∂z n-1 ∂z n-2 ... ∂z 1 ∂z 0 = ∂g n (z n-1 ) ∂z n-1 ∂g n-1 (z n-2 ) ∂z n-2 ... ∂g 1 (z 0 ) ∂z 0 . (2)
We can choose each g i to infinitesimally modify its input, i.e., g i (z i-1 ) = z i-1 +ϵf i (z i-1 , t i-1 , θ), so that the chain in Expression 1 transforms z 0 continuously. It is clear that z(t) = z(0) + t 0 f (z(t), t, θ)dτ is the limit of the previous iterative definition when ϵ → 0. Then the expression dzn dz0 in Equation (2) converges to dz(t) dz(0) , which as we show in Appendix B, satisfies the differential equation below:

d( dz(t) dz(0) ) dt = ∂f (z(t)) ∂z(t) dz(t) dz(0) . (3)
As this expression gives the generalisation of the chain rule in the continuous sense, we refer to it as the cable rule. We notice that Equation (3) gives the dynamics of the Jacobian of the state with respect to the initial condition z(0). The cable rule is therefore analogous to the forward sensitivity formula for ODEs which provides the dynamics of Jacobian of the state with respect to the parameters θ of the flow (Zhao & Mousseau, 2013). This relation is highlighted and explained in more detail in Appendix B, where in addition, the solution to Equation (3) is provided in Equation ( 38) by utilizing the Magnus Expansion, [START_REF] Magnus | On the exponential solution of differential equations for a linear operator[END_REF][START_REF] Blanes | The magnus expansion and some of its applications[END_REF]. Furthermore, we provide the continuous generalization of the total derivative decomposition, and we derive the instantaneous change of variables from the cable rule respectively in Appendix C and E.

AFFJORD, a Special Case of Augmented Neural ODE Flows. Since Neural Ordinary Differential Equations (NODEs) learn representations that preserve the topology of the input space, [START_REF] Dupont | Augmented neural odes[END_REF] introduce the Augmented Neural ODEs whose dynamics can be described as follows:

w(T ) = z(T ) z * (T ) = z(0) z * (0) + T 0 f (z(t), z * (t), θ) g(z * (t), z(t), ϕ) dt, (4)
where z * are the augmented dimensions. In this general form, the model is unsuitable to be used in the context of normalizing flows for two reasons:

1) The transformation of z(0) to z(T ) is not necessarily injective. Indeed, the transformation from [z(0), z * (0)] to [z(T ), z * (T )] is injective due to the Picard-Lindelöf Theorem, however, for two data points z ′ (0) and z ′′ (0), their images z ′ (T ), z ′′ (T ) might be identical as long as their respective augmented dimensions z ′ * (T ), z ′′ * (T ) differ.

2) The Jacobian determinant of this general transformation is computationally intractable. We can use the chain rule to express the Jacobian determinant of this transformation as

det dz(T ) dz(0) = det dz(T ) d [z(T ), z * (T )] d [z(T ), z * (T )] d [z(0), z * (0)] d [z(0), z * (0)] dz(0) . ( 5
)
and further develop the middle term in Expression 5 via the cable rule (Equation ( 38), Appendix B) to give the expression of the determinant of the Jacobian of general augmented neural ODE flows.

Unfortunately however, computing this expression is not feasible in general.

In order to mitigate these issues, we set z * (0) to be the zero vector and define g to be independent of the original dimensions. Therefore, [z(0), z * (0)] is transformed by h = (f, g), as follows:

w(T ) = z(T ) z * (T ) = z(0) z * (0) + T 0 f (z(t), z * (t), θ) g(z * (t), ϕ) dt. (6)
We name this model AFFJORD (Augmented FFJORD). We are not interested in z * (T ) as our interest lies on the transformed z(0), that is z(T ). Since by definition such coupled dynamics are contained in the formulation of neural ODEs, this implies that the instantaneous change of formula still holds, and the transformation is injective. A more detailed explanation proving that such a choice alleviates the problems discussed above is provided in Appendix F. The augmented dimensions z * (t) can also be seen as time dependent weights of the non-autonomous f whose evolution is determined by the autonomous ODE g, hence giving f greater flexibility in time (Zhang et al., 2019). In the case that the base distribution is multivariate normal, as shown in Appendix F.2, the expression of the loss function is:

L = ||Z(T )|| 2 2 - T 0 tr ∂f (z(t), z * (t)) ∂z(t) dt (7)
Multiscale Architecture in Augmented Neural ODE Flows. In the implementation of the model we modify the multiscale architecture in order to accommodate the augmented dimensions. We provide a more detailed description and an illustration of this modification in Appendix I.

Experiments

We compare the performance of AFFJORD with respect to the base FFJORD on 2D data of toy distributions, as well as on standard benchmark datasets such as MNIST, CIFAR-10 and CelebA(32 × 32). The goal is to train a neural network to generate a vector field h = (f, g) (as in Eq. ( 6)) that transforms our data distribution into a standard multivariate normal one. The augmented dimensions are fed to a hypernet (denoted as hyp) in order to generate the weights of the field of the main dimensions. The expression of the vector field to be learnt is the following:

dz dt = f [z(t), z * (t)], θ(t) = hyp(z * (t), w)
, where dz * (t) dt = g(z * (t), ϕ). Thus, the learnable parameters are w and ϕ. In all cases, the architecture of the main field f (in Eq. ( 6)) and the training time in AFFJORD are identical to that of FFJORD for the sake of fairness. Detailed information about the experiments such as hyper-parameters and network architectures can be found in Appendix I. Limitations and future prospects are discussed in Appendix K.

Toy 2D Datasets

In order to visualise the performance of the model, we first test FFJORD and AFFJORD on 2D data of toy distributions, depicted in Figure 1. In both cases we use the Runge-Kutta 4 solver with 40 time steps (160 function evaluations). While both FFJORD and AFFJORD are capable of modelling multi-modal and discontinuous distributions, Figure 1 shows that AFFJORD has higher flexibility in modeling the complex data distributions considered, in comparison to FFJORD. We can notice in Table 1 that the loss of our model is roughly two standard deviations lower than the one of FFJORD. For each model the experiment was repeated 30 times. The experiments provided here show that AFFJORD is characterized by high flexibility of the vector field.

Image Datasets

We show that AFFJORD outperforms FFJORD on MNIST, CIFAR-10 and CelebA (32 × 32). There are several architectures of FFJORD that can be used for this application, hence we first optimize the architecture of FFJORD and then add the augmented structure in AFFJORD. Out of the architectures introduced in (Grathwohl et al., 2019) that we tested, the one that performed best was the multiscale one, with three convolutional layers with 64 channels each. The number of CNF blocks was 1, and time was implemented by simply concatenating it as a channel into the data. When time was implemented via a hypernet, we observed that the training time increased and performance decreased, especially in the case of CIFAR-10. For AFFJORD we use the exact same base architecture, however we enable the evolution of the vector fields g(z * (t), ϕ) through time via a fully connected network with one hidden layer, which does not take as input all the dimensions of z * (t) but merely 20 of them. Number 20 was chosen as during fine-tuning, the best performance was reached in this setting. The width of the hidden layer is also 20, as it is the output. We fix 10 of these 20 dimensions and feed them to a linear hypernet with weight matrix shape [10, p] to output the p weights for the main component. As we show in Table 1, AFFJORD outperforms FFJORD on MNIST, CIFAR-10 and CelebA (32 × 32). Based on the conducted experiments the farther FFJORD is from optimal performance, the greater the improvements brought by AFFJORD are. The calculation of results is done as in [START_REF] Grathwohl | Ffjord: Freeform continuous dynamics for scalable reversible generative models[END_REF], where for each run the best evaluation result over epochs is taken. After 5 runs, for each model, the scores are averaged and reported in Table 1. In the case of MNIST, both models were trained for roughly 9 days, while in the case of CIFAR-10 and CelebA (32 × 32), they were trained for approximately 14 days. The results corresponding to the Real NVP and Glow models, are taken from the original papers: [START_REF] Dinh | Density estimation using real nvp[END_REF] and [START_REF] Kingma | Glow: Generative flow with invertible 1x1 convolutions[END_REF]. Like its predecessor FFJORD, AFFJORD can generate samples by backintegrating. The generative procedure is explained in more details in Appendix I. Generated samples from both AFFJORD and FFJORD can be found in Appendix J.

Conclusion

We presented the generalization of the total derivative decomposition in the continuous sense as well as the continuous generalization of the chain rule. Motivated by the latter, we propose a new type of continuous normalizing flow, namely AFFJORD which outperforms FFJORD in the experiments we conducted.

Pontrjagin, L., Boltyanskii, V., Gamkrelidze, R., Mishchenko, E., & Brown, D. (1962). The mathematical theory of optimal processes. In International series of monographs in pure and applied mathematics.

Rezende, D. J., & Mohamed, S. ( 2015). Variational inference with normalizing flows. In International Conference on Machine Learning.

Tabak, E. G., & Turner, C. V. (2013). A family of nonparametric density estimation algorithms.

Communications on Pure and Applied Mathematics.

Tabak, E. G., & Vanden-Eijnden, E. ( 2010). Density estimation by dual ascent of the log-likelihood.

Communications in Mathematical Sciences.

Zhang, T., Yao, Z., Gholami, A., Gonzalez, J. E., Keutzer, K., Mahoney, M. W., & Biros, G. ( 2019). Anodev2: A coupled neural ode framework. In Advances in Neural Information Processing Systems.

Zhao, H., & Mousseau, V. A. ( 2013). Extended forward sensitivity analysis for uncertainty quantification. Nuclear Technology.

A Background and Related Work

Normalizing Flows: The Normalizing Flow framework was previously defined in (Tabak & Vanden-Eijnden, 2010;Tabak & Turner, 2013), and was popularised by (Rezende & Mohamed, 2015) and by [START_REF] Dinh | Nice: Non-linear independent components estimation[END_REF] respectively in the context of variational inference and density estimation.

A Normalizing Flow is a transformation defined by a sequence of invertible and differentiable functions mapping a simple base probability distribution (e.g., a standard normal) into a more complex one. Let Z and X = g(Z) be random variables where g is a diffeomorphism with inverse h. If we denote their probability density functions by f Z and f X , based on the change of variable theorem we get:

f X (x) = f Z (z) det dz dx = f Z (h(x)) det dh(x) dx . (8)
In general, we want to optimize the parameters of h such that we maximize the likelihood of sampled points x 1 , ..., x n . Once these parameters are optimized, then we can give as input any test point x on the right hand side of Equation ( 8), and calculate its likelihood. For the generative task, being able to easily recover g from h is essential, as the generated point x g will take form x g =g(z s ), where z s is a sampled point from the base distribution f Z .

For increased modeling flexibility, we can use a chain (flow) of transformations,

z i-1 = g i (z i ), i ∈ [n], that is z i = h i (z i-1 ), i ∈ [n].
In this case due to chain rule we have:

f Z0 (z 0 ) = f Zn (z n ) det dz n dz 0 = f Zn (h n (...h 1 (z 0 ))) n i=1 det dh i (z i-1 ) dz i-1 , (9)
where z 0 is a data point. The interested reader can find a more in-depth review of normalizing flows in [START_REF] Kobyzev | Normalizing flows: An introduction and review of current methods[END_REF] and [START_REF] Papamakarios | Normalizing flows for probabilistic modeling and inference[END_REF].

Neural ODE Flows: Neural ODEs, [START_REF] Chen | Neural ordinary differential equations[END_REF], are continuous generalizations of residual networks:

z ti+1 = z i + ϵf (z ti , t i , θ) → z(t) = z(0) + t 0 f (z(τ ), τ, θ)dτ , as ϵ → 0. (10)
In (Pontrjagin et al., 1962;[START_REF] Chen | Neural ordinary differential equations[END_REF], it is shown that the gradients of neural ODEs can be computed via the adjoint method, with constant memory cost with regards to "depth":

dL dθ = T 0 ∂L ∂z(t) ∂f (z(t), t, θ) ∂θ dt, (11)
where ∂L ∂z(t) can be calculated simultaneously by

∂L ∂z(t) = ∂L ∂z(T ) - t T ∂L ∂z(τ ) ∂f (z(τ ), τ, θ(τ )) ∂z(τ ) dτ. (12)
In addition, they also derive the expression for the instantaneous change of variable, which enables one to train continuous normalizing flows:

log p(z(0)) = log p(z(T )) + T 0 tr ∂f (z(t), t, θ) ∂z(t) dt, (13)
where z(0) represents a sample from the data. A surprising benefit is that one does not need to calculate the determinant of the Jacobian of the transformation anymore, but simply the trace of a matrix. These models are collectively called Neural ODE flows (NODEFs) or simply Continuous Normalizing Flows (CNFs). In [START_REF] Grathwohl | Ffjord: Freeform continuous dynamics for scalable reversible generative models[END_REF], these ideas are further explored and computational simplifications are introduced, notably the use of Hutchinson's trace estimator, [START_REF] Hutchinson | A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines[END_REF][START_REF] Adams | Estimating the spectral density of large implicit matrices[END_REF], as an unbiased stochastic estimator of the trace in the likelihood expression in Equation ( 13). The resulting model is named FFJORD. , where c is the number of channels and s is the height and width of the image. Effectively this operation trades spatial size for additional channels, and after each transformation a normalizing flow is applied on half the channels, while the other half are saved in an array. This process can be repeated as many times as the width and height of the transformed image remain even numbers. In the end, all saved channels are concatenated to construct an image with the original dimensions. A visual description of this process can be found in Figure 2.

Augmented Neural ODEs: Considering that transformations by neural ODEs are diffeomorphisms (hence homeomorphisms), [START_REF] Dupont | Augmented neural odes[END_REF] show that Neural Ordinary Differential Equations (NODEs) learn representations that preserve the topology of the input space, and prove that this implies the existence of functions that Neural ODEs cannot represent. To address these limitations, they introduce the Augmented Neural ODEs:

d dt z(t) z * (t) = h( z(t) z * (t) , t) = f (z(t), z * (t), θ) g(z(t), z * (t), θ) for z(0) z * (0) = x 0 , (14)
where z * (t) is the augmented component, and h = [f, g] is the vector field to be learnt.

In addition to being more expressive models, [START_REF] Dupont | Augmented neural odes[END_REF] show that augmented neural ODEs are empirically more stable, generalize better and have a lower computational cost than Neural ODEs. A schematic of their architecture in the discrete case can be found in Figure 3.

B Cable Rule: The Continuous Generalization of the Chain Rule

We will first assume that z is one dimensional. If z i = f i (z i-1 ) for i ∈ {1, ..., n} then by the chain rule we have

dz n dz 0 = ∂z n ∂z n-1 dz n-1 dz 0 = ∂z n ∂z n-1 ∂z n-1 ∂z n-2 ... ∂z 1 ∂z 0 = ∂f n (z n-1 ) ∂z n-1 ∂f n-1 (z n-2 ) ∂z n-2 ... ∂f 1 (z 0 ) ∂z 0 (15)
Now, if we assume that z i = z(t i ) is transformed more gradually as in z i+1 = z i + ϵf (z i ), and that t i+1 = t i + ϵ, we get that

dz n dz 0 = ∂z n ∂z n-1 ∂z n-1 ∂z n-2 ... ∂z 1 ∂z 0 = (I + ϵ ∂f (z n-1 ) ∂z n-1 )(I + ϵ ∂f (z n-2 ) ∂z n-2 )...(I + ϵ ∂f (z 0 ) ∂z 0 ) (16)
We see that z(t) = z(0) + t 0 f (z(τ ))dτ is the limit of the previous iterative definition z i+1 = z i + ϵf (z i , θ(t i )), when ϵ → 0. For simplicity, we have written f (z i ) = f (z i , θ, t i ). If we decide to expand equation 16, we obtain

dz n dz 0 = I + n-1 i=0 ∂f (z i ) ∂z i ϵ + n-1 i=0 j<i ∂f (z i ) ∂z i ∂f (z j ) ∂z j ϵ 2 + ... + ∂f (z 0 ) ∂z 0 ... ∂f (z n ) ∂z n ϵ n (17) = I + S n 1 + S n 2 + ... + S n n+1 , (18)
where

S n 2 = n-1 i=0 j<i ∂f (z i ) ∂z i ∂f (z j ) ∂z j ϵ 2 , S n 3 = n-1 i=0 j<i k<j ∂f (z i ) ∂z i ∂f (z j ) ∂z j ∂f (z k ) ∂z k ϵ 3 , ... ( 19
)
We will focus on the sum S n 2 for a moment. Let us define

g n 2 (x, y) = ∂f (z ti ) ∂z ti ∂f (z tj ) ∂z tj , for x ∈ [t i , t i+1 ], y ∈ [t j , t j+1 ]. ( 20
)
It is clear that g n 2 is the discretization of

g 2 (t, u) = ∂f (z t ) ∂z t ∂f (z u ) ∂z u , for t ∈ [0, T ], u ∈ [0, T ]. (21)
We can notice that

S n 2 = n-1 i=0 j<i g n 2 (t i , t j )ϵ 2 and that g n 2 (t i , t j ) = g n 2 (t j , t i ) is symmetric, but in S n
2 , only takes values in the rectangles under the diagonal as illustrated in Figure 7. If we decide to expand the sum S n 2 on the rectangles above the diagonal as in Figure 8 and denote it as Sn 2 =

n-1 i=0 j̸ =i g n 2 (t i , t j )ϵ 2 , then due to the symmetry of g n 2 (t i , t j ) = g n 2 (t j , t i ), we deduce that S n 2 = 1 2 Sn 2 . The only rectangles missing in Sn 2 , regarding the discretization of [0, T ] × [0, T ], are the ones corresponding to the cases when i = j, which are the rectangles in the diagonal. It is important to emphasize here that S n 2 is the sum of C 2 n terms, while Sn

2 is made out of V 2 n = 2!C 2 n
terms. This is especially apparent when we notice the discretization of [0, T ] × [0, T ] is composed of n 2 rectangles, while the diagonal is composed of n rectangles, hence Sn 2 is the sum of n 2 -n = V 2 n terms. The ratio of the collective mass of the rectangles in the diagonal with respect to the entire [0, T ] × [0, T ] goes to zero as n → ∞, hence we conclude that

Sn 2 = n-1 i=0 j̸ =i ∂f (z i ) ∂z i ∂f (z j ) ∂z j ϵ 2 → [0,T ]×[0,T ] g 2 (t, u)d(t, u) = ( T 0 ∂f (z t ) ∂z t dt) 2 , (22)
thus,

S n 2 = 1 2! Sn 2 → 1 2! ( T 0 ∂f (z t ) ∂z t dt) 2 . ( 23
)
Similarly for

S n 3 = n-1 i=0 j<i k<j ∂f (z i ) ∂z i ∂f (z j ) ∂z j ∂f (z k ) ∂z k ϵ 3 ,
we can define

g n 3 (x, y, z) = ∂f (z ti ) ∂z ti ∂f (z tj ) ∂z tj ∂f (z t k ) ∂z t k , for x ∈ [t i , t i+1 ], y ∈ [t j , t j+1 ], z ∈ [t k , t k+1 ], (24)
and

g 3 (t, u, v) = ∂f (z t ) ∂z t ∂f (z u ) ∂z u ∂f (z v ) ∂z v , for t ∈ [0, T ], u ∈ [0, T ], v ∈ [0, T ]. (25)
In this case:

g n 3 (x, y, z) = g n 3 (x, y, z) = g n 3 (x, z, y) = g n 3 (y, x, z) = = g n 3 (y, z, x) = g n 3 (z, x, y) = g n 3 (z, y, x), (26)
where each equality corresponds to one of the 6 = 3! permutations of (x, y, z). As before we can expand the domain of S n 3 , by adding the rectangles in the discretization of [0, T ] × [0, T ] × [0, T ] such that i might be smaller than j, as well as that j could be smaller than k. We denote this expanded sum as Sn 3 , and by the symmetry of g n 3 , we notice that S n 3 = 1 3! Sn 3 . This implies that the rectangles participating in Sn 3 , now cover most of [0, T ] × [0, T ] × [0, T ], where the only exceptions are the ones when i = j or j = k (or both). The number of rectangles participating in Sn

3 is V 3 n = 3!C 3 n = n 3 -3n 2 + 2n
, and since the number of all rectangles in [0, T ] × [0, T ] × [0, T ] is n 3 , this implies that 3n 2 -2n rectangles are missing in sum Sn 3 from the cases when i = j or j = k (or both). The ratio of the collective mass of such rectangles with respect to the entire [0, T ] × [0, T ] × [0, T ] goes to zero as n → ∞, hence as before:

Sn 3 = n-1 i=0 j̸ =i i̸ =k̸ =j ∂f (z i ) ∂z i ∂f (z j ) ∂z j ∂f (z k ) ∂z k ϵ 3 → [0,T ]×[0,T ]×[0,T ] g 3 (t, u, v)d(t, u, v) (27) = ( T 0 ∂f (z t ) ∂z t dt) 3 , (28)
implying

S n 3 = 1 3! Sn 3 → 1 3! ( T 0 ∂f (z t ) ∂z t dt) 3 . ( 29
)
In a similar fashion we can prove that

S n k → 1 k! ( T 0 ∂f (zt)
∂zt dt) k . Thus we conclude that:

dz(T ) dz(0) = 1 0! I + 1 1! T 0 ∂f (z t ) ∂z t dt + 1 2! ( T 0 ∂f (z t ) ∂z t dt) 2 + 1 3! ( T 0 ∂f (z t ) ∂z t dt) 3 + ... (30) dz(T ) dz(0) = e T 0 ∂f (z t ) ∂z t dt . (31)
Unfortunately, this result does not hold when the dimensionality of z(t) is larger than one. Indeed, in this case,

∂f (zt i ) ∂zt i is a matrix, hence g n 2 (x, y) = ∂f (zt i ) ∂zt i ∂f (zt j ) ∂zt j is not necessarily symmetric, as the commutator [ ∂f (zt i ) ∂zt i , ∂f (zt j ) ∂zt j
] is not necessarily zero. For this reason, inspired by the previous result we try a different approach. Indeed, we can see from Equation 31that dz(T ) dz( 0) is the solution of the following ODE:

d( dz(t) dz(0) ) dt = ∂f (z(t)) ∂z(t) dz(t) dz(0) (32)
as the initial condition dz(t=0) dz(0) = I. Hence we wish to prove that dz(t) dz(0) satisfies the same ODE in higher dimensions as well.

We notice that we can write:

z(t) = z(0) + t 0 f (z(τ ))dτ = g(z(0), t), (33) therefore
, d( dz(t) dz(0) ) dt = ∂ 2 g(z(0), t) ∂t∂z(0) = ∂ 2 g(z(0), t) ∂z(0)∂t = = d( dg(z(0),t) dt ) dz(0) = df (z(t)) dz(0) = df (z(t)) dz(t) dz(t) dz(0) . (34)
We pause for a moment, in order to highlight the similarity of expression 34 and the forward sensitivity:

d( dz(t) dθ ) dt = df (z(t), t, θ) dθ = ∂f (z(t), t, θ) ∂z(t) dz(t) dθ + ∂f (z(t), t, θ) ∂θ . ( 35
)
Now from Expression 34, we infer that dz(0) is the solution of the following linear ODE:

d( dz(t) dz(0) ) dt = ∂f (z(t)) ∂z(t) dz(t) dz(0) . (36)
If we write Y (t) = dz(t) dz(0) , then the equation above becomes:

dY (t) dt = ∂f (z(t)) ∂z(t) Y (t). ( 37
)
The general solution of first-order homogeneous linear ODEs is given in [START_REF] Magnus | On the exponential solution of differential equations for a linear operator[END_REF][START_REF] Blanes | The magnus expansion and some of its applications[END_REF], and in our case can be written as follows:

dz(t) dz(0) = e Ω(t) dz(0) dz(0) = e Ω(t) , for Ω(t) = ∞ k=1 Ω k (t), (38)
where

Ω 1 (t) = t 0 A(t 1 ) dt 1 , (39)
Ω 2 (t) = 1 2 t 0 dt 1 t1 0 dt 2 [A(t 1 ), A(t 2 )], (40)
Ω 3 (t) = (41) = 1 6 t 0 dt 1 t1 0 dt 2 t2 0 dt 3 A(t 1 ), [A(t 2 ), A(t 3 )] + A(t 3 ), [A(t 2 ), A(t 1 )] , (42)
and so on for k > 3, where:

A(t) = ∂f (z(t), t, θ) ∂z(t) . ( 43
)
We notice that if z(t) is one dimensional, then this result agrees with the one in the previous approach. To conclude, we have proven the following theorem:

Theorem 1. Let dz(t) dt = f (z(t), θ, t)
, where f is continuous in t and Lipschitz continuous in z(t), furthermore let z(t) = z t (z(0), t) satisfy the conditions of Clairaut's theorem. Then the following holds:

d( dz(t) dz(0) ) dt = ∂f (z(t)) ∂z(t) dz(t) dz(0) . ( 44
)
Proof. f is continuous in t and Lipschitz continuous in z then due to Picard-Lindelöf theorem, z(T ) exists and is unique. Furthermore, since f is Lipschitz continuous in z(t), then ∂f (zt) ∂zt exists almost everywhere. Based on the previous derivations we reach the desired conclusion.

C The Continuous Generalization of the Total Derivative Decomposition

In the case that f = f (x(θ), y(θ)), then df dθ = ∂f ∂x dx dθ + ∂f ∂y dy dθ , as θ contributes to both x and y, which in turn determine f . If z(T ) = z(0) + T 0 f (z(t), θ(t), t)dt, then θ controls the vector field at each time point during integration, hence we expect that these infinitesimal contributions of the transformation from z(0) to z(T ) should be integrated.

In Appendix B, we assumed that z(t) = z(0) + t 0 f (z(τ ), θ, τ )dτ , which was the the limit of ϵ → 0 of the previous iterative definition of z i+1 = z i + ϵf (z i , θ, t i ). Now, we assume that the set of parameters in f is different at each discrete time, that is z i+1 = z i + ϵf i (z i , θ i , t i ), for independent sets θ i . First, we notice that z n = z n (z(t 0 ), θ(t 0 ), θ(t 1 ), ..., θ(t n-1 )) = z n (z(t 1 ), θ(t 1 ), θ(t 2 ), ..., θ(t n-1 )) = ...z n (z(t k+1 ), θ(t k+1 ), θ(t k+2 ), ..., θ(t n-1 )) = ...

... = z n (z(t n-1 ), θ(t n-1 )) = z(t n = T -ϵ).
This can be seen from Figure 9. Thus Figure 9: The dependencies in the discretisation of z(t) t = f (z(t), θ(t), t). Variable z t k+1 completely absorbs the contribution of θ t k to z T , hence z T is a function of z t k+1 and the following sets of parameters θ t k+1 , θ t k+2 , ..., θ tn .

dz n dθ(t k ) = dz n (z(t k + ϵ), θ(t k + ϵ), θ(t k+2 ), ..., θ(t n )) dθ(t k ) (45) dz n dθ(t k ) = ∂z n ∂z(t k + ϵ) dz(t k + ϵ) dθ(t k ) + 0 + 0 + ... + 0 = ∂z n ∂z(t k + ϵ) ∂z(t k + ϵ) ∂θ(t k ) . ( 46
)
We now focus on analysing ∂zn ∂z(t k +ϵ) and ∂z(t k +ϵ) ∂θ(t k ) . First we notice that from:

∂z n ∂z(t k ) = ∂z n ∂z(t k + ϵ) ∂z(t k + ϵ) ∂z(t k ) = ∂z n ∂z(t k + ϵ) (I + ϵ ∂f (z(t k ), θ(t k )) ∂z(t k ) + O(ϵ 2 )) (47) we get ∂z n ∂z(t k + ϵ) = ∂z n ∂z(t k ) -ϵ ∂z n ∂z(t k + ϵ) ∂f (z(t k ), θ(t k )) ∂z(t k ) + O(ϵ 2 ). ( 48
)
Regarding ∂z(t k +ϵ) ∂θ(t k ) the following holds:

∂z(t k + ϵ) ∂θ(t k ) = ∂(z(t k ) + f (z(t k ), θ(t k ))ϵ + O(ϵ 2 )) ∂θ(t k ) = 0 + ∂f (z(t k ), θ(t k )) ∂θ(t k ) ϵ + O(ϵ 2 ). ( 49
)
After combining both Equation 48and Equation 49in expression 46, we deduce that:

dz n dθ(t k ) = ∂z n ∂z(t k ) -ϵ ∂z n ∂z(t k + ϵ) ∂f (z(t k ), θ(t k )) ∂z(t k ) + O(ϵ 2 ) ∂f (z(t k ), θ(t k )) ∂θ(t k ) ϵ + O(ϵ 2 ) , ( 50
) thus ∂z n ∂θ(t k ) = dz n dθ(t k ) = ∂z n ∂z(t k ) ∂f (z(t k ), θ(t k )) ∂θ(t k ) ϵ + O(ϵ 2 ) (51)
Now assuming that θ(t k ) = g(t k , θ), we can write the total derivative of z n with respect to parameters θ :

dz n dθ = n k=0 ∂z n ∂θ(t k ) dθ(t k ) dθ = n k=0 ∂z n ∂z(t k ) ∂f (z(t k ), θ(t k )) ∂θ(t k ) dθ(t k ) dθ ϵ + O(ϵ 2 ) (52)
hence taking ϵ → 0, we conclude that

dz(T ) dθ = T 0 ∂z(T ) ∂z(t) ∂f (z(t), θ(t)) ∂θ(t) ∂θ(t) ∂θ dt. (53)
We can see that we are integrating the infinitesimal contributions of parameters θ, at each time t. In case that θ(t) = g(t, θ) = θ, then we have

dz(T ) dθ = T 0 ∂z(T ) ∂z(t) ∂f (z(t), θ) ∂θ ∂θ ∂θ dt = T 0 ∂z(T ) ∂z(t) ∂f (z(t), θ) ∂θ Idt = = T 0 ∂z(T ) ∂z(t) ∂f (z(t), θ) ∂θ dt = - 0 T ∂z(T ) ∂z(t) ∂f (z(t), θ) ∂θ dt. ( 54
) Theorem 2. Let dz(t) dt = f (z(t), θ(t), t)
, where f is continuous in t and Lipschitz continuous in z(t) and θ(t). Then the following holds:

dz(T ) dθ = T 0 ∂z(T ) ∂z(t) ∂f (z(t), θ(t), t) ∂θ(t) ∂θ(t) ∂θ dt. (55)
Proof. As before, since f is continuous in t and Lipschitz continuous in z then due to Picard-Lindelöf theorem, z(T ) exists and is unique. Furthermore, since f is Lipschitz continuous in θ(t), then ∂f (θ(t)) ∂θ(t) exists almost everywhere. Based on the previous derivations we reach the desired conclusion.

D Generalization of Continuous Backpropagation into Piecewise Continuous Backpropagation

We define z(t) as before, with the only difference being that its derivative is discontinuous at a point (say T 2 ):

z(t) = z(0) + t 0 f (z(τ ), θ(τ ))dτ, t ∈ [0, T 2 ] z(0) + T 2 0 f (z(τ ), θ(τ ))dτ + t T 2 g(z(τ ), ϕ(τ ))dτ, t ∈ ( T 2 , T ] (56)
As in [START_REF] Chen | Neural ordinary differential equations[END_REF], we can get

d( ∂L ∂z(t) ) dt = -∂L ∂z(t) ∂f (z(t),θ(t)) ∂z(t) , t ∈ [0, T 2 ] -∂L ∂z(t) ∂g(z(t),ϕ(t)) ∂z(t) , t ∈ ( T 2 , T ], (57)
thus

∂L ∂z(t) = ∂L ∂z(T ) - t T ∂L ∂z(τ ) ∂g(z(τ ),ϕ(τ )) ∂z(τ ) dτ, t ∈ ( T 2 , T ] ∂L ∂z(T ) - t T 2 ∂L ∂z(τ ) ∂f (z(τ ),θ(τ )) ∂z(τ ) dτ - T 2 T ∂L ∂z(τ ) ∂g(z(τ )ϕ(τ )) ∂z(τ ) dτ, t ∈ (0, T 2 ]. (58)
The approach developed in Appendix C can be used to generalize continuous backpropagation into piecewise continuous backpropagation. Indeed, identically as before, we have the first order appoximation:

dL ϵ dθ = n k=0 ∂L ϵ ∂z(t k ) ∂f (z(t k ), θ(t k )) ∂θ(t k ) ∂θ(t k ) ∂θ ϵ + n k=0 ∂L ϵ ∂z(t k ) ∂g(z(t k ), ϕ(t k )) ∂ϕ(t k ) ∂ϕ(t k ) ∂θ ϵ (59)
Taking the limit we get:

dL dθ = T 2 0 ∂L ∂z(t) ∂f (z(t), θ(t)) ∂θ(t) ∂θ(t) ∂θ dt + T T 2 ∂L ∂z(t) ∂g(z(t), ϕ(t)) ∂θ(t) ∂ϕ(t) ∂θ dt (60)
and in the same manner we get:

dL dϕ = T 2 0 ∂L ∂z(t) ∂f (z(t), θ(t)) ∂θ(t) ∂θ(t) ∂ϕ dt + T T 2 ∂L ∂z(t) ∂g(z(t), ϕ(t)) ∂θ(t) ∂ϕ(t) ∂ϕ dt, (61)
If we set θ(t) = θ and ϕ(t) = ϕ, then we get:

dL dθ = T 2 0 ∂L ∂z(t) ∂f (z(t), θ) ∂θ dt, (62)
and

dL dϕ = T T 2 ∂L ∂z(t) ∂g(z(t), ϕ) ∂ϕ dt. (63)

E Deriving the Instantaneous Change of Variable via the Cable Rule

The trace of a commutator [X, Y ] = XY -Y X is always zero, hence we prove below that for all terms Ω k>1 (t) given in Equations ( 39), ( 40) and ( 41) in Appendix B, we have tr(Ω k (t)) = 0.

Indeed, since the trace and the integral are interchangeable we have:

tr(Ω 2 (t)) = tr 1 2 t 0 dt 1 t1 0 dt 2 [A(t 1 ), A(t 2 )] = 1 2 t 0 dt 1 t1 0 dt 2 tr[A(t 1 ), A(t 2 )] = 0. ( 64
)
tr(Ω 3 (t)) = 1 6 t 0 dt 1 t1 0 dt 2 t2 0 dt 3 tr A(t 1 ), [A(t 2 ), A(t 3 )] + +tr A(t 3 ), [A(t 2 ), A(t 1 )] (65) = 1 6 t 0 dt 1 t1 0 (0 + 0)dt 2 = 0, (66)
and so on for k > 3.

The only exception is the case when k = 1:

tr(Ω 1 (t)) = tr t 0 A(t 1 )dt 1 = t 0 tr ∂f (z(t 1 ), t 1 , θ) ∂z(t 1 ) dt 1 . (67)
Finally, using Jacobi's formula, we conclude that:

log det dz(T ) dz(0) = log det e Ω(T ) = log e tr(Ω(T )) = T 0 tr ∂f (z(t), t, θ) ∂z(t) dt. (68)
F Simplifications and the Derivation of the Loss in Section 2

Hence,

Ω 1 (t) = t 0 dt 1 ∂f (z(t1),z * (t1),θ) ∂z(t1) ∂f (z(t1),z * (t1),θ) ∂z * (t1) 0 ∂g(z * (t1),ϕ) ∂z * (t1) = t 0 dt 1 ∂f (z(t1),z * (t1),θ) ∂z(t1) B 1 (t) 0 D 1 (t) (72) Ω 2 (t) = t 0 dt 1 t1 0 dt 2 ∂f (z(t2),z * (t2),θ) ∂z(t2) ∂f (z(t2),z * (t2),θ) ∂z * (t2) 0 ∂g(z * (t2),ϕ) ∂z * (t2) ∂f (z(t1),z * (t1),θ) ∂z(t1) ∂f (z(t1),z * (t1),θ) ∂z * (t1) 0 ∂g(z * (t1),ϕ) ∂z * (t1) - t 0 dt 1 t1 0 dt 2 ∂f (z(t1),z * (t1),θ) ∂z(t1) ∂f (z(t1),z * (t1),θ) ∂z * (t1) 0 ∂g(z * (t1),ϕ) ∂z * (t1) ∂f (z(t2),z * (t2),θ) ∂z(t2) ∂f (z(t2),z * (t2),θ) ∂z * (t2) 0 ∂g(z * (t2),ϕ) ∂z * (t2) = t 0 dt 1 t1 0 dt 2 [ ∂f (z(t1),z * (t1),θ) ∂z(t1)
, ∂f (z(t2),z * (t2),θ)

∂z(t2) ] B 2 (t) 0 D 2 (t) = Ω [z] 2 (t) B 2 (t) 0 D 2 (t) .
In this way we can prove that

Ω(t) = ∞ k=1 Ω(t) k = ∞ k=1 Ω [z] k (t) ∞ k=1 B k (t) 0 ∞ k=1 D k (t) = Ω [z] (t) B(t) 0 D(t) .
Considering that the exponential of a matrix whose lower left block is zero, will still have a zero lower left block, we have: Regarding issue 1), we have the following: Proposition 1. The architecture of AFFJORD ensures that the transformation is injective.

dz(t) dz(0) = e Ω(t) = e Ω [z] (t) B(t) 0 D(t) .
Proof. Indeed, as z * (t) is not dependent on external factors, and since z * (0) is constant regarding z(0), the end result z * (T ) will always be the same. Hence, for z ′ (0) and z ′′ (0) their images z ′ (T ) and z ′′ (T ) must be different, since their equality would imply [z ′ (T ), z * (T )] = [z ′′ (T ), z * (T )] contradicting the Picard-Lindelöf Theorem.

Issue 2) is mitigated since as we proved above, dz(T ) dz(0) simplifies to we derive the following loss function:

[I, 0] e T 0 ∂f (z(t),z * (t),θ) ∂z(t) dt+Ω [z] 2 (T )+... B(T ) 0 D(T ) I 0 , (74)
L = ||Z(T )|| 2 2 - T 0 tr ∂f (z(t), z * (t), θ) ∂z(t) dt. (75)
G Cable Rule Derived via Equation ( 12)

Differentiating Equation ( 12) in Appendix B, we get:

d dL dz(t) dt = - dL dz(t) ∂f (z(t), t, θ) ∂z(t) . ( 76
)
Choosing L to be z(0), we have

d dz(0) dz(t) dt = - dz(0) dz(t) ∂f (z(t), t, θ) ∂z(t) . (77)
We define A(t) := dz(0) dz(t) , and B(t) := A(t) -1 = dz(t) dz(0) , so that the equation above becomes

dA(t) dt = -A(t) ∂f (z(t), t, θ) ∂z(t) , (78) thus
-B(t) dA(t) dt = ∂f (z(t), t, θ) ∂z(t) . (79)
Then from

B(t)A(t) = I, (80)
we have

-B(t) dA(t) dt = dB(t) dt A(t), (81)
thus, using this expression in Equation ( 79) we get

dB(t) dt A(t) = ∂f (z(t), t, θ) ∂z(t) . (82)
Finally, we get the desired result by multiplying both sides from the right with B(t) .

H Equivalence Between Continuous Total Derivative Decomposition and the Continuous Backpropagation

In Appendix C. we derived the expression of continuous backpropagation from the continuous total derivative decomposition. However, we can also derive the the formula of the continuous total derivative decomposition (Equation ( 53)) from continuous backpropagation (Equation ( 11), Appendix A). Indeed, for a function f = f (z(t), θ(t), t), where θ(t) = g(θ, t), we can see f as h(z(t), θ, t) = f (z(t), θ(t), t). Hence,

∂L ∂θ = T 0 ∂L ∂z(t) dh(z(t), θ, t) dθ dt = T 0 ∂L ∂z(t) df (z(t), θ(t), t) dθ dt, (83)
therefore

∂L ∂θ = T 0 ∂L ∂z(t) ∂f (z(t), θ(t), t) ∂θ(t) dθ(t) dθ dt. (84)
Setting L = z(T ), gives the desired result.

I Additional Details About the Experiments

In the case of the 2D data, we use the implementation of CNFs provided in [START_REF] Chen | Neural ordinary differential equations[END_REF], since using the Hutchinson's trace estimator and GPUs as in FFJORD provides no computational benefits in low dimensions. In this case we also use the non-adaptive Runge-Kutta 4 ODE solver, while for image data we use Dopri5, as well as the FFJORD implementation of [START_REF] Grathwohl | Ffjord: Freeform continuous dynamics for scalable reversible generative models[END_REF]. We use a batch size of 200 for image data and a batch size of 512 for the toy datasets. In the case of image data, we use a learning rate of 6 × 10 -4 , while for toy data the learning rate is set to 10 -3 . All experiments were performed on a single GPU.

As mentioned in the main paper, we optimized the architecture of FFJORD first and tuned its hyperparameters. This architecture remains unchanged in the AFFJORD model for fair comparison, and we only finetuned the augmented structure and dimension in addition. Indeed, as it can be seen, the results we report for FFJORD (0.96 MNIST, 3.37 CIFAR) are better than those reported on the original paper (0.99 MNIST, 3.40 CIFAR). The aforementioned improvements were a result of reducing the number of parameters during our tuning process, by reducing the number of CNF blocks from 2 to 1. The total number of parameters for different architectures of FFJORD and AFFJORD can be found on Table 2. It should be emphasized that in the hypernet architecture, the parameters of AFFJORD are contained inside the parameters in the main architecture of FFJORD, so a bigger model is not being used, simply the flexibility of evolution of the main parameters in time is being increased. More generally, our experiments show that augmented architectures generally improve the performance of the standard non-augmented FFJORD counterparts, independently from the baseline architecture used. Furthermore, the farther the performance of FFJORD is from being optimal, the greater are the improvements when using AFFJORD.

I.1 Multiscale Architecture in Augmented Neural ODE Flows

With reference to Figure 10, the multiscale architecture in the augmented case performs an Augmented Continuous Flow on the data as described in Section 2, then squeezes the channels (the augmented as well as the original channels) as in the original multiscale architecture. However after the second transformation the augmented channels are removed and stored in a separated array (Array A). The data channels are treated as before, that is, half of them are saved (Array B), and the other half are squeezed again. After this process is finished we add new augmented channels, to repeat the cycle. Note that in order to generate data by the inverse transformation, we need to retrieve the transformed augmented dimensions previously stored (i.e., Array A). The original multiscale architecture is presented in Appendix A for comparison. As we mentioned earlier, AFFJORD retains the ability to generate samples from the learnt distribution, by simply integrating in the opposite direction. Indeed, for a given sample z s (T ) from the base distribution, we augment it to z * (T ), as z * (T ) is the same for all data points. Then, we can simply integrate backwards this concatenated vector [z s (T ), z * (T )] to [z s (0), z * (0)], drop the generated augmented dimensions z * (0), and simply keep the generated data point z s (0). Due to the use of the augmented multiscale architecture, where for each cycle we replace the augmented dimensions, these replaced augmented dimensions must be saved in an array for the backward generative pass. This is why in Figure 10 at the end of each scale i we save the augmented dimensions z * (T ) i . Data Dependent Augmented Dimensions As described in Section 2, several simplifications are made in the architecture of the general augmented neural ODE flows, in order to ensure immediate bijectivity and reduce the computational complexity. However, other possible architectures exist, where for example z * (0) is dependent on z(0), and g(z * (t)) = -z * (t). This would ensure that the augmented dimensions converge to zero, providing bijectivity. Since all augmented dimensions would be different during training, this would imply that the data is lifted to a higher plane, enabling richer transformations. Finding an approximation of the loss in Equation ( 5) remains a challenge for the future however.

Jacobian Regularization Theoretically speaking, all performance enhancing modifications that can be applied to FFJORD are also applicable to AFFJORD. Such a modification which reduces the training time of FFJORD is presented in [START_REF] Finlay | How to train your neural ode[END_REF], where both the vector field and its Jacobian are regularized. Thus, an interesting research direction in the future would be to test how the performance of AFFJORD is affected by such amendments.

Figure 1 :

 1 Figure 1: Probability density modeling capabilities of AFFJORD (second column) and FFJORD (third column) on 2D data of toy distributions.

Figure 2 :

 2 Figure 2: Example of the multiscale architecture on MNIST dataset

Figure 3 :

 3 Figure 3: Architecture of discretized augmented neural ODEs.

Figure 4 :

 4 Figure 4: The discretized augmented neural ODE flows implemented in AFFJORD.

Figure 5 :Figure 7 :Figure 8 :

 578 Figure 5: Rectangle (i, j) in discretized [0, T ] × [0, T ] Figure 6: Symmetry of g n 2

  2

  Alleviates the Issues Listed in Section 2

  and = det e Ω [z] (t) = e T 0 tr ∂f (z(t),z * (t),θ) ∂z(t) dt . In case that the base distribution is multivariate normal, from -log p(z(0)) = -log p(z(T )) det dz(T ) dz(0) = -log p(z(T )) det e T 0 ∂f (z(t),z * (t),θ) ∂z(t) dt+...

Figure 10 :

 10 Figure 10: Example of the augmented multiscale architecture on MNIST dataset

I. 3

 3 Validation Results During Training Additional results illustrating the training and validation loss through-out the training procedure are provided in Figures 11, 12 and 13.

Figure 11 :

 11 Figure 11: Performance of FFJORD and AFFJORD on the Hash-Gaussian toy 2D dataset.

Figure 12 :

 12 Figure 12: Performance of FFJORD and AF-FJORD on the TY toy 2D dataset.

Figure 13 :

 13 Figure 13: Graphs of training and evaluation losses of FFJORD and AFFJORD on CIFAR-10 as well as CelebA (32 × 32). We run the experiments 5 times. Lower is better.

Figure 20 :

 20 Figure 20: Number of function evaluations of FFJORD and AFFJORD when the 'concatenate' architecture is used.

Table 1 :

 1 Experimental Results for Density Estimation Models, in Bits/Dim for MNIST, CIFAR-10 and CelebA (32 × 32), and in NLL for the 2D data. Lower Is Better. The Multiscale Architecture Is Used in All Image Data Cases. AFFJORD 0.95 ± .01 3.32 ± .01 3.23 ± .00 1.035 ± .008 0.636 ± .02

	Model	MNIST	CIFAR10	CelebA	HG (2D)	TY (2D)
	Real NVP 1.06	3.49	-	-	-
	Glow	1.05	3.35	-	-	-
	FFJORD	0.96 ±.00	3.37 ±.00	3.28±.00	1.052 ±.007	0.674 ±.017

Table 2 :

 2 Number of Total Parameters of FFJORD and AFFJORD (Concat and Hypernet Architecture) for MNIST, CIFAR-10 and CelebA. The Multiscale Architecture Is Used in All Cases.

	Dataset	Concat FFJORD AFFJORD
	MNIST	400323	417629
	CIFAR10	679441	820815
	CelebA (32 × 32) 679441	820815
	Dataset	Hypernet FFJORD AFFJORD
	MNIST	783019	4556169
	CIFAR10	1332361 8976115
	CelebA (32 × 32) 1332361 8976115

Acknowledgments and Disclosure of Funding

This work has been supported by the French government, through the 3IA Côte d'Azur Investments in the Future project managed by the National Research Agency (ANR) with the reference number ANR-19-P3IA-0002. The authors are grateful to the OPAL infrastructure from Université Côte d'Azur for providing resources and support.

J Generated Samples

Below in Figure 14, Figure 15 Interestingly, if the concatenation architecture is used in AFFJORD, that is, we only replace the concatenated time channel in FFJORD with the augmented channel in AFFJORD, then the number of function evaluations decreases significantly. This aspect is illustrated in Figure 20. However, the hypernet architecture of AFFJORD is prone to the issue of the number of function evaluations. Indeed, forward passes in AFFJORD-hypernet can be challenging, as the learnt vector field becomes too stiff with an increasing number of integration steps.

Addition of Self-Attention [START_REF] Ho | Flow++: Improving flow-based generative models with variational dequantization and architecture design[END_REF] describe three modeling inefficiencies in prior work on flow models. One such factor is the lack of expressive power of convolutional layers used in normalizing flow models. Considering the performance improvements demonstrated in [START_REF] Ho | Flow++: Improving flow-based generative models with variational dequantization and architecture design[END_REF], in the future we intend to test the improvement in performance brought by the addition of self-attention in AFFJORD.


Mai-Linh Doan
email: mai-linh.doan@univ-grenoble-alpes.fr

Baumann, T.S Michael J Heap

Maude Schell

Ulrich Kueppers

M Komorowski

High strain rate damage in porous andesite

Keywords: High strain rate damage, Volcán de Colima, andesite, X-ray computed tomography, P-wave velocity, Split Hopkinson pressure bar

Deformation rates at volcanoes vary enormously, potentially altering the way in which volcanic rock accommodates strain. However, relationships between deformation rate, damage patterns, and physical properties are poorly understood. Laboratory deformation experiments are typically restricted to low strain rates (≪ 1 s -1 ). Here, we deformed samples of porous andesite in compression at strain rates from 10 -6 to ∼ 370 s -1 . We show, using Xray computed microtomography, that the failure mode changes as a function of strain rate: macroscopic fractures form at low strain rates and samples deformed at high strain rate (≥ 100 s -1 ) contain macroscopic fractures and collapsed pores. In general, high strain rate deformation results in more pervasively damaged samples. Deformation at high strain rate also results in larger decreases and increases to P-wave velocity and permeability, respectively. Strikingly, the change in the P-wave velocity of samples deformed at high strain rate is very large (decreases of up to 50%). We hypothesize

Introduction

1 Deformation rates at volcanoes vary tremendously, from 10 -14 -10 -12 s -1 2 [START_REF] Owen | Rapid deformation of 443 the south flank of Kilauea volcano[END_REF]Wadge et al., 2006;Moretti et al., 2020) to strain rates 3 sufficient to fragment melt (Dingwell, 1996;Wadsworth et al., 2018) and 4 form frictional melts [START_REF] Kendrick | Extreme frictional processes in the volcanic conduit of mount 401 st. helens (usa) during the 2004-2008 eruption[END_REF][START_REF] Hughes | Shear localisation, strain partitioning and frictional 397 melting in a debris avalanche generated by volcanic flank collapse[END_REF]. However, 5 the majority of laboratory deformation experiments on volcanic rocks have 6 been performed at low compressive strain rates of 10 -5 s -1 (Heap and Violay, 7 2021). Few experimental studies have focused on the strain rate dependence 8 of mechanical behavior at higher strain rates up to 10 -1 s -1 (Schaefer et al., 9 2015;Coats et al., 2018;[START_REF] Lavallée | The fragility of volcán de colima-a material constraint, in: Varley, 409 tion[END_REF][START_REF] Heap | The mechanical behaviour and failure modes 393 of volcanic rocks: a review[END_REF]. These 10 studies have shown that the strength of volcanic rocks increases as the strain 11 rate is increased. For example, the strength of porous dacite from Mt. Unzen 12 (Japan) increased from ∼20 to ∼28 MPa as the strain rate was increased from 13 10 -5 to 10 -1 s -1 (Coats et al., 2018).

14

Experiments on granite, marble, and porous sandstone have shown that 15 than forming the few macroscopic fractures observed at low strain rate (≤ 17 10 -1 s -1 ), very high strain rates (≥ 100 s -1 ) pulverize low-porosity granite 18 (Doan and Gary, 2009) and marble (Doan and Billi, 2011) and compact 19 porous sandstone (Aben et al., 2017b).

20

It is not clear, however, whether the failure mode of volcanic rock is 21 similarly changed at high strain rates. For example, while [START_REF] Olsson | The compressive strength of tuff as a function of strain 439 rate from 10 -6 to 10 3 /s[END_REF] 22 remarked that the porous tuff samples were "reduced to dust" at strain rates 23 of 10 3 s -1 , the failure of the low-porosity basalt was found to be largely 24 independent of strain rate (Lindholm et al., 1974). Indeed, there are very few 25 studies that have investigated the influence of high strain rates on the damage 26 evolution and failure mode of volcanic rocks. Further, while studies exist that show that low strain rate deformation in the brittle regime (≤ 10 -1 s -1 ) 28 can decrease P-wave velocities [START_REF] Stanchits | Ultrasonic velocities, acous-453 tic emission characteristics and crack damage of basalt and granite[END_REF]Fortin et al., 2011) 29 and increase permeability of volcanic rock (Fortin et al., 2011;Farquharson 30 et al., 2016), corresponding data for high strain rates-important for the 31 monitoring and understanding of volcanic systems-do not exist.

32

Here, we present the results of an experimental study on the influence of 33 strain rate on the mechanical behavior of porous andesite. Uniaxial defor-

σ i = A b A s E (ε i + ε r ) (1) σ o = A b A s E ε t (2) ∂ε ∂t = c L (ε i -ε r -ε t ) (3)
where E, c are the Young modulus and the elastic bar wave speed of the bar during loading, which can be altered for strain rates or samples that are too 84 high or too long, respectively. Once this quality control has been performed,

85

the output stress is chosen as the stress acting on the sample during the test.

86

Strain is computed by time integrating the strain rate. Full stress-strain 87 curves can then be reliably retrieved (gray curves on Fig. 2).

88

After the first loading, the stress waves reflect perfectly at the edges of Damage was also assessed using petrophysical data. The P-wave velocity 107 and permeability of select samples (deformed at low and high strain rates)

108 was measured at ITES before and after deformation using a digital oscillo- Violay, 2021). There are no qualitative differences between the stress-strain 126 curves at strain rates between 10 -6 to 10 (axial strains up to 3%, compared to 0.5-1% for the quasi-static tests, see 130 Fig. 5b). strain rates (Fig. 3).

142

For a similar initial porosity (20-22%), although the strength at low-143 strain-rate (blue squares in Fig. 3) is lower than the strength at high strain 144 rate (colored circles in Fig. 3), the influence of strain rate on strength plays 145 a second order role compared to porosity. An increase of strength as a func- The damage classification for each sample is presented in Fig. 5a, which 155

shows that the type of damage depends on the strain rate. All of the quasi-156 statically deformed samples (10 -6 and 10 -3 s -1 ) formed a macroscopic frac-157 ture, a failure mode that characterizes brittle deformation. However, ap-158 parent compaction occurred in addition to fracturing at strain rates above 159 120 s -1 and, at the highest strain rates, two of the samples macroscopi- dynamically to an axial strain of 0.6% appears to be localized on a plane sub-178 perpendicular to the maximum principal stress, similar to the compaction 179 bands seen previously to develop in porous volcanic rocks (Heap et al., 2015, 180 2020). The damage accumulated in the samples deformed dynamically to 181 higher axial strains (1.1 and 2.6%) is similar to that described for the quasi-182 static sample to an axial strain of 0.6%, but much more pervasive (Fig. 4).

183

Samples characterised by both fracturing and compaction have a higher 184 strength than those that simply fractured (Fig. 3). This can be explained does not depend on the initial sample porosity (Fig. 7a), but increases as a

197 function of increasing strain rate (Fig. 7b). P-wave velocity decrease is ≤ 198 16 10% for the low strain samples (Fig. 7c), which also dissipated more energy 199 (Figs. 6 and7d). But the P-wave velocity decrease can exceed 50% at high 200 strain (Fig. 7). It should be noted that such strains could only be achieved 201 for the damage pattern "compacted/fractured", which is the most diffuse 202 and occurs only for the dynamic tests (Fig. 5).

203

The pre-and post-deformation permeability was also measured for sam-204 ples deformed at a strain rates of 10 -4 , 49, and 171 s -1 , respectively, and

205
showed a positive correlation of permeability increase with strain rate (Tab. 206 1). The increase in permeability following deformation is therefore higher

207

(increase of about a factor of three) for the dynamic tests than for the quasi-208 static tests (increase of about 40%). Our data show samples accommodate strain differently at strain rates 212 above 100 -150 s -1 , from "fractured" to "fractured with apparent com-213 paction" (Fig. 4). The occurrence of compaction bands-planes of collapsed 214 pores connected by microcracks that formed sub-perpendicular to the max-215 imum principal stress-has been observed during high-pressure quasi-static 216 deformation experiments on porous andesites (Heap et al., 2015(Heap et al., , 2020)). Typ- showed that compaction was also observed in porous sandstone deformed Fortin et al., 2011;Heap et al., 2015;238 Farquharson et al., 2016), and that cataclastic pore collapse decreases the 239 permeability of volcanic rock (Heap et al., 2015;Farquharson et al., 2017; 240 Heap et al., 2020). Our data show, however, that the permeability of samples 241 deformed at high strain rates that are "fractured with apparent compaction"

242 increased by about a factor of three (Tab. 1). In other words, the observed 243 compaction (pores infilled with fragments and possible compaction bands;

244 Fig. 4) did not result in a decrease in permeability and is likely the result 245 of presence of axially orientated fractures (Fig. 4), not present in samples 246 deformed quasi-statically at high pressure (Heap et al., 2015;Farquharson 247 increasing strain rate is likely the result of the more pervasive damage in the 249 samples deformed at high strain rates (Fig. 4).

250

A decrease in P-wave velocity following quasi-static loading to failure (Fig. 251 7) has been observed previously in volcanic rocks and is due to the formation 252 of dilatant microcracks [START_REF] Stanchits | Ultrasonic velocities, acous-453 tic emission characteristics and crack damage of basalt and granite[END_REF]Fortin et al., 2011) The decrease in P-wave velocities after high-strain-rate loading is spec-266 tacular, especially considering they were obtained on laboratory samples, 267 which tend to have higher P-wave velocities than those observed in the field 268 (Lesage et al., 2018b). Although our porous andesite samples were not pul-269 verized at high strain rates (as was the case for low-porosity granites; Doan 270 and Gary ( 2009)), the significant decrease in P-wave velocity (Fig. 7) and 271 the pervasively damaged and fragmented nature of the samples deformed at 272 rated into the gas-particle jet of a subsequent high-energy explosive eruption.

301

Finally, increasing the permeability (Tab. 1) and the surface area available 302 for fluid-rock interactions (Fig. 4) could encourage efficient hydrothermal al-303 teration, a process that could also further reduce the stability of the volcanic 304 flank or lava dome [START_REF] Heap | The mechanical behaviour and failure modes 393 of volcanic rocks: a review[END_REF].

Figure 1 :

 1 Figure 1: Split Hopkinson Pressure bars (SHPB) at ISTerre (Grenoble, France). (a) Schematics of the SHPB device and relation for the samples. (b) The four strain gauges glued along the bars (at 120 cm and 74.4 cm from the right end of the input bar, and at 29.7 cm from the left end of the output bar) record the strain wave. The wave propagation can be also modeled using 1D stress propagation. Sample is at location 0.

89

  both bars: the bars have free edges and the sample is not reloaded. The 90 stress-strain curves of the first loading reflects the whole loading history 91 on the sample. By varying the striker length, the striker speed, the pulse 92 shaper, and the material of the bars, a range of strain rates and stresses can 93 be achieved.

Figure 2 :

 2 Figure 2: Strain-stress curves for andesite deformed quasi-statically (10 -6 to 10 -3 s -1 ; black curves) and dynamically (up to ∼ 370 s -1 ; gray curves).

109

  scope and a waveform pulse generator (see schematic provided in Heap et al. 110 (2014)) and a benchtop nitrogen permeameter (see schematic provided in 111 Heap and Kennedy (2016)), respectively. Both P-wave velocity and perme-112 ability were measured parallel to the sample axis and at ambient laboratory 113 temperatures. Volumetric flow rates were measured for different pressure 114 differentials to calculate permeability using Darcy's law and to check for 115 Klinkenberg and Forchheimer corrections. Due to the high permeability of 116 the studied andesite, a Forchheimer correction was applied in all cases. P-117 wave velocity was measured at ambient pressure and permeability was mea-Mechanical behavior at different strain rates 121 A total of 35 experiments were performed: 15 quasi-static and 20 dynamic 122 experiments. The stress-strain curves for all experiments are provided in Fig.1232. The stress-strain curves for the quasi-static tests (black curves) are similar 124 to those for porous volcanic rock deformed in uniaxial compression(Heap and 125

Figure 3 :

 3 Figure 3: Influence of porosity on the peak stress of porous andesite deformed over a wide range of strain rate (10 -6 to 10 -3 s -1 ).

131Fig. 3

 3 Fig.3shows the peak stress as a function of porosity for all samples (if

160

  Figure 4: X-ray computed tomographic slices (black and gray represent the porosity and rock, respectively) showing an intact sample (a) of the studied andesite and the failure modes at high (b-d) and low (e) strain rate.

Figure 5 :

 5 Figure5: Influence of strain and strain-rate (panels a and b) on the peak stress of porous andesite deformed over a wide range of strain rate (10 -6 to 10 -3 s -1 ). The damage classification for the samples tested at high strain rate was determined in two phases (visual macroscopic inspection and optional µCT imaging). All samples tested at low strain rate were fractured. Gray shade delimits the strain rate transition between fractured samples and compacted/fractured samples.

185Figure 6 :

 6 Figure 6: Increase of dissipated energy during loading as a function of strain. The samples that experience compaction/fracturing dissipate more energy than the samples who just endured fracturing.

  Change in failure mode at higher strain rates 211

Figure 7 :

 7 Figure 7: Relative change in P wave velocity as a function of (a) porosity, (b) strain rate, (c) strain, (d) dissipated energy.

224dynamically.

  These authors provided modeling to explain how the competi-225 tion between intragranular fracturing (i.e. grain crushing) and intergranular 226 fracture is resolved by kinetic effects. The stress intensity factor increases 227 as the axial stress increases, and fracturing can occur when the stress in-228tensity factor is grea ation process for intragranular processes, but at higher 229 stresses. Intergranular fracturing (i.e. macroscopic fracture formation) is 230 therefore favored thermodynamically but is inhibited kinetically, explaining 231 why compaction can occur at higher strain rates under uniaxial conditions 232(Aben et al., 2017b).

  shown that changes to permeability and P-wave velocity 235 following macroscopic failure are larger at higher strain rates (Tab. 1; Fig. 2367). Previous studies have shown that brittle failure in compression increases 237 the permeability of volcanic rock (

  acknowledges support from the Institut Universitaire de France 307 (IUF). T. Reuschlé and B. Renaudié are thanked for experimental assistance.

308N.

  Varley and J. Farquharson are thanked for their help in collecting the 309 experimental material.

310

  coseismic loadings leads to pulverization in active fault zones. Journal 322 of Geophysical Research: Solid Earth 121, 2338-2360. doi:Earth, 121,

Table 1 :

 1 Summary of the experimental results. "Compacted/Fractured" damage is labeled as "C/F". Some data were not recorded and the corresponding cells are left blank.

			Dimensions (mm)	Porosity	Vp (m/s)	Permeability (10 -13 m 2 )	Max Strain rate	Max strain	Max stress	Damage
		Sample
			Diameter	Height	(%)	Pre	Post	Pre	Post	(/s)	(%)	(MPa)	Macroscopic	X-ray CT scan
		CLB1	20.06	20.47	18.23	3529	1712			196	1.13	42.6	C/F	C/F
		CLB2	20.09	20.21	18.23	4150	3853			58	0.38	31.5	Fractured
		CLB3	20.07	17.63	19.22	3236	3181			121	0..5	16.7	Intact
		CLB4	20.09	18.89	22.12	3072	1660			183	1.25	32.4	C/F	C/F
		CLB5	20.1	17.7	22.12	3106	1846			277	2.02	32.9	C/F	C/F
		CLB6	20.08	20.74	22.36	2861	2816			138	0.59	15.2	Intact
		CLB7	20.07	19.72	21.31	2594	1814			259	1.58	37.7	C/F	C/F
		CLB8	20.09	19.54	17.76	2832				224	1.59	52.5	Fragmented
		CLB9	20.08	18.61	20.49	2978	1505			369	3.23	30	C/F	C/F
	SHPB	CLB10 CLB11	20.11 20.1	20 20.53	18.86 22.01	3252 2793	2967 1251			201 213	0.78 2.56	36.8 32.8	C/F C/F	C/F C/F
		CLB12	20.12	18.78	21.2	2867	2834			57	0.51	25.1	Intact	Fractured
		CLB13	20.09	18.88	19.76	3120	2648			190	1.95	34.8	C/F	C/F
		CLB14	20.1	19.46	19.22	3050	2039			259	3	40.6	C/F
		CLB15	20.1	20.32	20.96	3040	2239	9.31	27.1	171	1.63	34.1	C/F	C/F
		CLB16	20.05	21.6	22.99	2980				233	3.5	22.4	Fragmented
		CLB17	20.11	20.12	17.37	2980	2668	2.76	6.22	153	0.79	38.6	Intact	C/F
		CLB18	20.13	18.34	18.91	3160	2810			67	0.58	30.7	Fractured
		CLB20	20.12	19.16	21.99	3250	2798			125	1.38	24	C/F
		CLB21	20.14	18.67	21.3	3120	2110			139	1.79	26	C/F
		LLB12	20.02	39.92	20.9			4.86	6.76	10 -4	0.74	26.6	Fractured	Fractured
		LLB15	20	39.75	24.15	2697	2583			10 -5	0.54	20.3	Fractured
		LLB18	20	40.01	21.37					10 -4	0.81	29	Fractured
		LLB19	20	39.99	21.07					10 -3	1.07	32.6	Fractured
		LLB21	20.02	39.87	20.2	2820	2509			10 -5	0.49	31.9	Fractured
		LLB22	19.99	39.68	21.65					10 -6	0.48	24.1	Fractured	Fractured
	Slow tests	LLB31 LLB34 LLB42	19.99 20.02 19.85	39.99 40.01 39.94	21.93 22.23 25.53	2656	2351			10 -6 10 -3 10 -5	0.46 0.96 0.44	22.9 28.3 18.3	Fractured Fractured Fractured	Fractured
		LLB44	20	40	20.87					10 -5	0.52	24.8	Fractured
		LLB46	20	40	21.72					10 -5	0.51	22.5	Fractured
		LLB48	19.88	39.95	26.76	2674	2579			10 -5	0.37	13.7	Fractured
		LLB51	20	40	21.97					10 -5	0.51	20.6	Fractured
		LLB52	20	40	22.4					10 -5	0.54	22.5	Fractured
		LLB53	19.86	39.97	25.38	2807	2531			10 -5	0.39	15	Fractured

-3 s -1 . The stress-strain curves for 127 the dynamic tests are similar to those for the quasi-static tests, although 128 the samples deformed dynamically were typically deformed to larger strains 129

  . The very

	253
	254	large decrease in P-wave velocity (a decrease of up to 50%) following dy-
	255	namic loading to failure, not observed previously, is likely consequence of the
	256	pervasively damaged and fragmented nature of the samples deformed at high
		strain rate (Fig. 4). A decrease of 50% is higher than previously recorded

257

reductions in P-wave velocity associated with high strain rate damage. For 258 crystalline rocks, Aben et al. (2015) recorded P-wave velocity reductions up 259 to 30% for the diffuse fracturing associated with high strain rate damage. 260 Rempe et al. (2013) showed through subsurface seismic studies that veloc-261 ities in pulverized zones of the San Andreas Fault Zone (USA) can be as 262 low as 1 km/s, a reduction to 20% of the expected values for the granitic 263 protolith. 264 4.3. Volcanological implications of intense high strain rate damage 265

in a permanent decohesion that could severely weaken the rock units of the 274 volcanic edifice or the magma, both inside the conduit and above (e.g., the 275 lava dome).

276 Field measurements at Volcán de Colima show that the center of the 277 volcano exhibits similar decrease in seismic velocity at the kilometric scale, 278 extending down to 30 km (Escudero and Bandy, 2017). A decrease in P-279 wave velocity is sometimes attributed to an increase in temperature, due