Fast Differentiable Matrix Square Root

1 INTRODUCTION . Consider a positive semi-definite matrix A . The principle square root A 1 2 and the inverse square root A− 1 2 ( often derived by calculating the inverse of A 1 2 ) are mathematically of practical interests , mainly because some desired spectral properties can be obtained by such transformations . An exemplary illustration is given in Fig . 1 ( a ) . As can be seen , the matrix square root can shrink/stretch the feature variances along with the direction of principle components , which is known as an effective spectral normalization for covariance matrices . The inverse square root , on the other hand , can be used to whiten the data , i.e. , make the data has a unit variance in each dimension . Due to the appealing spectral properties , computing the matrix square root or its inverse in a differentiable manner arises in a wide range of computer vision applications , including covariance pooling ( Lin & Maji , 2017 ; Li et al. , 2018 ; Song et al. , 2021 ) , decorrelated batch normalization ( Huang et al. , 2018 ; 2019 ; 2020 ) , and Whitening and Coloring Transform ( WCT ) ( Li et al. , 2017b ; Cho et al. , 2019 ; Choi et al. , 2021 ) . To compute the matrix square root , the standard method is via Singular Value Decomposition ( SVD ) . Given the real Hermitian matrix A , its matrix square root is computed as : A 1 2 = ( UΛUT ) 1 2 = UΛ 1 2UT ( 1 ) where U is the eigenvector matrix , and Λ is the diagonal eigenvalue matrix . As derived by Ionescu et al . ( 2015b ) , the partial derivative of the eigendecomposition is calculated as : ∂l ∂A = U ( KT ( UT ∂l ∂U ) + ( ∂l ∂Λ ) diag ) UT ( 2 ) where l is the loss function , denotes the element-wise product , and ( ) diag represents the operation of setting the off-diagonal entries to zero . Despite the long-studied theories and well-developed algorithms of SVD , there exist two obstacles when integrated into deep learning frameworks . One issue is the back-propagation instability . For the matrix K defined in Eq . ( 2 ) , its off-diagonal entry is Kij=1/ ( λi−λj ) , where λi and λj are involved eigenvalues . When the two eigenvalues are close and small , the gradient is very likely to explode , i.e. , Kij→∞ . This issue has been solved by some methods that use approximation techniques to estimate the gradients ( Wang et al. , 2019 ; 2021 ; Song et al. , 2021 ) . The other problem is the expensive time cost of the forward eigendecomposition . As the SVD is not supported well by GPUs , performing the eigendecomposition on the deep learning platforms is rather time-consuming . Incorporating the SVD with deep models could add extra burdens to the training process . Particularly for batched matrices , modern deep learning frameworks , such as Tensorflow and Pytorch , give limited optimization for the matrix decomposition within the mini-batch . A ( parallel ) for-loop is inevitable for conducting the SVD one matrix by another . However , how to efficiently perform the SVD in the context of deep learning has not been touched . To avoid explicit eigendecomposition , one commonly used alternative is the Newton-Schulz iteration ( NS iteration ) ( Schulz , 1933 ; Higham , 2008 ) . It modifies the ordinary Newton iteration by replacing the matrix inverse but preserves the quadratic convergence . Compared with SVD , the NS iteration is rich in matrix multiplication and more GPU-friendly . Thus , this technique has been widely used to approximate the matrix square root in different applications ( Lin & Maji , 2017 ; Li et al. , 2018 ; Huang et al. , 2019 ) . The forward computation relies on the following coupled iterations : Yk+1 = 1 2 Yk ( 3I −ZkYk ) , Zk+1 = 1 2 ( 3I −ZkYk ) Zk ( 3 ) where Yk and Zk converge to the matrix square root A 1 2 and the inverse square root A− 1 2 , respectively . Since the NS iteration only converges locally , we need to pre-normalize the initial matrix and post-compensate the resultant approximation as : Y0 = 1 ||A||F A , A 1 2 = √ ||A||FYk . ( 4 ) Each forward iteration involves 3 matrix multiplications , which is more efficient than the forward pass of SVD . The backward pass of the NS iteration is calculated as : ∂l ∂Yk = 1 2 ( ∂l ∂Yk+1 ( 3I − YkZk ) −Zk ∂l ∂Zk+1 Zk −ZkYk ∂l ∂Yk+1 ) , ∂l ∂Zk = 1 2 ( ( 3I − YkZk ) ∂l ∂Zk − Yk ∂l ∂Yk+1 Yk − ∂l ∂Zk+1 ZkYk ) ( 5 ) where the backward pass needs 10 matrix multiplications for each iteration . The backward gradients for the post-compensation and pre-normalization steps are computed as : ∂l ∂A ∣∣∣ post = 1 2||A||3/2F tr ( ( ∂l ∂Yk ) TYk ) A , ∂l ∂A ∣∣∣ pre = − 1 ||A||3F tr ( ( ∂l ∂Y0 ) TA ) A+ 1 ||A||F ∂l ∂Y0 + ∂l ∂A ∣∣∣ post . ( 6 ) These two steps take 4 matrix multiplications in total . Consider the fact that the NS iteration often takes 5 iterations to achieve reasonable performances ( Li et al. , 2018 ; Huang et al. , 2019 ) . The backward pass is much more time-costing than the backward algorithm of SVD . As seen from Fig . 1 ( b ) and ( c ) , although the NS iteration outperforms SVD by 123 % in the speed of forward pass , its overall time consumption only improves that of SVD by 17 % . The speed improvement could be larger if a more efficient backward algorithm is developed . To address the drawbacks of SVD and NS iteration , i.e . the low efficiency in either the forward or backward pass , we derive two methods that are efficient in both forward and backward propagation to compute the differentiable matrix square root . In the forward pass , we propose to use Matrix Taylor Polynomial ( MTP ) and Matrix Padé Approximants ( MPA ) to approximate the matrix square root . The former approach is slightly faster but the latter is more numerically accurate ( see Fig . 1 ( c ) ) . Both methods yield considerable speed-up compared with the SVD or the NS iteration in the forward computation . For the backward pass , we consider the gradient function as a Lyapunov equation and propose an iterative solution using the matrix sign function . The backward pass costs fewer matrix multiplications and is more computationally efficient than the NS iteration ( see Fig . 1 ( d ) ) . Through a series of numerical tests , we show that the proposed MTP-Lya and MPA-Lya deliver consistent speed improvement for different batch sizes , matrix dimensions , and some hyperparameters ( e.g. , degrees of power series to match and iteration times ) . Moreover , our proposed MPA-Lya consistently gives a better approximation of the matrix square root than the NS iteration . Besides the numerical tests , experiments on the decorrelated batch normalization and second-order vision transformer also demonstrate that our methods can achieve competitive and even better performances against the SVD and the NS iteration with the least amount of time overhead . Finally , to promote the accessibility of our methods , the Pytorch implementation of all the utilized methods will be released upon acceptance . 2 DIFFERENTIABLE MATRIX SQUARE ROOT IN DEEP LEARNING . In this section , we first recap the previous approaches that compute the differentiable matrix square root and then discuss its usages in some applications of deep learning . 2.1 COMPUTATION METHODS . Ionescu et al . ( 2015b ; a ) first formulate the theory of matrix back-propagation , making it possible to integrate a spectral meta-layer into neural networks . Existing approaches that compute the differentiable matrix square root are mainly based on the SVD or NS iteration . The SVD calculates the accurate matrix square root but suffers from backward instability and expensive time cost , whereas the NS iteration computes the approximate solution but is more GPU-friendly . For the backward algorithm of SVD , several methods have been proposed to resolve this issue of gradient explosion ( Wang et al. , 2019 ; Dang et al. , 2018 ; 2020 ; Wang et al. , 2021 ; Song et al. , 2021 ) . Wang et al . ( 2019 ) propose to apply Power Iteration ( PI ) to approximate the SVD gradient . Recently , Song et al . ( 2021 ) propose to rely on Padé approximants to closely estimate the backward gradient of SVD . To avoid explicit eigendecomposition , Lin & Maji ( 2017 ) propose to substitute SVD with the NS iteration . Following this work , Li et al . ( 2017a ) ; Huang et al . ( 2018 ) adopt the NS iteration in the task of global covariance pooling and decorrelated batch normalization , respectively . For the backward pass of the differentiable matrix square root , Lin & Maji ( 2017 ) also suggest viewing the gradient function as a Lyapunov equation . However , their proposed exact solution is infeasible to compute practically , and the suggested Bartels-Steward algorithm ( Bartels & Stewart , 1972 ) requires explicit eigendecomposition or Schur decomposition , which is again not GPU-friendly . By contrast , our proposed iterative solution using the matrix sign function is more computationally efficient and achieves comparable performances against the Bartels-Steward algorithm ( see the Appendix ) . 2.2 APPLICATIONS . One successful application of the differentiable matrix square root is the Global Covariance Pooling ( GCP ) , which is a meta-layer inserted before the FC layer of deep models to compute the matrix square root of the feature covariance . Equipped with the GCP meta-layers , existing deep mod- els have achieved state-of-the-art performances on both generic and fine-grained visual recognition ( Lin et al. , 2015 ; Li et al. , 2017a ; Lin & Maji , 2017 ; Li et al. , 2018 ; Wang et al. , 2020a ; Song et al. , 2021 ) . More recently , Xie et al . ( 2021 ) integrate the GCP meta-layer into the vision transformer ( Dosovitskiy et al. , 2020 ) to exploit the second-order statistics of the high-level visual tokens , which solves the issue that vision transformers need pre-training on ultra-large-scale datasets . Another line of research proposes to use ZCA whitening , which applies the inverse square root of the covariance to whiten the feature , as an alternative scheme for the standard batch normalization ( Ioffe & Szegedy , 2015 ) . The whitening procedure , a.k.a decorrelated batch normalization , does not only standardize the feature but also eliminates the data correlation . The decorrelated batch normalization can improve both the optimization efficiency and generalization ability of deep neural networks ( Huang et al. , 2018 ; Siarohin et al. , 2018 ; Huang et al. , 2019 ; Pan et al. , 2019 ; Huang et al. , 2020 ; 2021 ; Ermolov et al. , 2021 ) . The WCT ( Li et al. , 2017b ) is also an active research field where the differentiable matrix square root and its inverse are widely used . In general , the WCT performs successively the whitening transform ( using inverse square root ) and the coloring transform ( using matrix square root ) on the multi-scale features to preserve the content of current image but carry the style of another image . During the past few years , the WCT methods have achieved remarkable progress in universal style transfer ( Li et al. , 2017b ; Wang et al. , 2020b ) , domain adaptation ( Abramov et al. , 2020 ; Choi et al. , 2021 ) , and image translation ( Ulyanov et al. , 2017 ; Cho et al. , 2019 ) . Besides the three main applications discussed above , there are still some minor applications , such as semantic segmentation ( Sun et al. , 2021 ) and super resolution ( Dai et al. , 2019 ) .

Summary: The work proposes a fast method to solve the matrix square root. The proposed approach is differentiable and is faster than SVD and NS iteration. Thus, the proposed method is very suitable for optimizing deep neural networks that involves matrix square root computation. Its forward pass uses matrix taylor polynomial or matrix pade approximants. Its backward propagation is done by solving continuous-time Lyapunov equation. Various experiments demonstrate its better performance in both accuracy and speed.