Year: 2020

Abstract:
Many types of neural network layers rely on matrix properties such as invertibility or orthogonality. Retaining such properties during optimization with gradientbased stochastic optimizers is a challenging task, which is usually addressed by either reparameterization of the affected parameters or by directly optimizing on the manifold. This work presents a novel approach for training invertible linear layers. In lieu of directly optimizing the network parameters, we train rank-one perturbations and add them to the actual weight matrices infrequently. This PInv update allows keeping track of inverses and determinants without ever explicitly computing them. We show how such invertible blocks improve the mixing and thus the mode separation of the resulting normalizing flows. Furthermore, we outline how the P concept can be utilized to retain properties other than invertibility.

text: 1 INTRODUCTION
Many deep learning applications depend critically on the neural network parameters having a certain mathematical structure. As an important example, reversible generative models rely on invertibility and, in the case of normalizing flows, efficient inversion and computation of the Jacobian determinant (Papamakarios et al., 2019).
Preserving parameter properties during training can be challenging and many approaches are currently in use. The most basic way of incorporating constraints is by network design. Many examples could be listed, like defining convolutional layers to obtain equivariances, constraining network outputs to certain intervals through bounded activation functions, Householder flows (Tomczak & Welling, 2016) to enforce layer-wise orthogonality, or coupling layers (Dinh et al., 2014; 2016) that enforce tractable inversion through their twochannel structure. A second approach concerns the optimizers used for training. Optimization routines have been tailored for example to maintain Lipschitz bounds (Yoshida & Miyato, 2017) or efficiently optimize orthogonal linear layers (Choromanski et al., 2020).
The present work introduces a novel algorithmic concept for training invertible linear layers and facilitate tractable inversion and determinant computation, see Figure 1. In lieu of directly changing the network parameters, the optimizer operates on perturbations to these parameters. The actual network parameters are frozen, while a parameterized perturbation (a rank-one update to the frozen parameters) serves as a proxy for optimization. Inputs are passed through the perturbed network
during training. In regular intervals, the perturbed parameters are merged into the actual network and the perturbation is reset to the identity. This stepwise optimization approach will be referred to as property-preserving parameter perturbation, or P4 update. A similar concept was introduced recently by Lezcano-Casado (2019), who used dynamic trivializations for optimization on manifolds.
In this work, we use P4 training to optimize invertible linear layers while keeping track of their inverses and determinants using rank-one updates. Previous work (see Section 2) has mostly focused on optimizing orthogonal matrices, which can be trivially inverted and have unity determinant. Only most recently, Gresele et al. (2020) presented a first method to optimize general invertible matrices implicitly using relative gradients, thereby providing greater flexibility and expressivity. While their scheme implicitly tracks the weight matrices’ determinants, it does not facilitate cheap inversion. In contrast, the present P4Inv layers are inverted at the cost of roughly three matrix-vector multiplications.
P4Inv layers can approximate arbitrary invertible matrices A ∈ GL(n). Interestingly, our stepwise perturbation even allows sign changes in the determinants and recovers the correct inverse after emerging from the ill-conditioned regime. Furthermore, it avoids any explicit computations of inverses or determinants. All operations occurring in optimization steps have complexity of at most O(n2). To our knowledge, the present method is the first to feature these desirable properties. We show how P4Inv blocks can be utilized in normalizing flows by combining them with nonlinear, bijective activation functions and with coupling layers. The resulting neural networks are validated for density estimation and as deep generative models. Finally, we outline potential applications of P4 training to network properties other than invertibility.

heading: 2 BACKGROUND AND RELATED WORK

heading: 2.1 RANK-ONE PERTURBATION
text: The P4Inv layers are based on rank-one updates, which are defined as transformations A 7→ A + uvT with u,v ∈ Rn. If A ∈ GL(n) and 1 + vTA−1u 6= 0, the updated matrix is also invertible and its inverse can be computed by the Sherman-Morrison formula
(A+ uvT )−1 = A−1 − 1 1 + vTA−1u A−1uvTA−1. (1)
Furthermore, the determinant is given by the matrix determinant lemma
det(A+ uvT ) = (1 + vTA−1u) det(A). (2)
Both these equations are widely used in numerical mathematics, since they sidestep the O(n3) cost and poor parallelization of both matrix inversion and determinant computation. The present work leverages these perturbation formulas to keep track of the inverses and determinants of weight matrices during training of invertible neural networks.

