Training Invertible Linear Layers through Rank-One Perturbations

1 INTRODUCTION Many deep learning applications depend critically on the neural network parameters having a certain mathematical structure . As an important example , reversible generative models rely on invertibility and , in the case of normalizing flows , efficient inversion and computation of the Jacobian determinant ( Papamakarios et al. , 2019 ) . Preserving parameter properties during training can be challenging and many approaches are currently in use . The most basic way of incorporating constraints is by network design . Many examples could be listed , like defining convolutional layers to obtain equivariances , constraining network outputs to certain intervals through bounded activation functions , Householder flows ( Tomczak & Welling , 2016 ) to enforce layer-wise orthogonality , or coupling layers ( Dinh et al. , 2014 ; 2016 ) that enforce tractable inversion through their twochannel structure . A second approach concerns the optimizers used for training . Optimization routines have been tailored for example to maintain Lipschitz bounds ( Yoshida & Miyato , 2017 ) or efficiently optimize orthogonal linear layers ( Choromanski et al. , 2020 ) . The present work introduces a novel algorithmic concept for training invertible linear layers and facilitate tractable inversion and determinant computation , see Figure 1 . In lieu of directly changing the network parameters , the optimizer operates on perturbations to these parameters . The actual network parameters are frozen , while a parameterized perturbation ( a rank-one update to the frozen parameters ) serves as a proxy for optimization . Inputs are passed through the perturbed network during training . In regular intervals , the perturbed parameters are merged into the actual network and the perturbation is reset to the identity . This stepwise optimization approach will be referred to as property-preserving parameter perturbation , or P4 update . A similar concept was introduced recently by Lezcano-Casado ( 2019 ) , who used dynamic trivializations for optimization on manifolds . In this work , we use P4 training to optimize invertible linear layers while keeping track of their inverses and determinants using rank-one updates . Previous work ( see Section 2 ) has mostly focused on optimizing orthogonal matrices , which can be trivially inverted and have unity determinant . Only most recently , Gresele et al . ( 2020 ) presented a first method to optimize general invertible matrices implicitly using relative gradients , thereby providing greater flexibility and expressivity . While their scheme implicitly tracks the weight matrices ’ determinants , it does not facilitate cheap inversion . In contrast , the present P4Inv layers are inverted at the cost of roughly three matrix-vector multiplications . P4Inv layers can approximate arbitrary invertible matrices A ∈ GL ( n ) . Interestingly , our stepwise perturbation even allows sign changes in the determinants and recovers the correct inverse after emerging from the ill-conditioned regime . Furthermore , it avoids any explicit computations of inverses or determinants . All operations occurring in optimization steps have complexity of at most O ( n2 ) . To our knowledge , the present method is the first to feature these desirable properties . We show how P4Inv blocks can be utilized in normalizing flows by combining them with nonlinear , bijective activation functions and with coupling layers . The resulting neural networks are validated for density estimation and as deep generative models . Finally , we outline potential applications of P4 training to network properties other than invertibility . 2 BACKGROUND AND RELATED WORK . 2.1 RANK-ONE PERTURBATION . The P4Inv layers are based on rank-one updates , which are defined as transformations A 7→ A + uvT with u , v ∈ Rn . If A ∈ GL ( n ) and 1 + vTA−1u 6= 0 , the updated matrix is also invertible and its inverse can be computed by the Sherman-Morrison formula ( A+ uvT ) −1 = A−1 − 1 1 + vTA−1u A−1uvTA−1 . ( 1 ) Furthermore , the determinant is given by the matrix determinant lemma det ( A+ uvT ) = ( 1 + vTA−1u ) det ( A ) . ( 2 ) Both these equations are widely used in numerical mathematics , since they sidestep the O ( n3 ) cost and poor parallelization of both matrix inversion and determinant computation . The present work leverages these perturbation formulas to keep track of the inverses and determinants of weight matrices during training of invertible neural networks . 2.2 EXISTING APPROACHES FOR TRAINING INVERTIBLE LINEAR LAYERS . Maintaining invertibility of linear layers has been studied in the context of convolution operators ( Kingma & Dhariwal , 2018 ; Karami et al. , 2019 ; Hoogeboom et al. , 2019 ; 2020 ) and using Sylvester ’ s theorem ( Van Den Berg et al. , 2018 ) . Those approaches often involve decompositions that include triangular matrices ( Papamakarios et al . ( 2019 ) ) . While inverting triangular matrices has quadratic computational complexity , it is inherently sequential and thus fairly inefficient on parallel computers ( see Section 4.1 ) . More closely related to our work , Gresele et al . ( 2020 ) introduced a relative gradient optimization scheme for invertible matrices . In contrast to this related work , our method facilitates a cheap inverse pass and allows sign changes in the determinant . On the contrary , their method operates in a higher-dimensional search space , which might speed up the optimization in tasks that do not involve inversion during training . 2.3 NORMALIZING FLOWS . Cheap inversion and determinant computation are specifically important in the context of normalizing flows , see Appendix C. They were introduced in Tabak et al . ( 2010 ) ; Tabak & Turner ( 2013 ) and are commonly used , either in variational inference ( Rezende & Mohamed , 2015 ; Tomczak & Welling , 2016 ; Louizos & Welling , 2017 ; Van Den Berg et al. , 2018 ) or for approximate sampling from distributions given by an energy function ( van den Oord et al. , 2018 ; Müller et al. , 2019 ; Noé et al. , 2019 ; Köhler et al. , 2020 ) . The most important normalizing flow architectures are coupling layers ( Dinh et al. , 2014 ; 2016 ; Kingma & Dhariwal , 2018 ; Müller et al. , 2019 ) , which are a subclass of autoregressive flows ( Germain et al. , 2015 ; Papamakarios et al. , 2017 ; Huang et al. , 2018 ; De Cao et al. , 2019 ) , and ( 2 ) residual flows ( Chen et al. , 2018 ; Zhang et al. , 2018 ; Grathwohl et al. , 2018 ; Behrmann et al. , 2019 ; Chen et al. , 2019 ) . A comprehensive survey can be found in Papamakarios et al . ( 2019 ) . 2.4 OPTIMIZATION UNDER CONSTRAINTS AND DYNAMIC TRIVIALIZATIONS . Constrained matrices can be optimized using Riemannian gradient descent on the manifold ( Absil et al . ( 2009 ) ) . A reparameterization trick for general Lie groups has been introduced in Falorsi et al . ( 2019 ) . For the unitary / orthogonal group there are multiple more specialized approaches , including using the Cayley transform ( Helfrich et al. , 2018 ) , Householder Reflections ( Mhammedi et al. , 2017 ; Meng et al. , 2020 ; Tomczak & Welling , 2016 ) , Givens rotations ( Shalit & Chechik , 2014 ; Pevny et al. , 2020 ) or the exponential map ( Lezcano-Casado & Martı́nez-Rubio , 2019 ; Golinski et al. , 2019 ) . Lezcano-Casado ( 2019 ) introduced the concept of dynamic trivializations . This method performs training on manifolds by combining ideas from Riemannian gradient descent and trivializations ( parameterizations of the manifold via an unconstrained space ) . Dynamic trivializations were derived in the general settings of Riemannian exponential maps and Lie groups . Convergence results were recently proven in follow-up work ( Lezcano-Casado ( 2020 ) ) . P4 training resembles dynamic trivializations in that both perform a number of iteration steps in a fixed basis and infrequently lift the optimization problem to a new basis . In contrast , the rank-one updates do not strictly parameterize GL ( n ) but instead can access all of Rn×n . This introduces the need for numerical stabilization , but enables efficient computation of the inverse and determinant through equation 1 and equation 2 , which is the method ’ s unique and most important aspect . 3 P4 UPDATES : PRESERVING PROPERTIES THROUGH PERTURBATIONS . 3.1 GENERAL CONCEPT . A deep neural network is a parameterized function MA : Rn → Rm with a high-dimensional parameter tensor A . Now , let S define the subset of feasible parameter tensors so that the network satisfies a certain desirable property . In many situations , generating elements of S from scratch is much harder than transforming any A ∈ S into other elements A′ ∈ S , i.e . to move within S. The efficiency of perturbative updates can be leveraged as an incremental approach to retain certain desirable properties of the network parameters during training . Rather than optimizing the parameter tensors directly , we instead use a transformation RB : S → S , which we call a property-preserving parameter perturbation ( P4 ) . A P4 transforms a given parameter tensor A ∈ S into another tensor with the desired property A′ ∈ S. The P4 itself is also parameterized , by a tensor B . We demand that the identity idS : A 7→ A be included in the set of these transformations , i.e . there exists a B0 such that RB0 = idS . During training , the network is evaluated using the perturbed parameters Ã = RB ( A ) . The parameter tensor of the perturbation , B , is trainable via gradient-based stochastic optimizers , while the actual parameters A are frozen . In regular intervals , every N iterations of the optimizer , the optimized parameters of the P4 , B , are merged into A as follows : Anew ← RB ( A ) , ( 3 ) Bnew ← B0 . ( 4 ) This update does not modify the effective ( perturbed ) parameters of the network Ã , since Ãnew = RBnew ( Anew ) = RB0 ( RB ( A ) ) = RB ( A ) = Ã . Hence , this procedure enables a steady , iterative transformation of the effective network parameters and stochastic gradient descent methods can be used for training without major modifications . Furthermore , given a reasonable P4 , the iterative update of A can produce increasingly non-trivial transformations , thereby enabling high expressivity of the resulting neural networks . This concept is summarized in Algorithm 1 . Further extensions to stabilize the merging step will be exemplified in Section 3.3 . Algorithm 1 : P4 Training Input : Model M , training data , loss function J , number of optimization steps Nsteps , merge interval N , perturbation R , optimizer OPT initialize A ∈ S ; initialize B : = B0 ; for i : = 1 . . . Nsteps do X , Y0 : = i-th batch from training data ; Ã : = RB ( A ) ; // perturb parameters Y : = MÃ ( X ) ; // evaluate perturbed model j : = J ( Y , Y0 ) ; // evaluate loss function gradient : = ∂j/∂B ; // backpropagation B : = OPT ( B , gradient ) ; // optimization step if i mod N = 0 then A : = RB ( A ) ; // merging step : update frozen parameters B : = B0 ; // merging step : reset perturbation end end 3.2 P4INV : INVERTIBLE LAYERS VIA RANK-ONE UPDATES Algorithm 2 : P4Inv Merging Step Input : Matrix A , Inverse Ainv , Determinant d det factor : = ( 1 + vTAinvu ) new det : = det factor · d ; if ln |det factor| and ln |new det| are sane then / * update frozen parameters ( equation 3 ) * / d : = new det ; A : = Ru , v ( A ) ; Ainv : = Ainv − 11+vTAinvuAinvuv TAinv ; / * reset perturbation ( equation 4 ) * / u : = 0 ; v : = N ( 0 , In ) ; // random reinitialization end The P4 algorithm can in principle be applied to properties concerning either individual blocks or the whole network . Here we train individual invertible linear layers via rank-one perturbations . Each of these P4Inv layers is an affine transformation Ax+b . In this context , the weight matrix A is handled by the P4 update and the bias b is optimized without perturbation . Without loss of generality , we present the method for layers Ax . We define S as the set of invertible matrices , for which we know the inverse and determinant . Then the rank-one update Ru , v ( A ) = A+ uv T ( 5 ) with B = ( u , v ) ∈ R2n is a P4 on S due to equations 1 and 2 , which also define the inverse pass and determinant computation of the perturbed layer , see Appendix B for details . The perturbation can be reset by setting u , v , or both to zero . In subsequent parameter studies , a favorable training efficiency was obtained by setting u to zero and reinitializing v from Gaussian noise . ( Using a unity standard deviation for the reinitialization ensures that gradient-based updates to u are on the same order of magnitude as updates to a standard linear layer so that learning rates are transferable . ) The inverse matrix Ainv and determinant d are stored in the P4 layer alongside A and updated according to the merging step in Algorithm 2 . Merges are skipped whenever the determinant update signals ill conditioning of the inversion . This is further explained in the following subsection .

Summary: This paper introduces an algorithm for training neural networks in a way that parameters preserve a given property. The optimization is based on using a transformation R that  perturbs parameters in a way that the desired property is preserved. Instead of directly optimizing the parameters of the network, the optimization is carried out on the parameters B of the auxiliary transformation R.