Training Invertible Linear Layers through Rank-One Perturbations
Authors: Andreas Krämer
Jonas Köhler
Frank Noe
Venue:  (2021)
Keywords: Parameter Perturbation
Reparameterization
Invertible Neural Networks
Normalizing Flows
Rank-one update

Abstract: Many types of neural network layers rely on matrix properties such as invertibility or orthogonality.
Retaining such properties during optimization with gradient-based 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 P$^{4}$Inv 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$^4$ concept can be utilized to retain properties other than invertibility.

Under review as a conference paper at ICLR 2021
TRAINING INVERTIBLE LINEAR LAYERS THROUGH RANK-ONE PERTURBATIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Many types of neural network layers rely on matrix properties such as invertibil- ity orthogonality. Retaining such properties during optimization with gradient- based stochastic optimizers is 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. lieu of directly optimizing the network parameters, we train rank-one perturbations and add them to the actual weight matrices infrequently. This pa Inv 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 out- line how the p concept be utilized to retain properties other than invertibility.
INTRODUCTION
Many deep learning applications depend criti- cally on the neural network parameters having certain mathematical structure. As impor- tant 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)
Standard Training
p4 Training
y
Optimizer
Pd
Biases
Perturbed Parameters
DNN
Preserving parameter properties during train- ing can be challenging and many approaches are currently in use. The most basic way of incorporating constraints is by network de- sign. Many examples could listed, like defining convolutional layers to obtain equiv- ariances, constraining network outputs to cer- tain intervals through bounded activation func- tions, Householder flows (Tomezak & Welling 2016) enforce layer-wise orthogonality, or coupling layers (Dinh et al. 2014: 2016) that enforce tractable inversion through their two- channel structure. A second approach concerns the optimizers used for training. Optimiza- tion routines have been tailored for example to maintain Lipschitz bounds (Yoshida & Miyato 2017) efficiently optimize orthogonal linear layers CChormaaski et al. 2020).
Weights
DNN
Blases
Weights
x
(merge)
Training of deep neural networks (DNN). Standard DNN transform inputs æ into outputs through activation functions and linear layers, which are tuned by an optimizer. In con- trast, p4 training operates perturbations to the parameters. Those are defined to retain certain network properties (here: invertibility well as tractable inversion and determinant computation). The perturbed parameters merged in regular intervals.
Figure 1:
The present work introduces a novel algorithmic concept for training invertible linear layers and facilitate tractable inversion and determinant computation, see Figure In lieu of directly changing the network parameters, the optimizer operates perturbations these parameters. The actual network parameters are frozen, while a parameterized perturbation (a rank-one update to the frozen parameters) serves proxy for optimization. Inputs passed through the perturbed network

---