heading: 2.2 EXISTING APPROACHES FOR TRAINING INVERTIBLE LINEAR LAYERS
text: Maintaining invertibility of linear layers has been studied in the context of convolution operators (Kingma & Dhariwal, 2018; Karami et al., 2019; Hoogeboom et al., 2019; 2020) and using Sylvester’s theorem (Van Den Berg et al., 2018). Those approaches often involve decompositions that include triangular matrices (Papamakarios et al. (2019)). While inverting triangular matrices has quadratic computational complexity, it is inherently sequential and thus fairly inefficient on parallel computers (see Section 4.1). More closely related to our work, Gresele et al. (2020) introduced a relative gradient optimization scheme for invertible matrices. In contrast to this related work, our method facilitates a cheap inverse pass and allows sign changes in the determinant. On the contrary, their method operates in a higher-dimensional search space, which might speed up the optimization in tasks that do not involve inversion during training.

heading: 2.3 NORMALIZING FLOWS
text: Cheap inversion and determinant computation are specifically important in the context of normalizing flows, see Appendix C. They were introduced in Tabak et al. (2010); Tabak & Turner (2013)
and are commonly used, either in variational inference (Rezende & Mohamed, 2015; Tomczak & Welling, 2016; Louizos & Welling, 2017; Van Den Berg et al., 2018) or for approximate sampling from distributions given by an energy function (van den Oord et al., 2018; Müller et al., 2019; Noé et al., 2019; Köhler et al., 2020). The most important normalizing flow architectures are coupling layers (Dinh et al., 2014; 2016; Kingma & Dhariwal, 2018; Müller et al., 2019), which are a subclass of autoregressive flows (Germain et al., 2015; Papamakarios et al., 2017; Huang et al., 2018; De Cao et al., 2019), and (2) residual flows (Chen et al., 2018; Zhang et al., 2018; Grathwohl et al., 2018; Behrmann et al., 2019; Chen et al., 2019). A comprehensive survey can be found in Papamakarios et al. (2019).

heading: 2.4 OPTIMIZATION UNDER CONSTRAINTS AND DYNAMIC TRIVIALIZATIONS
text: Constrained matrices can be optimized using Riemannian gradient descent on the manifold (Absil et al. (2009)). A reparameterization trick for general Lie groups has been introduced in Falorsi et al. (2019). For the unitary / orthogonal group there are multiple more specialized approaches, including using the Cayley transform (Helfrich et al., 2018), Householder Reflections (Mhammedi et al., 2017; Meng et al., 2020; Tomczak & Welling, 2016), Givens rotations (Shalit & Chechik, 2014; Pevny et al., 2020) or the exponential map (Lezcano-Casado & Martı́nez-Rubio, 2019; Golinski et al., 2019).
Lezcano-Casado (2019) introduced the concept of dynamic trivializations. This method performs training on manifolds by combining ideas from Riemannian gradient descent and trivializations (parameterizations of the manifold via an unconstrained space). Dynamic trivializations were derived in the general settings of Riemannian exponential maps and Lie groups. Convergence results were recently proven in follow-up work (Lezcano-Casado (2020)). P4 training resembles dynamic trivializations in that both perform a number of iteration steps in a fixed basis and infrequently lift the optimization problem to a new basis. In contrast, the rank-one updates do not strictly parameterize GL(n) but instead can access all of Rn×n. This introduces the need for numerical stabilization, but enables efficient computation of the inverse and determinant through equation 1 and equation 2, which is the method’s unique and most important aspect.

heading: 3 P4 UPDATES: PRESERVING PROPERTIES THROUGH PERTURBATIONS

