One Reflection Suffice

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 reflections . The only practical drawback is that many reflections cause low GPU utilization . We mitigate this final drawback by proving that one reflection is sufficient , if the reflection is computed by an auxiliary neural network . 1 INTRODUCTION . Orthogonal matrices have shown several benefits 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 reflections ( Mhammedi et al. , 2017 ) . The only practical drawback is low GPU utilization , which happens because the d reflections needs to be evaluated sequentially ( Mathiasen et al. , 2020 ) . Previous work often increases GPU utilization by using k d reflections ( Tomczak & Welling , 2016 ; Mhammedi et al. , 2017 ; Zhang et al. , 2018 ; Berg et al. , 2018 ) . Using fewer reflections limits the orthogonal transformations the reflections 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 sacrificing computation time ? We answer this question with a surprising “ yes. ” The key idea is to use an auxiliary neural network to compute a different reflection for each input . In theory , we prove that one such “ auxiliary reflection ” can represent any number of normal reflections . In practice , we demonstrate that one auxiliary reflection attains similar validation error to models with d normal reflections , 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 reflections train between 2 and 6 times faster for Fully Connected Neural Networks with orthogonal weight matrices ( see Section 3 ) . 1.1 OUR RESULTS . The Householder reflection of x ∈ Rd around v ∈ Rd can be represented by a matrixH ( v ) ∈ Rd×d . H ( v ) x = ( I − 2 vv T ||v||2 ) x . An auxiliary reflection uses a Householder matrix H ( v ) with v = n ( x ) for a neural network n. f ( x ) = H ( n ( x ) ) x = ( I − 2n ( x ) n ( x ) T ||n ( x ) ||2 ) x . One auxiliary reflection can represent any composition of Householder reflections . We prove this claim even when we restrict the neural network n ( x ) to have a single linear layer n ( x ) = Wx for W ∈ Rd×d such that f ( x ) = H ( Wx ) x. Theorem 1 . For any k Householder reflections U = H ( v1 ) · · ·H ( vk ) there exists a neural network n ( x ) = Wx with W ∈ Rd×d such that f ( x ) = H ( Wx ) x = Ux for all x ∈ Rd\ { 0 } . Previous work ( Mhammedi et al. , 2017 ; Zhang et al. , 2018 ) often employ k d reflections and compute Ux as k sequential Householder reflectionsH ( v1 ) · · ·H ( vk ) ·xwith weights V = ( v1 · · · vk ) . It is the evaluation of these sequential Householder reflection 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 sufficient to evaluate a single auxiliary reflection H ( Wx ) x instead of k reflections H ( v1 ) · · ·H ( vk ) · x , thereby gaining high GPU utilization while retaining the full representational power of any number of reflections . In practice , we demonstrate that d reflections can be substituted with a single auxiliary reflection 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 reflections is straightforward for Fully Connected Neural Networks and Recurrent Neural Networks , we needed additional ideas to support auxiliary reflections 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 Rd with d ≥ 2 if W = WT and has eigenvalues which satisfy 3/2 · λmin ( W ) > λmax ( W ) . Finally , we present a matrix formula for the Jacobian of the auxiliary reflection 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 : J = H ( Wx ) A− 2Wxx TW ||Wx||2 where A = I − 2x TWTx ||Wx||2 W. We prove Theorem 1 in Appendix A.1.1 while Theorems 2 and 3 are proved in Section 2 . 2 NORMALIZING FLOWS . 2.1 BACKGROUND . Let z ∼ N ( 0 , 1 ) d and f be an invertible neural network . Then f−1 ( z ) ∼ Pmodel defines a model distribution for which we can compute likelihood of x ∼ Pdata ( Dinh et al. , 2015 ) . log pmodel ( x ) = log pz ( f ( x ) ) + log ∣∣∣∣det ( ∂f ( x ) ∂x ) ∣∣∣∣ ( 1 ) This allows us to train invertible neural network as generative models by maximum likelihood . Previous work demonstrate how to construct invertible neural networks and efficiently compute the log jacobian determinant ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ; Ho et al. , 2019 ) . 2.2 INVERTIBILITY AND JACOBIAN DETERMINANT ( PROOF SKETCH ) . To use auxiliary reflections in Normalizing Flows we need invertibility . That is , for every y ∈ Rd there must exist a unique x ∈ Rd so f ( x ) = H ( Wx ) x = y.1 We find that f is invertible if its Jacobian determinant is non-zero for all x in Sd−1 = { x ∈ Rd | ‖x‖ = 1 } . Theorem 4 . Let f ( x ) = H ( Wx ) x with f ( 0 ) : = 0 , then f is invertible on Rd with d ≥ 2 if the Jacobian determinant of f is non-zero for all x ∈ Sd−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 : −det ( A ) ( 1 + 2 vTA−1u ||u||2 ) where vT = xTW , u = Wx and A = I − 2x TWTx ||Wx||2 W. It is then sufficient that det ( A ) 6= 0 and 1 + 2vTA−1u/||u||2 6= 0 . We prove that this happens if W = WT 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 Rd with d ≥ 2 if W = WT 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 reflections using Newtons method , see Appendix A.2.1 for details . 2.3 PROOFS . 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 first shows f is invertible if it has non-zero Jacobian determinant . Section 2.3.2 then present an expression for the Jacobian determinant , Lemma 1 , and prove the expression is non-zero if W = WT and 3/2 · λmin ( W ) > λmin ( W ) . 2.3.1 NON-ZERO JACOBIAN DETERMINANT IMPLIES INVERTIBILITY . In this section , we prove that f ( x ) = H ( Wx ) x is invertible on Rd if f has non-zero Jacobian determinant . To simplify matters , we first prove that invertibility on Sd−1 implies invertibility on Rd . Informally , invertibility on Sd−1 is sufficient because H ( Wx ) is scale invariant , i.e. , H ( c ·Wx ) = H ( Wx ) for all c 6= 0 . This is formalized by Lemma 2 . Lemma 2 . If f ( x ) = H ( Wx ) x is invertible on Sd−1 it is also invertible on Rd\ { 0 } . Proof . Assume that f ( x ) is invertible on Sd−1 . Pick any y′ ∈ Rd such that ||y′|| = c for any c > 0 . Our goal is to compute x′ such that H ( Wx′ ) x′ = y′ . By normalizing , we see y′/‖y′‖ ∈ Sd−1 . We can then use the inverse f−1 on y′/‖y′‖ to find x such that H ( Wx ) x = y′/‖y‖ . The result is then x′ = x‖y‖ since H ( Wx′ ) x′ = H ( Wx ) x||y|| = y due to scale invariance of H ( Wx ) . 1Note that we do not know H ( Wx ) so we can not trivially compute x = H ( Wx ) −1y = H ( Wx ) y . The main theorem we use to prove invertibiliy on Sd−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 satisfies 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 ) M1 = M2 = Sd−1 . Theorem 5 . ( Krantz & Parks , 2012 , 6.2.8 ) Let M1 and M2 be smooth , connected N -dimensional manifolds and let f : M1 → M2 be continuously differentiable . If ( 1 ) f is proper , ( 2 ) the Jacobian of f is non-zero , and ( 3 ) M2 is simple connected , then f is invertible . For M1 = M2 = Sd−1 the additional conditions are met if f is continuously differentiable . Corollary 1 . Let f : Sd−1 → Sd−1 with d ≥ 2 be continuously differentiable with non-zero Jacobian determinant , then f is invertible . Proof . Note that Sd−1 is smooth and simply connected if d ≥ 2 ( Lee , 2013 ) . Continuously functions on Sd−1 are proper . We conclude f is invertible on Sd−1 by Theorem 5 . We now show that f ( x ) = H ( Wx ) x is continuously differentiable on Sd−1 . Lemma 3 . The function f ( x ) = H ( Wx ) x is continuously differentiable on Sd−1 ifW 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|| = 0 iff x = 0 . But 0 /∈ Sd−1 and we conclude f is continuously differentiable on Sd−1 . Theorem 4 . Let f ( x ) = H ( Wx ) x with f ( 0 ) : = 0 , then f is invertible on Rd with d ≥ 2 if the Jacobian determinant of f is non-zero for all x ∈ Sd−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 Sd−1 if f has non-zero Jacobian determinant on Sd−1 . By Lemma 2 , we get that f is invertible on Rd if it has non-zero Jacobian on Sd−1 .

Summary: The authors present a way to learn the action of an arbitrary orthogonal matrix on a vector via a map from $\mathbb{R}^{n\times n}$ onto $\operatorname{O}(n)$. They show that the map is surjective, and give conditions under which they can invert this action. They then compare against previous proposed schemes in one task and show the performance of their models in other two.