Training Invertible Linear Layers through Rank-One Perturbations

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 P4Inv 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 P4 concept can be utilized to retain properties other than invertibility.

Reviews and author responses:

### Update: 
I'd like to thank the authors for carefully addressing my concerns.. 

Reviewer 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. 

I hope this does not discourage the authors, and wish them good luck with future research. 

___

### Summary 
**Objective:** train neural networks while retaining properties like invertibility or orthogonality of layers. 

**Approach:** instead of optimizing normal weights optimize an rank 1 update. Occasionally, the rank 1 updates are merged with the network parameters.  

### Strengths
**[+]** 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.  

**[+]** The article is very clearly written, it seems to me that the authors spent a great deal of time polishing the paper. 

### Weaknesses
**[-]** I have a minor concern wrt. optimization of rank 1 updates which I elaborate below. 

### Recommendation: 6
**[+]** The paper presents a novel approach for preserving invertibility, which could benefit many deep learning researchers. 

**[+]** The paper demonstrates how rank 1 updates can be used in deep learning, which I believe will inspire further research into this interesting direction. 

I 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. 

### Questions 
The following question was already answered by the authors before the submission of this review. I repeat the question here for completeness. 

---

Recent 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?

---

### Additional Feedback
I didn't find any typos, the article seems to be very polished. 

Merges are rejected whenever the rank-one update is ill-conditioned (Section 3.3). It is shown throughout the paper that this stabilizes the update so that the actual (frozen) weight matrix remains on the manifold. This is proven in the experiments. In fact, this stabilization step is one of the central aspects of the paper and one of the main reasons to choose a stepwise reparameterization approach. The ability to leave (and reenter) the set of feasible parameters between merging steps is in contrast to existing methods, particularly dynamic trivializations.

***Summary and general comments:***

This 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].

The authors then instantiate these ideas in the setting of normalizing flows via the Sherman–Morrison inversion formula.

***Questions:***

1) 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.

Why is that the case? Could you please justify the first assertion?

What 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.

2) "In contrast to this related work our method reparameterizes the update step during gradient descent rather than the parameter matrix itself"

This 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
$$\min_{x \in \operatorname{GL}(n)} f(x)$$
to
$$\min_{u,v \in \mathbb{R}^{2n}} f(X_0 + uv^\intercal)$$
for 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)

3) 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?

4) 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.

5) 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?


***Citations:***
- 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].

- 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.

- 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.

Minor:
- "which is a R-diffeomorphism" -> "an"

***Conclusion:***

I 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.

That 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.


[1] M. Lezcano-Casado. “Trivializations for gradient-based optimization on manifolds”. NeurIPS, 2019

[2] M. Lezcano-Casado. “Curvature-Dependant Global Convergence Rates for Optimization on Manifolds of Bounded Geometry”. https://arxiv.org/abs/2008.02517

[3] U. Shalit, G. Chechik. “Efficient coordinate-descent for orthogonal matrices through Givens rotations”. ICML, 2014

[4] P.-A. Absil, R. Mahony, and R. Sepulchre.Optimization algorithms on matrix manifolds. PrincetonUniversity Press, 2009- .

According to the whole paper, the authors argue that they study property preserving transformations, and they choose $\operatorname{GL}(n)$ to be their set, but now you have confirmed that the image of the low-rank transformation is not even in $\operatorname{GL}(n)$, so that it is not "property preserving" not even in this __very__ weak sense. As such, I think that these statements in the paper that eq. (5) is $P^4$, that is, image in $\operatorname{GL}(n)$ is wrong. This being the cornerstone of  the paper, I think that it disqualifies the paper for being accepted.

I prompt the authors to rethink the properties of their optimisation method, and try to match it with previous literature in the field, as the idea of parametrising sets of matrices to perform optimisation (what they call $P^4$ framework) is by no means new. Even more, I believe that they should rethink what is the image of the map that they use since, as it stands, it is not even contained in the set that they want to perform optimisation on. 

# Summary
This 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.

# Major comments
## Pros:
This 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. 
At 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.

## Cons:
1) 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.
2) 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.
3) A serious discussion about the computational cost of using perturbation instead of another method for learning an invertible matrix is lacking.
4) 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. 

## What could be done to address my comments
1) 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.
2) Add a comparison with the referenced paper. And discuss the potential advantage of each method.
3) Discuss the computational cost. As an example how long is the training of P4 inv in comparison to linear in fig 3.
4) 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.

If 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.

# Minor comments
Pages 2: B_0=id_S -> R_{B_0} = id_S ?
Page 3: Could clarify/elaborate on the last sentence just before algorithm 1.
page 5: I would avoid classifying NFs into two categories as you do.
page 6: Beginning, "quadratic matrix" -> "square matrix" ?!
4.3: "since it with" ?
Why not just using elu, sigmoid, or than instead of Bent identities? Could you explain why if it is important?

I thank the authors for their response, as promised I have raised my score to 6 thanks to the modifications made to the manuscript.

We thank the reviewer for those insightful comments. We are especially grateful for having been pointed to the dynamic trivializations introduced by Lezecano-Casado, which we had missed in our literature review. Since we agree that the P4 update scheme is closely related to dynamic trivializations, we have thoroughly revised the manuscript to take into account previous work by Lezcano-Casado. This is now mentioned in the introduction and discussed in more depth in Section 2.4. Note that due to the comments by all reviewers, the focus of the paper has shifted towards optimization of invertible matrices rather than the general P4 algorithm, which also underlines its novelty with respect to dynamic trivializations. Please find the responses to the detailed comments below.

