projUNN: efficient method for training deep networks with unitary matrices
Anonymous Author(s)
AffiliationAddress
Abstract
In learning with recurrent or very deep feed-forward networks, employing unitarymatrices in each layer can be very effective at maintaining long-range stability.
However, restricting network parameters to be unitary typically comes at the costof expensive parameterizations or increased training runtime.
We propose insteadan efficient method based on rank- k updates – or their rank- k approximation – thatmaintains performance at a nearly optimal training runtime.
We introduce twovariants of this method, named Direct (projUNN-D) and Tangent (projUNN-T) projected Unitary Neural Networks, that can parameterize full N -dimensional unitaryor orthogonal matrices with a training runtime scaling as O ( kN 2 ) .
Our methodeither projects low-rank gradients onto the closest unitary matrix (projUNN-T) ortransports 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 unitaryparameters 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 exploreprojUNN in training orthogonal convolutional neural networks, which are currentlyunable to outperform state of the art models but can potentially enhance stabilityand robustness at large depth.
Anonymized Google Colab RNN demo and CNNdemo.
Introduction
Learning in neural networks can often be unstable when networks are very deep or inputs are longsequences of data [5, 83].
For example, vanilla recurrent neural networks (RNNs) have recurrentstates that are evolved via repeated application of a linear transformation followed by a pointwisenonlinearity, which can become unstable when eigenvalues of the linear transformation are not ofmagnitude one.
Unitary matrices, which have eigenvalues of magnitude one, can naturally avoid thisissue and have been used as a means to overcome these so-called vanishing and exploding gradients[5, 44].
More recently, unitary convolutional layers have been similarly constructed to help buildmore stable deep networks that are norm-preserving in their transformations [58, 72].
In the RNN setting, prior algorithms to apply n × n unitary matrices in RNNs have parameterizedmatrices into layers of unitary or orthogonal transformations or parameterized the Lie algebra ofthe unitary or orthogonal group (see Table 1).
In the layer-wise setting, unitarity is enforced forall values of parameters, but many layers are required to form a composition that can recreate anydesired unitary, i.e., fully parameterizing an n × n unitary requires O ( n ) layers.
By parameterizingthe Lie algebra [56, 41], algorithms perform better on common benchmarks but have the drawbackthat performing gradient optimization on an n × n unitary requires O ( n 3 ) operations generically per [Table]
step. Though not an issue with the small to medium sized models used today, this O ( n 3 ) is still O ( n ) [Table] slower than standard methods of forward- and back-propogation in RNNs.
[Table]
Motivated by the feature that gradients in neural networks are typically approximately low rank
[Table]
weshow that gradient updates to unitary/orthogonal matrices can be efficiently performed in low ranksettings.
We propose a new model called PROJ UNN where matrices are first updated via gradientbased optimization and then projected back onto the closest unitary ( PROJ UNN-D) or transportedin the direction of the gradient ( PROJ UNN-T).
PROJ UNN has near-optimal runtime complexityunlike other existing algorithms for unitary RNNs (Table 1) and is especially effective even in themost extreme case where gradients are approximated by rank one matrices.
In RNN learning tasks,
PROJ UNN matches or exceeds benchmarks of state-of-the-art unitary neural network algorithms.
Though we present our model first in the RNN setting, we show that there is a direct extensionof PROJ UNN to the case of orthogonal/unitary convolution which we explore further.
Here, weperform unitary/orthogonal convolution in the Fourier domain as inspired by [76].
Our algorithmruns efficiently in the convolutional setting especially for filters of large size and many channels (see
Appendix G for more details).
Related works
Maintaining stability in neural networks via orthogonal or unitary matrices has a rich history of studyin machine learning, both from an applied and theoretical perspective.
Here, we briefly mentionthe most related works and algorithms we use in comparison to our PROJ UNN.
For a more holisticreview of prior work in unitary neural networks and other related topics, please see Appendix B.
Unitary neural networks were first designed to address the issue of vanishing and exploding gradientsin RNNs while learning information in very long sequences of data more efficiently than existingparameterizations such as the long-short term memory unit (LSTM) [38].
Early algorithms [5, 66]maintained unitarity by constructing a series of parameterized unitary transformations.
Perhaps themost effective of these methods is the efficient unitary recurrent neural network (EUNN) [44] whichparameterized unitary matrices by composing layers of Givens rotations, Fourier transforms, andother unitary transformations.
The unitary RNN (uRNN) of [82] and the Cayley parameterization(scoRNN) of [35] parameterized the full unitary space and maintained unitarity by performing a
Cayley transformation.
Later, [56] introduced the exponential RNN (expRNN) which parameterizedunitary matrices in the Lie algebra of the orthogonal/unitary group.
Though the uRNN, scoRNN, andexpRNN perform well on benchmarks, their algorithms require matrix inversion or SVD steps whichare time-consuming in high dimensions.
For convolutional neural networks, [72] showed how to efficiently calculate the singular valuesof a linear convolution and proposed an algorithm for projecting convolutions onto an operatornorm ball which relied on a series of costly projection steps.
[58] introduced a block convolutionalorthogonal parameterization (BCOP) which was faster and more efficient than the methods in[72], but required extra parameters in its parameterization and only parameterized a subset of thespace of orthogonal convolutions.
Most recently, [73] implemented orthogonal convolutions byparameterizing the Lie algebra of the orthogonal group via their skew orthogonal convolution (SOC)algorithm which approximates orthogonal convolutions especially well for small filter sizes.
Finally,
[Figure]
𝒅𝑾 = −𝜂∇ 𝑾 ℒ projUNN−D projUNN−T
[Figure]
[76] performs convolutions in the Fourier domain via application of the Cayley transform.
Ourorthogonal/unitary convolutional parameterization is inspired by their approach and improves theirruntime for convolutions over many channels
Notation and background
Vectors and matrices are denoted with bold lower-case and upper-case script, v and V , respectively.
Scalars are denoted by regular script e and tensors are denoted by bold text T .
The complex conjugateof a complex-valued input · is denoted by · ∗ (ignored when real-valued).
The transpose of a matrix
M is denoted by M ⊺ and the conjugate transpose of a matrix is denoted by M † .
We denote the
Frobenius norm of a matrix by ∥ · ∥ F and the spectral norm of a matrix by ∥ · ∥ 2 .
Here, we provide a brief overview of the unitary/orthogonal groups and refer readers to Appendix Afor a more detailed mathematical background.
The set of n × n orthogonal O ( n ) and unitary U ( n )matrices are both Lie groups defined as
[Equation]
Constraining matrices in O ( n ) and U ( n ) to have determinant equal to one constructs the specialorthogonal SO ( n ) and unitary SU ( n ) groups respectively.
The Lie algebra or tangent space of theidentity of O ( n ) and U ( n ) are the set of skew symmetric o ( n ) and skew Hermitian u ( n ) matrices,
[Equation]
The matrix exponential exp( · ) is a map from the Lie algebra to the associated Lie group.
The map issurjective if the Lie group is compact and connected – a property which holds for the unitary andspecial orthogonal groups but not the orthogonal group.
Projected unitary networks
Our PROJ UNN algorithm is motivated by the simple observation that most of the “information" of atypical gradient in a deep learning task is captured in a low rank subspace of the complete gradient.
Fig.
15 illustrates this feature when training our PROJ UNN convolutional network on CIFAR10.
We include further analysis and justification of this low rank behavior in Appendix F. As we will show,we can perform updates on the low rank subspace of the gradient efficiently by approximating thegradient with a low rank matrix and performing projections of parameters onto that low rank subspace.
Our experiments show that this methodology, even with rank one approximations, is effective atlearning and empirically introduces a form of “beneficial" stochasticity during gradient descent.
Based on how the projection is performed, our PROJ UNN algorithm takes two forms illustrated in
Fig.
1b.
The directly projected unitary neural network ( PROJ UNN-D) projects an update onto theclosest unitary/orthogonal matrix in Frobenius norm.
The tangent projected unitary neural network( PROJ UNN-T) projects gradients onto the tangent space and transports parameters in that direction.
4. PROJ UNN-D
PROJ UNN-D takes advantage of the fact that the polar transformation returns the closest unitary ororthogonal matrix in the Frobenius norm to a given matrix (not necessarily unitary or orthogonal):
Lemma 4.1 (Projection onto unitary manifold [46]) .
Given a matrix A ∈ C n × n :
[Equation]
(3) where U ( n ) indicates the set of n × n unitary matrices.
Note, that if the matrix A is real, then the projection above will be onto an orthogonal matrix.
Given
Lemma 4.1, PROJ UNN-D performs optimization in two steps, which are illustrated in Fig.
1b.
First,matrix entries are updated via a standard learning step as in gradient descent, constructing a newmatrix that is generally no longer unitary.
In the second step, PROJ UNN-D returns the unitary ororthogonal matrix closest in the Frobenius norm to the inputted matrix using Lemma 4.1.
At firstsight, the second step would require O ( n 3 ) time to perform, but we can take advantage of the factthat gradient updates are typically approximately low rank (see Appendix F).
Efficient low rankapproximations can be obtained using sampling methods detailed in Section 4.3.
With this in mind,we show that rank k updates can be performed in O ( kn 2 ) time when k ≪ n .
Theorem 4.2 (Low rank unitary projection) .
Let U be an n × n orthogonal/unitary matrix perturbedby G k , a rank k matrix.
Then the projection onto the closest orthogonal/unitary matrix defined belowcan be performed in O ( k ( n 2 + nk + k 2 )) steps.
[Equation]
To achieve this runtime, we perform updates completely in an O ( k ) subspace of the full vectorspace.
The operation ( U + G k )
[( U + G k ) † ( U + G k )] − 1 / 2 can be decomposed into a series of O ( k )matrix-vector operations and an eigendecomposition of a 2 k × 2 k sub-matrix.
The complete proof anddetails are deferred to Appendix C. One limitation of the above is that the eigendecomposition andinversion of a low rank matrix can cause numerical instability after many update steps.
We discussthis further in Appendix H.3 where we also provide options to alleviate this instability.
PROJ UNN-T,which we discuss next, does not require matrix inversion and is thus empirically more stable.
4. PROJ UNN-T
PROJ UNN-T maintains unitarity of matrices by orthogonally projecting gradient updates onto thetangent space and then performing a rotation in the direction of the projection ( i.e., along the geodesic).
As in PROJ UNN-D, there is a closed form for the orthogonal projection:
Lemma 4.3 (Tangent space projection [82]) .
Given the tangent space T U U ( n ) of an orthogonal/unitary matrix U , the orthogonal projection Π T U with respect to the canonical metric⟨ X , Y ⟩ =
Re (cid:0)
Tr[ X † Y ] (cid:1) is
[Equation]
Similar to Lemma 4.1, this projection also returns the closest matrix in Frobenius norm to X in thetangent space,
[Equation]
Similar to PROJ UNN-D, PROJ UNN-T performs learning in two steps.
First, a gradient update G isprojected onto the tangent space using Lemma 4.3.
Then, the orthogonal/unitary matrix is transportedor rotated in the direction of the projection by application of the exponential map via the update rule[56, 82],
[Equation]
where η denotes the learning rate.
This update rule is an example of Riemannian gradient descentwhere we use the exponential map to transport gradient updates along the unitary/orthogonal manifold[15].
Here, we transport the matrix U along the geodesic in the direction of Π T U ( G ) .
This can berelated to the update of PROJ UNN-D which is an example of a retraction or an approximation to theexponential map of PROJ UNN-T (see Appendix C.3).
The update rule above requires matrix exponentiation and multiplication, both costly steps which canbe sped up when G is a low rank matrix.
Namely, to perform a rank k gradient update, we obtain anequivalent runtime scaling of O ( kn 2 ) for the PROJ UNN-D when k ≪ n .
Theorem 4.4 (Low rank tangent transport) .
Let U be an n × n orthogonal/unitary matrix perturbedby G k , a rank k matrix.
Then projecting G k onto the tangent space and performing a rotation inthat direction as defined in Eq.
(7) can be performed in O ( k ( n 2 + nk + k 2 )) steps.
As with the P ROJ UNN-D, we achieve this runtime by performing the update above completely inan O ( k ) subspace of the full vector space.
The update via the exponential map can similarly bedecomposed into a series of O ( k ) matrix-vector operations and an eigendecomposition of a 2 k × 2 ksub-matrix.
Proper manipulations of the eigenvalues of the sub-matrix implement updates via theexponential map.
The complete proof and details are deferred to Appendix C.
4.3 Sampling methods
[Figure]
Commonly, gradients can have large rank but have still have many small singular values ( e.g., see
Fig. 1a).
Here, a matrix A is deemed approximately low rank (see more details in Appendix F), andone can obtain a rank k approximation
A k of A by sampling from rows and columns of A .
Weuse two sampling algorithms.
The LSI sampling algorithm [68] obtains a rank k approximationto an n × n matrix A in time
O ( kn 2 log n ) .
The algorithm projects the matrix A onto a randomorthogonal subspace and then applies SVD based methods to the projected matrix to obtain the lowrank approximation to that matrix.
This algorithm features low approximation errors even for small kand is used extensively in our implementation.
The column sampling (linear time SVD) algorithm[23] samples from the columns of an n × n matrix A to obtain a rank k approximation in O ( c 2 n + c 3 )time, where c is a hyperparameter indicating the number of columns sampled.
Typically, c is chosenas a multiple of k
so the runtime is O ( k 2 n + k 3 ) .
In implementing this algorithm, we calculate theright singular vectors via matrix multiplication of the left singular vectors so the total runtime is
O ( kn 2 + k 2 n + k 3 ) .
We note that the two procedures described above, though sufficient for our purposes, can be furtheroptimized in their asymptotic runtime.
For sake of completeness, we discuss two of these othersampling algorithms in Appendix F.
4.4 Extension to unitary or orthogonal convolution
Unitary/orthogonal convolutions are linear convolution operations that also preserve the 2 -norm (isometric).
Restricting convolutions to be unitary/orthogonal typically results in a drop in performanceon standard imaging tasks when used in isolation, but prior work has explored unitary/orthogonalconvolutions to potentially improve algorithmic stability and robustness (see Appendix B.1 for morebackground)
We describe here how PROJ UNN can be used to implement unitary/orthogonalconvolutions in potentially a more efficient manner.
Given input tensor X ∈ C M × N × C where C is the number of channels of an M × N input, linearconvolution (or technically cross-correlation) with a filter W ∈ C M × N × C × C is defined as
[Equation]
where the indexing above is assumed to be cyclic (taken modulus the corresponding dimension)[55, 28].
Orthogonal/unitary convolutions form a subset of filters that preserve norms, i.e., filters W such that ∥ conv W ( X ) ∥ =
∥ X ∥ .
Equivalently, conv W ( · ) is orthogonal/unitary if the Jacobian of thetransformation is also orthogonal/unitary.
To maintain unitarity/orthogonality, we set the dimensionsof the filter W above such that it returns an output Y of the same dimension as the input X .
One canalso perform semi-orthogonal or semi-unitary convolution by appropriately zero-padding an input ortruncating from dimensions in the output.
Standard convolutional filters are typically supported over a sparse set of local elements, but performing orthogonal/unitary convolution generally requires implementing convolutions with filterssupported over all elements resulting in slower runtimes.
One can locally parameterize convolutionalfilters in the Lie algebra of the orthogonal/unitary group; nevertheless the exponential map into the
Lie group expands the support of the filter:
[Equation]
Thus, enforcing unitarity in convolutions generally requires additional overhead over the traditionalsetting of locally supported filters, but by performing convolution in the Fourier domain, runtimes forfull-width filters can be optimally improved to O ( N 2 C log( N ) + N 2 C 2 ) [64]:
[Equation]
(10) where (cid:98) W i,j, : , : is the value of the (cid:98) r and (cid:98) s frequency of W across all channels in the Fourier domainand FFT is the 2-dimensional fast Fourier transformation.
Our method is inspired by that of [76] which transformed W into Fourier space and performed a
Cayley transformation (approximation to the exponential map into the Lie group) over the matricesindexed by (cid:98) W (cid:98) r, (cid:98) s, : , : which requires O ( N 2 C 2 log( N ) + N 2 C 3 ) operations.
For our algorithm, weparameterize W in the Fourier domain and only manipulate (cid:98) W (see Appendix B.1 for a depictionof our parameterization).
By parameterizing (cid:98) W directly and performing rank k updates using our
PROJ UNN, this runtime can be improved to O ( N 2 C log( N ) + kN 2 C 2 ) which is optimal whenk ≪ N .
Our procedure for performing unitary/orthogonal convolution on an input X with filter
W essentially follows the steps in Eq.
(10): perform an FFT on X , block-multiply this by (cid:98) W , andperform an inverse FFT on the output to obtain the final result.
Limitations Unitary/orthogonal convolutions are implemented in a cyclic fashion ( i.e., indicesare taken modulus the dimension) which is not the standard approach but has been used before toaccelerate convolutional operations [64].
Additionally, we parameterize convolution filters to havesupport over all possible elements (full-width), which can be expensive in memory.
One can restrictthe convolution to local terms in the Lie algebra, but this would not improve runtime as our algorithmruns in the Fourier space.
To target local terms in a convolution, we instead propose for future workto implement a regularizer which has a specified support and penalizes the norm of the filter outsidethat support.
Finally, the space of orthogonal convolutions has multiple disconnected components,which can present challenges for gradient based learning [58].
However, we can avoid this drawbackby implementing PROJ UNN using fully supported filters in the space of unitary convolutions whichis connected (proof deferred to Appendix C.4).
Theorem 4.5 (Unitary convolutional manifold is connected) .
The space of unitary convolutions withfilters of full support has a single connected component.
4.5 Runtime comparisons
PROJ UNN has a nearly optimal asymptotic runtime scaling which offers practical benefits in highdimensions.
In the RNN setting, Fig.
2a shows that the low rank version of PROJ UNN has a runtimethat scales at the same rate as that of a vanilla RNN albeit with increased overhead.
Updating theunitary matrix of PROJ UNN takes O ( kn 2 ) time for performing updates of rank k ≪ n , only a factork more than a vanilla RNN which performs updates in O ( n 2 ) time.
Note, that exact (full rank)updates to the n × n unitary matrices of a PROJ UNN take roughly O ( n 3 ) time corresponding to theruntime of an SVD and equivalent to the runtime of expRNN and scoRNN [56, 35].
In the convolutional setting, PROJ UNN offers the most benefit when there are many channels, filterswith large support (very wide), or a need for exact unitary/orthogonal operations (in contrast with an [Figure] approximate method like [73]).
Given an N × N input with C channels a and backward pass of PROJ UNN runs in time O ( N 2 C log( N ) + kN 2 C 2 ) when performing rank k updates.
This is afactor of C faster than the Cayley implementation [76] which runs in time O ( N 2 C 2 log( N )+ N 2 C 3 ) .
For a more complete analysis of the asymptotic and empirical runtimes of various models includingmany not listed here, please see Appendix G.
Experiments
We propose in this section a variety of benchmarked experiments to validate the efficiency andperformance of the proposed PROJ UNN method focusing mostly on RNN tasks.
We include furtherdetails of the experiments in Appendix E including a preliminary empirical analysis of PROJ UNN inconvolutional tasks.
Toy model: learning random unitary
To study the learning trajectories of PROJ UNN, we considera simple toy model aimed at learning a target random unitary.
More specifically, we parameterize alarge unitary matrix U ∈ C 2048 × 2048 to learn a Haar random target unitary U tar ∈ C 2048 × 2048 givena dataset { x i , y i
= U tar x i } 4096 i
=1 of size 4096 where x i ∈ C 2048 has entries drawn i.i.d.
randomnormal.
U is initialized as a random unitary matrix, and each step, we perform vanilla gradientdescent over a batch of 16 training points using mean-squared error loss ℓ ( x i , y i ) =
∥
Ux i − y i ∥ 22 .
Approximations of rank k to the gradient are obtained using the column sampling algorithm.
Fig. 2b, which plots the Frobenius error ∥
U − U tar ∥ 2 F , shows that PROJ UNN-T equipped with thecolumn sampling approximator is able to learn the random target unitary even when k
= 1 (see
Appendix E.1 for plots with PROJ UNN-D).
Furthermore, for a fixed learning rate, learning requiresfewer steps with larger k up to k = 16 , the maximum rank of the gradient (note that ∇ U ℓ ( x i , y i ) isrank 1).
Therefore, approximating the gradient via low rank approximations can significantly speedup learning in this task (see Appendix E.1 for further details).
Adding task In the adding task, an RNN must learn to add two numbers in a long sequence.
Weconsider a variant of the adding task studied in [5], where the input consists of two data sequences oflength T .
The first is a list of T numbers sampled uniformly from [0 , 1] , and the second is a list ofbinary digits set to zero except for two locations (those which must be summed) set to one locateduniformly at random within the intervals [1 , T/ 2) and [ T/ 2 , T ) respectively. [Figure] ss [Figure]
ss [Figure] [Figure] Consistent with [35], we train our PROJ UNN-T using an RNN with hidden dimension of 170 and the
RMSprop optimizer to reduce the mean-squared error of the output with respect to the target.
Naivelypredicting the average value of one for a random input achieves mean-squared error of approximately0.167.
As shown in Fig.
3, PROJ UNN-T is able to learn the target function even with rank k = 1approximations.
Surprisingly, for a fixed learning rate and scheduler, convergence to the true solutionis almost equally fast for k = 1 , k = 4 , and k = 16 .
Further details are provided in Appendix E.2.
Copy memory task
The copying memory task is a common benchmark for RNNs [38, 5, 36],where the aim is to memorize input data by ignoring a long sequence of void tokens.
Given analphabet of n + 2 symbols { a i } n +2 i =1 , n of which represent data (sequence of letters A, B, . . .
andadditional void (-) and start recall (:) tokens, the RNN must output the first K input tokens as the last
K output tokens and void otherwise.
An example input/output for M = 6 with n = 4 is
Input:
ABCDAD---···-------:-----
Output: ---------···-------ABCDAD
Here, T = 1000 or T = 2000 so the network must memorize data over a very long sequence of voidtokens.
As in [44], we consider n = 8 and input length K = 10 and train networks with batch size128 using the RMSProp algorithm.
Naively predicting T + K void tokens followed by K randomselections of the n possible tokens achieves a baseline loss of K log( n ) /
( T + 2 K ) .
PROJ UNN-Tis able to learn the copy task efficiently as shown in Fig.
In fact, for fixed learning rates, rankone approximations using the column sampling algorithm provide the fastest convergence to thetrue solution in comparison to higher rank approximations.
Networks were initialized using Henaffinitialization (see Appendix H.4) and the learning rate for unitary parameters was set to 32 times lessthan that of regular parameters (see Appendix E.3 for more details).
Permuted MNIST Another challenging long-term memory task we consider is the permutedpixel-by-pixel MNIST dataset.
Here, MNIST images are flattened, and pixels are randomly shuffledand placed in a sequence thereby creating some non-local dependencies.
MNIST images have 28 × 28resolution, so the pixel-by-pixel sequences have length T = 784 .
The task is digit classification (10classes) as in standard MNIST models.
We employ the same data processing, shuffle permutation,and formatting as that in prior works [56].
We perform cross-validation over different learning rates [Table] -D with different low-rank values k ∈ { 1 , 2 , 4 , 8 , 16 } .
The final test accuracy is shown in Table 2.
As observed in the copy and adding tasks, we find thatusing k > 1 does not lead to improved performances.
In fact, we provide the evolution of the testset accuracy during training in Fig. 13 and note that as the number of updates is large (hundreds perepoch), even rank k = 1 update are able to move the model’s parameters to their local optimum.
CNN experiments To explore the performance of our PROJ UNN training algorithm for convolutional layers, we first analyzed its performance on CIFAR10 classification using a Resnet architecture[34].
Our aim was not to “beat" benchmarks but to provide an honest comparison of the performanceof PROJ UNN to existing methods.
In fact, as noted earlier, enforcing unitarity generically resultsin a drop in accuracy for commonly used architectures.
Consistent with prior work [76] we employdata-augmentation of random translations and left-right flips.
Previous analysis in the RNN settingshowed that rank k = 1 is sufficient for convergence so we always set k = 1 when using PROJ UNNin the convolutional setting.
For Resnet9 trained using the RMSprop optimizer, PROJ UNN-T and
PROJ UNN-D reached 80 .
75% and 82 .
06% accuracy respectively, matching or outperforming reported results from existing unitary CNN models which achieved accuracies of 80 .
72% for BCOP[58] and 81 .
70% for Cayley [76] (further details in Appendix E.5).
Note, that all of these methodsresulted in a performance drop compared to the standard model (without unitary constraints) whichachieved accuracy of 92 . 26% .
Hence, we believe that there remain a large potential for unitary modelsto close this gap.
Separate from just performance and to motivate the use of unitary parameterization,we provide in Fig.
5, test accuracy results from a simple CNN model with progressively increasingdepth trained with and without unitary parameterization on MNIST data.
We observe that unitaryweights might provide benefits for vanilla CNN architectures that have not been designed to handlevery deep settings.
Of course, various techniques and tricks have been designed to enable CNNs tobe trainable at large depths [83, 34, 14].
Unitary convolutions, which are simple and theoreticallymotivated, can potentially be used either separately or in-tandem with these other techniques.
Discussion
Our PROJ UNN shows that one need not sacrifice performance or runtime in training unitary neuralnetwork architectures.
Our results broadly take advantage of the approximate low rank structureof parameter gradients to perform updates at nearly optimal runtime.
Looking beyond the settingstudied here, it is an interesting question how our framework can be applied to other neural networkarchitectures or parameter manifolds.
Group convolutional neural networks and Riemannian gradientdescent offer two promising avenues for further application of our techniques.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes] Contributions and scope of the paper are specified inthe abstract and introduction.(b) Did you describe the limitations of your work?
[Yes] Limitations are addressedthroughout this work, including after the descriptions of the algorithm in Section 4.1and Section 4.4, in Appendix H.3, in our comparison to related work in Appendix B, inthe runtime comparison of Appendix G, and various other places.(c)
Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes] We confirm that we have read the ethics guideline to ensure our paper
conforms to it.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] All assumptionsare given in the sections themselves or part of the statements of the relevant theorems,
propositions, and lemmas.(b) Did you include complete proofs of all theoretical results?
[Yes] Proofs are included for all results and deferred to the appendix.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] An anonymized
demo is included in this anonymous manuscript.
Full code will be made available upon
publication.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] Relevant details are included in Section 5 and Appendix E.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] we employed the same experiment setup as prior work thatconsists in reporting the best test performance based on validation set performances
after cross-validation.
As we report such numbers for a variety of settings, error barscould be derived from our tables.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] Relevant details are included inAppendix H.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] Please refer to
Appendix H.
(b) Did you mention the license of the assets?
[Yes] Please refer to Appendix H.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
A Mathematical backgroundA.1 Lie groups and Lie algebras
Here, we provide a brief mathematical background to Lie groups and Lie algebras with a particular focus on the unitaryand orthogonal groups.
For a more comprehensive overview, we recommend [32].
Though Lie groups are typically defined with respect to differentiable manifolds, we restrict ourselves here to the subsetof matrix Lie groups which is less general but allows for a more concise and simpler theoretical overview.
Informally,
Lie groups are groups whose elements are specified by a set of parameters that vary smoothly and continuously, i.e.,the group is also a differentiable manifold.
Specific to matrices, Lie groups are commonly defined with respect to thegeneral linear group GL ( n, C ) denoting the set of invertible n × n matrices with complex valued entries [32].
Liegroups are closed subgroups of GL ( n, C ) that have the following smoothness property.
Definition A.1 (Matrix Lie groups [32]) .
A matrix Lie group is any subgroup of GL ( n, C ) with the property that anysequence of matrices A m ∈ C n × n in the subgroup converge to a matrix A that is either an element of the subgroup ornot invertible ( i.e., not in GL ( n, C ) ).
Two important Lie groups studied in this work are the unitary and orthogonal groups whose definitions are copiedbelow.
[Equation]
The Lie algebra is the tangent space of a Lie group at the identity element.
To observe this, we introduce the matrixexponential map which is central to the connection between Lie groups and their corresponding Lie algebras.
[Equation]
For compact groups, the exponential map is a smooth map whose image is the connected component to the identity ofthe Lie group [52, 32].
The special orthogonal and unitary groups are both compact and connected so the exponentialmap is surjective for these groups ( i.e., for every group element, there exists an element of the Lie algebra whoseexponential is equal to that group element).
However, the orthogonal group has two connected components, i.e.,elements with positive and negative determinant, and the image of exponential map are only orthogonal matrices withpositive determinant.
Since the matrix exponential map is a smooth function, we can take the derivative of the exponential map with respectto a parameter as below.
[Equation]
and thus,
[Equation]
The above gives us the Lie algebra to a given group.
Definition A.2 (Lie algebra [32]) .
Given a Lie group G , the Lie algebra g of G is the set of matrices X such thate t X ∈ G for all t ∈ R .
Typically, Lie algebras of a Lie group are denoted with Gothic or Fraktur font.
Using the above definition, one canconstruct the corresponding Lie algebras.
As an example, consider the unitary group where given a matrix U ∈ U ( n )and X ∈ u ( n ) ,
[Equation]
and since the above holds for all t ∈ R , we can differentiate the above at t = 0 , obtaining the property of elements ofthe unitary Lie algebra that − X
= X † as seen in the main text:
[Equation]
Proceeding in a similar fashion with the orthogonal group, we obtain the following result copied from the main text,
[Equation]
A.2 Projecting onto the group or algebra
Projections onto the orthogonal/unitary groups or their tangent spaces are central to optimizing over the space oforthogonal/unitary matrices.
We focus the discussion here to the case of unitary matrices, but note that all of thefollowing statements apply to orthogonal matrices as well by simple adjustments such as replacing the conjugatetranspose ( † ) with the transpose ( ⊺ ) .
The tangent space T U to a matrix U ∈ U ( n ) is equal to
[Equation]
[Equation]
Given the canonical inner product ⟨ A , B ⟩ = Re (cid:0)
Tr[ A † B ] (cid:1) , we can show that the orthogonal projection onto thetangent space T
U U ( n ) is equal to that given by Lemma 4.3 copied below.
Lemma 4.3 (Tangent space projection [82]) .
Given the tangent space T U U ( n ) of an orthogonal/unitary matrix U , theorthogonal projection Π T U with respect to the canonical metric ⟨ X , Y ⟩ =
Re (cid:0)
Tr[ X † Y ] (cid:1) is
[Equation]
Similar to Lemma 4.1, this projection also returns the closest matrix in Frobenius norm to X in the tangent space,
[Equation]
Proof.
By the definition of an orthogonal projection, we need that Π T U (Π T U ( · ))
=
Π T U ( · ) and for all X
, Y ∈ T U U ( n )Π T U ( X ) =
X and ⟨ Π T U ( X ) , X − Π T U ( X ) ⟩ = 0 .
The first two properties are straightforward to check.
The lastproperty can be shown as below using the definition of Π T U and cyclic property of trace:⟨ Π T U ( X ) , X − Π T U (
[Equation]
Furthermore, we have for all Y ∈ T U U ( n ) using triangle inequality and unitary invariance of the Frobenius norm:
[Equation]
which proves that Π T U ( · ) projects onto the closest matrix Y ∈ T
U U ( n ) in Frobenius norm.
Since the set of unitary/orthogonal matrices does not form a vector space, an orthogonal projection is not a well definedoperation in this space.
However, it is still valid to ask what is the “closest" unitary/orthogonal matrix to a given matrixin a given norm.
This is exactly what is stated in Lemma 4.1 copied from the main text and proven below.
Lemma 4.1 (Projection onto unitary manifold [46]) .
Given a matrix A ∈ C n × n :
[Equation]
where U ( n ) indicates the set of n × n unitary matrices.
Proof.
We follow the approach of [46] to prove this result.
To shorten our notation, let V
=
A ( A † A ) − 1 / 2 .
Given anyunitary U ∈ U ( n ) , let U =
M + V for the properly chosen M ∈ C n × n .
From unitarity of U and V , we have
[Equation]
Then,
[Equation]
and since from Eq.
(A.13) we have that MV † + V M † = − MM † ,
[Equation]
The second term above is non-negative since Tr[( A † A )
[Equation]
] =
Tr[ M † ( A † A ) 1 / 2 M ] and ( A † A ) 1 / 2 is positivesemi-definite.
Thus, for all U ∈ U ( n ) ,
[Equation]
which proves the result.
B Review of previous unitary neural network techniques
The integration of orthogonal/unitary matrices into neural networks is broadly aimed at maintaining stability in neuralnetworks with many layers.
For vanilla recurrent neural networks, repetitive application of the hidden-to-hiddentransformation matrix exponentially amplifies or decays the eigenvalues of the transformation, thus resulting inexponentially large or small gradients.
Similarly, in deep network architectures where weight matrices are drawnrandomly, the norms of hidden states can similarly grow or decay exponentially with added layers.
Enforcing unitarityor orthogonality of neural network layer transformations offers a straightforward method to address these issues ofinstability since orthogonal/unitary matrices have eigenvalues of unity.
To establish notation and provide motivation for later analysis, consider a RNN whose input is a sequence of vectorsx ( t ) with hidden layer h ( t ) updated according to the following rule:
[Equation]
In the unitary or orthogonal formulation for RNNs, the matrix W in Eq.
(B.1) is replaced by a unitary or orthogonalmatrix U .
During optimization of a unitary RNN, one must enforce the unitarity or orthogonality of U during training.
Existing methods to parameterize and enforce unitarity/orthogonality in neural network layers can be separated intothree categories depending on the method of parameterization.
We discuss each of these in detail below:
• Layer-wise transformations: among the first methods employed in this line of work, these methods parameterize orthogonal/unitary matrices in a layer-wise fashion where each layer is a parameterized orthogonal/unitarymatrix.
Example parameterizations include Givens rotations and Householder reflections.
These methods areefficient when there are not many layers included in the parameterization and typically not employed whenfull access to the unitary/orthogonal group is required as full parameterization of an n × n unitary/orthogonalmatrix requires O ( n ) layers.• Lie algebra parameterization: motivated by the fact that staying on the manifold of the Lie algebra isoften easier than staying on the manifold of a Lie group, these methods parameterize the Lie algebra of theunitary/orthogonal groups and later obtain the actual unitary/orthogonal matrix by implementing the matrixexponential map.
This matrix exponential map is typically the most costly step in these methods as one caneither perform it directly ( e.g., using an SVD) or approximate it via Taylor series or Padé approximationswhich require repetitive application of matrices when applying it to an input.
Though we adapt techniques
from these methods in our work to efficiently perform gradient updates, we do not parameterize matrices inthe Lie algebra.
• Matrix entry parameterization: as in typical neural network architectures, this method parameterizes anorthogonal/unitary matrix by directly parameterizing the entries of the matrix.
This method is optimal in the“forward" direction since performing the transformation on an input simply requires matrix multiplication (asin vanilla architectures).
However, updates to the matrix will no longer maintain unitarity or orthogonality, andone must employ methods to project these updates back onto the unitary/orthogonal manifold.
This method isemployed in our work.
We now present each of the above methods in the order given.
Layer-wise transformations Early algorithms [44, 5, 66] maintained unitarity by parameterizing a matrix U as aseries or layered set of k parameterized unitary transformations:
[Equation]
where F ( i ) uni
( θ i ) indicates a transformation that maps parameters θ i into a unitary matrix.
These transformations includeparameterized Givens rotations or Householder reflections.
As a concrete example, consider the Givens rotationparameterization which is an orthogonal matrix that performs a rotation of the i, j -th dimensions by an amount θ :
G ( i, j, θ ) =   1 · · · · · · · · ·. ..
0 · · · cos θ · · · − sin θ · · ·... 0 · · · sin θ · · · cos θ · · ·...
0 · · · · · · · · ·  ,
[Equation]
B.
[Equation]
or in other words, entries G ii
=
G jj = cos θ , G ij = − G ji = sin θ , and G kl =
δ kl for all other entries.
In general, at least k = O ( n ) layers are needed to achieve a full parameterization of the n × n unitary matrices.
Training is performed by updating the parameters θ i within each layer.
However, for large matrices, due to the factthat O ( n ) layers are required to parameterize the full space of transformations, these algorithms are only efficientwhen parameterizations over a subset of the space of unitary/orthogonal matrices suffices.
Achieving this balance ofparameterization versus performance is challenging as prior work – especially work studying the learnability of unitarymatrices in quantum computation – has shown that loss landscapes over the unitary/orthogonal manifold may containmany bad local minima [50, 26, 4, 22].
Lie algebra parameterization As a reminder, the Lie algebra of the orthogonal and unitary groups are the set ofskew symmetric ( o ( n ) ) and skew Hermitian ( u ( n ) ) matrices,
[Equation]
Transformations from the Lie algebra to the Lie group are performed using the exponential map which is surjectiveonto the connected components of the identity:
[Equation]
Thus, these methods parameterize the full space of unitary or special orthogonal matrices.
However, the orthogonalgroup has two connected components (matrices with determinant equal to one and negative one) so the exponentialmaps only onto the positive determinant matrices.
Note that the Lie algebra is a vector space so the sum of two matrices in the Lie algebra is also in the algebra.
Thus,gradient updates can typically be very easily performed.
However, the exponential map is often expensive to compute,and much prior work has focused on implementing this transformation efficiently.
[35] and [73] employ approximationsto the matrix exponential via Padé or Taylor series approximations.
For example, in the Taylor series approximation,one simply truncates the Taylor series to K entries:
[Equation]
where the norm ∥ · ∥ 2 indicates the spectral norm ( i.e., largest singular value).
This application of the matrix exponentialcomes at an added cost both when calculating derivatives and when applying the matrix in the forward direction ( e.g.,when one simply desires the output of a network).
Applying the approximation above to the input of a layer requires
O ( K ) applications of the matrix.
The matrix exp( X ) can be explicitly constructed to avoid this added cost; however,obtaining this matrix requires a one time cost of O ( K ) matrix-matrix multiplication operations which can be costly forlarge matrices.
The Padé approximants are more typically used in approximating the exponential map since they can guarantee that theapproximated matrix is actually a unitary/orthogonal matrix (as opposed to a simple Taylor series approximation whichdoes not provide this guarantee).
A Padé approximant is an optimal rational function approximation to a given function.
The Padé approximant to the exponential map takes the form exp( A )
≈ r mn
( A ) =
p m ( A ) q n ( A ) − 1 where p m ( · ) andq n ( · ) are order m and n polynomials respectively which have closed forms:
[Equation]
The error of the approximation scales as O ( ∥
A ∥ m + n +1 )
For unitary/orthogonal matrices, this approximation hasthe feature that for order m = n , applying the approximation to an element of the Lie algebra of the orthogonal/unitarymatrices outputs a orthogonal/unitary matrix.
Notably, setting m = n = 1 obtains the Cayley transform which wasused in [35].
Finally, we note that approximations to the exponential function often can bias the gradients so in the RNN setting, [56]actually perform an SVD to obtain the actual output of the exponential map and analytically calculate gradients.
Matrix entry parameterization Parameterizing the entries of a matrix U directly is an obvious and simple means ofconstructing a unitary or orthogonal matrix.
However, the set of unitary/orthogonal matrices do not form an algebraso one cannot simply update these matrices by simply adding a gradient update to a given matrix.
Instead, updates tothese matrices must be performed by projecting the updated matrix back onto the set of unitary/orthogonal matrices.
Obviously, U must be initialized to be unitary/orthogonal in these methods.
Given a gradient update G , the matrix U + G is typically no longer unitary.
Methods of Riemannian optimization areemployed to update the matrix in a unitary/orthogonal fashion.
One means of updating the matrix is via the Cayleytransform which computes a parametric curve in the direction, employed in [57, 82, 61, 76].
The Cayley transform“transports" U in the direction of the projection of the gradient in the tangent space Π( G ) .
i.e., let U (0) be the initialunitary matrix, then one can transport the matrix in the direction of Π( G ) (where Π( · ) is the projection onto the tangentspace) by a “length" of α via the Cayley transform [67],
[Equation]
[Equation]
The above update formula requires a matrix inversion step which is the most costly step for large matrices.
[57, 61]approximate the transformation via a fixed point interation which avoids having to invert a matrix but still requiresmatrix-matrix multiplication.
More generally, the Cayley transform is a first order Padé approximant to the exponentialmap which provides a connection between the matrix entry parameterization and the Lie algebra parameterizationdiscussed previously [37], i.e., setting m = n = 1 in Eq.
(B.8) obtains the Cayley transform.
Our algorithms described in the main text directly parameterize matrix entries and thus follow this parameterizationmethod.
In performing the exponential map, we do not resort to any approximations since the exponential is efficient toperform in low rank settings.
Updates to unitary matrices are either projected directly onto the closest unitary matrix in
Frobenius norm ( PROJ UNN-D) or transported along the geodesic in the direction of the projection of the gradient ontothe tangent space ( PROJ UNN-T).
We refer the reader to the main text for a full description of our methodology.
B.1 Orthogonal or unitary convolution
Since convolutions are linear operators, one can perform orthogonal or unitary convolutions using similar techniquesas in the general case studied above.
As a reminder, for 2-D convolution, given input tensor X ∈ C M × N × C where Cdenotes the number of channels of the M × N input, linear convolution (or technically cross-correlation) with a filter
W ∈ C M × N × C × C takes the form
[Equation] [Equation]
where the indexing above is assumed to be cyclic (taken modulus the corresponding dimension).
Unitary or orthogonalconvolutions form a subset of filters W which preserve the norm, i.e., ∥ conv W ( X ) ∥ =
∥ X ∥ .
Equivalently, conv W ( · ) isorthogonal/unitary if the Jacobian of the transformation is also orthogonal/unitary.
The first method for orthogonalizing convolutions in neural networks was proposed in [72].
Their algorithm showedhow to calculate the singular values of a linear convolution operation and then performed a series of projections ona given convolutional filter to project it onto an operator norm ball.
Their algorithm made use of the convolutiontheorem which shows that linear convolution is diagonalized by Fourier transformations.
[59] also proposed a method toorthogonalize convolutions by performing an SVD of the linearization of a convolution and then bounding the singularvalues accordingly.
This method outputs linear transformations that are close to a convolution operation but no longernecessarily maintaining the equivariance of a standard convolution.
Also, their method is expensive for large matricesas it requires implenting an SVD.
More recent methods implement orthogonal convolutions via approximations to the exponential map [76, 73].
Thesemethods first form convolution filters W whose Jacobians J ( W ) satisfy the skew symmetry property of the orthogonalgroup:
[Equation]
where conv-transpose is the equivalent transposition operation in the filter space defined as [73][conv-transpose( L )] m,n,c,d
[Equation]
for a filter L ∈ R M × N × D × C .
A given filter W can be transformed into a skew symmetric convolution filter by simply applying L = W −conv-transpose( W ) .
As in the general case, one can apply the exponential map to the convolution filter to perform orthogonal convolution.
[Equation]
where ∗ p indicates the convolution is applied p times, e.g., L ∗ 2 X =
L ∗ L ∗ X .
Since the above operation iscomputationally expensive [73] implement a k -th order Taylor approximation to the exponential map:
[Equation]
Given 2-D inputs of dimension N × N × C , applying convolution with filters with support over W elements in eachdimension has runtime O ( pC 2 W 2 N 2 ) .
The added factor p in the runtime is required both upon training and evaluationof the network.
Though this number is held constant, this adds a multiplicative overhead to running the algorithmboth when evaluating an input and when training the algorithm.
[73] use the above method to implement orthogonalconvolution in their skew orthogonal convolution (SOC) algorithm.
Similar techniques have been used to performinvertible (not necessarily orthogonal or unitary) convolutions in [39].
There are two key considerations or drawbacks associated with using the Taylor expansion to perform unitary/orthogonalconvolution as in [73].
First, approximations via the Taylor approximation necessarily include an error that can beaccounted for by scaling the factor p .
However, unlike the Padé approximants discussed earlier such as the Cayleyapproximant, the Taylor approximation to the exponential map does not return a unitary/orthogonal operator.
Thus, thevalue p needed to bound the error from any unitary/orthogonal operator also scales logarithmically with the desired error;however, this factor can also grow with the size of the filter, the number of channels, size of the input to a convolutionoperation, and especially the depth of the network.
Second, Taylor approximations to a function can significantlybias the gradient.
As noted in [56], their expRNN algorithm avoided using the Taylor or Padé approximation inthe implementation of the unitary/orthogonal operation for this reason.
As an example of how approximations canfail to represent the gradient, they provide the set of functions f n ( x ) = sin(2 πnx ) /n approximating f = 0 , wheref n → f uniformly but the derivatives do not converge to zero.
Especially when constructing very deep networks, suchinstabilities can potentially cause issues in training.
Convolution operations over filters with a large support can be performed more efficiently in the Fourier domain asexplored in various prior work [64, 11].
Specific to orthogonal convolution, [76] use the 2-D fast Fourier transform tomore efficiently perform orthogonal convolution.
For single channel inputs and outputs, convolution in the Fourierregime corresponds to pointwise multiplication over Fourier bases.
Given multi-channel inputs and outputs, convolutionin the Fourier regime corresponds to matrix multiplication over the blocks indexed by channels where each blockcorresponds to a specific Fourier basis, i.e., for a filter W and input X :
[Equation]
B.
[Equation]
where (cid:98) W (cid:98) r, (cid:98) s, : , : is the representation of the filter in the Fourier regime.
If the filter and and input are flattened in theirspatial dimensions, we can represent the operation above as block-diagonal multiplication.
For example, one canvisually input the Fourier multiplication on an 2-channel input with spatial dimension [2 , 2] .[FFT
[Equation]

[Equation]
(cid:98)
[Equation]
[Table]
[Table]
,(cid:98) W
[Equation]
[Table] [Table] [Table]
,(cid:98) W
[Table] [Table]
(cid:98) W
[Table]
(cid:98) W
[Table]
(cid:98) W
[Table]
(cid:98) W
[Table]
W
[Table]
(cid:98) W
[Table]
(cid:98) W
[Table]
(cid:98) W
[Table]

   (cid:98) X 1 ,
[Table] [Table]
(cid:98) X 1 ,
[Table] [Table]
(cid:98) X
[Table] [Table]
[Table]
(cid:98) X
[Table]
(cid:98) X
[Table] [Table] [Table]
(cid:98) X
[Table]
[Table]
[Table]
(cid:98) X
[Table]
[Table]
(cid:98) X
[Table]
[Table] ,  [76] use the above technique, parameterizing filters in the Lie algebra of the orthogonal group.
Filters are mapped intothe orthogonal group using the Cayley transform (see Eq.
(B.9) for its form).
Since the Cayley transform requires matrixinversion over the matrices above, their total runtime scales as O ( N 2 C 2 log( N ) + N 2 C 3 ) operations for convolutionover N × N images with C channels.
Our methodology is adapted from the techniques in [76] and scales moreefficiently as O ( N 2 C log( N ) + kN 2 C 2 ) time when the rank of gradient updates k ≪ N .
We instead parameterizethe convolution filter directly in the Fourier regime, by parameterizing the individual blocks ( (cid:98) W (cid:98) r, (cid:98) s, : , : ) of the block 813 diagonal matrix above.
One drawback of our method compared to [76] is that we cannot specify the support of theconvolution in the real space of the Lie algebra as we parameterize matrices directly in the Fourier space.
Of course,one can perform projections of the convolutions in Fourier space onto the bases spanned by the support of the elementsin the Lie algebra; however, we opt instead to implement a regularizer which biases the filter towards elements of thedesired space ( e.g., local elements).
Finally, [58] performed orthogonal convolution by parameterizing a block of matrices and corresponding projectorsin their Block Convolutional Orthogonal Parametrization (BCOP) algorithm.
Their method preceded those of [76]and [73].
However, their method has two key drawbacks: it only parameterizes a subset of the space of orthogonalconvolutions and is slower than other methods as it adds additional parameters to connect the various components in thespace of orthogonal convolution.
Finally, methods have been proposed to iteratively or approximately maintain orthogonality.
[9] propose a regularizerfor convolutional layers which penalize matrices having singular values far from unity.
[40] propose using a Newton’smethod iteration to update linear transformations to be closer to orthogonal/unitary.
Their method requires iterativematrix-matrix multiplication over the unrolled weight matrix which can become expensive for large images.
B.2 Other related works
In the context of recurrent neural networks, designing RNNs to learn long sequences of data has a rich history of study.
Some of the first and most celebrated algorithms include the long short-term memory networks (LSTM) [38] and gatedreccurent unit (GRU) networks [20].
Since these works, various techniques have been used to more optimally avoidissues with learning long-sequence data.
Beyond the unitary RNNs discussed earlier, some work has explored usingbi-directional RNNs [33, 53], including those that have a more biologically inspired design [12].
These networksperform well on the copy task.
Another line of research studies the stability properties of continuous state space models which can be converted into a
RNN formulation [24].
Some algorithms construct continuous state space models whose attractors are stable points inthe dynamical system [19, 25].
More recently, continuous state space models have been designed with hidden statetransformations that are customized to memorize data by limiting the learning over time to a subset of orthogonalpolynomials [29, 30, 77].
Here, hidden states are in a sense parameterized over a set of polynomial coefficients[77].
These algorithms perform very well on tasks such as the copy task or Permuted MNIST but do not includeunitary/orthogonal transformations in their network.
In fact, more recent models [30, 77] achieve slightly higher scoreson the TIMIT and permuted MNIST benchmarks compared to the unitary RNN formulations studied here.
In the convolutional setting, [83] form very deep convolutional neural networks by initializing parameters to constructnorm-preserving orthogonal transformations.
Orthogonality is not preserved during training however.
[80] biasfilters towards orthogonality by implementing a regularizer that penalizes the weights when norms of outputs arelarger or smaller than norms of inputs.
The networks used in [83, 80] do not necessarily preserve the orthogonalityproperty of their convolutional transformations during training.
[45] study orthogonal convolutions from the basis ofunitary/orthogonal wavelets showing how to represent convolution in terms of these wavelets.
Low rank approximations have been used in prior work in deep learning to prune neural network models [84, 74] andaccelerate convolutions [43, 75, 42].
More related to our work, recent research has compressed gradients efficientlyusing low rank compression methods [78].
Although their focus was in sharing compressed information acrosscomputing units, their work lends support to the notion that the information of a gradient can be effectively andefficiently stored in low rank components.
Research in learning theory has also noted connections between the stablerank of neural network parameters and generalization.
[7] prove a generalization bound by compressing the modelsin the hypothesis class based on the stable rank of individual layers.
[63] study the phases of learning from a randommatrix theory setting showing that the stable rank of a matrix tends to decay over training.
Various works have studiedthe implicit bias induced by optimization algorithms such as gradient descent showing that in many cases, the implicitbias is towards low rank solutions [31, 21].
Such bias towards low rank solutions has been explicitly proven in thesetting of 2-layer matrix factorization [60], deep matrix factorization [6], linear group convolutional networks [54],and matrix recovery from Pauli measurements [62].
We note that relating low-rankness to the generalization ability oflearning algorithms is a richly studied topic, and there are numerous papers that we did not mention here.
Finally, a wide range of work in quantum computation and quantum machine learning studies unitary learning algorithmsin the context of quantum systems [13].
In fact, since the state space of closed quantum systems is transformed by unitaryoperators, quantum computers offer a unique platform for performing machine learning on the unitary manifold.
This isan active area of research and existing methods for performing quantum machine learning on quantum architecturesinclude variational algorithms which parameterize a quantum circuit and update the parameters via classical optimizationmethods [17, 81, 85, 48, 49, 79], quantum neural networks which design analogues to classical deep neural networks[51, 10, 71], and direct implementations of classical deep learning algorithms on quantum architectures [47, 16, 3].
We stress that these quantum algorithms are inherently different in nature than their classical counterparts and therestill exist significant challenges that must be surmounted before they become practically feasible [17, 13, 65].
Forexample, the loss landscapes of quantum algorithms often have many more poor local minima in comparison to classicalcounterparts [50, 4] and training of quantum architectures requires sampling of outputs which is challenging whenderivatives decay with the size of a model – a phenomenon described as “barren plateaus" [65, 18].
Furthermore, evenif algorithms can be efficiently run on a quantum computer, preparing data for use in a quantum computer and readingout information from the quantum computer are challenging tasks which are not guaranteed to be efficient [1].
C Deferred proofs
C.1 Proof of Theorem 4.2
Recall Theorem 4.2:
Theorem 4.2 (Low rank unitary projection) .
Let U be an n × n orthogonal/unitary matrix perturbed by G k , arank k matrix.
Then the projection onto the closest orthogonal/unitary matrix defined below can be performed in
O ( k ( n 2 + nk + k 2 )) steps.
[Equation]
Proof.
We proceed to prove the above statement by first analyzing the case
where k
= 1 and then generalizing to higherrank k .
Recall from Lemma 4.1 that we would like to perform the following update:
[Equation]
C [Equation] [Equation]
Let the rank one vector components of G k = ab † and define
[Equation]
With the above, we can rewrite ( ˜ M † ˜ M ) − 12 in Eq.
(C.1) as:
[Equation]
where ˆ a = ˜ U † a and c a = a † a .
Eq.
(C.3) is the Identity matrix plus the update of a rank two matrix.
To see this, we decompose ˆ a and b into orthogonalcomponents using Gram Schmidt:
[Equation]
In this new basis:
[Equation]
and
[Equation]
Performing an eigendecomposition of the above 2 × 2 matrix into a diagonal eigenvalue matrix S and eigenvectormatrix C , we can rewrite ˜ M † ˜ M in a convenient form:
[Equation]
Taking the inverse square root of the above can be performed by manipulating singular values:
[Equation]
Finally, we multiply the above on the left by ˜ M :
[Equation]
The above requires performing an eigendecomposition of a 2 × 2 matrix and a series of matrix-vector multiplication,matrix additions, and vector-vector outer products – in total scaling as O ( n 2 ) time.
Rank k updates Note that the above method can be extended to low rank updates, running in O ( kn 2 ) time when therank of the update k ≪ n .
Specifically, now our update is:
[Equation]
Following the same steps would ultimately require performing an eigendecomposition of a 2 k × 2 k matrix which takes
O ( k 3 ) time.
Additionally, one must perform Gram-Schmidt decomposition on a set of k vectors of length n which takes
O ( k 2 n ) time.
Finally, a series of O ( k ) matrix-vector multiplications and vector-vector outer products is performedresulting in a total runtime of O ( k ( n
+ nk + k 2 )) time.
In cases where k ≪ n , the time to perform the Gram-Schmidtdecomposition and the eigendecomposition is negligible and an overall runtime of O ( kn 2 ) time is achieved.
Even incases where updates are not low rank, one can apply efficient sampling procedures (see Section 4.3) to find low rankapproximations to the update matrix and apply the methods above while maintaining runtimes.
C.2 Proof of Theorem 4.4
Recall Theorem 4.4:
Theorem 4.4 (Low rank tangent transport) .
Let U be an n × n orthogonal/unitary matrix perturbed by G k , a rank kmatrix.
Then projecting G k onto the tangent space and performing a rotation in that direction as defined in Eq.
(7) canbe performed in O ( k ( n 2 + nk + k 2 )) steps.
Proof.
We would like to efficiently perform the update below:
[Equation]
As before, we proceed to prove the above statement by first analyzing the case where k = 1 and then generalizing tohigher rank k .
Let the rank one vector components of G k = ab † .
Then
[Equation]
where ˆ a =
U † a .
The above is a rank 2 matrix.
As before, we now proceed to perform an eigendecomposition in thelow rank subspace of the above matrix.
Using Gram Schmidt, we have
[Equation]
In this new basis:
[Equation]
and
[Equation]
Performing an eigendecomposition of the above 2 × 2 matrix into a diagonal eigenvalue matrix S and eigenvectormatrix C , we can rewrite ˜ M † ˜ M in a convenient form:
[Equation]
We apply the exponential map scaled by η to obtain
[Equation]
Finally, we multiply the above on the left by U to obtain the final result
[Equation]
The above requires performing an eigendecomposition of a 2 × 2 matrix and a series of matrix-vector multiplication,matrix additions, and vector-vector outer products – in total scaling as O ( n 2 ) time.
Rank k updates Note that as before, the above method can be extended to low rank updates, running in O ( kn 2 ) timewhen the rank of the update k ≪ n .
Specifically, now our update is:
[Equation]
Following the same steps, the projection onto the tangent space would be a rank 2 k matrix.
The steps that follow wouldultimately require performing an eigendecomposition of a 2 k × 2 k matrix which takes O ( k 3 ) time.
Additionally, one must perform Gram-Schmidt decomposition on a set of k vectors of length n which takes O ( k 2 n ) time.
Finally, aseries of O ( k ) matrix-vector multiplications and vector-vector outer products is performed resulting in a total runtimeof O ( k ( n 2 + nk + k 2 )) time.
In cases where k ≪ n , the time to perform the Gram-Schmidt decomposition and theeigendecomposition is negligible and an overall runtime of O ( kn 2 ) time is achieved.
Even in cases where updates arenot low rank, one can apply efficient sampling procedures (see Section 4.3) to find low rank approximations to theupdate matrix and apply the methods above while maintaining runtimes.
C.3 First order equivalence to PROJ UNN-D
Though PROJ UNN-D and PROJ UNN-T perform different updates, we can show that up to first order, the updates are infact equivalent.
Furthermore, as one may expect, this first order update is equal to the projection of the gradient onto thetangent space (see Lemma 4.3).
In the case of PROJ UNN-D, this shows that the update step is a retraction or first orderapproximation to the matrix exponential implemented in PROJ UNN-T [15].
Proposition C.1 (First order equivalence) .
For an n × n unitary/orthogonal matrix U perturbed by ∆ U , gradientupdates applied by algorithms PROJ UNN-D
(Eq.
(4) ) and PROJ UNN-T (Eq.
(7) ) are, up to first order, equal to
U + Π T U (∆ U ) , i.e.,
[Equation]
Proof.
We first show that the above formula holds for PROJ UNN-D and then show the same for PROJ UNN-T. Recallthe update formula for PROJ UNN-D:
[Equation]
Expanding the above up to first order and applying the first order Taylor expansion of ( · ) − 1 / 2 , we have:
[Equation]
Similarly, for PROJ UNN-T, recall the update formula (ignoring − η term for learning rate):
[Equation]
[Equation]
Expanding the above and applying the first order approximation of the matrix exponential,
[Equation]
C.4 Unitary convolutional manifold is connected
Before proceeding to prove that the space of unitary convolutions is connected, we first provide a version of theconvolution theorem which shows that convolution in the Fourier domain corresponds to block multiplication overchannels.
This is a classic result also contained in various prior works [72, 76], and we provide a short proof here forcompleteness.
Lemma C.2.
Given convolution filter W ∈ C M × N × C × C and input X ∈ C M × N × C , recall the definition of cyclicconvolution (or technically cross-correlation) of W and X :
[Equation]
Let FFT be the two dimensional fast Fourier transform, then convolution in the Fourier domain corresponds toblock-wise multiplication over channels:
[Equation]
Proof.
Let the roots of unity be denoted as ω M = e 2 πi/M and ω N = e 2 πi/N .
Then,[FFT conv [Equation] X )]
[Equation]
As an aside, the complex conjugation of the filter above is due to the fact that convolution in neural networkscorresponds to the more commonly used term cross-correlation in mathematics.
If the convolution operation wereredefined as what is more commonly known as convolution in mathematics, i.e., define conv ′ as [conv ′ W
( X )] p,q,d = [Equation]
n [Equation] q
[Equation]
n [Equation] c , then the complex conjugate would no longer appear on the filter term.
From here, we simply show that each block in the above is connected which allows us to prove that the unitary manifoldis connected.
Theorem 4.5 (Unitary convolutional manifold is connected) .
The space of unitary convolutions with filters of fullsupport has a single connected component.
Proof.
Since the fourier transform is unitary and invertible, we can represent every filter W in the fourier domain as(cid:98) W ∈ C M × N × C × C and vice-versa.
In the Fourier domain, given Lemma C.2, we have that convolution corresponds toblock-wise multiplication over channels.
For unitary convolution, each block (cid:99) W (cid:98) r, (cid:98) s, : , : indexed by frequencies (cid:98) r and (cid:98) s 972 must be a unitary matrix, so we now analyze the set of filters whose blocks are unitary matrices in the Fourier domain.
The space of unitary matrices U ( C ) is connected [32].
Therefore for every block, there exists a connected path betweenany two unitary matrices (cid:99) W (1) (cid:98) r, (cid:98) s, : , : , (cid:99) W (2) (cid:98) r, (cid:98) s, : , : ∈ U ( C ) .
The full space of unitary convolutions is parameterized by the MNtimes direct product of groups
U ( C ) , i.e., U ( C ) ( × MN ) .
Since the direct product of finitely many connected spaces isalso connected, then the space of unitary convolutions is connected.
D Pseudocode for performing projUNN updates
Pseudocode for performing an update step on a unitary matrix U ∈ C N × N with a gradient update of ∆ U is shownin Algorithm 1.
The steps are the same if U is an orthogonal real-valued matrix; the output in this case will alsobe real-valued.
As a cautionary note, especially in the last step of Algorithm 1, where there is a composition ofmultiple matrix-vector multiplications, the order of these multiplications must be chosen to only perform matrix-vectoroperations to ensure optimal runtime. In other words, two N × N matrices should never be multiplied by each other atany point in this algorithm. [Figure] rank k= k= k= k= k= exact
In convolutional settings, the steps in Algorithm 1 are applied across blocks of the convolution in Fourier space.
Asstated in the main text and detailed in Appendix B.1, we parameterize filters as blocks in the Fourier regime of theconvolution (see Eq.
(B.15)).
Given gradient updates for each block, operations in Algorithm 1 are applied to eachblock simultaneously which can be done in parallel.
Algorithm 1 PROJ UNN update step
Require: unitary matrix U ∈ C N × N or orthogonal matrix U ∈ R N × N , gradient update ∆ U ∈ C N × N ( ∈ R N × N )
Require: hyperparameter k corresponding to rank of approximation 1: Obtain rank k approximation to ∆ U with output (cid:80) ki
=1 a i b † i ≈ ∆ U (see Section 4.3) 2: Follow steps in Theorem 4.2 ( PROJ UNN-D) or Theorem 4.4 ( PROJ UNN-T) in Appendix C: 3:
Perform Gram-Schmidt (via QR decomposition) on concatenation of vectors U † a i and b i for all i ∈
[ k ] : output Q ∈ C N × k as semi-orthogonal matrix containing basis after Gram-Schmidt 4:
Form matrix K ∈ C 2 k × 2 k below:
[Equation]
[Equation]
Q see Eq.
(C.1) to Eq.
(C.6)
PROJ UNN-T: K = 12 (cid:104)(cid:80) ki
=1 Q † U † a i b † i Q − Q † b i
[Equation]
† i UQ (cid:105) see Eq.
(C.12) to Eq.
(C.15) 5:
Find eigenvalues s 1 , . .
, s 2 k and eigenvectors u 1 , . .
, u 2 k of K 6:
Perform update step by applying eigenvalue function:
[Equation]
see Eq.
(C.8) and Eq.
(C.9), ϵ added to ensure stability when s j is close to − 1 (we set ϵ = 10 − 8 )
PROJ UNN-T: U ←
[Equation]
(cid:104)
[Equation]
(exp
[Equation]
ηs j )
[Equation]
1) u j u
[Equation]
j (cid:105) where η is the learning rate see Eq.
(C.17) and Eq.
(C.18)
E Analysis on various benchmarked tasks
E.1 Learning random unitary
In addition to the analysis shown in the main text, here we include a plot showing the value of the Frobenius error as afunction of the runtime (Fig. 6).
Using PROJ UNN-T to peform learning via low rank approximations to the gradientsignificantly speed up learning.
Optimization was performed with vanilla gradient descent over batch sizes of 16 withlearning rates of 0.5 and 0.33 for PROJ UNN-T and PROJ UNN-D respectively.
Since the learning rate was fixed {+acrossall k for each of these experiments, the norm of the update was smaller for lower values of k . Scaling up the learningrate based on the value of k could make runtimes even faster for lower values k . step F r o b e n i u s E rr o r rank k= k= k= k= k= exact    step 
+}
F r o b e n i u s E rr o r rank k= k= k= k= k= exact
[Figure]
o b e n i u s E rr o r
[Figure]
k= k= k= k= k=     step
[Figure]
[Figure]
[Figure]
[Figure]
E rr o
[Figure]
rank k= k= k= k= k=
[Figure]
[Figure]
{+F r o b e n i u s E rr o r rank k= k= k= k= k= exact    step 
 F r o b e n i u s E rr o r rank k= k= k= k=+} k= exact [Figure] r o