heading: 3.1 GENERAL CONCEPT
text: A deep neural network is a parameterized function MA : Rn → Rm with a high-dimensional parameter tensor A. Now, let S define the subset of feasible parameter tensors so that the network satisfies a certain desirable property. In many situations, generating elements of S from scratch is much harder than transforming any A ∈ S into other elements A′ ∈ S, i.e. to move within S. The efficiency of perturbative updates can be leveraged as an incremental approach to retain certain desirable properties of the network parameters during training. Rather than optimizing the parameter tensors directly, we instead use a transformation RB : S → S, which we call a property-preserving parameter perturbation (P4). A P4 transforms a given parameter tensor A ∈ S into another tensor with the desired property A′ ∈ S. The P4 itself is also parameterized, by a tensor B. We demand that the identity idS : A 7→ A be included in the set of these transformations, i.e. there exists a B0 such that RB0 = idS.
During training, the network is evaluated using the perturbed parameters Ã = RB(A). The parameter tensor of the perturbation, B, is trainable via gradient-based stochastic optimizers, while the actual parameters A are frozen. In regular intervals, every N iterations of the optimizer, the optimized parameters of the P4, B, are merged into A as follows:
Anew ← RB(A), (3) Bnew ← B0. (4)
This update does not modify the effective (perturbed) parameters of the network Ã, since
Ãnew = RBnew(Anew) = RB0(RB(A)) = RB(A) = Ã.
Hence, this procedure enables a steady, iterative transformation of the effective network parameters and stochastic gradient descent methods can be used for training without major modifications. Furthermore, given a reasonable P4, the iterative update of A can produce increasingly non-trivial transformations, thereby enabling high expressivity of the resulting neural networks. This concept is summarized in Algorithm 1. Further extensions to stabilize the merging step will be exemplified in Section 3.3.
Algorithm 1: P4 Training Input : Model M , training data, loss function J , number of optimization steps Nsteps, merge interval N , perturbation R, optimizer OPT initialize A ∈ S; initialize B := B0; for i := 1 . . . Nsteps do
X,Y0 := i-th batch from training data; Ã := RB(A) ; // perturb parameters Y := MÃ(X) ; // evaluate perturbed model j := J(Y ,Y0) ; // evaluate loss function gradient := ∂j/∂B ; // backpropagation B := OPT(B, gradient) ; // optimization step if i mod N = 0 then
A := RB(A) ; // merging step: update frozen parameters B := B0 ; // merging step: reset perturbation
end end
3.2 P4INV: INVERTIBLE LAYERS VIA RANK-ONE UPDATES
Algorithm 2: P4Inv Merging Step Input : Matrix A, Inverse Ainv, Determinant d det factor := (1 + vTAinvu) new det := det factor · d; if ln |det factor| and ln |new det| are sane then
/* update frozen parameters (equation 3) */ d := new det; A := Ru,v(A); Ainv := Ainv − 11+vTAinvuAinvuv
TAinv; /* reset perturbation (equation 4) */ u := 0; v := N (0, In) ; // random reinitialization
end
The P4 algorithm can in principle be applied to properties concerning either individual blocks or the whole network. Here we train individual invertible linear layers via rank-one perturbations. Each of these P4Inv layers is an affine transformation Ax+b. In this context, the weight matrix A is handled by the P4 update and the bias b is optimized without perturbation. Without loss of generality, we present the method for layers Ax.
We define S as the set of invertible matrices, for which we know the inverse and determinant. Then the rank-one update
Ru,v(A) = A+ uv T (5)
with B = (u,v) ∈ R2n is a P4 on S due to equations 1 and 2, which also define the inverse pass and determinant computation of the perturbed layer, see Appendix B for details. The perturbation can be reset by setting u, v, or both to zero. In subsequent parameter studies, a favorable training efficiency was obtained by setting u to zero and reinitializing v from Gaussian noise. (Using a unity standard deviation for the reinitialization ensures that gradient-based updates to u are on the same order of magnitude as updates to a standard linear layer so that learning rates are transferable.) The
inverse matrix Ainv and determinant d are stored in the P4 layer alongside A and updated according to the merging step in Algorithm 2. Merges are skipped whenever the determinant update signals ill conditioning of the inversion. This is further explained in the following subsection.

