Title: projUNN: efficient method for training deep networks with unitary matrices
Authors: Bobak Kiani
Randall Balestriero
Yann LeCun
Seth Lloyd
Venue: NeurIPS 2022

Abstract: In learning with recurrent or very deep feed-forward networks, employing unitary matrices in each layer can be very effective at maintaining long-range stability. However, restricting network parameters to be unitary typically comes at the cost of expensive parameterizations or increased training runtime. We propose instead an efficient method based on rank-$k$ updates -- or their rank-$k$ approximation -- that maintains performance at a nearly optimal training runtime. We introduce two variants of this method, named Direct (projUNN-D) and Tangent (projUNN-T) projected Unitary Neural Networks, that can parameterize full $N$-dimensional unitary or orthogonal matrices with a training runtime scaling as $O(kN^2)$. Our method either projects low-rank gradients onto the closest unitary matrix (projUNN-T) or transports unitary matrices in the direction of the low-rank gradient (projUNN-D). Even in the fastest setting ($k=1$), projUNN is able to train a model's unitary parameters to reach comparable performances against baseline implementations. In recurrent neural network settings, projUNN closely matches or exceeds benchmarked results from prior unitary neural networks. Finally, we preliminarily explore projUNN in training orthogonal convolutional neural networks, which are currently unable to outperform state of the art models but can potentially enhance stability and robustness at large depth.

TL;DR: We propose an efficient method to train neural networks with layers whose linear transformations are strictly unitary or orthogonal.