[Figure]
e n i u s E rr o r
[Figure]
k
[Figure]
k
[Figure]
k
[Figure]
k
[Figure]
k= [Figure] e n i u s E rr o r rank k= k= k= k= k=
[Figure]
epoch [Figure] t e s t l o ss T= k= k= k= baseline
[Figure] epoch 0. t e s t l o ss T= k= k= k= baseline
For sake of completeness, we also include plots (in both regular and logarithmic scaling axis) showing the learningtrajectory of the various combinations of samplers and PROJ UNN architectures in the random unitary task.
See Fig.
7for PROJ UNN-T with column sampling (left hand side repeated from main text), Fig. 8 for PROJ UNN-D with columnsampling, Fig. 9 for PROJ UNN-T with LSI sampling, and Fig. 10 for PROJ UNN-D with LSI sampling.
Fig. 8 and Fig.
10 also highlight potential instabilities in training PROJ UNN-D over long periods of time, a featurenoted in the main text and discussed further in Appendix H.3.
To alleviate this, for PROJ UNN-D architectures, we setthe learning rate to slightly lower at 0.33 and projected parameter matrices onto the closest unitary using Lemma 4.1every 2048 steps.
Note, that since this projection step was performed sparingly every n steps (where n is the dimensionof the matrix), we did not find any significant decrease in runtime. [Figure]
E.2 Adding task
[-F r o b e n i u s E rr o r rank k= k= k= k= k=     step 
-]
[Figure]
[-F r o b e n i u s E rr o r rank k= k= k= k= k= exact    step 
 F r o b e n i u s E rr o r rank k= k= k= k= k= exact step F r o b e n i u s E rr o r rank k= k= k= k= k=     step 
-]As shown in the main text, our PROJ UNN is able to learn the adding task even for long sequence lengths.
Rankk gradient approximations were obtained using the column sampling approximation.
RMSprop was used to trainthe RNN in this task.
The learning rate was initialized to 0 .
001 for non-orthogonal parameters and 0 .
001 / 32 for allorthogonal parameters.
Each epoch, the learning rate was reduced by multiplying it by 0 .
We found that initializingthe orthogonal matrix as the identity matrix worked well for this task.
This is in distinction with e.g., [35] whichinitialized using the Cayley initialization (block diagonal).
For reference, we also include here an enhanced version ofthe plot in the main text in Fig.
11 and a similar plot obtained using PROJ UNN with LSI sampling in Fig.
E.3 Copy task
Fig.
4 in the main text shows that our PROJ UNN-T can efficiently learn the copy task.
For this task, we set the learningrate of the RMSprop optimizer to 7 e − 4 .
For orthogonal parameters, this learning rate was divided by 32.
Orthogonalmatrices were intialized using Henaff initialization [36] as described in Appendix H.4.
We found that this initializationscheme worked best in comparison to other methods.
E.4 Pixel permuted MNIST
Fig.
13 charts the trajectory of learning on the permuted MNIST task.
As is evident in the figure, rank k = 1 is sufficientto guarantee optimal or nearly optimal convergence in all settings. epoch 0.
t e s t l o ss
T= k= k= k= baseline     epoch 0.
t e s t l o ss
T= k= k= k= baseline
[Figure] E.5 CIFAR10 CNN experiments
To explore the performance of our PROJ UNN training algorithm for convolutional layers, we analyze its performanceon CIFAR10 classification.
Here, we provide further details to the results in the main text as well as some furtherpreliminary experiments on unitary CNNs in resnet architectures.
As in prior work [76] we employ the usual dataaugmentation of random translations and left-right flips.
We leverage the Resnet [34] architecture.
As our previousanalysis in the RNN setting has shown that rank k = 1 is sufficient for convergence, we always set k = 1 when using [Table]
PROJ UNN in the convolutional setting.
We leverage the RMSprop optimizer and perform cross-validation on thelearning rate.
We present our results in Table 3 and make two observations.
First, the smaller model (Resnet9) is able toreach or slightly outperform the alternative exact orthogonal constraints.
Second, the larger model (Resnet18) fallsbehind the approximate orthogonal constraint method.
This result is perhaps expected as our convolutional layers arefull-width, i.e., they allow for much greater degree of over-fitting.
This was not detrimental in the small model withfewer convolutional layers.
As a result, although we validate the ability of PROJ UNN to produce state-of-the-art smallconvolutional networks, there remains open avenues of research to extend the method to larger models.
As an aside, theresnet architecture is, in a sense, a very stable architecture since it is precisely designed to be able to incorporate manylayers.
To provide an honest comparison to prior work, we analyzed the performance of PROJ UNN with respect to thisresnet architecture, but note that more “vanilla" CNN architectures may be better targets for unitary constraints.
F Analysis of low rank updates
It is typically the case that gradients of matrices with respect to a loss function are not exactly low rank but approximatelylow rank.
More specifically, a matrix A is approximately low rank if there exists a rank k matrix A k such that therelative error of the approximation E rel is small:
[Equation]
where ∥ · ∥ F denotes the Frobenius norm of a matrix.
We note that there is a connection between the above and thestable rank of a matrix A defined as ∥ A ∥ 2 F / ∥
A ∥ 22 .
The stable rank is upper bounded by the exact or hard rank of amatrix A .
In RNN settings, Fig. 14 which plots 1 − E rel shows that low rank approximations to a gradient can typically be veryclose to the true gradient even as the hidden size of the RNN is large.
In Fig.
14, the RNN was given an input of batchsize 64 where each input is a sequence of length 20 and each element of the sequence is a vector in R 200 with elementsdrawn from the standard normal distribution.
The RNN outputs logits to learn random target integers ranging from 1 to10 .
Gradients were calculated for a single optimization step over a single batch of the random input and output data.
As observed in the main text, such low rank behavior is often more evident in real-world data.
In fact, as seen in Fig.
1a,virtually all the information of the gradient in a convolutional architecture is captured in the first few singular vectors.
Gradients here are shown for a single C × C block in the Fourier regime of the convolution filter parameterized viaour orthogonal PROJ UNN convolution (see Section 4.4 and Appendix B.1 for form of parameterization).
This filter iscontained in the last residual block of the Resnet-9 network [34] and has 512 channels so the gradient is a 512 × 512matrix.
Since networks were trained with a batch size of 128, the maximum rank of this gradient is actually 128 (seeshort proof below in Proposition F.2).
Since the weight matrix in hidden layers of RNNs is repetitively applied, gradients tend to be of a higher stable rank inthese settings.
Nevertheless, as evident in Fig.
15, gradients over the course of an epoch of training an RNN of hiddendimension 512 on the sequential MNIST task still exhibit clear low rank behavior, albeit not to the same degree asconvolutional architectures.
In constructing these plots, a batch size of 128 was used and networks were trained usingthe RMSprop algorithm.
Throughout our experience with training the PROJ UNN on various tasks, we observe thatsetting the rank of an approximation to even just one is effective at ensuring convergence of training to a good solution.
We furthermore note that the exact rank of a gradient in neural network architectures can typically be bounded bythe dimensions of the input.
For convolutional networks, this exact rank is bounded by the batch size and for vanillarecurrent neural networks, this exact rank is bounded by the batch size times the sequence length.
We provide formalstatements and short proofs of these propositions below.
[Figure]
p e r c e n t o f t o t a l g r a d i e n t n o r m
[Figure]
a
[Figure]
n [Figure] [Figure] Proposition F.1.
Given a loss function ℓ : R × R
→ R taking in two real numbers and outputting a real number, letf ( t ) denote the t -th sequential output of a vanilla RNN defined as
[Equation]
where t ∈ [ T ] indexes the sequence of inputs, inputs x
( t ) ∈ R d ( x denotes the concatenation of all inputs), inputtransformation M ∈ R h × d , hidden transformation W ∈ R h × h , output transformation V ∈ R h × d , and σ and τ arepointwise non-linearities.
Let ∇ W L indicate the gradient of the hidden transformation matrix W with respect to theloss function L = (cid:80) bi =1 ℓ ( f ( T ) ( x i ) , y i ) over a batch of inputs { x i , y i } bi =1 .
The rank of ∇ W L is bounded above bybT .
Proof.
Define z ( t ) =
M x ( t ) + W h ( t − 1) ( x ) .
Calculating the gradient of the loss function with respect to the parameter
W , we have:
[Equation]
The above is the sum of bT rank one matrices which is at most rank bT concluding the proof.
Proposition F.2.
Given a loss function ℓ : R × R → R taking in two real numbers and outputting a real number andconvolution filter W ∈ R N × M × C × C , let f ( X ) denote the convolutional neural network defined as
[Equation]
where input X ∈ R N in × M in × C in and g pre and g post are functions which apply the layers before and after theconvolution with filter W .
Let (cid:98) W (cid:98) r, (cid:98) s, : , : denote the C × C matrix storing the values of the convolution filter in the Fourierregime for frequencies (cid:98) r and (cid:98) s (see Eq.
(B.15) ).
Let ∇ (cid:98) W (cid:98) r, (cid:98) s, : , : L indicate the gradient of (cid:98) W (cid:98) r, (cid:98) s, : , : with respect to the lossfunction L = (cid:80) bi =1 ℓ ( f ( X i ) , y i ) over a batch of inputs { X i , y i } bi =1 .
Then, for all (cid:98) r and (cid:98) s in the support of the filter, therank of ∇ (cid:98) W (cid:98) r, (cid:98) s, : , : L is bounded above by b .
Proof.
We have that
[Equation]
We now proceed to show that the rank of ∇ (cid:98) W [Equation]
[Equation]
r
[Equation]
s ℓ ( f ( X i ) , y i ) is equal to one which completes the proof.
As a reminder,we have via the convolution theorem that
[Equation]
Let (cid:98) Y ( X ) =
FFT conv W ( X ) , then we have
[Equation]
The above is a rank one matrix which concludes the proof. [Table]
F.1 Other sampling algorithms
[Figure] The prior analysis showed that gradients are typically low rank, and in the main text, we listed a couple efficientalgorithms that allow one to approximate a full but approximately low rank matrix with an explicitly low rank matrix.
Enhancements to these algorithms exist that may provide further improvements in very large dimensions.
For example,the FKV sampling [27] and constant time SVD [23] algorithms provide runtimes that are often logarithmic in thedimension of the matrix for sampling entries from the low rank approximation to a matrix.
Nevertheless, explicitlyconstructing all entries of a rank k approximation to an n × n matrix A requires time at least O ( k 2 n ) .
Improvements inruntime are achieved by applying random projections to both the column and row subspaces of the matrix A to performthe final approximation.
In practice, the runtime of these algorithms depends on other factors that limit the applicabilityof the algorithm even for relatively large matrices.
Numerical simulations show that FKV sampling achieves a practical [Table] speedup for matrices of dimension approximately 10 6 or higher [8].
Since matrices in our PROJ UNN were of muchsmaller dimension, we did not use these methods.
G Runtime comparisons
Table 1 (see main text) and Table 4 lists the asymptotic runtime complexities of unitary/orthogonal models in the
RNN and convolutional setting respectively.
For RNNs, PROJ UNN has the best runtime scaling of all models whichparameterize the full orthogonal/unitary space with just one layer.
In fact, apart from the rank approximation factor k ,the runtime of PROJ UNN is optimal in the RNN setting.
In the convolutional setting, PROJ UNN scales most efficientlywhen there are many channels or the filter size W scales faster than the logarithm of the input dimension log N .
Forsmall filter sizes, BCOP and SOC scale relatively efficiently.
In fact, SOC is nearly optimal for W ≪ log N up to theapproximation factor p in the Taylor expansion which is required both upon evaluation and training of the network.
Runtimes for RNN implementations in practice Enforcing unitarity in PROJ UNN incurs a computational overheadassociated to performing eigendecompositions and QR decompositions entailed in updating the gradient.
Even thoughsuch computations are performed on a small subspace of the total dimension of the matrix, such computations maynot effect increase training times by a constant factor as evident in Fig.
2a.
Empirically, we observed that much of theincreased overhead was due to performing eigendecompositions and QR decompositions on a GPU, both tasks whichare challenging to parallelize on GPU architectures.
We note that similar issues may be one reason why the expRNNruntimes appear much slower in our simulations as shown in Fig.
2a.
Recent updates to pytorch [70] and tensorflow [2]have focused on improving the runtimes of these common linear algebraic routines, and we expect these runtimes tocontinue to improve in the future.
Runtimes for convolutional network implementations in practice Fig.
16, Fig. 17, and Fig.
18 compare runtimesof orthogonal convolution algorithms when varying the number of channels, input size, and filter size respectively.
Fig.
19 plots the runtime when the filter and image size are set equal to each other and varied together.
In summary,we find that the skew orthogonal convolution (SOC) [73], the Cayley model [76], and our PROJ UNN are all relativelyefficient for small to medium sized models.
PROJ UNN empirically performs best in comparison to other models whenthere are many channels or the filter size is large.
SOC [73] is the fastest when the filter size is small.
As a reminder,
SOC [73] employs the Taylor series approximation to form a unitary/orthogonal matrix.
We did not change the numberof terms in the approximation used even though increasing the size of a filter, dimension of an input, or the number ofchannels will likely require more terms to maintain the same error in the approximation.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
The Cayley model [76] has a runtime that is empirically similar to PROJ UNN in most of our experiments despite
PROJ UNN having an improved asymptotic scaling with regards to the number of channels (see Table 4).
Runtimes hereare calculated by averaging across the first 10 steps of gradient descent on random data for initialized networks.
Closeto initialization, matrices in the convolution filter are generally well conditioned and this may be one reason
why the
Cayley model showed good scaling in Fig.
16 since the matrix inversion for the well conditioned matrices runs quickly.
We also suspect that for the number of channels considered, the added costs of performing sampling and projections in
PROJ UNN do not materialize into a very clear runtime benefit.
Interestingly, the exact version of PROJ UNN (no lowrank approximation) runs in similar time to the Cayley method.
Finally, we note that in comparison to RNN implementations, there is increased overhead associated to performingeigendecompositions and QR decompositions with our projUNN convolutional network implementations since unitarymatrices are batched across the elements of the convolutional filter.
As mentioned before, much of the increasedoverhead was due to performing eigendecompositions and QR decompositions on a GPU, both tasks which can bechallenging to parallelize on GPU based architectures.
Future updates to Pytorch [70] and Tensorflow [2] may helpimprove runtimes of these operations.
H Network architectures and training details
ExpRNN trivializations https://github.com/Lezcano/expRNN
MIT scoRNN https://github.com/SpartinStuff/scoRNN
NA scuRNN https://github.com/Gayan225/scuRNN
NA skewed orthogonal convolutions (SOC)
https://github.com/singlasahil14/SOC
NA
Cayley convolution https://github.com/locuslab/orthogonal-convolutions
MIT
BCOP github.com/ColinQiyangLi/LConvNet
MIT
H.1 Handling complex numbers
Since matrices in the unitary group are complex valued, care must be taken to ensure that the neural network can handlesuch complex numbers.
In these settings, standard activation functions need to be adapted for complex numbers.
It isoften advantageous to have a nonlinearity which does not change the phase of its input.
For these reasons, the activationfunction used is typically the modRELU activation shown below:
[Equation]
where b is a bias term (trainable).
In calculations of the modulus | z | , a small additive constant is often added to avoidstability issues when | z | is small.
Another challenge that arises with performing orthogonal convolution is that the representation of a real-valuedconvolutional filter in the Fourier domain will be complex-valued.
More specifically, the fast Fourier transform of areal signal is Hermitian-symmetric.
For example, for a one dimensional vector x ∈ R N where (cid:98) x = FFT x ∈ C N and(cid:98) x i = (cid:98) x − i mod N .
Therefore, when initializing weight filters in our orthogonal convolution operations, one must becareful to ensure the Hermitian symmetric property holds.
Similarly, when performing convolutions in the Fourierspace, care must be taken in converting data types to and from complex and real space.
In our implementation, we usedthe pytorch FFT.RFFT2 and FFT.IRFFT2 commands to implement these operations efficiently [69].
H.2 Optimizers
For training PROJ UNN architectures, a separate optimizer is used for unitary/orthogonal parameters which must beprojected after updates and all other parameters.
Learning rates for the unitary/orthogonal parameters are typicallyset to less than that of the other parameters.
In our implementations, we found that a learning rate for the unitary ororthogonal parameters of one-tenth or one-twentieth of that of the other parameters works well in practice.
Furthermore,when using optimizers with momentum terms, we constructed a variant of standard optimizers to apply changes tothe momentum terms after projecting gradients onto the unitary/orthogonal manifold or tangent space.
This optimizerperformed well in our experiments though sometimes added instability.
Therefore, unless otherwise stated, standardoptimizers were used.
H.3 Numerical stability
Though updates via PROJ UNN mathematically guarantee that parameters remain unitary/orthogonal, performing asignificant number of sequential updates to a matrix can add numerical errors over time slowly drifting parameters awayfrom the unitary/orthogonal manifold.
This is especially true in the case of PROJ UNN-D where updating matricesrequires division of eigenvalues in the eigendecompositon of the rank k subspace (see Appendix C and Eq.
(C.16)).
In contrast, we empirically find that PROJ UNN-T maintains stability throughout optimization even when requiringmany updates.
To actively address potential instabilities, one can perform eigendecompositions with higher floatingpoint precision or sporadically project unitary/orthogonal parameters onto the closest unitary/orthogonal matrix via
Lemma 4.1.
Though this projection step requires O ( n 3 ) time for an n × n matrix, performing this projection only every
O ( n ) gradient updates can still preserve efficient runtimes of on average O ( kn 2 ) time per gradient update.
Unlessotherwise stated, we do not perform any additional steps for maintaining stability.
H.4 Initialization of unitary/orthogonal parameters
Empirically, we found that initializing unitary/orthogonal matrices to be Haar random or close to Haar
random doesnot perform well.
This is also an observation noted in prior works [36, 35, 56].
Instead, we used one of two differentinitialization schemes.
The first initializes unitary/orthogonal parameters as the identity matrix.
The second initializesunitary/orthogonal matrices using variants of the so-called Henaff initialization [36] where 2 × 2 diagonal blocks of thematrices are initialized as samples from the below
[Equation]
where s i is sampled uniformly from [ − π, π ] [36].
In other words, an n × n matrix U is initialized as
[Equation]
(cid:21)(cid:19) ·
[Equation]
exp (cid:18)(cid:20)
[Equation]
(cid:21)(cid:19)
[Equation] [Equation]
(cid:18)(cid:20)
[Equation]
[Equation]
(cid:21)(cid:19) 
 .
(H.3)
Note, that since we parameterize matrices in the Lie group instead of the Lie algebra, we include the exponential mapin the parameterization above.
Other variants of the above have been used in prior works.
For example, the Cayleyinitialization samples s i = (cid:113) 1 − cos t i 1+cos t i where t i is sampled uniformly from [0 , π/ 2] [35].
H.5 Computational details
We employed the latest version of PyTorch (1.10.1+cu102) with Python3.8 and all the pre-requisite libraries comingalong.
The hardware leverages Quadro GP100 GPU and Intel(R)
Xeon(R) CPU E5-2698 v4 @
2.20GHz.