heading: 3.3 NUMERICAL STABILIZATION
text: The update to the inverse and determinant can become ill-conditioned if the denominator in equation 1 is close to zero. Thankfully, the determinant lemma from equation 2 provides an indicator for illconditioned updates (if absolute determinants become very small or very large). This indicator in combination with the stepwise merging approach can be used to tame potential numerical issues. Concretely, the following additional strategies are applied to ensure stable optimizations.
• Skip Merges: Merges are skipped whenever the determinant update falls out of predefined bounds, see Appendix A for details. This allows the optimization to continue without propagating numerical errors into the actual weight matrix A. Note that numerical errors in the perturbed parameters Ã are instantaneous and vanish when the optimization leaves the ill-conditioned regime. As shown in our experiments in Section 4.2, merging steps that occur relatively infrequently without drastically hurting the efficiency of the optimization.
• Penalization: The objective function can be augmented by a penalty function g(u,v) in order to prevent entering the ill-conditioned regime {(u,v) : det (Ru,v(A)) = 0} , see Appendix A.
• Iterative Inversion: In order to maintain a small error of the inverse throughout training, the inverse is corrected after every Ncorrect-th merging step by one iteration of an iterative matrix inversion (Soleymani, 2014). This operation is O(n3) yet is highly parallel.

heading: 3.4 USE IN INVERTIBLE NETWORKS
text: Our invertible linear layers can be employed in normalizing flows (Appendix C) thanks to having access to the determinant at each update step. We tested them in two different application scenarios:
P4Inv Swaps In a first attempt, we integrate P4Inv layers with RealNVP coupling layers by substituting the simple coordinate swaps with general linear layers (see Figure 9 in Appendix H). Fixed coordinate swaps span a tiny subset of O(n). In contrast, P4Inv can express all of GL(n). We thus expect more expressivity with the help of better mixing. The parameter matrix A is initialized with a permutation matrix. Note that the P4 training is applied exclusively to the P4Inv layers rather than all parameters.
Nonlinear invertible layer In a second attempt, we follow the approach of Gresele et al. (2020) and stack P4Inv layers with intermediate bijective nonlinear activation functions. Here we use the elementwise Bent identity
B(x) =
√ x2 + 1− 1
2 + x.
In contrast to more standard activation functions like sigmoids or ReLU variants, the Bent identity is an R-diffeomorphism. It thereby provides smooth gradients and is invertible over all of R.

heading: 4 EXPERIMENTS
text: P4Inv updates are demonstrated in three steps. After a runtime comparison, single P4Inv layers are first fit to linear problems to explore their general capabilities and limitations. Second, to show their performance in deep architectures, P4Inv blocks are used in combination with the Bent identity to perform density estimation of common two-dimensional distributions. Third, to study the generative performance of normalizing flows that use P4Inv blocks, we train a RealNVP normalizing flow with P4 swaps as a Boltzmann generator (Noé et al., 2019). One important feature of this test problem is the availability of a ground truth energy function that is highly sensitive to any numerical problems in the network inversion.

heading: 4.1 COMPUTATIONAL COST
text: P4Inv training facilitates cheap inversion and determinant computation. To demonstrate those benefits, the computational cost of computing the forward and inverse KL divergence in a normalizing flow framework was compared with standard linear layers and an LU decomposition. Importantly, the KL divergence includes a network pass and the Jacobian determinant.
Figure 2 shows the wall-clock times per evaluation on an NVIDIA GeForce GTX 1080 card with a batch size of 32. As the matrix dimension grows, standard linear layers become increasingly infeasible due to the O(n3) cost of both determinant computation and inversion. The LU decomposition is fast for forward evaluation since the determinant is just the product of diagonal entries. However, the inversion does not parallelize well so that inverse pass of a 4096-dimensional matrix was almost as slow as a naive inversion. Note that this poor performance transfers to other decompositions involving triangular matrices, such as the Cholesky decomposition. In contrast, the P4Inv layers performed well for both the forward and inverse evaluation. Due to their O(n2) scaling, they outperformed the two other methods by two orders of magnitude on the 4096-dimensional problem.
This comparison shows that P4Inv layers are especially useful in the context of normalizing flows whose forward and inverse have to be computed during training. This situation occurs when flows are trained through a combination of density estimation and sampling.