Review discussion:
Strength
- The authors propose two novel projected gradient step procedures for training constraint to the space of unitary matrices.
In contrast to exisitng approaches, we achieve more computational efficiency by exploiting the low-rank structure of gradients that is present in minibatch training.
- Experimental evaluations show that the low-rank approximation even allows ranks substantially smaller than the batch size ( times sequence length for RNN), increasing the computational complexity even more.
- Discussion of related work places this work among existing work, clearly stating advances and differences.
- The extension to convolutional layers by exploiting the relation between the Fourier transformation and the discrete convolution operation enables the use of unitary operators in the image domain. This may allow users more control about properties of the linear operators that have been linked to generalization in deep networks like spectral norm, stable rank or eigenvalue decay.
- Source code for recurrent networks is available which facilitates widespread use. Unfortunately the CNN code was not available in time for the review
Weaknesses
- The experiments mostly show technical advantages like increases in computational performance, not benefits in applications where this technique is strictly required
-What are applications beside the RNN toy problems where unitary operators may be needed?
The authors have adequately addressed the limitations. I cannot think of any potential negative societal impact
No
We thank the reviewer for their feedback and their positive comments about our work.
**Reviewer**: “Unfortunately the CNN code was not available in time for the review” \
**Response**: For reviews, we are including an example simulation of the unitary CNN in an anonymous Google Colab (https://colab.research.google.com/drive/1FCfRKxrbhelP1mcLuiq6PxnCt1BG0wNp?usp=sharing) where we learn MNIST data. This illustrates the use of projUNN in a simple CNN setting. Our full code base will be released as a github link upon publication.
**Reviewer**: “What are applications beside the RNN toy problems where unitary operators may be needed?”\
**Response**: We agree with the reviewer that it would be useful to find further applications where strictly enforcing unitarity/orthogonality enhances performance for practical applications. Looking beyond the RNN experiments performed here, possible avenues of exploration include enforcing isometry as an inductive bias of the architecture (e.g., see [1] for inspiration), constructing models whose layers are more robust to adversarial attacks (prior works explored this in [2],[3] among others), developing stable invertible layers for normalizing flows (e.g., see [4]), and potentially learning quantum data [5]. In many practical settings, simply regularizing weights to be close to unitary can be effective, so we are searching for settings where strict unitarity is actually necessary.
**References**\
[1] Baddoo, P. J., Herrmann, B., McKeon, B. J., Kutz, J. N., & Brunton, S. L. (2021). Physics-informed dynamic mode decomposition (piDMD). arXiv preprint arXiv:2112.04307.\
[2] Singla, S., & Feizi, S. (2021, July). Skew orthogonal convolutions. In International Conference on Machine Learning (pp. 9756-9766). PMLR.\
[3] Trockman, A., & Kolter, J. Z. (2021). Orthogonalizing convolutional layers with the cayley transform. arXiv preprint arXiv:2104.07167.\
[4] Golinski, A., Lezcano-Casado, M., & Rainforth, T. (2019). Improving normalizing flows via better orthogonal parameterizations. In ICML Workshop on Invertible Neural Networks and Normalizing Flows. \
[5] Biamonte, J., Wittek, P., Pancotti, N. et al. Quantum machine learning. Nature 549, 195–202 (2017). https://doi.org/10.1038/nature23474
Strengths:
- Very compelling result, might enable making URNNs practical
- Very good exposition
- Extensive experiments
Weaknesses:
- Both methods seem fairly complicated to implement, full peudocode algorithm descriptions (including the low-rank gradient approximation) would be appreciated. Hopefully code will be released.
None
None
No
We thank the reviewer for their feedback and nice comments about the exposition and presentation of our paper.
We agree with the reviewer that pseudocode would be helpful and we have added this in Appendix D. Space-permitting given the added page in the full publication, we will use the remaining space to add the pseudo-code to the main text.
Furthermore,  we have included another google colab link (https://colab.research.google.com/drive/1FCfRKxrbhelP1mcLuiq6PxnCt1BG0wNp?usp=sharing) which has an example of projUNN used in a convolutional network trained on MNIST. Our complete code base will be released as a github link upon publication of the manuscript.
The paper is well-motivated and clearly-written. The proposed methods are simple-yet-effective, and the theoretical results are strong (which I have checked in details). Comparison with prior works looks comprehensive. Due to my limited knowledge in related works, I will leave the evaluation of novelty as well as experimental results to other reviewers.
The authors propose two methods (projUNN-D and projUNN-T). Is there any difference between these two methods other than the numerical issues? I would like to see some comparisons (e.g., the pros and cons) between them.
Yes. The authors discussed the numerical instability of the direct method as well as the memory concern when applied to orthogonal convolutions.
No
We thank the reviewer for their feedback and for their overall positive assessment of our work. We have copied their question below followed by our answer.
**Reviewer**: “The authors propose two methods (projUNN-D and projUNN-T). Is there any difference between these two methods other than the numerical issues?”\
**Response**: As projUNN-D and projUNN-T are both first-order gradient based methods, their performance is very similar.
However, projUNN-D closely follows methods of projected gradient descent whereas projUNN-T follows methods of Riemannian gradient descent. They both agree up to first order (i.e. are essentially equivalent for infinitesimally small learning rate) and are a retraction of the exponential map. Generally, we have found that Riemannian gradient descent methods are more stable and faster as stated in the paper, but we also find projUNN-D to be effective in many cases and also potentially simpler to implement. We hope future work can be inspired by either of these methods to extend the speed-ups found in unitary/orthogonal matrices to other matrix manifolds.
Overall, the paper is clear and well written; the appendix provides a nice mathematical background/review, many interesting remarks/extensions/experiments. The main strength of the paper arises from the observation that low rank matrices can accurately approximate gradients (where fidelity is measured by residual frobenius norm error). This approximation allows the proposed method to be essentially a factor of O(n) faster than its competitors.
One potential downside is that the proposed method (projUNN) and some of the baselines use eigen decomposition/QR decompositions which suffer some overhead due to parallelization issues in the GPU, causing some methods to appear slower in the experiments. Authors mention that this issue could be potentially fixed in the future with updates to DL libraries.
It would be interesting to compare train/test loss evolution over iterations between projUNN and the other baselines (i.e. testing convergence in terms of iterations rather than time) as well as total time to achieve similar loss (as opposed to single iteration cost).
Minor Remarks:
Typo: pg. 23 $$\hat a = \tilde{M}^{\dagger} a \rightarrow \hat a = U^{\dagger} a$$
Yes
No
We thank the reviewer for their feedback and nice comments about the presentation and relevance of our manuscript. We also thank the reviewer for pointing out the typo in pg. 23 which we have corrected.
**Reviewer:** “It would be interesting to compare train/test loss evolution over iterations between projUNN and the other baselines (i.e. testing convergence in terms of iterations rather than time) as well as total time to achieve similar loss (as opposed to single iteration cost).”  \
**Response:** For most tasks considered here, we benchmarked performance with respect to reported results. Most of these tasks were learned with RNNs whose matrices are small enough to efficiently perform operations like matrix inversion and SVD, and thus runtimes, in most of these settings, are roughly equivalent for projUNN and other models. Since our goal was to address poor runtime scalings of the prior methods, we did not directly compare runtimes for the relatively small models considered. Even with that caveat, comparisons can be made by comparing convergence times in our plots to those of previous models.
For example expRNN [1] and scoRNN [2] report learning the copy task (T=2000) in about 2000 and 4000 iterations (times taken from plots in the respective papers) respectively whereas projUNN learns the same task in about 3000-5000 iterations depending on the rank approximation in our experiment. In general, we find that since we are approximating gradients at each step with low rank matrices, projUNN needs more iterations to learn but this overhead is roughly a small constant multiple of comparable models. Looking ahead, we would like to analyze potential areas where the scaling benefits of projUNN have practical applications.
**References**\
[1] Lezcano-Casado, M., & Martınez-Rubio, D. (2019, May). Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In International Conference on Machine Learning (pp. 3794-3803). PMLR.\
[2] ​​Helfrich, K., Willmott, D., & Ye, Q. (2018, July). Orthogonal recurrent neural networks with scaled cayley transform. In International Conference on Machine Learning (pp. 1969-1978). PMLR.
Thank you for the detailed responses.
Indeed, it will be interesting to analyze and identify potential practical applications.

Meta-review:
role: AC
data: metareview: This paper provides two routines to replace gradient updates with low-rank unitary updates, and provides extensive technical discussion and experiments.  Reviewers are uniformly positive, and I also voice similar praises, e.g., I too appreciate the extensive discussion in appendices A and B, and the detailed experiments in the later appendices.  As such, it is easy to recommend acceptance, and I will push for this to receive at least a spotlight.  Even so, I urge the authors to make careful revisions for remaining issues raised by the reviewers, and to perform a full pass of their own.
role: PC

Decision: Accept