Under review as a conference paper at ICLR 2021
during training. In regular intervals, the perturbed parameters merged into the actual network and the perturbation is reset to the identity. This stepwise optimization approach will be referred to property property-preserving preserving parameter perturbation, or p update. A similar concept was introduced (2019), who used dynamic trivializations for optimization on manifolds.
cently by [ezcano-
In this work, use Pa training to optimize invertible linear layers while keeping track of their in- verses 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 gradhents. 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 p Inv layers are inverted the cost of roughly three matrix-vector multiplica- tions.
I'Inv layers can approximate arbitrary invertible matrices A E 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, avoids any explicit computations verses or determinants. All operations occurring in optimization steps have complexity of at most On?)). To our knowledge, the present method is the first to feature these desirable properties.
in-
We show how I'Inv 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 deep generative models. Finally, outline potential applications of pt training to network properties other than invertibility.
BACKGROUND AND RELATED WORK
2.1 RANK-ON PERTURBATION
The I'Inv layers are based on rank-one updates, which are defined transformations A A we" with u, v E " A E GL(n) and 1 + v² -uuu f O, the updated matrix is also invertible and its inverse can be computed by the
Morrison formula
(A + uv
A-1 1 TT A
A uv A''uo" A-1
4 2
Furthermore, the determinant is given by the matrix determinant lemma
det (A + ue²) (1 - o" A-'u)) (A).
(2)
Both these equations are widely used in numerical mathematics, since they sidestep the O(n²)
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.
2.2 EXISTING APPROACHES FOR TRAINING INVERTIBLE LINEAR LAYERS
Maintaining invertibility of linear layers has been studied in the context of convolution opera- tors (Kingma & Dhariwal 2018 Karami et al. 2019: Hoogeboom et al.. 2019 2020) and using
Sylvester's 's theorem (Van Den Berg et 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 [..). More closely related to our work. Gresele et al.] (2020) introduced relative gradient optimization scheme for invertible matrices. In contrast to this related work, our method facilitates cheap inverse pass and allows sign changes in the determinant. On the contrary, their method operates in a nigher-dimensionee search space, which might speed up the optimization in tasks that do not involve inversion during training.
2.3 NORMALIZING FLOWSS
Cheap inversion and determinant computation are specifically important in the context of normal- izing flows, Appendix C They were introduced in Tabak al. (2010); Tabak & Turner (2013)
2

---

Under review as a conference paper at ICLR 2021
and are commonly used, either in variational inference (Rezende & Mohamed 2015 Tomezak & 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 Müller et al. 2019 Noč et al. 2019 Köhler et al., 2020). The most important normalizing flow architectures are coupling layers (Dinh et 2014 2014|2206 Kingma & Dhariwal [2018] Müller et al. 2019), which are subclass of autoregressive flows (Germain et 2015) Papamakarios et al. [2017 Huang et al. |0018: De Cao al. 2019), and (2) residual flows (Chen 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).
2.4 OPTIMIZATION UNDER CONSTRAINTS AND DYNAMIC TRIVIALIZATIONS
Constrained matrices can be optimized using Riemannian gradient descent on the manifold (Absil al. ()009) A reparameterization trick for general Lie groups has been introduced Falorsi et al.
(2019). For the unitary/ / orthogonal group there are multiple more specialized approaches, inc Including using the Cayley transform (Helfrich et al. 2018). Houscholder Reflections (Mhammedi et al. 2017
Meng et al. 2020 Tomczak & Welling 2016). Givens rotations (Shalit & Chechik, [0144] Pevny
et al
2020 or the exponential map eezcano-Casade & Martinez-Rubil 2019 Golinski et al.
2019
Lezcano-Casado (2019) introduced the concept of dynamic trivializations. This method performs
training manifolds by combining ideas from Riemannian gradient descent and trivializations (pa- rameterizations 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 Sezcano-Casade (2020). p training resembles dynamic trivi- alizations in that both perform number of iteration steps in a fixed basis and infrequently lift the optimization problem to new basis. In contrast, the rank-one updates do not strictly parameterize GL(n) but instead can access all of Rax* This introduces the need for numerical stabilization, but enables efficient computation of the inverse and determinant through equation | which is the method's ' unique and most important aspect.
and equation
3 DA UPDATES: PRESERVING PROPERTIES THROUGH PERTURBATIONS
3.1 GENERAL CONCEPT
A deep neural network is parameterized function MA R RT with a high-dimensional parameter tensor A. Now, let S define the subset of feasible parameter tensors so that the network satisfies certain desirable property. In many situations, generating elements of S from scratch is much harder than transforming any A E S into other elements A' E S, i.c. to move within S.
The efficiency of perturbative updates can be leveraged an incremental approach retain certain desirable properties of the network parameters during training. Rather than optimizing the parameter tensors directly, we instead use transformation Rg S S, which we call a ceopertyy-prseeriing parameter perturbation (P1). A pa transforms given parameter tensor A E S into another tensor with the desired property A' E The p1 itself is also parameterized, by a tensor B. We demand that the identity idg A - A be included in the set of these transformations, i.e. there exists a Bo such that Rs. idg.
During training, the network is evaluated using the perturbed parameters Ā = 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 Pa B, are merged into as follows:
Anew f- RB (A).),
(3)
B new By.
(4)
This update does not modify the effective (perturbed) parameters of the network A, since
À Ånew
RB RB..... (Anew) RB,
( (Ro (Rs(A)) Rs (A) A.
3

---

Under review as a conference paper at ICLR 2021
Hence, this procedure enables steady, iterative transformation of the effective network parame- ters and stochastic gradient descent methods can be used for training without major modifications. Furthermore, given reasonable Pt, the iterative update of A can produce increasingly non-trivial transformations, thereby enabling high expressivity the resulting neural networks. This concept summarized in Algorithm Further extensions 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 == Bo:
for i = 1 Steeps do
X, Yo == i-th batch from training data;
= (B)
// perturb parameters // evaluate perturbed model
Y = M, (X)
j= J(Y (Y)
// evaluate loss function
gradient := B
// backprepagati // optimization step
= OPT(B, gradient))
if' mod N then
= RB(An
// merging step: update frozen parameters merging step: reset perturbation
B B
end
end
3.2 P'INV: INVERTIBLE LAYERS VIA RANK-ONE UPDATES
Algorithm 2: P'Inv Merging Step
Input : Matrix A, Inverse Aivy, Determinant d det factor (1 + TT Auy
new det == det det_ factor d;
if In [det_factor] and In |new_det]
then update frozen parameters (equation
3
net_
A = Ru,v (A);
A = Ainv T4 A uv Amvi
/* reset perturbation (equation
+/
= 0;
(0, I,);
// random reinitialization
end
The pd 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 I'I layers is an affine transformation Az+ 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 Az.
We define & as the set of invertible matrices, for which we know the inverse and determinant. Then the rank-one update
R (A) = + uo
(5)
with = (u, v) E 2 and determinant computation of the perturbed layer, 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 standard deviation for the reinitialization ensures that muddeent-based updates to u are on the same order of magnitude updates standard linear layer that learning rates transferable.) The
a pa on S due to equations 1 and which also define the inverse pass
zero and reinitializing from Gaussian noise. (Using unity

---

Under review as a conference paper at ICLR 2021
inverse matrix Anv and determinant d are stored in the p4 layer alongside A and updated according to the merging step in Algorithm? 2 conditioning of the inversion. This is further explained in the following subsection.
Merges are skipped whenever the determinant update signals ill
3.3 NUMERICAL STABILIZATION
The update to the inverse and determinant can become ill-conditioned if the denominator in equation is close to zero. Thankfully, the determinant lemma from equation 2 provides indicator for ill- conditioned updates (if absolute determinants become very small or very large). This indicator in combination with the stepwise merging approach be used to tame potential numerical issues. Concretely, the following additional strategies are applied to ensure stable optimizations.
Skip Merges: Merges 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 erros 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, (A)) = 0) see
Appendix
Iterative Inversion: In order to maintain small eπror of the inverse throughout training, the inverse is corrected after every N Neorreet-th merging step by one iteration of iterative matrix inversion (Soleymani 2014). This operation is (n²² yet is highly parallel.
3.4 USE IN INVERTIBLE NETWORKS
Our invertible linear layers can be employed in normalizing flows (Appendix C thanks having access to the determinant at each update step. We tested them in two different application scenarios:
I'Iny Swaps In first attempt, we integrate I'Inv layers with RealNVP coupling layers by sub- stituting the simple coordinate swaps with general linear layers (see Figure In Appendix H). Fixed coordinate swaps span a tiny subset of O(n). In contrast, P'Iny can express all of GT.(n). We thus expect more expressivity with the help of better mixing. The parameter matrix A is initialized with permutation matrix. Note that the pd training is applied exclusively to the I'Inv layers rather than all parameters.
Nonlinear invertible layer In second attempt, we follow the approach of Gresele et al. (2020) and stack Inv layers with intermediate bijective nonlinear activation functions. Here the elementwise Bent identity
B(x)
In contrast more standard activation functions like sigmoids ReLU variants, the Bent identity an R-diffeommrphiss. It thereby provides smooth gradients and invertible over all of R.
EXPERIMENTS
I'Iny updates are demonstrated in three steps. After a runtime comparison, single P'[ny layers are first fit to linear problems to explore their general capabilities and limitations. Second, to show their performance in deep architectures, P"Iny 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 p* Inv blacks, we train a RealNVP normalizing flow with p swaps Boltzmann generator (Noé 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.
5

---

Under review as a conference paper at ICLR 2021
dim 64
512
dim 4096
Standard
30
WU
300
Time [ms]
Pfinv
20
200
100
Forward
(Backward)
Inverse
(Backward)
Forward
(Backward)
Inverse
(Backward)
Forward
(Backward)
Inverse
(Backward)
Figure 2: Wall-clock times of forward and backward pass of linear normalizing flows including determinant computation. Three methods are compared: (a) standard linear layers ("standard"), where the inverses and determinants are computed through PyTorch's ' inverse and der functions; (b) LU decompositions ("LU"", where the determinants are products over the diagonal entries and the matrices are inverted through triangular -solee;, (c) I'Inv updates that keep track of inverses and determinants through rank-one updates. Timings are compared for square matrices of dimension 64, 512, and 4096.
4.1 COMPUTATIONAL COST
I'Inv training facilitates cheap inversion and determinant computation. To demonstrate those ben- efits, the computational cost of computing the forward and inverse KL divergence a normalizing flow framework
compared with standard linear layers and LU decomposition. Importantly, the KI. divergence includes a network pass and the Jacobian determinant.
Figure 2 shows the wall-clock times per evaluation batch size of 32. As the matrix dimension grows, standard linear layers become increasingly infeasi- ble due to the (n²) 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 di9e-ddmnns matrix was almost as slow naive inversion. Note that this poor performance transfers to other decompositions involv- ing triangular matrices, such as the Cholesky decomposition. In contrast, the I'Inv layers performed well for both the forward and inverse evaluation. Due to their O(n²) scaling, they outperformed the two other methods by two orders of magnitude on the 4096-dimensional problem.
NVIDIA GeForce GTX 1080 card with
This comparison shows that I'Iny 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 trained through combination of density estimation and sampling.
4.2 LINEAR PROBLEMS
Fitting linear layers to linear training data trivial in principle using basic linear algebra methods. However, the optimization with stochastic gradient descent small learning rate will help illumi- nate some important capabilities and limitations of I'Inv layers. It will also help answer the open question if gradient-based optimization of invertible matrix A allows crossing the ill-conditioned regime {A E Æexx. det A 0}. Furthermore, the training efficiency of perturbation updates can be directly compared arbitrary linear layers that are optimized without perturbations.
Specifically, each target problem defined by square matrix T. The training data is generated by sampling random vectors and computing targets y Tw. Linear layers are then initialized as the identity matrix A = I and the loss function J(A) = A y/1² is minimized in three ways:
1. by directly updating the matrix elements (standard training of linear layers),
2. through InIny updates, and
3. through the inverses of I'v updates, i.e., by training A through the updates in equation]

---

Under review as a conference paper at ICLR 2021
10-4
Tinear
Linear
Play
Pla via Inverse
P'ny (N-1)
B
Loss
10-6
Pinny Inv. (N=1)
10-10
Plly (N-10)
000
X
P Iny (N=50)
15000
30000
0.9
1.0
0. 1.0
11
0.9 10
Iseration
Re
Re
X
Re A
Figure 3: Training towards 32-dimensional positive definite target matrix T. Left: Losses during training. Right: Eigenvalues during training. Final eigenvalues are shown as red crosses. Eigenval- ues of the target matrix are shown as black squares.
The first linear problem is a 32-dimensional positive definite matrix with cigenvalues close to 1. Figure shows the evolution of eigenvalues and losses during training. All three methods of opti- mization successfully recovered the target matrix. While training I'Inv via the inverse led to slower convergence, the forward training of P'Inv converged in the same number of iterations constrained linear layer for a merge interval = 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.
un-
Linear
P-Inv
P'lny via Inverse
Liucar phh (N=1)
P'lay via Iuv. (N=1)
Im
10-3
25000
50000
Iseration
Re X
Re
Re A
Figure 4: Training towards an orthogonal target matrix T E SO(128). Left: Losses during training. Right: Eigenvalues during training. Final eigenvalues are shown as red crosses. Eigenvalues of the target matrix are shown as black squares.
The second target matrix was 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 I'Iny update avoided the region around the origin. This detour led to slower convergence in the initial phase of the optimization. Notably, the inverse stayed accurate up to decimals throughout training. Training inverse P'Inv was not successful for this example. This shows that the inverse P'Iny update can easily get stuck in local minima. This is not surprising as the elements of the inverse (equation [] are parameterized by 12"=dimmessoon rational quadratic functions. When training against linear training data with unique global optimum, the multi- modality 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 parameterizzatins.
The final target matrix was T = 1001 matrix with determinant -1. In order to iteratively converge the target matrix, the of singular matrices has to be crossed. As expected, using nonzero penalty parameter prevented the P'Inv 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 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 con- siderably reduced search space dimension. In fact, 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

---

Under review as a conference paper at ICLR 2021
4.3 2D DISTRIBUTIONS
The next step was to assess the effectiveness of P'Inv layers in deep networks. This was particularly important rule out a potentially harmful accumulation of rounding errors. Density estimation of common 2D toy distributions was performed by stacking I' layers with Bent identities and their inverses. For comparison, an RNVP flow was constructed with the same number of tunable parameters as the P'Inv flow, see Appendix G for details.
Data
RNVP
P*Inv
Linçur
pllnv
10a
10-4
||Ag A- Talla
10-1
10
ddd
10-26 26
10-1
30
10-7
10000
20000
Iseration
Figure 5: Training towards the matrix T 110 using no penalty. Residue of inversion (black line) and absolute determinants the standard linear and P'Inv layer. Both converge the target in a similar number iterations (dashed line).
Figure 6: Density estimation for wwo-dimensional distribu- tions from RealNVP (RNVP) and P'Inv networks with sim- ilar numbers of tunable parameters.
Figure 6 compares the generated distributions from the two models. The samples from the P'Inv model aligned favorably with the ground truth. In particular, they reproduced the multimodality of the data. In contrast to RNVP, p'Iny cleanly separated the modes, which underlines the favorable mixing achieved by general linear layers with elementwise nonlinear activations.
4.4 BOLTZMANN GENERATORS OF ALANINE DIPEPTIDE
Boltzmann generators (Noé al., 2019) combine normalizing flows with statistical mechanics in order to draw direct samples from a given target density, e.g. given by many-body physics system. This setup ideally suited to assess the inversion of normalizing flows since the given physical po- tential energy defines the target density and thereby provides quantitative measure for the sample quality. In molecular examples, specifically, the target densities are multimodal, contain singulari- ties, and are highly sensitive to small perturbations in the atomic positions. Therefore, the generation of the 66- 56-dimensional alanine dipeptide conformations is highly nontrivial test for generative mod- els.
The training efficiency and expressiveness of Boltzmann Generators (see Appendix F for details) were compared between pure RNVP baseline models as used in Noé et al. and models augmented by P'Inv swaps (see Section 3.4). The deep neural network architecture and training strategy are described in Appendix H Both flows had 25 blocks from Figure in the appendix, resulting in 735,050 RNVP parameters. In contrast, the I'Iny blocks had only 9,000 tunable parameters. Due to this discrepancy and the depth of the network, adding I'Inv swaps. However, sigmificant numerical errors in the inversion would definitely show in such setup due to the highly sensitive potential energy.
cannot expect dramatic improvements from
8

---

Under review as a conference paper at ICLR 2021
Data
RNVP
p*Inv
Patential Energies
Low-energy fraction: 70.1%
97.4%
99.3%
100.0%
0%
-10
-2
0.0
2.5
2.5 -2.5
2.5
Data +0.004
RNVP
P*inv
Figure 7: Left: Energy distributions of generated samples in dimensionless units of kgT; the second (orange) violin plot shows energies when the training data was perturbed by normal distributed random noise with 0.004 nm standard deviation. The low-energy fraction for each column denotes the fraction of samples that had potential energy lower than the maximum energy from the training set (Ne 20 kg?"). Right: Joint marginal distribution of the backbone torsions and eh: training data compared to samples from RealNVP Boltzmann generators with and without P'Inv swaps (denoted p² In and RNVP, respectively).
Figure (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 I % of the total length of the molecule) and energies were evaluated for the perturbed data set. As consequence, the mean of the potential energy distribution increased by 13 kgT.
In comparison, the Boltzmann generators produced much more accurate samples. The energy distri- butions from RNVP and I'Inv blocks were only shifted upward by = 2.6 kg and rarely generated samples with infeasibly large energies. The performance of both models was comparable with slight advantages for models with P'Inv swaps. This shows that the P'Inv 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 I'v layers provided a cleaner separation of modes.
OTHER POTENTIAL APPLICATIONS OF Pd UPDATES
Perturbation theorems are ubiquitous in mathematics and physics so that p updates will likely prove useful for retaining other properties of individual layers or neural networks as a whole. To this end, the pd scheme Section 3.1 formulated in general terms. Orthogonal matrices may be parameterized in similar manner to I'Inv through Givens rotations or double-House Houble-Householer updates. Optimizers that constrain a joint property 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 physics often rely on networks that obey the relevant physical invariances and equivariances (e.g. Köhler al. (2020): Boyda et al. (2020); Kanwar et al. (2020); Hermann et al. (2020); Pfau et al.] (2020); Rezende et al. (2019)). These propertics might also be amenable to pd training if suitable property-presentey perturbations be defined.
CONCLUSIONS
We have introduced P'Iny updates, novel algorithmic concept to preserve tractable inversion and determinant computation linear layers using parameterized perturbations. Applications to nor- malizing flows proved the efficiency and accuracy of the inverses and determinants during training. crucial aspect of the p4 method is its decoupled merging step, which allows stable and efficient updates. As a consequence, the invertible linear I' layers 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 pt concept presents an efficient and widely applicable way of preserving desirable network properties during training.

---

Under review as a conference paper at ICLR 2021
ACKNOWLEDGEMENTS
We thank the anonymous reviewers for their valuable suggestions that helped a lot in improving the manuscript.
REFERENCES
P. A. Absil, R. Mahony, and R. Sepulchre. Optimization algorithms matrix manifolds. Princeton University Press, 2009. ISBN 7780661132998 doi: 10. 10.15155/9781400830444
Jens Behrmann, Will Grathwohl, Ricky T. Q. Chen, David Duvenaud, and Joern-Henrik Jacob- sen. Invertible residual networks. In Proceedings of the 36th International Conference on Ma- chine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 573-582, Long Beach, California, USA, 2019. PMLR. URL http::/prooceeiing// http:///.......r..//e//s. behrmann html
Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez. Rezende, Michael S Albergo, Kyle Cranmer, Daniel C Hackett, and Phiala Shanahan. Sampling using su(n) gauge equivari- ant flows. arXiv preprint 2000:200000045566 2020.
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differ- ential equations. In Advances in neural information processing systems, pp. 6571-6583 2018.
Tian Qi Chen, Jens Behrmann, David Duvenaud, and Menn-Kenrik Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, pp- 9913-9923, 2019.
Krzysztof Choromanski, David Cheikhi, Jared Davis, Valerii Likhosherstov, Achille Nazaret, Achraf Bahamou, Xingyou Song, Mrugank Akarte, Jack Parker-Holder, Jacob Bergquist, Yuan Gao, Aldo Pacchiano, Tamas Sarlos, Adrian Weller, and Vikas Sindhwani. Stochastic flows and geometric optimization on the orthogonal group. In 37th International Conference on Machine Learning (ICMI. 2020), 2020. URI https: / arxivv... /atpg///////aivvooor////////00 .0t.//300031.1333333 13563
Nicola De Cao, Ivan Titov, and Wilker Aziz. Block neural autoregressive flow. In Conference on Uncertainty in Artificial Intelligence (UAI), 2019. UKLhhctpss//axxi..rrr///////1
04676
Manuel Dibak, Leon Klein, and Frank Noé. Temperature--seerabb flows. In preparation. 2020.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components esti- mation. arXiv preprint arXiv:1410.8516, 2014.
Laurent Dinh, Jascha Sohk-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv: 2016.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems, pp. 7509-7520, 2019.
Luca Falorsi, Pim de Haan, Tim R. Davidson, and Patrick Forré. Reparameterizing distributions on lie groups In Proceedings of Machine Learning Research, volume 89 of Proceedings of Machine Learning Research, pp. 3254-3253, PMLR, 2019. URL httt::/pprcceeiinngg mlr. wttpppessss///slssssse.o 19a. html
Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. Made: Masked autoencoder for distribution estimation. In International Conference on Machine Learning, pp. 881-889, 2015.
Adam Golinski, Mario Lezzano--asadee and Tom Rainforth. Improving normalizing flows via better orthogonal parametcrizations. In ICML Workshop on Invertible Neural Networks and Normalizing Flows, 2019.
Henry Gouk, Eibe Frank, Bernhard Pfahringer, and Michael Cree. Regularisation of Neural Net- works by Enforcing Lipschitz Continuity. arXiy preprint arXiv:1804.04368, apr 2018. URL
10

---

Under review as a conference paper at ICLR 2021
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. In 7th International Conference on Learning Representations (ICLR), 2018. htttt:t/taaaastt org/ abs/
01367
Luigi Gresele, Giancarlo Fissore, Adrián Javaloy, Bernhard Schölkopf, and Aapo Hyvärinen. Rel- ative gradient optimization of the jacobian term unsupervised deep learning. 34th Con- ference on Neural Information Processing Systems (NeurIPS 2020), 2020. URL https:
//arxiv.org/ abs/ 2006. 100000
Kyle Helfrich, Devin Willmott, and Qiang Ye. Orthogonal recurrent neural networks with scaled cayley transform. In International Conference Machine Learning, pp. 1969-1978. PMLR, 2018.
Jan Hermann, Zeno Schätzle, and Frank Noé. Deep-neural-networkk solution of the electronic Schrödinger equation. Nat. Chem., 12(10):891-88) 897, sep 2020. ISSN 17554349. doi: 10.1038/ s41557-020-0544-y
Emiel Hoogeboom, Rianne van den Berg, and Max Welling, Emerging convolutions for generative normalizing flows. International conference on machine learning, 2019. URL https://
arxiv. abs http:////b///91.....1
Emiel Hoogeboom, Victor Garcia Satorras, Jakub M Tomezak, and Max Welling. The convolution
exponential and generalized sylvester flows. In 34th Conference on Neural Information Pro- cessing Systems (NeurIPS 2020), 2020. URL https: /tatttsss//paprrss..te /wwpppaappppeeepppeeeeeeeeeeeeee wwwperrssspiss....e.eee..ee.eeee cc/paper//02// fata://300e2222faaaaaaao000000
pdf
Chin-We Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. arXiv preprint arXiv: 1804.00790000779, 2018
Gurtej Kanwar, Michael S. Albergo, Denis Boyda, Kyle Cranmer, Daniel C. Hackett, Sébastien Racanière, Danilo Jimenez Rezende, and Phiala E. Shanahan. Equivariant flow-based sampling for lattice gauge theory. Phys. Rev. Lett., 125:121601, Sep 2020. doi: 10.1 10. 125.
URL https: aps. wwt.///////oooooooo/ http://........//eeeeeeee wtt.:///yeeeeeeeetttt.oom 125.
121
Mahdi Karami, Dale Schuurmans, Jascha Sohl-Dickstein, Laurent Dinh, and Daniel Duckworth. Invertible convolutional flow. In Advances Neural Information Processing Systems, pp. 5635- 2019.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions, In Advances in neural information processing systems, pp. 10215-10224, 2018.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. arXiv preprint arXiv:2006.02425 2020.
Mario Lezcano-Casado. Trivializations for puudient-based optimization on manifolds. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Garnett (eds.), Ad- vances in Neural Information PwceiennPPooooo Processing Systems, volume 32, pp. 9157-9168. Curran Asso- ciates, Inc., 2019. URL http:s::/proceeediiigg nww..ips........................ cc cc/ppaper/2119//ieeee.. 11338856222244445399999933332223 Paper. pdf
Mario Lezcano-Casado, Curvature-dependant global convergence rates for optimization on mani- folds of bounded geometry. arXiv preprint 17X1::2008.02517,222 2020.
Mario Lezcano-Casado and David Martínez-Rubio Cheap orthogonal constraints in neural net- works: A simple parametrization of the orthogonal and unitary group. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pp. 3794-3803. PMLR, 2019. URL. http://proceed ngs .m pww.srsss/eoe wttss///991111cc//o1a/ casado" 9a html
Christos Louizos and Max Welling. Multiplicative normalizing flows for variational bayesian neural networks. In Proceedings of the 34th International Conference on Machine Learning Volume 70, pp. 2218-2227, 2017.
11

---

Under review as a conference paper at ICLR 2021
Chenlin Meng, Yang Song, Jiaming Song, and Stefano Ermon. Gaussianization flows. arXiv preprint arXiv:2003.01941, 2020.
Zakaria Mhammedi, Andrew Hellicar, Ashfaqur Rahman, and James Bailey. Efficient orthogonal parametrisation of recurrent neural networks using householder reflections. In International Con- ference Machine Leurning, pp. 2401-2409. PMLR, 2017.
Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus Gross, and Jan Novák. Neural
importance sampling. ACM Trans. Graph., 38(5), 2019. ISSN 0730-0301. URL https:
//doi. 11../3341133155
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu, Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457):eaaw 1147, 2019.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems, pp. 2338-2347, 2017.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762, 2019.
Tomas Pevny, Vasek Smidl, Martin Trapp, Ondrej Polacek, and Tomas Oberhuber. Sum-product transform networks: Exploiting symmetries using invertible transformations. arXiv preprint arXiv:2005.01297, 2020.
David Pfau, James S. Spencer, Alexander G. D. G. Matthews, and W. M. C. Foulkes. Ab ini-
tio solution of the many-electron schrödinger equation with deep neural networks. Phys. Rev.
Research, 2:033429, Sep 2020. doi: 10. 2o00:phssseeeecrrrh2.........4
URL https:
//link. aps. ott.pppp.../ooooo/..ooo http:////.....o.eeeeeeeeeeeeeeea 033429
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. Proceedings
of the 32nd International Conference Machine Learning, volume 37, pp. 1530-1538, 2015. URL http:://ppppt http::///occeeeedeeeeeee m
html
Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant hamilto- nian flows. arXiv preprint arXiv: 13739, 2019.
Uri Shalit and Gal Chechik. Coordinate-descent for learning orthogonal matrices through Givens rotations. In 31st Int. Conf. Mach. Learn. ICMI 2014, volume 1, pp. 833 845, dec 2014. ISBN 9781634393973. URL htr.:tttp://aaaa...rrr/aaa/////s
0624
Fazlollah Soleymani. A fast convergent iterative solver for approximate inverse of matrices. Numer- ical Linear Algebra with Applications, 21(3):439-452, 2014.
Esteban G Tabak and Cristina V Turner. A family of nonparametric density estimation algorithms.
Communications on Pure and Applied Mathematics, 66(2):145-164, 2013.
Esteban G Tabak, Eric Vanden -ijnden,, et al. Density estimation by dual ascent of the log-
likelihood. Communications in Mathematical Sciences, (()::217-233, 2010. 2010
Jakub M Tomezak and Max Welling. Improving variational auto-encoders using householder flow.
arXiv preprint 1rXiv:161..09630 2016.
Rianne Van Den Berg, Leonard Hasenclever, Jakub M. Tomezak, and Max Welling. Sylvester normalizing flows for variational inference. In 34th Conference on Uncertainty in Artificial In- telligence 2018, UAI 2018, 34th Conference Uncertainty in Artificial Intelligence 2018, UAI 2018, pp. 393-402, 2018. 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018 Conference date: 06-08-2013 Through 10-08-2018.
Aaron van den Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George den Driessche. Edward Lockhart, Luis Cobo. Florian Stimberg, Nor- man Casagrande, Dominik Grewe, Seb Noury, Sander Dieleman, Erich Elsen, Nal Kalchbrenner, Heiga Zen, Alex Graves, Helen King, Tom Walters, Dan Belov, and Demis Hassabis. Parallel
12

---

Under review as a conference paper at ICLR 2021
Fast high-fidelity speech synthesis. In Proceedings of the 35th International Con- ference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp- 3918-3926, 2018. URLpttt:://proceeiiii
WaveNet:
wwr.
Hao Wu, Jonas Köhler, and Frank Noé. arXiv:2002.06707, 2020.
Stochastic normalizing flows. arXiv preprint
Yuichi Yoshida and Takeru Miyato. Spectral norm regularization for improving the generalizability of deep learning, 2017.
Linfeng Zhang, Lei Wang. et al. Monge-ampera flow for generative modeling. arXiv preprint 1809: 10188, 2018.
13

---

Under review as a conference paper at ICLR 2021
A SANITY CHECK FOR THE RANK-ONE UPDATE
Based on the matrix determinant lemma (equation 2 rank-one updates are ill-conditioned if the term G := 1 v" Ai vanishes. If such 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
S In | S C and
C S In det A c(
The constants Cmin and Cmax regularize the matrix A and its inverse Alav, respectively, since In det A| = det Aivl-
The penalty function is also based on these constraints:
g(u, v) Cp (ReLU? (In |G| - Cmax) + ReLU? (Cmin In (Gl)
+ReLU² (In |G det A| - Cunx) ReLU (Cmin - InIInnC IG det Al)
with a penalty parameter Cp. For the experiments in this work, we used Ciin -2, Cmax 15, (000) min -6 c(0) inf, c. 2.5, and (1)) 15.5.
B IMPLEMENTATION OF P'INV LAYERS
In practice, the I'v layer stores the current inverse Ainv A-1 and determinant alongside A. The forward pass of an input vector x can computed efficiently by first computing u² and then adding vu x to Ax. The inverse pass can be similarly structured 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 the loss function J/ is composed of elements from aJ/0B and aJ/0A. 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 and equation can additionally be augmented rectify numerical errors made during optimization.
In order to allow crossings of otherwise inaccessible regions of the parameter space the merging step accepted every Nioree merges, even if the determinant poorly conditioned. If u or ever contain non-numeric values, merging steps were rejected and the perturbation reset without performing merge.
C NORMALIZING FLOWS
A Normalizing flow is learnable bijection f: R" - R" which transforms 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 has the density:
q(x) = p (f-)(()) J.-
Given a target distribution p(x), this tractable density allows minimizing the reverse Kullback- Leibler (KL) divergence DEL [g(x)]l c.g., if p(x) is known up to normalizing constant, or the forward KL divergence DKL (((x)]l(()) having access to samples from p(x).
D COUPLING LAYERS
Maintaining invertibility and computing the inverse and Jacobian is challenge, if f could be an arbitrary function. Thus, it is common to decompose f in sequence of coupling layers
/ (= ) 8(1) o...g(k) S(K).
Each (k) constrained the form (x) (() (Tk) (x), x2) 0æz, where 2 @32, 3i R and J2 n- Here (k); R x nm Rmm is point-wiss transformation, which is invertible in
14

---

Under review as a conference paper at ICLR 2021
its first component given fixed 22. Possible implementations include simple affine transformations (Dinh et al. 2014; 2016) as well as complex transformations based on splines (Müller et al. 2019
Durkan et 2019 Each (4) thus has block-triangular Jacobian matrix
J J
(
(
Imxm
where JTa is (n m) x (n m) diagonal matrix. The layers S(A) take care of achieving mixing between coordinates and are usually represented as simple swaps
S
I 1n-mxn-m
=
I Imxm
The total log Jacobian of Ja is then given by
K
log det Jo
Σ
tr [log (Jroo)] + log det ((k)
where log det (k) when ((k) is given by the simple swaps above.
E BOLTZMANN GENERATORS
Boltzmann Generators (Noé al., 2019) are generative neural networks that sample with respect to known target energy, as for example given by molecular force fields. The potential energy u S - IR of such systems a function of the atom positions x. The corresponding probability of each configuration is given by the Boltzmann distribution (x(x) exp(-3u(x))/2 /ZZ /Z, where = 1/(kgT) is inverse temperature with the Boltzmann constant kg. 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-hased training. Concretely, the following loss function is minimized
J(A) (i (A)),
where + = denote weighting factors between density estimation and energy-based training. The maximum likelihood and KI. divergence equation are defined respectively as
[ Fozz Fox (x; A)l1 In Rxz(x; A)
Ji Ex
and
Ezup. (F (z; A)) In Rux (z; A)]
As example, train model for alanine dipeptide, molecule with 22 atoms, in water. Water is represented by an implicit solvent model. This system was previously used in WWu et al. (2020). Training data was generated using MDD simulations at 600K to sample all metastable regions.
F TRAINING OF LINEAR TOY PROBLEMS
The I'Inv layers were trained using a stochastic gradient descent optimizer with a learning rate of 10 -2 and the hyperparameters from Table I
The matrices were initialized with the identity.
Table 1: Hyperparameters for linear toy problems
Problem (Dimension)
N
Ntoree
Neorrect
Penalty Co
Positive Definite (32)
1, 10, 50
10
50
0.1
Orthogonal (128)
1
10
50
0.1
-I (101)
10
50
0.0
15

---

Under review as a conference paper at ICLR 2021
G TRAINING OF 2D DISTRIBUTIONS
The P'lny layers used for 2D distributions were composed of blocks containing
1. P'Inv layer with bias (2D),
2. an elementwise Bent identity,
3. another I'Inv layer with bias (2I),, and
4. an inverse elementwise Bent identity.
100 of these blocks stacked resulting in 1200 tunable parameters (counting elements of and v). p'In training was performed with = 10, Ntorce = 10 and Ncorrect 50. No penalty was used. Matrices were initialized with the identity Iz.
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 distributions. Adam optimization was performed for 8 epochs of 20000 steps and with a batch size of 200. The initial learning rate 5 10-3 and decreased by factor of 0.5 in each epoch. After each merging step, the metaparameters of the Adam optimizer were reset to their initial state.
two-dimensional standard normal
Figure 8:
Samples from P'Iny training via the inverse.
Figure showing samples from network with only inverse P'Iny blocks. While the samples worse than with forward blocks, the distributions still well represented. This result indicates that the premature convergence encountered for linear test problems is lesser problem in nonlinear problems and deep architectures.
H TRAINING OF BOLTZMANN GENERATORS
The normalizing flows in Boltzmann generators were composed of the blocks shown in Figure and mixed coordinate transform as defined in Noć et al. (2019). The test problem was taken from Dibak 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 P'Inv 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 P'Iny layers are affected by the pt updates. Merging was performed every N 100 steps with Nioree 10 and Neorrect 50. No penalty was used, i.c. Co = 0.0. The P*Inv matrices were initialized with the reverse permutation, i.e. Aij 5((n-j)-
Density estimation with Adam was performed for 40,000 optimization steps with a batch size of 256 and learning rate of 10-3, A short energy-based training epoch appended for 2000 steps with a learning rate 10 and w. /w, = 0.05. After each merging step. the metaparameters the Adam optimizer reset their initial state for all I'Iny parameters.
16

---

Under review as a conference paper at ICLR 2021
Baseline Block
Płinv Block
RNVP
cat
P4Inv
swap
51
RNVP
RNVP
swap
P4Inv
RNVP
RNVP
split
Figure 9: Neural network blocks used in the Boltzmann generator application. The baseline archi- tecture is a normalizing flow composed of Real NVP (RNVP) coupling blocks. RNVP uses input from two channels æ1 and 22. The input from the first channel is left untouched, 311 21, while the output 1/2 from the second channel is conditioned on the first channel through two neural networks t and s. Each block of the baseline model contains are bracketed by splitting and concatenation (cat) of the data. Instead of the swapping steps, the I'Inv model uses invertible linear layers that are trained through p updates.
RNVP blocks and two swapping steps that
I MNIST CLASSIFICATION VIA RANK-ONE UPDATES
Compared to fully-connected multi-layer perceptrons (MLP), rank-one updates reduce the number of independent trainable parameters per layer from m n m + n where m and are the input and output dimension, respectively. It is therefore useful to study how the reduced search space dimension affects the training efficiency a nonlinear setup. To this end, non-invertible classifier MLPs were trained on MNIST with unrestricted search space well through rank- one updates. The original network was composed two convolutional layers, two dropout layers, and two linear layers. In pa training, the two linear layers (dimensions 9216x 128 and 128x 10) were trained through rank-one updates which were merged in every iteration (N = 1). Naturally, this training did not involve any inverse determinant updates, A vanilla SGD optimizer was used with various learning rates.
100
Learning Rate
0.02
Test Accuracy [%6]
5e-03 (Standard MLP)
Training Loss
95
5e-03 (P*)
5e-04 (Standard MLP)
0.01
90
5e-04 (P*)
5e-05 (Standard MLP)
0.00
85
5e-05 (P4)
0
20
40
0
20
40
Epoch
Epoch
Figure 10: Training loss and test accuracy during MNIST training with a vanilla SGD optimizer averaged over ten replicas. Standard multilayer perceptrons (MLP) are compared with rank-one updates.
Figure 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 mank-nn updates became less efficient.
17


Reviews:
[{"review_id": "vVG46522n9c", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.  \n\nThe method is (only) exemplified with the particular case where one needs to optimize a network with the property of having invertible layers (which is an important use case for example for normalizing flows, and invertible networks). In this particular case, the paper shows that the transformation of the parameters can be cast as a perturbation of rank one matrices using closed-form formulas that can be used to check and guarantee the invertibility. The parameters are updated periodically, after a series of perturbations which helps to stabilize the optimization. The paper shows three experiments on two synthetic datasets (linear data and 2d manifold) and  one on Boltzmann generators to generate samples of Alanine Dipeptide.\n\nThe paper presents an interesting idea and some empirical evidence showing promising results. My main concern with the paper is that the experimental evidence is quite limited so it is hard to judge the real contribution of the method. Additionally, the paper could improve the organization.\n\nIn particular, Section 2 needs a better organization. The title mentions a general idea, but in practice only the case of invertible matrices is analyzed and discussed. Section 2.2. is rather disconnected from Section 1.1. There's no motivation why this is there. (Maybe this should be moved to a subpart of Section 2.3). Same applies to Section 2.4 (implementation details?). It might be better to specifically focus and present the method for the the case mentioned in Section 2.5. Also, it seems more reasonable to discuss the related work (Section 3) before jumping into the algorithm details (part of Section 2).\n\nSome questions:\n\n-- Update u and v. The optimization in u and v has too many degrees of freedom. For example, you could constrain to always have ||v||=1 without losing anything. Will this help to avoid local minima?\n\n-- In the first experiment (linear data) it is mentioned that when training against linear training data P^4Inv can get stuck into local minima, and that this does not happen when using more complex data. Would you mind elaborating a little more on this observation?\n\n-- What can be concluded from Figure 4? It seems to me that all the tested methods did a good job here.\n\nTypos and minor comments:\n--Pag2 - B_0 = ids. I think this should be R_{B_0} = ids\n\n--Pag6 -  B=-I_{101}, I assume this should be A=-I_{101}.\n\n--Pag7 - \"since it with a given\" (writing)\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ypay74WVDiD", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Update: \nI'd like to thank the authors for carefully addressing my concerns.. \n\nReviewer 2 claims $P^4$ training is a special case of previous work. It seems the authors agree at least to some extent *\"we agree that the P4 update scheme is closely related to dynamic trivializations\"*. It is not clear to me how harshly this should be penalized. In some cases, it is interesting to research special cases of known general results. Unfortunately, it is hard for me to judge if this is the case here, as I know of the previous articles reviewer 2 refer to, but I have not studied them carefully. As a result, I am decreasing my confidence from 3 to 1 and my score from 7 to 6. \n\nI hope this does not discourage the authors, and wish them good luck with future research. \n\n___\n\n### Summary \n**Objective:** train neural networks while retaining properties like invertibility or orthogonality of layers. \n\n**Approach:** instead of optimizing normal weights optimize an rank 1 update. Occasionally, the rank 1 updates are merged with the network parameters.  \n\n### Strengths\n**[+]** Preserving properties like invertibility is important and has been studied by much previous work. The authors present a novel approach based on rank 1 updates, which, to the best of my knowledge, is completely novel.  \n\n**[+]** The article is very clearly written, it seems to me that the authors spent a great deal of time polishing the paper. \n\n### Weaknesses\n**[-]** I have a minor concern wrt. optimization of rank 1 updates which I elaborate below. \n\n### Recommendation: 6\n**[+]** The paper presents a novel approach for preserving invertibility, which could benefit many deep learning researchers. \n\n**[+]** The paper demonstrates how rank 1 updates can be used in deep learning, which I believe will inspire further research into this interesting direction. \n\nI condition my recommendation on an MLP experiment I already discussed with the authors before submitting this review. Furthermore, I'll re-evaluate my conviction after reading the comments by the other reviewers. \n\n### Questions \nThe following question was already answered by the authors before the submission of this review. I repeat the question here for completeness. \n\n---\n\nRecent research explore the idea that SGD variants perform better with networks that has more parameters. Informally, it is argued it is harder to get stuck at local minima when SGD can move in more directions. If one optimize rank 1 updates, SGD would have $2d$ directions to move in instead of $d^2$ directions. I am concerned this might impacts the performance of SGD. This concern would be address by the following experiment: train two MLPs on MNIST, one with SGD and one with $P^4$, do they attain similar loss curves?\n\n---\n\n### Additional Feedback\nI didn't find any typos, the article seems to be very polished. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "FOwNrGzxJ4z", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Summary\nThis work proposes parameter perturbations for the training of neural networks while preserving some properties. It particularly focuses on the case of invertible layers and demonstrates improvements over classical architectures for density estimation with normalizing flows. It is empirically shown that parameter perturbations allow for sign changes in the determinant.\n\n# Major comments\n## Pros:\nThis work is original and tackles an important problem, that is learning invertible matrices. I think the approach is original and the fact that it could be applied for other properties makes the development and the proof of work very interesting. \nAt some point you discuss some weaknesses of your work, I really liked that as it will definitely help people using your method. The paper is clearly written and so easy to follow. The experiments are insightful and help the reader to get a better feeling of what is happening at training time. Some results suggest that using your method clearly helps for certain density estimation tasks. Additional experiments could maybe show that P4 + Bent leads to efficient and expressive normalizing flows.\n\n## Cons:\n1) The focus of the paper is not really on parameter perturbations for training neural networks in general. Most of the paper discusses how perturbations can be applied for learning invertible matrices. Introducing the paper as if it could be applied successfully for any kind of properties does not make a lot of sense to me as only experiments on invertible matrices is presented.\n2) The paper discusses some concurrent methods for learning invertible matrices such as https://arxiv.org/pdf/2006.15090.pdf. However, it is not clear which method is preferable in which situation. What is even more embarrassing is that no comparison at all is performed against this work.\n3) A serious discussion about the computational cost of using perturbation instead of another method for learning an invertible matrix is lacking.\n4) The benchmark you are using to assess the performance gain on the density estimation task is a bit odd to me and does not allow the reader to really assess the usefulness of your method in such a setting. \n\n## What could be done to address my comments\n1) Either perform additional experiments showing the applicability of your method for a large bunch of interesting properties (e.g. orthogonality). Instead, I would suggest clarifying the title, the abstract, and the intro. Adding a discussion about other interesting properties that could be kept with your method between the experiments and the conclusion would make sense to me. As is I feel like you're selling something more generic than what you really show with your experiments whereas I believe learning invertible matrices is already an important and interesting achievement.\n2) Add a comparison with the referenced paper. And discuss the potential advantage of each method.\n3) Discuss the computational cost. As an example how long is the training of P4 inv in comparison to linear in fig 3.\n4) Try to discuss the more general advantage of using learned invertible linear layers in flows instead of predefined PLU factorization. I don't think benchmarking your \"flow\" with respect to the UCI datasets will really improve the quality of your manuscript but it would help the reader to get a sense of some possible additional expressivity gained with your method.\n\nIf you address comments 1, 2, and 3 (or convince these are not important), I will raise my score. Addressing 4 is less important but it would definitely improve the manuscript's quality.\n\n# Minor comments\nPages 2: B_0=id_S -> R_{B_0} = id_S ?\nPage 3: Could clarify/elaborate on the last sentence just before algorithm 1.\npage 5: I would avoid classifying NFs into two categories as you do.\npage 6: Beginning, \"quadratic matrix\" -> \"square matrix\" ?!\n4.3: \"since it with\" ?\nWhy not just using elu, sigmoid, or than instead of Bent identities? Could you explain why if it is important?", "extra_scores": {}, "extra_text": {}}, {"review_id": "-1RKMsN6j77", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 2, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "***Summary and general comments:***\n\nThis paper presents a method to parametrize a set of constraints via a linear space that may change. This method was already studied in much more depth in [1], where this idea is explored through the lens of vector bundles and retractions on them, with a convergence result appearing in the follow-up work [2].\n\nThe authors then instantiate these ideas in the setting of normalizing flows via the Sherman–Morrison inversion formula.\n\n***Questions:***\n\n1) For instance, parameterizing the whole orthogonal group (via the Cayley parameterization or matrix exponentials of skew-symmetric matrices) is computationally expensive. In contrast, orthogonal matrices can be cheaply perturbed into other orthogonal matrices using double Householder transforms or Givens rotations.\n\nWhy is that the case? Could you please justify the first assertion?\n\nWhat do you mean by \"it can be cheaply perturbed into an orthogonal matrix\"? The best orthogonal approximation to a given matrix is the orthogonal projection given by the polar decomposition. That is expensive.\n\n2) \"In contrast to this related work our method reparameterizes the update step during gradient descent rather than the parameter matrix itself\"\n\nThis is not really different as explained in [1]. Using a dynamic trivialization changes the metric that you are working with, but you can still use it as a parametrization. In your example, it would account for changing the problem\n$$\\min_{x \\in \\operatorname{GL}(n)} f(x)$$\nto\n$$\\min_{u,v \\in \\mathbb{R}^{2n}} f(X_0 + uv^\\intercal)$$\nfor a fixed $X_0$ and changing the problem via the dynamic trivialization framework every $N$ steps. Using this formulation, you can let the optimizer do all the heavy lifting for you. See also 5)\n\n3) Trivializations are known to have a problem with the momentum and adaptive term. I see that you use Adam in the training of the 2D distributions. How do you solve the problem of the incorrect momentum and adaptive terms when you change the base?\n\n4) A general question. I do not understand what is the motivation behind using the low-rank update besides the fact that it allows for an (amortized) low cost of the inversion.\n\n5) A follow-up of the previous question: How is $R_{X}(u,v) = X + uv^\\intercal$ a sensible map to use? This map does not have its image on $\\operatorname{GL}(n)$ but on all $\\mathbb{R}^{n\\times n}$! For example, $R_{-uv^\\intercal}(u,v) = 0$ which is clearly not invertible. As such, I do not understand the claims in section 2.3 about this being a \"fully flexible invertible layer\". This is the reason behind having to use Algorithm 2 in P⁴Inv. Given that Algorithm 2 is used, what is the reason behind using this parametrization at all over just doing unconstrained optimization?\n\n\n***Citations:***\n- Lezcano-Casado & Martínez-Rubio do not use Givens rotations but the exponential. The first to use Givens rotations in the context of ML was [3].\n\n- When it comes to Riemannian gradient descent, it is probably better to cite Absil's book [4] as a general reference rather than two recent papers, as this is a well studied topic.\n\n- I would recommend to clean-up the bibliography, as there are many citations that point to the arXiv when the articles have indeed been published in peer-reviewed venues.\n\nMinor:\n- \"which is a R-diffeomorphism\" -> \"an\"\n\n***Conclusion:***\n\nI really like the first experiment for its simplicity, trying to elucidate the behavior of the layer. It is also nice to see the improvement that this idea gives over RNVP. At the same time, when it comes to the experiments, I believe that it would have been of interest to compare this approach with other known ways to parametrize invertible linear layers, such as those that use QR, SVD or Choleski factorizations.\n\nThat being said, as mentioned in 4) and 5) I do not see the reason for this being a good way to obtain invertible layers, given that even the image of the parametrization does not lie in GL(n). Furthermore, the ideas behind the framework presented in this paper were already studied in previous papers in much more depth.\n\n\n[1] M. Lezcano-Casado. “Trivializations for gradient-based optimization on manifolds”. NeurIPS, 2019\n\n[2] M. Lezcano-Casado. “Curvature-Dependant Global Convergence Rates for Optimization on Manifolds of Bounded Geometry”. https://arxiv.org/abs/2008.02517\n\n[3] U. Shalit, G. Chechik. “Efficient coordinate-descent for orthogonal matrices through Givens rotations”. ICML, 2014\n\n[4] P.-A. Absil, R. Mahony, and R. Sepulchre.Optimization algorithms on matrix manifolds. PrincetonUniversity Press, 2009- .", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper discusses a method to update/optimise invertible matrices via low-rank updates. The key property of the proposed method is that it keeps track of the matrix inverse and its determinant through the optimisation (with updates that are much cheaper to compute than a direct inversion/determinant computation).

While the method of performing low-rank updates for invertible matrices itself has already been extensively studied in the literature as pointed out by reviews, this work focuses (after extensive revision) on the properties of this update method.
 Since the updates may leave the manifold of invertible matrices, a numerical stabilisation step was introduce whereby updates that produce ill-conditioned matrices are rejected during optimisation.

Rank-one updates allow for fast update of matrix inverse and determinants. So this is particularly interesting when applied to normalising flows, as it allows for cheaper computation of the log-det-Jacobian terms.  

The novelty of this approach is rather limited (as also pointed out by R2). The experiments and, in particular, the application to normalising flows are interesting, well-executed. It is not clear if there are advantages of the method in other domains where log-det-Jacobians are not necessary relative to existing literature.

Decision: Reject