heading: 4.2 LINEAR PROBLEMS
text: Fitting linear layers to linear training data is trivial in principle using basic linear algebra methods. However, the optimization with stochastic gradient descent at a small learning rate will help illuminate some important capabilities and limitations of P4Inv layers. It will also help answer the open question if gradient-based optimization of an invertible matrix A allows crossing the ill-conditioned regime {A ∈ Rn×n : detA = 0}. Furthermore, the training efficiency of perturbation updates can be directly compared to arbitrary linear layers that are optimized without perturbations.
Specifically, each target problem is defined by a square matrix T . The training data is generated by sampling random vectors x and computing targets y = Tx. Linear layers are then initialized as the identity matrix A := I and the loss function J(A) = ‖Ax− y‖2 is minimized in three ways:
1. by directly updating the matrix elements (standard training of linear layers),
2. through P4Inv updates, and
3. through the inverses of P4Inv updates, i.e., by training A through the updates in equation 1.
The first linear problem is a 32-dimensional positive definite matrix with eigenvalues close to 1. Figure 3 shows the evolution of eigenvalues and losses during training. All three methods of optimization successfully recovered the target matrix. While training P4Inv via the inverse led to slower convergence, the forward training of P4Inv converged in the same number of iterations as an unconstrained linear layer for a merge interval N = 1. Increasing the merge interval to N = 10 only affected the convergence minimally. Even for N = 50, the optimizer took only twice as many iterations as for an unconstrained linear layer.
The second target matrix was a 128-dimensional special orthogonal matrix. As shown in Figure 4, the direct optimization converged to the target matrix in a linear fashion. In contrast, the matrices generated by the P4Inv update avoided the region around the origin. This detour led to a slower convergence in the initial phase of the optimization. Notably, the inverse stayed accurate up to 5 decimals throughout training. Training an inverse P4Inv was not successful for this example. This shows that the inverse P4Inv update can easily get stuck in local minima. This is not surprising as the elements of the inverse (equation 1) are parameterized by R2n-dimensional rational quadratic functions. When training against linear training data with a unique global optimum, the multimodality can prevent convergence. When training against more complex target data, the premature convergence was mitigated, see Appendix G. However, this result suggests that the efficiency of the optimization may be perturbed by very complex nonlinear parameterizations.
The final target matrix was T = −I101, a matrix with determinant -1. In order to iteratively converge to the target matrix, the set of singular matrices has to be crossed. As expected, using a nonzero penalty parameter prevented the P4Inv update from converging to the target. However, when no penalty was applied, the matrix converged almost as fast as the usual linear training, see Figure 5. When the determinant approached zero, inversion became ill-conditioned and residues increased. However, after reaching the other side, the inverse was quickly recovered up to 5 decimal digits. Notably, the determinant also converged to the correct value despite never being explicitly corrected.
The favorable training efficiency encountered in those linear problems is surprising given the considerably reduced search space dimension. In fact, a subsequent rank-one parameterization of an MNIST classification task suggests that applications in nonlinear settings also converge as fast as standard MLPs in the initial phase, but slow down when approaching the optimum, see Appendix I.

heading: 4.3 2D DISTRIBUTIONS
text: The next step was to assess the effectiveness of P4Inv layers in deep networks. This was particularly important to rule out a potentially harmful accumulation of rounding errors. Density estimation of common 2D toy distributions was performed by stacking P4Inv layers with Bent identities and their inverses. For comparison, an RNVP flow was constructed with the same number of tunable parameters as the P4Inv flow, see Appendix G for details.
Figure 5: Training towards the matrix T = −I101 using no penalty. Residue of inversion (black line) and absolute determinants of the standard linear and P4Inv layer. Both converge to the target in a similar number iterations (dashed line).
Figure 6: Density estimation for two-dimensional distributions from RealNVP (RNVP) and P4Inv networks with similar numbers of tunable parameters.
Figure 6 compares the generated distributions from the two models. The samples from the P4Inv model aligned favorably with the ground truth. In particular, they reproduced the multimodality of the data. In contrast to RNVP, P4Inv cleanly separated the modes, which underlines the favorable mixing achieved by general linear layers with elementwise nonlinear activations.

