O NE R EFLECTION S UFFICE
Anonymous authors
Paper under double-blind review
A BSTRACT
Orthogonal weight matrices are used in many areas of deep learning.
Much previous work attempt to alleviate the additional computational resources it requires to constrain weight matrices to be orthogonal.
One popular approach utilizes many
Householder reﬂections.
The only practical drawback is that many reﬂections cause low GPU utilization.
We mitigate this ﬁnal drawback by proving that one reﬂection is sufﬁcient, if the reﬂection is computed by an auxiliary neural network.
I NTRODUCTION
Orthogonal matrices have shown several beneﬁts in deep learning, with successful applications in
Recurrent Neural Networks, Convolutional Neural Networks and Normalizing Flows.
One popular approach can represent any d × d orthogonal matrix using d Householder reﬂections (Mhammedi et al., 2017).
The only practical drawback is low GPU utilization, which happens because the d reﬂections needs to be evaluated sequentially (Mathiasen et al., 2020).
Previous work often increases GPU utilization by using k (cid:28) d reﬂections (Tomczak & Welling, 2016; Mhammedi et al., 2017; Zhang et al., 2018; Berg et al., 2018).
Using fewer reﬂections limits the orthogonal transformations the reﬂections can represent, yielding a trade-off between representational power and computation time.
This raises an intriguing question: can we circumvent the trade-off and attain full representational power without sacriﬁcing computation time?
We answer this question with a surprising “yes.”
The key idea is to use an auxiliary neural network to compute a different reﬂection for each input.
In theory, we prove that one such “auxiliary reﬂection” can represent any number of normal reﬂections.
In practice, we demonstrate that one auxiliary reﬂection attains similar validation error to models with d normal reﬂections, when training Fully
Connected Neural Networks (Figure 1 left), Recurrent Neural Networks (Figure 1 center) and convolutions in Normalizing Flows (Figure 1 right).
Notably, auxiliary reﬂections train between 2 andtimes faster for Fully Connected Neural Networks with orthogonal weight matrices (see Section 3)
[Figure]
1. O UR R ESULTS
The Householder reﬂection of x ∈ R d around v ∈ R d can be represented by a matrix H ( v ) ∈ R d × d .
[Equation]
An auxiliary reﬂection uses a Householder matrix H ( v ) with v = n ( x ) for a neural network n .
[Equation]
One auxiliary reﬂection can represent any composition of Householder reﬂections.
We prove this claim even when we restrict the neural network n ( x ) to have a single linear layer n ( x ) =
Wx for
W ∈ R d × d such that f ( x ) =
H ( Wx ) x .
Theorem 1.
For any k Householder reﬂections U = H ( v 1 ) · · · H ( v k ) there exists a neural network n ( x ) =
Wx with W ∈ R
d
× d such that f ( x ) =
H ( Wx ) x
= Ux for all x ∈ R d \{ 0 } .
Previous work (Mhammedi et al., 2017; Zhang et al., 2018) often employ k (cid:28) d reﬂections and compute
Ux as k sequential Householder reﬂections H ( v 1 ) · · · H ( v k ) · x with weights V = ( v 1 · · · v k ) .
It is the evaluation of these sequential Householder reﬂection that cause low GPU utilization (Mathiasen et al., 2020), so lower values of k increase GPU utilization but decrease representational power.
Theorem 1 states that it is sufﬁcient to evaluate a single auxiliary reﬂection H ( Wx ) x instead of k reﬂections H ( v 1 ) · · · H ( v k ) · x , thereby gaining high GPU utilization while retaining the full representational power of any number of reﬂections.
In practice , we demonstrate that d reﬂections can be substituted with a single auxiliary reﬂection without decreasing validation error, when training Fully Connected Neural Networks (Section 3.1),
Recurrent Neural Networks (Section 3.2) and Normalizing Flows (Section 3.3).
While the use of auxiliary reﬂections is straightforward for Fully Connected Neural Networks and Recurrent Neural
Networks, we needed additional ideas to support auxiliary reﬂections in Normalizing Flows.
In particular, we developed further theory concerning the inverse and Jacobian of f ( x ) =
H ( Wx ) x .
Note that f is invertible if there exists a unique x given y
= H
( Wx ) x and W .
Theorem 2.
Let f ( x ) =
H ( Wx ) x with f (0) :
= 0 , then f is invertible on R d with d ≥ 2 if
W = W T and has eigenvalues which satisfy 3 / 2 · λ min ( W ) >
λ max ( W ) .
Finally, we present a matrix formula for the Jacobian of the auxiliary reﬂection f ( x ) =
H ( Wx ) x .
This matrix formula is used in our proof of Theorem 2, but it also allows us simplify the Jacobian determinant (Lemma 1) which is needed when training Normalizing Flows.
Theorem 3.
The Jacobian of f ( x ) =
H ( Wx ) x is:
[Equation]
We prove Theorem 1 in Appendix A.1.1 while Theorems 2 and 3 are proved in Section 2.
N ORMALIZING F LOWS 2. B ACKGROUND
Let z ∼ N (0 , 1) d and f be an invertible neural network.
Then f − 1 ( z ) ∼ P model deﬁnes a model distribution for which we can compute likelihood of x ∼ P data (Dinh et al., 2015).
[Equation]
This allows us to train invertible neural network as generative models by maximum likelihood.
Previous work demonstrate how to construct invertible neural networks and efﬁciently compute the log jacobian determinant (Dinh et al., 2017; Kingma & Dhariwal, 2018;
Ho et al., 2019).
2. I NVERTIBILITY AND J ACOBIAN D ETERMINANT (P ROOF S KETCH )
To use auxiliary reﬂections in Normalizing Flows we need invertibility.
That is, for every y ∈ R d there must exist a unique x ∈ R d so f ( x )
=
H ( Wx ) x = y .
We ﬁnd that f is invertible if its
Jacobian determinant is non-zero for all x in S d − 1 = { x ∈ R d | (cid:107) x
(cid:107) = 1 } .
Theorem 4.
Let f ( x ) =
H ( Wx ) x with f (0) :
= 0 , then f is invertible on R d with d ≥ 2 if the
Jacobian determinant of f is non-zero for all x ∈ S
d − 1 and W is invertible.
The Jacobian determinant of H ( Wx ) x takes the following form.
Lemma 1.
The Jacobian determinant of f ( x ) =
H ( Wx ) x is:
[Equation]
||
[Equation]
|(cid:19) where v T = x T W, u
=
Wx and A
= I − 2 x T W T x
||
Wx ||
W.
It is then sufﬁcient that det( A ) (cid:54) = 0 and 1 + 2 v T A − 1 u/ || u
|| 2 (cid:54) = 0 .
We prove that this happens if
W = W T with eigenvalues 3 / 2 · λ min ( W ) >
λ max ( W ) .
This can be achieved with W =
I + V V T if we guarantee σ max ( V V T ) <
1 / 2 by spectral normalization (Miyato et al., 2018).
Combining these results yields Theorem 2.
Theorem 2.
Let f ( x ) =
H ( Wx ) x with f (0) :
= 0 , then f is invertible on R d with d ≥ 2 if
W = W T and has eigenvalues which satisfy 3 / 2 · λ min ( W ) >
λ max ( W ) .
Computing the Inverse.
In practice, we use Newtons method to compute x so H ( Wx ) x = y .
Figure 2 show reconstructions n − 1 ( n ( x ))
=
x for an invertible neural network n with auxiliary reﬂections using Newtons method, see Appendix A.2.1 for details.
[Figure]
2. P ROOFS
The goal of this section is to prove
that f ( x ) =
H ( Wx ) x is invertible.
Our proof strategy has two parts.
Section 2.3.1 ﬁrst shows f is invertible if it has non-zero Jacobian determinant.
Section 2.3.
then present an expression for the Jacobian determinant, Lemma 1, and prove the expression is non-zero if W = W T and 3 / 2 · λ min ( W )
> λ min ( W ) .
2.3. N ON -Z ERO J ACOBIAN D ETERMINANT I MPLIES I NVERTIBILITY
In this section, we prove that f ( x ) =
H ( Wx ) x is invertible on R d if f has non-zero Jacobian determinant.
To simplify matters, we ﬁrst prove that invertibility on S d − 1 implies invertibility on R d .
Informally, invertibility on S d − 1 is sufﬁcient because H ( Wx ) is scale invariant, i.e., H ( c · Wx ) =
H ( Wx ) for all c (cid:54) = 0 .
This is formalized by Lemma 2.
Lemma 2.
If f ( x ) = H ( Wx ) x is invertible on S d − 1 it is also invertible on R d \{ 0 } .
Proof.
Assume that f ( x ) is invertible on S d − 1 .
Pick any y (cid:48)
∈ R d such that || y (cid:48)
||
=
c for any c > 0 .
Our goal is to compute x (cid:48) such that H ( Wx (cid:48) )
x (cid:48) =
y (cid:48) .
By normalizing, we see y (cid:48) / (cid:107) y (cid:48) (cid:107) ∈
S
d − 1 .
We can then use the inverse f − 1 on y (cid:48) / (cid:107) y (cid:48)
(cid:107) to ﬁnd x such that H ( Wx ) x = y (cid:48) / (cid:107) y (cid:107) .
The result is then x (cid:48)
=
x (cid:107) y (cid:107) since H ( Wx (cid:48) )
x (cid:48)
=
H ( Wx ) x
|| y
||
=
y due to scale invariance of H ( Wx ) .
The main theorem we use to prove invertibiliy on S d − 1 is a variant of Hadamards global function inverse theorem from (Krantz & Parks, 2012).
On a high-level, Hadamard’s theorem says that a function is invertible if it has non-zero Jacobian determinant and satisﬁes a few additional conditions.
It turns out that these additional conditions are meet by any continuously differentiable function f ( x ) when (in the notation of Theorem 5) M 1 =
M 2 = S
d − 1 .
Theorem 5.
(Krantz & Parks, 2012, 6.2.8) Let M 1 and M 2 be smooth, connected N -dimensional manifolds and let f : M 1 → M 2 be continuously differentiable.
If (1) f is proper, (2) the Jacobian of f is non-zero, and (3) M 2 is simple connected, then f is invertible.
For M 1 = M 2 = S d − 1 the additional conditions are met if f is continuously differentiable.
Corollary 1.
Let f : S d − 1 → S d − 1 with d ≥ 2 be continuously differentiable with non-zero
Jacobian determinant, then f is invertible.
Proof.
Note that S d − 1 is smooth and simply connected if d ≥ 2 (Lee, 2013).
Continuously functions on S d − 1 are proper.
We conclude f is invertible on S d − 1 by Theorem 5.
We now show that f ( x ) =
H ( Wx ) x is continuously differentiable on S d − 1 .
Lemma 3.
The function f ( x ) =
H ( Wx ) x is continuously differentiable on S d − 1 if W is invertible.
Proof.
Compositions of continuously differentiable functions are continuously differentiable by the chain rule.
All the functions used to construct H ( Wx ) x are continuously differentiable, except the division.
However, the only case where division is not continously differentiable is when ||
Wx || = 0 .
Since W is invertible,
||
Wx || =
iff x = 0 .
But 0 / ∈ S
d − 1 and we conclude f is continuously differentiable on S d − 1 .
Theorem 4.
Let f ( x ) =
H ( Wx ) x with f (0) :
= 0 , then f is invertible on R d with d ≥ 2 if the
Jacobian determinant of f is non-zero for all x ∈ S
d − 1 and W is invertible.
Proof.
By Lemma 3, we see f is continuously differentiable since W is invertible, which by Corollary 1 means f is invertible on S d − 1 if f has non-zero Jacobian determinant on S d − 1 .
By Lemma 2, we get that f is invertible on R d if it has non-zero Jacobian on S d − 1 .
2.3. E NFORCING N ON -Z ERO J ACOBIAN D ETERMINANT
The goal of this section is to present conditions on W that ensures the Jacobian determinant of f ( x ) is non-zero for all x ∈ S d − 1 .
We ﬁrst present a matrix formula for the Jacobian of f in Theorem 3.
By using the matrix determinant lemma , we get a formula for the Jacobian determinant in Lemma 1.
By investigating when this expression can be zero, we ﬁnally arive at Lemma 4 which states that
the
Jacobian determinant is non-zero (and f thus invertible) if W = W T and 3 / 2 · λ min > λ max .
Theorem 3.
The Jacobian of f ( x ) =
H ( Wx ) x is:
[Equation]
See Appendix A.2.2 for PyTorch implementation of J and a test case against PyTorch autograd.
Proof.
The ( i, j ) ’
th entry of the Jacobian determinant is, by deﬁnition,
[Equation]
Then, by the product rule, we get
[Equation]
[Equation]
First, (Petersen & Pedersen, 2012) equation (81) gives ∂x T W T x ∂x j
[Equation]
x
[Equation]
j .
Second ||
Wx || − 2 can be found using the chain rule:
[Equation]
Petersen & Pedersen, 2012, equ. 81)
[Equation]
Combining everything we get
[Equation]
In matrix notation
[Equation]
this translates into the following, if we let
A =
I − 2 · x T W T x
||
Wx || 2 W .
[Equation]
This concludes the proof.
Theorem 3 allows us to write J as a rank one update M + ab T for a, b ∈ R d , which can be used to simplify det( J ) as stated in the following lemma.
Lemma 1.
The Jacobian determinant of f ( x ) =
H ( Wx ) x is: − det( A
[Equation]
|| u ||(cid:19) where v T = x T W, u
=
Wx and A
= I − 2 x T W T x
||
Wx ||
W.
Proof.
The matrix determinant lemma allows us to write det( M + ab T )
= det( M )(1 + b T M − 1 a ) .
Let M = H ( Wx ) A and b T = − 2 · x T W/
||
Wx || 2 and a =
Wx .
The Jacobian J from Theoremis then J = M + ab T .
The determinant of J is then:
[Equation]
| [Equation] This is true because H ( Wx ) − 1 =
H
( Wx ) , H ( Wx ) ·
Wx = − Wx and det( H ( Wx ))
= − 1 .
We can now use Lemma 1 to investigate when the Jacobian determinant is non-zero.
In particular, the Jacobian determinant must be non-zero if both det( A ) (cid:54) =
and 1+2 v T A − 1 u/ || u
|| 2 (cid:54) = 0 .
In the following lemma, we prove that both are non-zero if W = W T and 3 / 2 · λ min > λ max .
Lemma 4.
Let W = W T and 3 / 2 · λ min > λ max
then λ i ( A − 1 ) < − 1 / 2 for A from Lemma 1.
These conditions imply that det( A ) (cid:54) = 0 and 1 + 2 v T A − 1 u/ || u
|| 2 (cid:54) =
with v T , u from Lemma
Proof.
We ﬁrst show that the inequality 3 / 2 · λ min ( W ) > λ max ( W ) implies λ i ( A − 1 ) <
− 1 / 2 .
[Equation]
−
[Equation]
If γ i := x T W T x
||
Wx || 2 · λ i ( W ) ∈ (1 / 2 , 3 / 2)
we get that 1 / (1 − 2 γ i ) ∈ ( −∞ , − 1 / 2)
so λ i
( A − 1 ) < − 1 / 2 .
If we let y :
=
Wx we get x T W T x
||
Wx || 2 =
y T W − 1 y
||
y
|| 2 .
This is the Rayleigh quotient of W − 1 at y , which for W = W T is within [ λ min ( W − 1 ) , λ max ( W − 1 )] .
Therefore γ i ∈ [ 1 λ max ( W ) , 1 λ min ( W ) ] · λ i ( W ) .
Note ﬁrst that γ min ≤ 1 and γ max ≥ 1 .
It is left to show that γ min ≥ λ min /λ max > 1 / 2 and γ max ≤ λ max /λ min < 3 / 2 .
Both conditions on eigenvalues are met if 3 / 2 · λ min > λ max .
We now want to show that det( A ) (cid:54) = 0 and 1 + 2 v T A − 1 u/ || u
|| 2 (cid:54) = 0 .
First, notice that det( A ) =
(cid:81) di
=
λ i ( A ) (cid:54)
=
since λ i ( A ) < − 1 / 2 .
Second, note that W = W T implies that the v T from
Lemma 1 can be written as v T =
x T W = x T W T = u T .
This means we only need to ensure u T A − 1 u/ || u
|| 2 , the Rayleigh quotient of A − 1 at u , is different to − 1 / 2 .
But W = W T implies
A =
A T because A = I − 2 x T W T
x/ ||
Wx || 2 · W .
The Rayleigh quotient is therefore bounded by [ λ min ( A − 1 ) , λ max ( A − 1 )] , which means it is less than − 1 / 2 since λ
i ( A − 1 ) < − 1 / 2 .
We can then conclude that also 1 + 2 v T A − 1 u/ || u
|| 2 = 1 + 2 u T A − 1 u/ || u
|| 2 < 1 + 2 · − 1 / 2 = 0 .
So det( J ) (cid:54) = 0 by Lemma 4 and Lemma 1, which by Theorem 4 implies invertibility (Theorem 2).
Remark.
Note that the constraints W = W T and 3 / 2 · λ min > λ max were introduced only to guarantee det( A ) (cid:54) = 0 and 1 + 2 v T A − 1 u/ || u
|| 2 (cid:54) = 0 .
Any argument or constraints on W that ensures det( A ) · (1 + v T A − 1 u/ || u || 2 ) (cid:54) = 0 are thus sufﬁcient to conclude f ( x ) is invertible.
E XPERIMENTS
We compare a single auxiliary reﬂections against d normal reﬂections when training Fully Connected Neural Networks ( d = 784 ), Recurrent Neural Networks ( d = 170) and Normalizing Flows ( d = 48 ).
The experiments demonstrate that neural networks with a single auxiliary reﬂections attain similar performance to neural networks with many normal reﬂections.
All plots show means and standard deviations over 3 runs.
See Appendix B for experimental details.
3. F ULLY C ONNECTED N EURAL N ETWORKS
We trained four different Fully Connected Neural Networks (FCNNs) for classiﬁcation on MNIST.
We compared a FCNN with 6 auxiliary reﬂections against a FCNN with 6 orthogonal matrices each represented by 784 normal reﬂections.
For completeness, we also trained two FCNNs where theorthogonal matrices where attained by the matrix exponential and Cayley map, respectively, as done in (Casado, 2019; Lezcano-Casado & Mart´ınez-Rubio, 2019).
The FCNN with auxiliary reﬂections attained slightly better validation error, see (Figure 3 left).
Furthermore, we found the auxiliary reﬂections were 2 to 6 times faster than competing methods, see (Figure 3 right).
This was true even though we used (Mathiasen et al., 2020) to speed up the sequential Householder reﬂections.
See
Appendix B.1 for further details.
3. R ECURRENT N EURAL N ETWORKS
We trained three Recurrent Neural Networks (RNNs) for classiﬁcation on MNIST as done in (Mhammedi et al., 2017).
The RNNs had a transition matrix represented by one auxiliary reﬂection, one normal reﬂection and 170 auxiliary reﬂections.
See (Figure 4 left) for a validation error during training, including the model from (Mhammedi et al., 2017).
As indicated by the red curve, using only one normal reﬂection severely limits the transition matrix.
In the right plot, we magnify the ﬁrst 20 epochs to improve readability.
The RNNs with 1 auxiliary reﬂection attains similar mean validation accuracy to the RNNs with 170 normal reﬂections.
See Appendix B.2 for further details.
[Figure]
[Figure] [Figure]
3. N ORMALIZING F LOWS AND C ONVOLUTIONS
We initially trained two Normalizing Flows (NFs) on CIFAR10.
Inspired by (Hoogeboom et al., 2019), we used reﬂections to parameterize the 1x1 convolutions of an NF called Glow (Kingma & Dhariwal, 2018), see Appendix B.3 for details.
We trained an NF with many reﬂections and an
NF with a single auxiliary reﬂection constrained to ensure invertible (see Section 2.2).
The single auxiliary reﬂection attained worse validation NLL compared to the model with 48 normal reﬂections.
We suspected the decrease in performance was caused by the restrictions put on the weight matrices
W i of the auxiliary reﬂections to enforce invertibility, i.e., W i =
W Ti and 3 / 2 · λ min ( W i ) >
λ max ( W i ) .
To investigate this suspicion, we trained a model with no constraints on W .
This improved performance to the point were one auxiliary reﬂections tied with many normal reﬂections (see Figure 5 left).
Even though the resulting auxiliary reﬂections are not provably invertible, we found that Newtons method consistently computed the correct inverse.
Based on this observation, we conjecture that the training dynamics caused the auxiliary reﬂections to remain invertible.
By this we mean that the auxiliary reﬂections were initialized with non-zero Jacobian determinants (see Appendix B.3) and the loss function (Equation (1)) encourages the auxiliary reﬂections to increase their Jacobian determinants during training.
Since Newtons method consistently computed the correct inverse, we were able to generate samples from all models, see (Figure 5 right).
[Figure]
R ELATED W ORK
Orthogonal Weight Matrices.
Orthogonal weight matrices have seen widespread use in deep learning.
For example, they have been used in Normalizing Flows (Hoogeboom et al., 2019), Variational Auto Encoders (Berg et al., 2018), Recurrent Neural Networks (Mhammedi et al., 2017) and Convolutional Neural Networks (Bansal et al., 2018).
Different Approaches.
There are several ways of constraining weight matrices to remain orthogonal.
For example, previous work have used Householder reﬂections (Mhammedi et al., 2017), the Cayley map (Lezcano-Casado & Mart ´ ınez-Rubio, 2019) and the matrix exponential (Casado, 2019).
These approaches are sometimes referred to as hard orthogonality constraints , as opposed to soft orthogonality constraints , which instead provide approximate orthogonality by using, e.g., regularizers like || WW T − I || F (see (Bansal et al., 2018) for a comprehensive review).
Reﬂection Based Approaches.
The reﬂection based approaches introduce sequential computations, which is, perhaps, their main limitation.
Authors often address this by reducing the number of reﬂections, as done in, e.g., (Tomczak & Welling, 2016; Mhammedi et al., 2017; Berg et al., 2018).
This is sometimes undesirable, as it limits the expressiveness of the orthogonal matrix.
This motivated previous work to construct algorithms that increase parallelization of Householder products, see, e.g., (Mathiasen et al., 2020; Likhosherstov et al., 2020).
Similar Ideas.
Normalizing Flows have been used for variational inference, see, e.g., (Tomczak & Welling, 2016;
Berg et al., 2018).
Their use of reﬂections is very similar to auxiliary reﬂections, however, there is a very subtle difference which has fundamental consequences.
For a full appreciation of this difference, the reader might want to consult the schematic in (Tomczak & Welling, 2016,
Figure 1), however, we hope that the text below clariﬁes the high-level difference.
Recall that auxiliary reﬂections compute H ( Wx ) x
so H ( Wx ) can depend on x .
In contrast, the previous work on variational inference instead compute H ( v ) z where v and z both depend on x .
This limits H ( v ) in that it can not explicitly depend on z .
While this difference is subtle, it means our proof of Theorem 1 does not hold for reﬂections as used in (Tomczak & Welling, 2016).
C ONCLUSION
In theory, we proved that a single auxiliary reﬂection is as expressive as any number of normal reﬂections.
In practice, we demonstrated that a single auxiliary reﬂection can attain similar performance to many normal reﬂections when training Fully Connected Neural Networks, Recurrent Neural Networks and Normalizing Flows.
For Fully Connected Neural Networks, we reduced training time by a factor between 2 and 6 by using auxiliary reﬂections instead of previous approaches to orthogonal matrices (Mhammedi et al., 2017; Lezcano-Casado & Mart´ınez-Rubio, 2019; Casado, 2019).
P ROOFS
A
T HEOREM
Our proof of Theorem 1 is an follows Lemma 5 which we state below.
Theorem 1.
For any k Householder reﬂections U = H ( v 1 ) · · · H ( v k ) there exists a neural network n ( x ) =
Wx with W ∈ R
d
× d such that f ( x ) =
H ( Wx ) x
= Ux for all x ∈ R d \{ 0 } .
Proof.
Let W =
I − U
then H
( Wx ) x
= H ( x − Ux ) x
= Ux for all x ∈ R d since ||
Ux ||
=
||
x
|| .
Lemma 5.
Let || x
||
=
||
y
|| then H ( x − y ) x
= y .
Proof.
The result is elementary, see, e.g., (Wang, 2015).
For completeness, we derive it below.
[Equation]
A. P Y T ORCH E XAMPLES AND T EST C ASES
To ease the workload on reviewers, we opted to use small code snippets that can be copied into www.colab.research.google.com and run in a few seconds without installing any dependencies.
Some PDF viewers do not copy line breaks, we found viewing the PDF in Google Chrome works.
A.2. T EST C ASE : I NVERSE USING N EWTONS M ETHOD
Given y we compute x such that H ( Wx ) x =
y using Newtons method.
To be concrete, the code below contains a toy example where x ∈ R 4 and W =
I + V V T / (2 · σ max ( V V T ))
∈ R 4 × 4 .
The particular choice of W makes H ( Wx ) x invertible, because λ i ( W ) =
+ λ i ( V V T )
= 1+ σ i ( V V T ) ∈ [1 , 3 / 2) because V V T is positive deﬁnite.
Any possible way of choosing the eigenvalues in the range [1 , 3 / 2) guarantees that 3 / 2 · λ min > λ max which implies invertibility by Theorem 2.
import torch print ("torch version: ", torch.__version__) torch.manual_seed(42)d=
# Create random test-case.
I = torch.eye(d)
V = torch.zeros((d, d)).uniform_() x
= torch.zeros((d, 1)).uniform_()
W =
I + V @ V.T / torch.svd(V @
V.T)[1].max() # Define the function f(x)=H(Wx)x.
def H(v): return torch.eye(d) - 2 * v @ v.T / (v.T @ v) def f(x): return H(W @ x ) @ x
# Print input and output print ("x\t\t", x.data.view(-1).numpy()) print ("f(x)\t", f(x).data.view(-1).numpy()) print ("")
# Use Newtons Method to compute inverse.
y = f(x) xi = y for i in range (10): print ("[ % .2i/
% .2i]" % (i+1, 10), xi.data.view(-1).numpy())
# Compute Jacobian using Theorem 3.
A = torch.eye(d) - 2* (xi.
T @ W.T @
xi) /
torch.norm(W @
xi)**
* W J = -2*W @
xi @
xi.
T @
W/torch.norm(W@xi)**2 + H(W @ xi)
A xi = xi - torch.inverse(J) @
(f(xi)- y) assert torch.allclose(xi, x, atol=10**(-7)) print ("The two vectors are torch.allclose") torch version: 1.6.0+cu101 x [0.
0.57390445 0.26658005 0.62744915] f(x)
[Table]
The two vectors are torch.allclose
A.2. T EST C ASE : J ACOBIAN AND A UTOGRAD
import torch print ("torch version: ", torch.__version__) torch.manual_seed(42)
# Create random test-case.
d =
W = torch.zeros((d, d)).uniform_(-1, 1) x
= torch.zeros((d, 1)).uniform_(-1, 1)
I = torch.eye(d) # Compute Jacobian using autograd.
def H(v): return I - 2 * v @ v.T / (v.T @ v) def f(x): return H(W @ x ) @
x J =
torch.autograd.functional.jacobian(f, x)[:, 0, :, 0] print (J) #
Compute Jacobian using Lemma 4.
A = [Equation]
torch. [Equation]
[Equation] [Equation] [Equation]
/torch.norm [Equation]
print (J_)
# Test the two matrices are close.
assert torch.allclose(J, [Table] torch version:
[Table]
B E XPERIMENTAL DETAILS
In this section, we specify the details of the three experiments presented in the Section 3.
The experiments were run on a single NVIDIA RTX 2080 Ti GPU and Intel Xeon Silver 4214 CPU @
2.20GHz.
B. F ULLY C ONNECTED N EURAL N ETWORKS
For the experiment in Section 3.1 we trained four Fully Connected Neural Networks (FCNNs) as
MNIST classiﬁers.
All FCNNs had the same structure which we now explain.
Inspired by (Zhang et al., 2018)
the layers of the FCNNs were parametrized in their Singular Value Decomposition (SVD).
This just means each layer consisted of two orthogonal matrices
U, V ∈ R 784 × 784 and a diagonal matrix Σ ∈ R 784 × 784 , so the forward pass computes y =
U Σ V T x .
The FCNNs had three such fully connected layers with relu non-linearity in between, and a ﬁnal linear layer of shape
W ∈ R 784 × 10 .
We used the Adam optimizer (Kingma & Ba, 2015) with default parameters 2 to minimize cross entropy.
To speed up the network with 784 normal reﬂections, we used the FastH algorithm from (Mathiasen et al., 2020).
For the network with auxiliary reﬂections, we had U, V be auxiliary reﬂections instead of orthogonal matrices.
In all experiments, we initialized the singular values Σ ij ∼ U (0 . 99 , 1 . 01) .
We used orthogonal matrices with reﬂections, the Cayley transform and the matrix exponential as done in (Mhammedi et al., 2017; Casado, 2019; Lezcano-Casado & Mart ´ ınez-Rubio, 2019), respectively.
The orthogonal matrices are constructed using a weight matrix W .
In all cases, we initialized
W ij ∼ U ( − 1 √ d , 1 √ d ) .
It is possible one could initialize W ij in a smarter way, which could change the validation error reported in Figure 3.
That said, we did try to initialize W using the Cayley initialization suggested by (Casado, 2019).
However, we did not ﬁnd it improved performance.
B. R ECURRENT N EURAL N ETWORKS
For the experiment in Section 3.2, we trained three Recurrent Neural Networks as MNIST classiﬁers as done in (Arjovsky et al., 2016; Mhammedi et al., 2017; Zhang et al., 2018; Casado, 2019).
We used the open-source implementation from (Casado, 2019).
They use a clever type of “Cayley initialization” to initialize the transition matrix U .
We found it worked very well, so we choose to initialize both the normal and auxiliary reﬂections so they initially represented the same transition matrix U .
For normal reﬂections, this can be done by computing v 1 , ..., v d so H 1 · · · H d = U by using the QR decomposition.
For the auxiliary reﬂection, this can be done using W =
I − U so
H ( Wx ) x
= Ux (see Theorem 4).
In (Casado, 2019), they use h 0 = 0 as initial state and report “ We choose as the initial vector h
=for simplicity, as we did not observe any empirical improvement when using the initialization given in (Arjovsky et al., 2016).
We sometimes encountered division by zero with auxiliary reﬂections when h 0 = 0 , so we used the initialization suggested by (Arjovsky et al., 2016) in all experiments.
The open-source implementation (Casado, 2019) use RMSProp (Hinton et al., 2012) with different learning rates for the transition matrix and the remaining weights.
This was implemented in PyTorch by using two RMSProp optimizers.
We found training auxiliary reﬂectons to be more stable with
Adam (Kingma & Ba, 2015).
We believe this happens because the “averaged gradients” v become very small due to the normalization term
||
Wx || 2 in H ( Wx ) x
= x −
Wxx T W T
x/ ||
Wx || 2 .
When v becomes small the scaling 1 / ( √ v + (cid:15) ) of RMSProp becomes very large.
We suspect the 1 / ( (cid:112) v/ (1 − β T 2 )+ (cid:15) ) scaling used by Adam ﬁxed the issue, which caused more stable training with
Adam.
This caused us to use Adam optimizer for the transition matrix instead of RMSProp for all the RNNs we trained.
B. N ORMALIZING F LOW
For the experiment in Section 3.3, we trained three Normalizing Flows as generative models on
CIFAR10 as done in (Dinh et al., 2015; 2017; Kingma & Dhariwal, 2018; Ho et al., 2019).
We used an open-source PyTorch implementation of Glow (Kingma & Dhariwal, 2018) 4 with default parameters, except for the number of channels “-C” and the number of steps “-K.” In particular, to decrease training time, we reduced “-C” from 512 to 64 and “-K” from 32 to 8 .
This caused an increase in validation NLL (worse performance) from 3 .
49 to 3 .
66 after 80 epochs.
Auxiliary Reﬂections for 1x1 Convolutions.
(Kingma & Dhariwal, 2018) suggested using invertible 1 × 1 convolutions for Normalizing Flows.
That is, given an input x ∈ R c × h × w and kernel
W ∈ R c × c they compute z : ,i,j =
Wx : ,i,j for all i, j .
The resulting function is invertible if W is, and it has Jacobian determinant hw det( W ) .
It was suggested by (Hoogeboom et al., 2019) to represent
W in its QR decomposition so det( W ) = det( QR ) = det( Q ) det( R ) = det( R ) =
(cid:81) i R ii .
To this end, they represent the orthogonal matrix Q as a product of reﬂections, in particular, they use c = 12 , 24 , 48 reﬂections at different places in the network.
The main goal of this experiment, was to compare c = 12 , 24 , 48 normal reﬂections against a single auxiliary reﬂection, which computes z : ,i,j = H
( Wx : ,i,j ) x : ,i,j instead of z : ,i,j =
Wx : ,i,j .
To isolate the difference in performance due to reﬂections, we further removed the rectangular matrix.
Provably Invertible.
One of the Normalizing Flows with auxiliary reﬂections had the weights of its auxiliary reﬂections constrained to ensure invertibility.
In particular, we let each weight matrix be
W =
I + V V T and used spectral normalization V V T / (2 σ max ( V V T )) to ensure σ i ( V V T ) < 1 / 2 .
The largest singular value can be computed efﬁciently using power iteration (Miyato et al., 2018).
For ease of implementation, we circumvented using power iteration due to a known open issue in the ofﬁcial PyTorch implementation.
We instead used TORCH .
SYMEIG to compute the largest singular value by computing the largest eigenvalue λ max ( V V T ) =
σ max ( V V T ) , which holds because V V T is positive deﬁnite for V ∈ R c × c .
This was only possible because the matrices where at most 48 × 48 , for larger problems one would be forced to use the power iteration.
Initialization.
The open-source implementation of Glow initializes W = Q with Q from
TORCH .
QR ( TORCH .
RANDN (( C , C )))
For the case with normal reﬂections, we computed v 1 , ..., v c such that H ( v 1 ) · · · H ( v c ) =
Q (Wang, 2015).
For the auxiliary reﬂection without constraints we let W
=
I − Q
such that H ( Wx ) x
= H ( x − Qx ) =
Qx by Lemma 5.
However, for the experiment with constraints on W , we could not initiallize W =
I − Q and instead used W
= I + V V T where (initially) V ij ∼ U ( − 1 √ c , 1 √ c ) .
This increased error at initialization from 6 . 5 to 8 .
We suspect this happens because the previous initialization of W has complex eigenvalues which W
= I + V V T does not (because it is symmetric).
In practice, we mitigate the poor initialization by using an additional ﬁxed matrix Q = Q T which does not change during training.
This is essentially the same as using a ﬁxed permutation as done in (Dinh et al., 2017), but, instead of using a ﬁxed permutation, we use a ﬁxed orthogonal matrix.
While using the ﬁxed Q , we found simply initializing V
= I worked sufﬁciently well.