1. This sentence was removed since it no longer fit in with the shifted focus of the paper. What we meant to say there was that rank-one or Householder perturbations only require a few matrix vector multiplications in the forward pass, which is computationally cheap. In contrast, exponential maps often involve matrix exponentials, which are comparably more expensive.
2. Thanks for pointing this out. We have generally augmented the comparison with this previous work by Gresele et al. In this context, this sentence was revised and now reads “In contrast to this related work, our method facilitates a cheap inverse pass and allows sign changes in the determinant.“ (See Section 2.2.)
3. We reset the metaparameters of the optimizer corresponding to the P4Inv parameters. This is now mentioned in Appendices G and H.
4. Facilitating cheap inversion and determinant computation is exactly the point of the method. This is particularly interesting for normalizing flows that are trained via a combination of the inverse and forward KL divergence (density estimation + sampling) as in our Section 4.4. This aspect is now highlighted in the introduction and amplified by the new Section 4.1 on computational cost.
5. The reviewer is right that rank-one perturbations are not sensible if the goal is just to retain invertibility. The main reasoning behind this choice is that it allows cheap inversion. This is the point, where our approach deviates from dynamic trivializations. While dynamic trivializations are primarily meant to stay on the manifold, our main concern is how to retain knowledge of the inverse and determinant. This point is now highlighted in the introduction and mentioned again throughout the manuscript.

Other decompositions: Due to this comment, we have compared the computational cost of P4Inv layers to LU factorizations in Section 4.1. Parameterizations based on triangular matrices are fairly inefficient on parallel computers, while the P4Inv layers are inverted at the cost of 3 matrix vector multiplications.

We have also revised our references to include the suggestions of the reviewer.

We hope that these revisions mitigate the concerns and that the reviewer deems the updated version of the manuscript acceptable for the conference.

We thank the reviewer for the positive evaluation and the valuable suggestion. Note that the manuscript underwent  a major revision to incorporate comments by the other reviewers. 

Regarding the reduced search space dimension, we have included the suggested experiment as Appendix I, which is briefly discussed in the main text in Section 4.2. Please note that we have to revise our initial answer to the comment. At first, we did not recognize that scaling the standard deviation of the reinitialization essentially amounts to a change in learning rate. This aspect is now mentioned in Section 3.2.


We thank the reviewer for those valuable suggestions. The manuscript underwent a major revision to address all comments.

1. As suggested, we shifted the focus of the paper so that the whole paper, including title, abstract, and introduction are now centered around invertible layers. Potential applications to other properties are briefly discussed before the conclusions.
2. The main advantage with respect to the paper by Gresele et al. is the cheap inversion during training. This aspect is now highlighted in the introduction and in Section 2.2. We also mention a potential drawback of our method induced by the lower-dimensional search space to provide a fair comparison.
3. A new subsection (4.1) has been added to the experiments discussing solely the computational cost. We felt that a comparison with a standard layer would not be fair, since the main advantage of the method is the cheap evaluation of determinants and inverses, which is not a part of standard training. As for this (unfair) comparison, the standard training was approximately a factor of 8 faster in all linear test problems. Instead we compared the computational cost of P4Inv layers with other methods to parameterize inverses and determinants.
4. The subsection on computational cost also includes LU factorizations.

The minor comments were also valid and helpful. We hope that the reviewer might reconsider his or her score and rate the revised version of the manuscript acceptable.

We thank the reviewer for those helpful comments. 
Due to comments by all reviewers, we have decided to explicitly put the focus of the manuscript on invertible layers rather than the general update scheme. This is also reflected in our revised title and abstract and generally helps the organization of the methods section. We have also added further experiments: (a) a runtime comparison between P4Inv layers, standard linear layers, and LU decompositions in Section 4.1 and (b) training of an MNIST classifier through rank-one updates in Appendix I. We are confident that these experiments further underline the capabilities and limitations of our approach. 

Regarding the reviewer's specific questions:
* Dimension of rank-one updates: We agree that each rank-one update has an intrinsic dimensionality of 2n - 1, while it is represented by 2n parameters in our reparameterization. We tested a version of the algorithm, where ||v|| was constrained to unity. This did not prevent premature convergence and also altered the optimal learning rate. When reinitializing v from Gaussian Noise with sigma=1 (as we did in our experiments), gradient updates to u are of the same magnitude as direct gradient updates to A so that the same learning rate can be used for standard training and P4Inv training. This is reflected by the fact that P4 training converges with the exact same rate as standard training on linear test problems (see Figure 3).
* Local minima: We have added the results of density estimation of 2D distributions through via the inverses of P4Inv layers to Appendix G. The networks are still capable of approximating the optimal solution, although the match is inferior to forward P4Inv layers.
* Figure 4: This Figure shows that sign changes in the determinant are possible in P$^4$Inv training without permanently hurting the inverse. As noted by Papamakarios et al. (arXiv:1912.02762), a continuous update rule cannot parameterize all of GL(n), since determinants can only change sign by passing through zero, where the matrix is singular. Figure 4 shows that passing through zero is possible in discrete steps.
* Thanks also for catching some minor inaccuracies and typos.

We hope that the paper is deemed acceptable with those changes.


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.  

The 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.

The 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.

In 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).

Some questions:

-- 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?

-- 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?

-- What can be concluded from Figure 4? It seems to me that all the tested methods did a good job here.

Typos and minor comments:
--Pag2 - B_0 = ids. I think this should be R_{B_0} = ids

--Pag6 -  B=-I_{101}, I assume this should be A=-I_{101}.

--Pag7 - "since it with a given" (writing)


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: withdrawn-rejected-submissions