heading: 4.4 BOLTZMANN GENERATORS OF ALANINE DIPEPTIDE
text: Boltzmann generators (Noé et al., 2019) combine normalizing flows with statistical mechanics in order to draw direct samples from a given target density, e.g. given by a many-body physics system. This setup is ideally suited to assess the inversion of normalizing flows since the given physical potential energy defines the target density and thereby provides a quantitative measure for the sample quality. In molecular examples, specifically, the target densities are multimodal, contain singularities, and are highly sensitive to small perturbations in the atomic positions. Therefore, the generation of the 66-dimensional alanine dipeptide conformations is a highly nontrivial test for generative models.
The training efficiency and expressiveness of Boltzmann Generators (see Appendix E for details) were compared between pure RNVP baseline models as used in Noé et al. and models augmented by P4Inv swaps (see Section 3.4). The deep neural network architecture and training strategy are described in Appendix H. Both flows had 25 blocks as from Figure 9 in the appendix, resulting in 735,050 RNVP parameters. In contrast, the P4Inv blocks had only 9,000 tunable parameters. Due to this discrepancy and the depth of the network, we cannot expect dramatic improvements from adding P4Inv swaps. However, significant numerical errors in the inversion would definitely show in such a setup due to the highly sensitive potential energy.
Figure 7 (left) shows the energy statistics of generated samples. To demonstrate the sensitivity of the potential energy, the training data was first perturbed by 0.004 nm (less than 1% of the total length of the molecule) and energies were evaluated for the perturbed data set. As a consequence, the mean of the potential energy distribution increased by 13 kBT .
In comparison, the Boltzmann generators produced much more accurate samples. The energy distributions from RNVP and P4Inv blocks were only shifted upward by ≈ 2.6 kBT and rarely generated samples with infeasibly large energies. The performance of both models was comparable with slight advantages for models with P4Inv swaps. This shows that the P4Inv inverses remained intact during training. Finally, Figure 7 (right) shows the joint distribution of the two backbone torsions. Both Boltzmann generators reproduced the most important local minima of the potential energy. As in the 2D toy problems, the P4Inv layers provided a cleaner separation of modes.

heading: 5 OTHER POTENTIAL APPLICATIONS OF P4 UPDATES
text: Perturbation theorems are ubiquitous in mathematics and physics so that P4 updates will likely prove useful for retaining other properties of individual layers or neural networks as a whole. To this end, the P4 scheme in Section 3.1 is formulated in general terms. Orthogonal matrices may be parameterized in a similar manner to P4Inv through Givens rotations or double-Householder updates. Optimizers that constrain a joint property of multiple layers have previously been used to enforce Lipschitz bounds (Gouk et al. (2018), Yoshida & Miyato (2017)) and could also benefit from the present work. Applications in physics often rely on networks that obey the relevant physical invariances and equivariances (e.g. Köhler et al. (2020); Boyda et al. (2020); Kanwar et al. (2020); Hermann et al. (2020); Pfau et al. (2020); Rezende et al. (2019)). These properties might also be amenable to P4 training if suitable property-preserving perturbations can be defined.

heading: 6 CONCLUSIONS
text: We have introduced P4Inv updates, a novel algorithmic concept to preserve tractable inversion and determinant computation of linear layers using parameterized perturbations. Applications to normalizing flows proved the efficiency and accuracy of the inverses and determinants during training. A crucial aspect of the P 4 method is its decoupled merging step, which allows stable and efficient updates. As a consequence, the invertible linear P4Inv layers can approximate any well-conditioned regular matrix. This feature might open up new avenues to parameterize useful subsets of GL(n) through penalty functions. Since perturbation theorems like the rank-one update exist for many classes of linear and nonlinear functions, we believe that the P4 concept presents an efficient and widely applicable way of preserving desirable network properties during training.

heading: ACKNOWLEDGEMENTS
text: We thank the anonymous reviewers for their valuable suggestions that helped a lot in improving the manuscript.

heading: A SANITY CHECK FOR THE RANK-ONE UPDATE
text: Based on the matrix determinant lemma (equation 2) rank-one updates are ill-conditioned if the term G := 1+vTAinvu vanishes. If such a perturbation is ever merged into the network parameters, the stored matrix determinant and inverse degrade and cannot be recovered. Therefore, merges are only accepted if the following conditions hold:
C (0) min ≤ ln |G| ≤ C (0) max and C (1) min ≤ ln |GdetA| ≤ C (1) max.
The constants Cmin and Cmax regularize the matrix A and its inverse Ainv, respectively, since ln |detA| = − ln |detAinv|. The penalty function is also based on these constraints:
g(u,v) = Cp · ( ReLU2 (ln |G| − Cmax) + ReLU2 (Cmin − ln |G|)
+ReLU2 (ln |GdetA| − Cmax) + ReLU2 (Cmin − ln |GdetA|) )
with a penalty parameter Cp. For the experiments in this work, we used Cmin = −2, Cmax = 15, C
(0) min = −6, C (0) max = inf , C (1) min = −2.5, and C (1) max = 15.5.
B IMPLEMENTATION OF P4INV LAYERS
In practice, the P4Inv layer stores the current inverse Ainv ≈ A−1 and determinant alongside A. The forward pass of an input vector x can be computed efficiently by first computing uTx and then adding vuTx to Ax. The inverse pass can be similarly structured to avoid any matrix multiplies.
Note that P4 training can straightforwardly be applied to only a part of the network parameters. In this case, all other parameters are directly optimized without the perturbation proxy and the gradient of the loss function J is composed of elements from ∂J/∂B and ∂J/∂A. Furthermore, the perturbation of parameters and evaluation of the perturbed model can sometimes be done more efficiently in one step. Also, the merging step from equation 3 and equation 4 can additionally be augmented to rectify numerical errors made during optimization.
In order to allow crossings of otherwise inaccessible regions of the parameter space the merging step was accepted every Nforce merges, even if the determinant was poorly conditioned. If u or v ever contain non-numeric values, merging steps were rejected and the perturbation is reset without performing a merge.

heading: C NORMALIZING FLOWS
text: A Normalizing flow is a learnable bijection f : Rn → Rn which transforms a simple base distribution p(z), by first sampling z ∼ p(z), and then transforming into x = f(z). According to the change of variables, the transformed sample x has the density:
q(x) = p ( f−1(x) ) ∣∣detJf−1(x)∣∣ . Given a target distribution ρ(x), this tractable density allows minimizing the reverse KullbackLeibler (KL) divergence DKL [q(x)‖ρ(x)] , e.g., if ρ(x) is known up to a normalizing constant, or the forward KL divergence DKL [ρ(x)‖q(x)] , if having access to samples from ρ(x).

heading: D COUPLING LAYERS
text: Maintaining invertibility and computing the inverse and its Jacobian is a challenge, if f could be an arbitrary function. Thus, it is common to decompose f in a sequence of coupling layers
f = g(1) ◦ S(1) ◦ . . . g(K) ◦ S(K).
Each g(k) is constrained to the form g(k)(x) = T (k)(x1, x2)⊕x2, where x = x1⊕x2, x1 ∈ Rm and x2 ∈ Rn−m. Here T (k) : Rm × Rn−m → Rm is a point-wise transformation, which is invertible in
its first component given a fixed x2. Possible implementations include simple affine transformations (Dinh et al., 2014; 2016) as well as complex transformations based on splines (Müller et al., 2019; Durkan et al., 2019). Each g(k) thus has a block-triangular Jacobian matrix
Jg(k) =
[ JT (k) M (k)
0 Im×m
] ,
where JT (k) is a (n − m) × (n − m) diagonal matrix. The layers S(k) take care of achieving a mixing between coordinates and are usually represented as simple swaps
S(k) = [
0 In−m×n−m Im×m 0
] .
The total log Jacobian of fθ is then given by
log detJfθ = K∑ k=1 tr [log (JT (k))] + log detS(k),
where log detS(k) = 0 when S(k) is given by the simple swaps above.

heading: E BOLTZMANN GENERATORS
text: Boltzmann Generators (Noé et al., 2019) are generative neural networks that sample with respect to a known target energy, as for example given by molecular force fields. The potential energy u : R3n → R of such systems is a function of the atom positions x. The corresponding probability of each configuration x is given by the Boltzmann distribution px(x) = exp(−βu(x))/Z, where β = 1/(kBT ) is inverse temperature with the Boltzmann constant kB . The normalization constant Z is generally not known.
The generation uses a normalizing flow and training is performed via a combination of density estimation and energy-based training. Concretely, the following loss function is minimized
J(A) = wlJl(A) + weJe(A), (6)
where wl+we = 1 denote weighting factors between density estimation and energy-based training. The maximum likelihood and KL divergence in equation 6 are defined respectively as
Jl = Ex∼px [ 1
2 ‖Fxz(x;A)‖2 − lnRxz(x;A)
] and
Je = Ez∼pz [u(Fzx(z;A))− lnRzx(z;A)] .
As an example, we train a model for alanine dipeptide, a molecule with 22 atoms, in water. Water is represented by an implicit solvent model. This system was previously used in Wu et al. (2020). Training data was generated using MD simulations at 600K to sample all metastable regions.

heading: F TRAINING OF LINEAR TOY PROBLEMS
text: The P4Inv layers were trained using a stochastic gradient descent optimizer with a learning rate of 10−2 and the hyperparameters from Table 1. The matrices were initialized with the identity.

heading: G TRAINING OF 2D DISTRIBUTIONS
text: The P4Inv layers used for 2D distributions were composed of blocks containing
1. a P4Inv layer with bias (2D),
2. an elementwise Bent identity,
3. another P4Inv layer with bias (2D), and
4. an inverse elementwise Bent identity.
100 of these blocks were stacked resulting in 1200 tunable parameters (counting elements of u and v). P4Inv training was performed with N = 10, Nforce = 10 and Ncorrect = 50. No penalty was used. Matrices were initialized with the identity I2.
The RealNVP network used for comparison was composed of five RealNVP layers. The additive and multiplicative conditioner networks used dense nets with two 6-dimensional hidden layers each and tanh activation functions, respectively. This resulted in a total of 1230 parameters.
The examples are taken from Grathwohl et al. (2018). Priors were two-dimensional standard normal distributions. Adam optimization was performed for 8 epochs of 20000 steps and with a batch size of 200. The initial learning rate was 5 · 10−3 and decreased by a factor of 0.5 in each epoch. After each merging step, the metaparameters of the Adam optimizer were reset to their initial state.
Figure 8 complements Figure 6 by showing samples from a network with only inverse P4Inv blocks. While the samples are worse than with forward blocks, the distributions are still well represented. This result indicates that the premature convergence encountered for linear test problems is a lesser problem in nonlinear problems and deep architectures.

heading: H TRAINING OF BOLTZMANN GENERATORS
text: The normalizing flows in Boltzmann generators were composed of the blocks shown in Figure 9 and a mixed coordinate transform as defined in Noé et al. (2019). The test problem was taken from Dibak et al. (2020). RNVP layers contained two 60-dimensional hidden layers each and ReLU and tanh activation functions for both t and s, respectively. The baseline model consisted of blocks of alternated RNVP blocks and swaps. The P4Inv model used invertible linear layers instead of the swapping of input channels in the baseline model. The computational overhead due to this change was negligible. RNVP parameters were optimized directly as usual and only the P4Inv layers are affected by the P4 updates. Merging was performed every N = 100 steps with Nforce = 10 and Ncorrect = 50. No penalty was used, i.e. C0 = 0.0. The P4Inv matrices were initialized with the reverse permutation, i.e. Aij = δi(n−j).
Density estimation with Adam was performed for 40,000 optimization steps with a batch size of 256 and a learning rate of 10−3. A short energy-based training epoch was appended for 2000 steps with a learning rate of 10−5 and we/wl = 0.05. After each merging step, the metaparameters of the Adam optimizer were reset to their initial state for all P4Inv parameters.

heading: I MNIST CLASSIFICATION VIA RANK-ONE UPDATES
text: Compared to fully-connected multi-layer perceptrons (MLP), rank-one updates reduce the number of independent trainable parameters per layer from m · n to m + n − 1, where m and n are the input and output dimension, respectively. It is therefore useful to study how the reduced search space dimension affects the training efficiency in a nonlinear setup. To this end, non-invertible classifier MLPs were trained on MNIST with an unrestricted search space as well as through rankone updates. The original network was composed of two convolutional layers, two dropout layers, and two linear layers. In P4 training, the two linear layers (dimensions 9216×128 and 128×10) were trained through rank-one updates which were merged in every iteration (N = 1). Naturally, this training did not involve any inverse or determinant updates. A vanilla SGD optimizer was used with various learning rates.
Figure 10 shows the training loss and test accuracy during training. As for the linear problems from the previous subsection, the training efficiency was virtually unaffected during the first phase of the optimization, i.e. when the descent direction did not change significantly between subsequent iterations. However, as the descent direction became more noisy in the vicinity of the optimum, the training with rank-one updates became less efficient.

Review summaries:
This paper introduces an algorithm for training neural networks in a way that parameters preserve a given property. The optimization is based on using a transformation R that  perturbs parameters in a way that the desired property is preserved. Instead of directly optimizing the parameters of the network, the optimization is carried out on the parameters B of the auxiliary transformation R.