Fast Differentiable Matrix Square Root
Authors: Yue Song
Nicu Sebe
Wei Wang
Venue: ICLR 2022 Poster (2022)
Keywords: Differentiabl Matrix Square Root
Differentiable Matrix Decomposition
Vision Transformers

Abstract: Computing the matrix square root or its inverse in a differentiable manner is important in a variety of computer vision tasks. Previous methods either adopt the Singular Value Decomposition (SVD) to explicitly factorize the matrix or use the Newton-Schulz iteration (NS iteration) to derive the approximate solution. However, both methods are not computationally efficient enough in either the forward pass or in the backward pass. In this paper, we propose two more efficient variants to compute the differentiable matrix square root. For the forward propagation, one method is to use Matrix Taylor Polynomial (MTP), and the other method is to use Matrix Pad\'e Approximants (MPA). The backward gradient is computed by iteratively solving the continuous-time Lyapunov equation using the matrix sign function. Both methods yield considerable speed-up compared with the SVD or the Newton-Schulz iteration. Experimental results on the de-correlated batch normalization and second-order vision transformer demonstrate that our methods can also achieve competitive and even slightly better performances. The code is available at \href{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}{https://github.com/KingJamesSong/FastDifferentiableMatSqrt}.

Published as a conference paper at ICLR 2022
FAST DIFFERENTIABLE MATRIX SQUARE RoOT
Yue Song, Nicu Sebe & Wei Wang
Department of Information Engineering and Computer Science (DISI)
University of Trento Trento, TN 38123, Italy {yue syngg @unitn.it
ABSTRACT
Computing the matrix square root or its inverse in a differentiable manner is im- portant in variety of computer vision tasks. Previous methods either adopt the Singular Value Decomposition (SVD) to explicitly factorize the matrix or use the Newton-Schulz iteration (NS iteration) derive the approximate solution. How- ever, both methods are not computationally efficient enough either the forward pass in the backward pass. In this paper, we propose two more efficient variants to compute the differentiable matrix square root. For the forward propagation, one method to use Matrix Padé Approximants (MPA). The backward gradient is computed by iteratively solving the continuous- -time Lyapunov equation using the matrix sign function. Both methods yield considerable speed-up compared with the SVD or the Sewton-Schulz iteration. Experimental results on the de-correlated batch nor- malization and second-order vision transformer demonstrate that our methods can also achieve competitive and even slightly better performances. The code is avail- able at wttp:///ghhllooooonnignnmmmsoo/s
use Matrix Taylor Polynomial (MTP), and the other method is
INTRODUCTION
Consider a positive semi-definite matrix A. The principle square root A and the inverse square root A (often derived by calculating the inverse of A are mathematically of practical interests, mainly because some desired spectral properties can be obtained by such transformations. An exem- plary illustration is given Fig. (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 appeal- ing 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 Li et al. 2018 Song et al. 2021), decorrelated batch normalization (Huang et al. 2018 2019 and Whitening and Coloring Transform (WCT) (Li et al. 2017b Cho et al. 2019 Choi al.
2017
2020
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:
At ( U
(1)
where U is the eigenvector matrix, and A is the diagonal eigenvalue matrix. As derived by IInescu et al. (2015b), the partial derivative of the eigendecomposition is calculated as:
0l
Öl OA
U K"
dding U
04
0U
where / is the loss function, C denotes the element-wise product, and ()diag represents the opera- tion 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- proppaaaagnt instability. For the matrix K defined in Eq. ( its off-diagonal entry is K,y=1/(,,A,, AAs,, where X and >, are involved eigenvalues. When the eigenvalues are close

---

Published as a conference paper at ICLR 2022
and small, the gradient is very likely to explode, i.e., K2y This issue has been solved by some methods that use approximation techniques to estimate the gradients (Wang et al. 2019 2021 Song et al.] (02I). The other problem is the expensive time cost of the forward eigendecompostiio.. As the SVD is not supported well by GPUs, performing the eigendeconposition on the deep learning platforms is rather time-consuming. Incorporating the SVDD with deep models could add extra bur- dens to the training process. Particularly for batched matrices, modern deep learning frameworks, such Tensorflow and Pytorch, give limited optimization for the matrix decomposition within the mini- A (parallel) for-loop is inevitable for conducting the SVD one matrix by another. How- ever, how to efficiently perform the SVD in the context of deep learning has not been touched.
PCA thehedng Broeeesee Spent
Perwart fror sme Seeee
Backwert Speed sampe Hemdoore
0:3-
Orginal Didbo
Time
Square Root
300
15
15
(c)
(d)
Figure 1: (a) Exemplary visualization of the matrix square root and its inverse. Given the origi- nal data X XeR2x", the matrix square root stretches the data along the axis of small variances and squeezes the data in the direction with large variances, serving as a spectral normalization for co- variance matrices. The inverse square root, on the other hand, can be used to transform the data into the uncorrelated structure, i.e., have the unit variance in each dimension. (b) The comparison of error and speed using the setting of our ZCA whitening experiment. Our MTP and MPA are faster than the SVD and the NS iteration, and our MPA more accurate than the NS iteration. comparison of speed and error in the forward pass (FP). The iteration times of the NS iteration range from to 7, while the degrees of our MTP and MPA vary from 6 to 18. Our MPA computes the more accurate and faster matrix square root than the NS iteration, and our MTP enjoys the fastest calculation speed. (d) The speed comparison in the backward pass (BP). Our Lyapunov solver more efficient than the NS iteration as fewer matrix multiplications involved.
(c) The
To avoid explicit eigendecomposstion, one commonly used alternative is the Newton-Schula iter- ation (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
ct al. 2018 Huang et al. 2019). The forward computation relics on the following coupled iterations:
k++1
Y3 Yk), Z (3I ZzYz)
where Y2 and Za converge to the matrix square root At and the inverse square root A A- respec- tively. Since the NS iteration only converges locally, we need to pre-normalize the initial matrix and post-compensate the resultant approximation as:
1
Yo=
A, A = VI Yt.
(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 calculated as:
Øl 1
il
Øl
oY, 2 0Y (3I
Z%
Z₁
Z₁
Y
0Y
(5)
Ø1 (3I Yn Zx)
Øl
a1
01
Yk
Yv
Ze Y
az₁
0Y
0Zz+1
where the backward pass needs 10 matrix multiplications for each iteration. The backward gradients for the post-compensation and pre- normaizatiinn steps are computed as:
Øl
1
Øl
tr
Y8
A
0A
post
2/1
A|
3/2
OY,
(6)
al
Ø1
Ø1
01
tr
T
A
A
+
+
0A
pre
A
3 F
oY
||A
oY,
0A
post.
2

---

Published as a conference paper at ICLR 2022
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. (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 prop- agation to compute the differentiable matrix square root. In the forward pass, we propose to use Matrix Taylor Polynomial (MTP) and Matrix Padé Approximants (MPA) to approximate the matrix square root. The former approach slightly faster but the latter is more numerically accurate (see Fig. (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 Lya- punov 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. (d). Through a series of numerical tests, we show that the proposed MTP-Lya and MPAA-Lya deliver consistent speed improvement for different batch sizes, matrix dimensions, and some hyper- parameters (e.g., degrees of power series to match and iteration times). Moreover, our proposed MPA-Lyy 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 per- formances 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 spectral meta-layer into neural networks. Existing approaches that compute the differ- entiable 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 explo- sion (Wang et al. 2019 Dang et al.] [018 2020 Wang et al. (2021: [Song et al.] (202T) Wang et al. 2019 propose to apply Power Iteration (PI) to approximate the SVD gradient. Recently, [Snn et al. 1021
propose to rely on Padé 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, [i et al. (2017a); Huang 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 (201)) 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-
3

---

Published as a conference paper at ICLR 2022
els have achieved state-of- totte-of-the-arre the--rt 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
eet al. 2021 More recently. Xie et al. (2021) integrate the GCP meta-layer into the vision trans- former (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 lltra-large-scale datasets.
Another line of research proposes to use ZCA whitening, which applies the inverse square root of the covariance whiten the feature, as an alternative scheme for the standard batch normaliza- tion (Toffe & 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 al. 2018 Huang et al. 2019 Pan al. 2019: Huang al. 2020 2021; Ermolov et al. (021).
The WCT eL al. 2017b) is also an active research field where the differentiable matrix square root and its inverse 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 al. 2020b), domain adaptation (Abramov et al. 2020 Choi et al. 2021), and
image translation [Ulannv 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).
3 FAST DIFFERENTIABLE MATRIX SQUARE ROOT
This section presents the forward propagation and the backward computation our differentiable matrix square root. For the inverse square root, it can be easily derived by computing the inverse of the matrix square root. Moreover, depending on the application, usually the operation of matrix inverse can be avoided by solving the linear system, i.e., computing A=B" 'C' is equivalent to solving BA=C. Thanks to the use of LU factorization, such solution is more numerically stable and usually much faster than the matrix inverse.
3.1 FORWARD PASS
For the forward pass, we derive two variants: the MTP and MPA. The former approach is slightly faster but the latter more accurate. Now we illustrate these two methods in detail.
3.1.1
MATRIX TAYLOR POLYNOMIAL
We begin with motivating the Taylor series for the scalar case. Consider the following power series:
(1 - = (1=
where the notion
denotes the binomial coefficients that involve fractions, and the series con-
verges when 2<1 according to the Cauchy root test. For the matrix case, the power series can be similarly defined by:
--- = 1 Σ
(8)
A=1
where I is the identity matrix. Substituting Z with (I (I-A) leads to:
A - r Σ
2
A)*
Similar with the scalar case, the power series converge only if ||(I - A) <1, where ||-1|p denotes any vector-induced matrix norms. To cireumvent this issue, can first pre-normalize the matrix
4

---

Published as a conference paper at ICLR 2022
Table 1: Comparison of forward operations.
Table 2: Comparison of backward operations.
Op.
MTP
MPA
NS iteration
Op.
Lya
NS iteration
Mat. Mul.
K-1
(K 1121//2
#iters
Mat. Mul.
x #iters
4 + 10 #iters
Mat. Inv. 0
Mat. Inv.
0
0
A by dividing
||
This can guarantee the convergence as || [
П| 1b <1 is always satisfied.
Afterwards, the matrix square root A is post-compensated by multiplying vI| [p. Integrated with these two operations, Eq. (9) can be re-formulated as:
A A F (r-C (--(()-)A)))>))))
A A
y*)
(10)
Truncating the series to a certain degree K yields the MTP approximation for the matrix square root. For the MTP of degree K, K-1 matrix multiplications needed.
3.1.2 MATRIX PADÉ APPROXIMANTS
The MTP enjoys the fast calculation, but it converges uniformly and sometimes suffers from the so-called "hump phenomenon", i.e., the intermediate terms of the series grow quickly but cancel each other in the summation, which results in a large approximation error. Expanding the series to a higher degree does not solve this issue either. The MPA, which adopts two polynomials of smaller degrees to construct rational approximation, is able to avoid this caveat. The coefficients of the MPA polynomials need to be pre-computed by matching to the corresponding Taylor series. Given the power series of scalar in Eq. 7 the coefficients of [M, N] scalar Padé approximant computed by matching to the truncated series of degree M M+N+ N++1:
1 1 Σ m=1 Pm2
M+N Σ
(11)
Ex=1 In
where Pm and In are the coefficients that also apply to the matrix case. This matching gives rise to system of linear equations, and solving these equations directly determines the coefficients. The numerator polynomial and denominator polynomial of MPA are given by:
A
A
Pa =
Σ Pm (I
T QN = I- Σ gn
)M
(12)
|
||||
Then the MPA for approximating the matrix square root is computed as:
A / Q PM.
(13)
Compared with the MTP, the MPA trades off half of the matrix multiplications with matrix inverse, which slightly increases the computational cost but converges more quickly and deliv- better approximation abilities. Moreover,
note that the matrix inverse can be avoided, as Eq. (13) can be more efficiently and numerically stably computed by solving the linear system According to Van Assche (2006). diagonal Padé approximants (i.e., PM and QN have the same degree) usually yleld better approxımation than the non-diagonal ones. There- fore, to match the MPA and MTP of the same degree, we set M=N-K-1 K-1
A [A
Table summarizes the forward computational complexity. As suggested in Li et al. 2018); Huang et al. (2019), the iteration times for NS iteration are often set as 5 such that reasonable performances can achieved. That is. consume the same complexity as the NS iteration does, our MTP and
MPA can match to the power series to degree 16. However, as depicted in Fig. (c), our MPA achieves the better accuracy than NS iteration even degree 8. This observation implies that our MPA is better option in terms of both accuracy and speed than the NS iteration.
3.2 BACKWARD PASS
Though one can manually derive the gradient of the MPA and MTP, their backward algorithms computationally expensive they involve the matrix power up degree K, where K can be
5

---

Published as a conference paper at ICLR 2022
arbitrarily large. Relying the AutoGrad package of deep learning frameworks can be both time and memory- To attain a more efficient backward algorithm, we propose to iteratively solve the gradient equation using the matrix sign function. Given the matrix A and its square root A since we have A A =A a perturbation on leads
A + d dA
(14)
Using the chain rule, the gradient function of the matrix square root satisfies:
al
01
A
(15)
0A
0A
dA+
actually defines the continuous-time Lyapunov equa- tion (BX + X special case of Sylvester equation (BX+X solution is given by:
As pointed out Lin & Maji (017). Eq. (15)
The closed-form
0l A & + I 39 A
vee( 0A
vea
(16)
0A
where vee(-) denotes unrolling matrix to vectors, and @ is the Kronecker product. Although the closed-form solution exists theoretically, it can not be computed in practice due to the huge memory consumption of the Kronecker product. Another approach to solve Eq. (15) is via Bartels-Stewart algorithm (Bartels Stewart, 1972). However, it requires explicit eigendecomposition or Schulz decomposition, which is not GPU-friendly and computationally expensive.
We propose to use the matrix sign function and iteratively solve the Lyapunov equation. Solving the Sylvester equation via matrix sign function has been long studied in the literature of numerical analysis [Roberts| 1980; Kenney & Laub 1995 Benner et al. (0066.. One notable line of research using the family Newton iterations. Consider the following continuous Lyapunov function:
BX+ XB =
(17)
and X denotes the seeking solution 04 where B refers to A in Eq. (15). C represents 8A3'
Eq. 17 can be represented by the following block using a Jordan decomposition:
B
C
I
X
B
0
I X
H
(18)
0
B
0
I
0
- B
0
I
To derive the solution of Lyapunov equation, we need utilize two important properties of matrix sign functions. Specifically, we have:
Lemma 1. For a given matrix H with no eigenvalues on the imaginary axis, its sign function has the following properties: 1) sign(H H I; 2) if H has the Jordan decomposition -=TMT
then its sign function satisfies sign(H) =
Lemma shows that sign( H) is the matrix square root of the identity matrix, which indicates the possibility of using Newton's 's root-finding method to derive the solution (Higham 2008). Here we iteration, the modified inverse-free and multiplication-rich Newton
also adopt the iteration, to iteratively compute sign H)) as:
{
1
3Bs 0
3C₁
B₁
C
B
BBC CCBr
H HI H
-3
0
BB
B²
]
B3I
B
0 3C,
B
(B, C,,
Cr
B3I - B
(19)
The equation above defines two coupled iterations for solving the Lyapunov equation. Since the NS
iteration converges only locally, i.e, converges when || [ here we divide Ho by | |r
to meet the convergence condition. This normalization defines the initialization Bo |
B
and
C Co= B|| and we have the following iterations:
B - BRI B
(20)
C++1 - 2 1- B-CCC + BBC.B Cx BB)

---

Published as a conference paper at ICLR 2022
Relying on Lemma 2 the sign function of Eq.
18
can be also calculated as:
B
X
B
X
-1
sign H) sign
sign
B
0
I
0
B
0
I
I
X
0
I
-X
(21)
0
I
0
-1
0
I
I
2X
0
-I
As indicated above, the iterations in Eq. 20 have the convergence:
B = I, lim Cs = 2X
(22)
10000
k-200
After iterating k times, we can get the resultant approximate solution X X= Ct. Instead of manually choosing the iteration times, one can also the termination criterion by checking the convergence ||B F where is the pre-defined tolerance.
Table compares the backward computation complexity of the iterative Lyapunov solver and NS iteration. Our proposed Lyapunov solver spends fewer matrix multiplications and is thus more efficient than the NS iteration. Even if we iterate the Lyapunov solver more times (e.g., or 8), still costs less time than the backward calculation of NS iteration that iterates 5 times.
4 EXPERIMENTS
In this section, we first perform series of numerical tests compare our proposed methods with the SVD and NS iteration. Subsequently, we validate the effectiveness of our MTP and MPA in two deep learning applications: ZCA whitening and covariance pooling. The ablation studies that illustrate the hyper-parameters are put in the Appendix.
4.1 NUMERICAL. TESTS
To comprehensively evaluate the numerical performance, we compare the speed and error for the input of different batch sizes, matrices in various dimensions, different iteration times of the back- ward pass, and different polynomial degrees of the forward pass. In each of the following tests, the comparison is based on 10, 000 random covariance matrices and the matrix size is consistently 64x64 unless explicitly specified. The error is measured by calculating the Mean Absolute Error (MAE) of the matrix square root computed by the approximate methods (NS iteration, MTP, and MPA) and accurate methods (SVD).
Speed verse Siee
Speed versa Batch
Speed veraa Dinanates
Enogr eurra Meire Dineraice
sark
5.50
NiT-us
0801 225
Datch 2014
Mall Dmwannndddn
(a)
(b)
Figure 2: The results the numerical tests. (a) The computation speed (FP+BP) of each method versus different batch sizes. (b) The speed comparison (FP+BP) of each method versus different matrix dimensions.
(c) The error comparison of each method versus different matrix dimensions. The hyper-parameters follow our experimental setting of ZCA whitening and covariance pooling.
4.1.1 FORWARD ERROR VERSUS SPEED
Both the NS iteration and our methods have a iteration times for NS iteration and polynomial degrees for our MPA and MTP. To validate the
to tune in the forward pass, i.e.,

---

Published as a conference paper at ICLR 2022
impact, we measure the speed and error for different apper-parameters.. The degrees of our MPA and MTP vary from 6 to 18, and the iteration times of NS iteration range from 3 to We give the preliminary analysis in Fig. M (c). As can be observed,
MTP has the least computational time, and our MPA consumes slightly more time than MTP but provides a closer approximation.
Moreover, the curve of our MPA consistently lies below that of the NS iteration, demonstrating our MPA is a better choice in terms of both speed and accuracy.
4.1.2 BACKWARD SPEED VERSUS ITERATION
Fig. (d) compares the speed of our backward Lyapunov solver and the NS iteration versus different iteration times. The result is coherent with the complexity analysis in Table our Lyapunov solver much more efficient than NS iteration. For the NS iteration of times, our Lyapunov solver still has an advantage even when we iterate 8 times.
4.1.3 SPEED VERSUS BATCH SIZE
In certain applications such as covariance pooling and instance whitening, the input could be batched matrices instead of single matrix. To compare the speed performance for batched input, we con- duct another numerical test. The hyper-parameter choices follow our experimental settings in ZCA whitening. As seen in Fig. (a), our MPA-Lya and MTP-Lya are consistently more efficient than the NS iteration and SVD. To give a concrete example, when the batch size is 64, our MPA-LLa is 2.58X faster than NS iteration and 27. 27.25X faster than SVD, while our MTP-Lya 5.82X faster than the NS iteration and 61 .32 faster than SVD.
As discussed before, the current SVD implementation adopts for-loop compute each matrix one by one within the mini-batch. This accounts for why the time consumption of SVD grows almost linearly with the batch size. For the NS iteration, the backward pass is not as batch-friendly as our Lyapunov solver. The gradient calculation defined in Eq. (6) requires measuring the trace and han- dling the multiplication for each matrix in the batch, which haichhhas   be accomplished ineluctably by for-loop" Our backward pass can be more efficiently implemented by batched matrix multiplication.
4.1.4 SPEED AND ERROR VERSUS MATRIX DIMENSION
In the last numerical test, we compare the speed and error for matrices in different dimensions. The hyper-parameter settings also follow our experiments of ZCA whitening. As seen from Fig. (b), our proposed MPA-Lya and MTP-Lya consistently outperform others terms of speed. In particular, when the matrix size is very small (<32), the NS iteration does not hold speed advantage over the SVD. By contrast, our proposed methods still have competitive speed against the SVD. Fig. (c) presents the approximation error. Our MPA-Lya always has better approximation than the NS iteration, whereas our MTP-Lya gives a worse estimation but takes the least time consumption, which can be considered trade-off between speed and accuracy.
4.2 ZCA WHITENING: DECORRELATED BATCH NORMALIZATION
Following Wang et al. (02T). we insert the decorrelated batch normalization layer after the first convolutional layer of ResNet (He et al. 2016). For our forward pass, match the MTP to the power series of degree 11 and set the degree for both numerator and denominator of our MPA as 5. We keep iterating times for our backward Lyapunov solver. The detailed architecture changes and the implementation details of other methods are kindly referred to the Appendix.
Table 3 displays the speed and validation Our MTP-Lya is 1.25X faster than NS iteration and 1.48X faster than SVD-Padé, and our MPA-Lya 1.17X and 1.34X faster. Furthermore, our MPA-Lya achieves state-of-the-an performances across datasets and models. Our MTP-Lya has comparable performances on ResNet- 18 but slightly falls behind on ResNet-50. We guess this mainly because the relatively large approximation error of MTP might affect little on the small model but can hurt the large model.
CIFAR10 and CIFAR100 (Krizhevsky 2009
See the code in the official Pytorch implementation of Li al. 2018 via this link.
8

---

Published as a conference paper at ICLR 2022
Table 3: Validation error of different ZCA whitening methods. The covariance matrix is of size 1x64x44. The time consumption measured for computing the matrix square rool (BP+FP) on workstation equipped with Tesla K40 GPU and 6-core Intel(R) Xeon(R) CPU @ 2.20GHz. For each method, we report the results based on five runs.
ResNet-18
ResNet-50
Methods
Time (ms)
CIFARIO
CIFAR100
CIFARIO0
mean=std
min
mtan-ssid
min
mean-tsid
min
SVD-PI
3.49
4.5910.09
4.44
21.3910:23
21.04
19.9410.44
19.28
SVD-Taylor
3,41
4.50:10.08
4.40
21.14=0.20
20.91
19.81
19.26
SVD-Padé
3.39
4.65=t0.11
4.50
21.41=0.115
21.26
20.25::::3
19.98
NS Iteration
2.96
4.57++0.15
4.37
21.24-00.20
21.01
19.39
19.01
Our MPA-Lya
2.52
4.39:10.09
4.25
21.11=0.12
20.95
19.55+0.20
19.24
Our MTP-Lya
2.36
4.49:
4.31
21.42+0.21
21.24
20.55
20.12
Table 4: Validation top-1/top-5 accuracy of the second-order vision transformer on ImageNet (Deng
al.) 2009). The covariance matrices of size 48x48x448 where 64 is the mini-batch size. The
time cost is measured for computing the matrix square root (BP+FP) on a workstation equipped with
Tesla 4C GPU and 6-core Intel(R) Xeon(R) CPU 2.20GHz. For the So- ViT- model, all the methods achieve similar performances but spend different epochs.
Architecture
Methods
Time (ms)
So-ViT-7
So-ViT-10
So-ViT ViT-14
Pl
1.84
75.93/93.04
77.96 94.18
82.16/ 22.02 (303 epoch)
SVD-PI
83.43
76.55 /93.42
78.53 94.40
82.16 /9.01 (278 epoch)
SVD-Taylor
83,29
76.66 99.52
78.64 / 94.49
82.15 / 99.02 (271 epoch)
SVD-Padé
83.25
76.71 /93.49
78.77 94.51
82.17 99.02 (265 epoch)
NS Iteration
10.38
76.50 / 93.44
78.50 94.44
82.16 / 96.01 (280 epoch)
Our MPA-Lya
3.25
76,84/ 93.46
78.83 94.58
82.17 /96.03 (254 epoch)
Our MTP-Lya
2.39
76.46 / 93.26
78.44 94.33
82.16 82.16/96.02 96.02 (279 epoch)
4.3 COVARIANCE POOLING: SECOND-ORDER VISION TRANSFORMER
Now we turn to the experiment on second-order vision transformer ViT). The implementation of our methods follows our ZCA whitening experiment, and other settings put in the Appendix.
Table compares the speed and performances on three So-Vi ViT architectures with different depths. Our proposed methods predominate the SVD and NS iteration in terms of speed. To be more specific, our MPA-Lya is 3.19X faster than the NS iteration and 25.63X faster than SVD-Padé, and our MTP- Lya is 4.34X faster than the NS iteration and 34.85X faster than SVD-Padé. For the So-ViT-7 and 1i-- our MPA-Lya achieves the best evaluation results and even slightly outperforms the SVD-based methods. Moreover, on the So- ViT-14 model where the performances are saturated, our method converges faster and spends fewer training epochs. The performance of our MTP-Lya is also on par with the other methods.
For the vision transformers, to accelerate the training and avoid gradient explosion, the mixed- precision techniques are often applied and the model weights are in half- calf-precision (i.e., float16). In the task of covariance pooling, the SVD often requires double precision (i.e., float64) to ensure the effective numerical representation of the eigenvalues (Song et al., (021). The SVD methods might not benefit from such low precision, as large round-off errors are likely to be triggered. We expect the performance of SVD-based methods could be improved when using a higher precision. The PI suggested in the ViT only computes the dominant eigenpair but neglects the rest. In spite of the fast speed, the performance is not comparable with other methods.
CONCLUSION
In this paper, we propose two fast methods to compute the differentiable matrix square root. In the forward pass, the MTP and MPA are applied approximate the matrix square root, while an iterative Lyapunov solver is proposed to solve the gradient function for sack-propagation Several numerical tests and computer vision applications demonstrate the effectiveness of our proposed model. In future work, we would like to extend our work to other applications differentiable matrix square root, such neural style transfer and covariance pooling for CNNs.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
This work was supported by EU H2020 SPRING No. 871245 and EU H2020 A14Media No. 951911 projects. We thank Professor Nicholas J. Higham for valuable suggestions.
REFERENCES
Alexey Abramov, Christopher Bayer, and Claudio Heller. Keep it simple: Image statistics matching for domain adaptation. arXiv preprint arXiv:2005.12551, 2020.
George A Baker. Defects and the convergence of padé approximants. Acta Applicandae Mathemat- ica, 61():37755, 2000.
Richard H. Bartels and George W Stewart. Solution of the matrix equation ax+ xb= c [f4]. Commu- nications of the ACM, 15(9):820 826, 1972.
Peter Benner, Enrique Quintana-Orti, and Gregorio Quintana-Ortí. Solving stable sylvester equa- tions via rational iterative schemes. Journal of Scientific Computing, 28(1):51-83, 2006.
Wonwoong Cho, Sungha Choi, David Keetae Park, Inkyu Shin, and Jacgul Choo. Image-to-image to-image translation via group-wise deep whitening-and-colrriin transformation. In CVPR, 2019.
Sungha Choi, Sanghun Jung, Huiwon Yun, Joanne Kim, Seungryong Kim, and Jaegul Choo. Robustnet: Improving domain generalization in urban-scene segmentation via instance selective whitening. In CVPR, 2021.
Tao Dai. Jianrui Cai, Yongbing Zhang, Shu-Tao Xia, and Lei Zhang. Second-order attention network for single image super-resolution. CVPR, 2019.
Z. Dang, K. M. Yi, Y. Hu, F. Wang, P. Fua, and M. Salzmann. Figendecomppostiion-Fccc Training of Deep Networks with Zero Eigenvalue-Band Baseed Losses. In ECCV, 2018.
Z. Dang, K.M. Yi, F. Wang, Y. Hu, P. Fua, and M. Salzmann. Eigendeccopposition Free Training of Deep Networks for Linear Least-Square Problems. TPAMI, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: hierarchical image database. In CVPR, 2009.
large-scale
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.
Aleksandr Ermolov, Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening for self- supervised representation learning. In ICML,, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, 2016.
Nicholas Higham. Functions of matrices: theory and computation. SIAM, 2008
Lei Huang, Dawei Yang, Bo Lang, and Jia Deng. Decorrelated batch normalization. In CVPR, 2018.
Lei Huang, Yi Zhou, Fan Zhu, Li Liu, and Ling Shao. Iterative normalization: Beyond standardiza- tion towards efficient whitening. In CVPR, 2019.
Lei Huang, Lei Zhao, Yi Zhou, Fan Zhu, Li Liu, and Ling Shao. An investigation into the stochas- ticity batch whitening. CVPR, 2020.
Lei Huang, Yi Zhou, Li Liu, Fan Zhu, and Ling Shao. Group whitening: Balancing learning effi- ciency and representational capacity. CVPR, 2021.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML. 2015.
10

---

Published as a conference paper at ICLR 2022
Catalin lonescu, Orestis Vantzos, and Cristian Sminchisescu. Matrix backpropagation for deep networks with structured layers. In ICCV, 2015a.
Catalin lonescu, Orestis Vantzos, and Cristian Sminchisescu. Training deep networks with structured layers by matrix backpropagation. arXiv preprint PXXiv:1509.0766 2015b.
Charles S Kenney and Alan J Laub. The matrix sign function. IEEE transactions on automatic control, 4038):1330-1348 1995.
A Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University of Tront, 2009.
Peihua Li, Jiangtao Xie, Qilong Wang, and Wangmeng Zuo. second-order information helpful for large-scale visual recognition? In ICCV, 2017a.
Peihua Li, Jiangtao Xie, Qilong Wang, and Zilin Gao. Towards faster training of global covariance pooling networks by iterative matrix square root normalization. CVPR, 2018.
Yijun Li, Chen Fang, Jimei Yang, Zhaowen Wang, Xin Lu, and Ming-Hsuan Hsuan Yang. Universal style transfer via feature transforms. In NewrIPS, 2017b.
Tsung- Yu Lin and Subhransu Maji. Improved bilinear pooling with cnns. BMVC, 2017.
Tsung-Y Yu Lin, Aruni RoyChowdhury, and Subhransu Maji. Bilinear enn models for fine-grained visual recognition. In ICCV, 2015.
Xingang Pan, Xiaohang Zhan, Jianping Shi, Xiaoou Tang, and Ping Luo. Switchable whitening for deep representation learning. ICCV, 2019.
John Douglas Roberts. Linear model reduction and solution of the algebraic riccati equation by use of the sign function. International Journal of Control, 32(4))677-687, 1980.
Günther Schulz. Iterative berechung der reziproken matrix. ZAMM-Journal of Applied Mathematics and Mechonics/Zeitsccriff für Angewandte Mathematik und Mechanik, 13(1):57-59, 1933.
Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening and coloring batch transform for gans. In ICLR, 2018
Yue Song, Nicu Sebe, and Wei Wang. Why approximate matrix square root outperforms accurate svd in global covariance pooling? In ICCV, 2021.
Herbert Stahl. Spurious poles in padé approximation. Journal of computational and applied math- ematics, (9(1-)):511-527 1998.
Qiule Sun, Zhimin Zhang, and Peihua Li. Second-order encoding networks for semantic segmenta- tion. Neurocomputing, 2021.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Improved texture networks: Maximizing quality and diversity in feed-forward stylization and texture synthesis. In CVPR, 2017.
Walter Van Assche. Padé and hermite-padé approximation and orthogonality. arXiv preprint math/0609094, 2006.
Qilong Wang, Jiangtao Xie, Wangmeng Zuo, Lei Zhang, and Peihua Li. Deep meet global covariance pooling: Better representation and generalization. TPAMI, 2020a.
Wei Wang, Zheng Dang, Yinlin Hu, Pascal Fua, and Mathicu Salzmann. Backpropagattion-fiee friendlly cigendecomppsstion. In NeurIPS, 2019.
Wei Wang, Zheng Dang, Yinlin Hu, Pascal Fua, and Mathieu Salzmann. Robust differentiable svd. TPAMI, 2021.
Zhizhong Wang, Lei Zhao, Haibo Chen, Lihong Qiu, Qihang Mo. Sihuan Lin, Wei Xing, and Dong- ming Lu. Diversified arbitrary style transfer via deep feature perturbation. In CVPR, 2020b,
Jiangtao Xie, Ruiren Zeng, Qilong Wang. Ziqi Zhou, and Peihua Li. So-vit: Mind visual tokens for vision transformer. arXiv preprint arXiv:2104. 10935, 2021.

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 SUMMARY OF OUR ALGORITHMS
Algorithm. [ and Algorithm. 2 summarize the forward pass (FP) and the backward pass (BP) of our proposed methods, respectively. The hyper-parameter K in Algorithm. D means the degrees of power series, and T in Algorithm. denotes the iteration times.
Algorithm 1: FP of our MTP and MPA.
Algorithm 2: BP of our Lyapunov solver.
Input: and
Input: 81
A and T
Output: A
0A7
if MTP then
Output: 8A
// FP method is MTP
Cot 01 it0;;
A 4 1--I-X0 MK
(I
A AAT
Normalize Bo 1
Co ||Bllp
C₁
k
while i < T do
else
// Coupled iteration
// FP method is MPA
B++1 2 Bx1 B
M+ K-1
2 Nt K-1,
PM I- Σ Pm (I || F
ym;
C++1+ B.C. + B,C.B,++
QN+-+ Σ (I [ Ir
C3I
- Pas:
iii+1 1;
end Post-compensate A || A
end
81
2 Ck:
DA
ResNet
Ours
Vision Transformer
Cass Prodicione
Strite
Stide
DN
ZCA whitering
Gorarience
Sat
Postin
Block
Blook 1; Sindee
Closs tokon
Vaual taker
(a)
(b)
Figure (a) The architecture changes of ResNet models in the experiment of ZCA whitening. Following Wang et al.] (2021), the decorrelated batch normalization layer is inserted after the first convolutional layer. The kernel sizes, the stride of the first convolution layer, and the stride of the first ResNet block changed correspondingly. (b) The scheme of So-ViT (Xie et al., (021). The covariance square root of the visual tokens are computed to assist the classification. In the original vision transformer (Dosovitskiy al., 2020), only the class token is utilized for class predictions.
A.2 ZCA WHITENING
Consider the reshaped feature map X CxBC The ZCA whitening procedure first computes its sample covariance as:
A-(X (XX)( (xx))*+el
(23)
where AERCXC (XX is the mean of X, and e a small constant to make the covariance strictly positive definite. Afterwards, the inverse square root is calculated whiten the feature map:
X Wohitand A-13
(24)
12

---

Published as a conference paper at ICLR 2022
By doing so, the eigenvalues of X are all ones, i.e., the feature is uncorrelated. During the training process, the training statistics are stored for the inference phase. Compared with the ordinary batch normalization that only standardizes the data, the ZCA whitening further eliminates the correlation of the data. The detailed architecture changes of ResNet are displayed in Fig. 3 (a). For the NS iteration and our methods, we first calculate the matrix square root A and compute X whitend by solving the linear system A X whitend==X
A.3 SECOND-ORDER VISION TRANSFORMER
For the task of image recognition, ordinary vision transformer (Dosovitskiy et al., 2020) attaches an empty class token to the sequence of visual tokens and only uses the class token for prediction, which may not exploit the rich semantics embedded in the visual tokens. Instead, So- ViT (Xie et al. 2027 proposes to leverage the high-level visual tokens assist the task of classification:
-FC(c) + FC ((xx²)+)
(25)
where e the output class token, X denotes the visual token, and y the combined class predic- tions. We show the model overview in Fig. 3 (b). Equipped with the covariance pooling layer, the second-order vision transformer removes the need for pre-training on the ultra-large-scale datasets and achieves state-of-the-art performance even when trained from scratch. To reduce the computa- tional budget, the So-ViT further proposes to use Power Iteration (PI) to approximate the dominant eigenvector and eigenvalue of the covariance X XXT..
A.4 BASELINES
In the experiment section, we compare our proposed two methods with the following baselines:
Power Iteration (PI). It suggested in the original So- SiT to compute only the dominant eigenpair.
SVD-PI (Wang et al. 2019) that uses PI to compute the gradients of SVD.
SVD-Taylor (Wang et al.
2021 Song et al. 2021) that applies the Taylor polynomial to
approximate the gradients.
SVD-Padé (Song et al. (021)) that proposes to closely approximate the SVD gradients using Padé approximants. Notice that our MTP/MPA used in the FP fundamentally different from the Taylor polynomial or Padé approximants used in the BP of SVD-Padé. For our method, we use Matrix Taylor Polynomial (MTP) and Matrix Padé Approximants (MPA) to derive the matrix square root in the FP. For the polynomial and scalar Padé approximants to approximate the gradient x, the BP. That is to say, their aim is to use the technique to compute the gradient and this will not involve the back-propagation of Taylor polynomial or Padé approximants.
they use scalar Taylor
NS iteration (Schulz 1933 Higham 2008) that uses the Newton-Schulle iteration to com- pute the matrix square root. has been widely applied in different tasks, including co- variance pooling (Li et al. 2018) and ZCA whitening (Huang et al. 2018). We note that although Huang et al. (2019) and Higham (2008) use different forms of NS iteration, the two representations are equivalent each other. Huang al. (2019) just replace Y,, with Z A and re-formulate the iteration using one variable. The computation complexity is still the same.
As the ordinary differentiable SVD suffers from the gradient explosion issue and easily causes the program fail, we not take it into account for comparison.
Unlike previous methods such as SVD and NS iteration, our MPA-Lya/MTP--Lya does not have consistent FP and BP algorithm. However, we do not think will bring any caveat to the stability or performance. Our ablation study in the Appendix A.7.2 implies that our BP Lyapunov solver approximates the real gradient very well (i.e., B <e - and ||0.5C our experiments on ZCA whitening and vision transformer demonstrate superior performances. In light these experimental results, we argue that as long as the BP algorithm is accurate enough, the inconsistency between the BP and FP is not an issue.
<><76--). Also,
13

---

Published as a conference paper at ICLR 2022
A.5 IMPLEMENTATION DETAILS
All the source codes are implemented in Pytorch. For the SVD methods, the forward eigendecom- position is performed using the official Pytorch function TORCH.SVD. which calls the LAPACK's routine gesdd that uses the Divide-and- and---uend- algorithm for the fast calculation.
All the numerical tests are conducted on single workstation equipped with a Tesla K40 GPU and 6-core Intel(R) Xeon(R) GPU @ 2.20GHz.
A.5.1 TRAINING SETTINGS OF ZCA WHITENING
Suggested by Wang et al. (2020a), we truncate the Taylor polynomial to degree 20 for SVD-Taylor. To make Padé approximant match the same degree with Taylor polynomial, we set the degree of both numerator and denominator to 10 for SVD-Padé. For SVD-PI, the iteration times are also set as 20. For the NS iteration, according to the setting in Li et al. (2018); [Huang et al. (2018), we set the iteration times to 5. The other experimental settings follow the implementation in Wang et al. 2021.. We the workstation equipped with Tesla K40 GPU and 6-core Intel(R) Xeon(R) GPU @ 2.20GHz for training.
A.5.2 TRAINING SETTINGS OF COVARIANCE POOLING
We use 8 Tesla G40 GPUs for distributed training and the NVIDIA Apex mixed-precision trainer used. Except that the spectral layer uses the single-precision (i.e. float32), other layers use the half-precision (i.e., f) to accelerate the training. Other implementation details follow the ex- perimental setting of the original So- ViT (Xie et al. 2021).
Following the experiment of covariance pooling for CNNs (Song et al. 2021), the degrees of Taylor polynomial are truncated to 100 for SVD-Taylor, and the degree of both the numerator and denom- inator of Padé approximants 100. In the experiment of covariance pooling, more terms of the Taylor series are used because the covariance pooling meta-layer requires more accurate gradient estimation (Song al. 202I).
set to 50 for SVD-Padé. The iteration times of SVD-PI
set to
For the SVD-based methods, usually the double-precision is required to ensure an effective numeri- cal representation of the eigenvalues. Using lower precision would make the model fail to converge at the begimming of the training (Song et al. [202]) To resolve this issue. we first apply the NS itera- tion train the network for epochs, then switch to the corresponding SVD method and continue the training till the end. This hybrid approach can avoid the non-convergence of the SVD methods the beginning of the training phase.
A.6 EXTENSION TO INVERSE SQUARE ROOT
Although we use the LU factorization to derive the inverse square root in the paper for comparison fairness, our proposed method naturally extend the inverse square root. As the matrix square root At of our MPA calculated as VI|AA|| QN PM, the inverse square root can be directly computed as VIAAAE Par QN.
A.7 ABLATION STUDIES
conduct three ablation studies to illustrate the impact degree of power series the forward pass, the termination criterion during the pacc-propagttion, and the possibility of combining our Lyapunov solver with the SVD and the NS iteration.
A.7.1 DEGREE OF POWER SERIES TO MATCH FOR FORWARD PASS
Table 5 displays the performance of our MPA-Lya for different degrees of power series. As use more terms of the power series, the approximation error gets smaller and the performance gets steady improvements from the degree [3. 3| to [5, 5]. When the degree of our MPA is increased from [5. 5] to [6, 6). there are only marginal improvements. We hence set the forward degrees as [5, 5] for MPA and as 11 for MTP trade-off between speed and accuracy.
14

---

Published as a conference paper at ICLR 2022
Table 5: Performance of MPA-Lya versus different degrees of power series to match.
ResNet-18
ResNet-50
Degrees
Time (ms)
CIFARTO
CIFARTO0
CIFARTO0
mcantstd
min
mcan
min
mcantstd
min
3, 3
0.80
0.64
4.54
2..85+0..18
21.20
20.14:10.43
19.56
4, 4
0.86
0.55
4.51
21.26=0.222
21.03
19.87=00229
19.64
6, 0.98
4.45:10.07
4.33
21.09
21.04
19.51:10.24
19.26
5,
0.93
4.39:0009
4.25
21.11=0.11
20.95
19.55=0.20
19.24
A.7.2 TERMINATION CRITERION FOR BACKWARD PASS
comppress the performance of backward algorithms with different termination criteria as well as the exact solution computed by the hartels-Stewary algorithm (BS algorithm) (Bartels & Stewart 1972). Since the NS iteration has the property of quadratic convergence, the errors B -[l and [0.5C& X|IF decrease at a larger rate for more iteration times. When we iterate more than times, the error becomes sufficiently neglectable, i.e., the NS iteration almost converges. Moreover, from 8 iterations to iterations, there are no obvious performance improvements. We thus terminate the iterations after iterating 8 times.
The exact gradient calculated by the BS algorithm does not yield the best results. Instead, it only achieves the least fluctuation ResNet-50 and other results inferior to our iterative solver. This because the formulation of our Lyapunov equation is based on the assumption that the accurate matrix square root is computed, but in practice only compute the approximate pass. In this case, calculating the accurate gradient of the approximate matrix square root might not necessarily work better than the approximate gradient of the approximate matrix square root.
in the forward
Table 6: Performance of our MPA-Lya versus different iteration times. The residual errors and 0.5C&
|| B 1
|||F measured based 10, 000 randomly sampled covariance matrices.
ResNet-18
ResNet-50
Methods
Time (ms)
| Flle
[00.CC X|l
CIFARIO
CIFARTOU
CIFAR TOO
mean-Fstd
mm
mean
min
mean-Fstd
min
BS algorithm
4.3
4.45
21.20:10:23
21.01
19.60=0.16
19,55
#iter 5
1.05
==0 35-11
==0 290.2019
4.48:00.13
4.31
21.15+0.24
20.84
20.03=0.19
19.78
#iter 6
1.23
=4000410
e00.0231
4.33-0.100
4.28
21.16:40.19
20.93
19.83+0.24
19.57
fiter 7
1.43
R=3.5c
4.45-00111
4.29
21.18:10.20
20.95
19.69=0.20
19.38
#iter 9
1.73
=:20 7
ss 7
4.40-00.07
4.28
21.08-+0.15
20.89
19.52-0.22
19.25
Riter 8
1.59
203e-1 7
767e-6
4.39:1.09
4.25
21.11:10.12
20.95
19.55 0.20
19.24
A.7.3 ITERATIVE LYAPUNOV SOLVER AS A GENERAL BACKWARD ALGORITIM
Notice that our proposed iterative Lyapunov solver is general backward algorithm for computing the matrix square root. That is to say, should be also compatible with the SVD and NS iteration as the forward pass. Table 7 compares the performance of different methods that use the Lyapunov solver as the backward algorithm. As can be seen, the SVD-Lya can achieve competitive perfor- mances compared with other differentiable SVD methods. However, the combination of Lyapunov solver with the NS iteration, i.e., the NS-Lya, cannot converge on any datasets. Although the NS iteration is applied different. For the FP algorithm, the NS iteration is two coupled iterations that use two variables Y2 and Z to compute the matrix square root. For the BP algorithm, the NS iteration defined to com- pute the matrix sign and only uses the variable Yt. The term Z is not involved in the BP and we have no control over the gradient propaapppggging through it. We conjecture this might introduce some instabilities to the training process. Despite the analysis, developing an effective remedy is direction of our future work.
both the FP and BP of the NS-Lya, the implementation and the usage are
A.8 COMPUTATION ANALYSIS ON MATRIX INVERSE VERSUS SOLVING LINEAR SYSTEM
Suppose we want to compute C-A-IB There are two options available. One is first computing the inverse of A and then calculating the matrix product A-1B. The other option is solve the
15

---

Published as a conference paper at ICLR 2022
Table 7: Performance comparison of SVD-Lya and NS-L.a.
ResNet-18
ResNet-50
Methods
Time (ms)
CIFARTO
CIFAR100
CIFARTOO
mean:istd
min
mtan=std
min
mean=std
min
SVD-Lya 4.47
4.4
4.20
21.24=0.24
21.02
19.41=0.11
19.26
SVD-PI 3,49
4.59::0.09
4.44
21.39=0.23
21.04
19.94=0.44
19.28
SVD- SVD-Taylor
3,41
4.50
4.40
21.14=0.20
20.91
19.81=00244
19.26
SVD-Padé
3.39
4.65++0.11
4.50
21.41=00155
21.26
20.25+:0.20
19.98
NS-Lya
2.87
NS Iteration
2.96
4.57:0.15
4.37
21.2411.20
21.01
19.3910.30
19.01
MPA-Lya
2.52
4.39.10.09
4.25
21.11:0.112
20.95
19.55:10.20
19.24
MTP-Lya 2.36
4.49:10.3
4.31
21.4240.21
21.24
20.55:10.77
20.12
linear system AC=B. This process involves first performing the GPU-friendly LU decomposition decompose and then conducting two substitutions obtain C.
Solving the linear system is often preferred over the matrix inverse for accuracy, stability, and speed reasons. When the matrix is ill-conditioned, the matrix inverse is very unstable because the eigen- value X3 would become X. after the inverse, which might introduce instability when ); is very small. For the LU-based linear system, the solution is still stable for ill coldciiitioned matrices. As for the ac- curacy aspect, according to Higham 2008), the errors of the two computation methods are bounded by:
B ACimo S a A
(26)
B ACw S a C
We roughly have the relation |A] |L||U]. So the difference is only related to |A- and |Crurl. When ill-conditioned, |A-1| could be very large and the error of B AC... will be significantly larger than the error of linear system. About the speed, the matrix inverse option takes about 14 flops, while the linear system consumes about 3 + 2n flops. Consider the fact that LU is much more GPU-friendly than the matrix inverse. The actual speed of the linear system could even be faster than the theoretical analysis.
For the above reasons, solving linear system is a better option than matrix inverse when need to compute
A-1 41B.
A.9 SPEED COMPARISON ON LARGER MATRICES
In the paper, we only compare the speed of each method till the matrix dimension 128. Here we add the comparison larger matrices till the dimension 1024. We choose the maximal dimension as 1024 because it should be large enough to cover the applications in deep learning. Tuble&ccompe the speed of our methods against the NS iteration.
Table 8: Time consumption (ms) for a single matrix whose dimension is larger than 128. The measurement based on 10, 000 randomly sampled matrices.
Matrix Dimension xx5 x x
x512x512|
1022
MTP-Lya 4.88
5,43
7.32
MPA-Lya 5.56
7.47
11.28
NS iteration 6.28
9.86
12.69
As can be observed, our MPA-Lya is consistently faster than the NS iteration. When the matrix dimension increases, the computation budget of solving the linear system for deriving our MPA increases, but the cost of matrix multiplication also increases. Consider the huge amount of matrix multiplications of NS iteration. The computational cost of the NS iteration grows at a larger speed compared with MPA-Lya.
16

---

Published as a conference paper at ICLR 2022
Table 9: Comparison of validation accuracy ImageNet (Deng al. 2009) and ResNet-50 (He
et al. 2016). The time consumption is measured for computing the matrix square root (FP+BP). We follow Song et al. (202T) for all the experimental settings.
Methods/Covariaan Size (BxCxC) Time (ms) top-1 acc (%) top-5 acc (%)
(ms)
MPAA-La
110.61
77.13
93.45
256x256x256
NS iteration
164.43
77.19
93.40
A. EXTRA EXPERIMENT ON GLOBAL COVARIANCE POOLING FOR CNNS
To evaluate the performance of our MPA-Lya on batched larger matrices, we conduct another ex- periment the task of Global Covariance Pooling (GCP) for CNNs (Li et al. 2017a 2018, Song et al. 2021). In the GCP model, the matrix square root of the covariance of the final convolutional feature is fed into the FC layer for exploiting the second-order statisties. Table ! presents the speed and performance comparison. As can be seen, the performance of our MPA-Lya is very competitive against the NS iteration, but our MPA-Lya is about 1.5X faster.
A.11 DEFECT OF PADÉ APPROXIMANTS
When there is the presence of spurious poles Stahl 1998 Baker 2000), the Padé approximants are very likely to suffer from the well-known defects of instability. The spurious poles mean that when the approximated function has very close poles and zeros, the corresponding Padé approximants will also have close poles and zeros. Consequently, the Padé approximants will become very unstable in the region of defects (i.e., when the input is in the neighhorhood of poles and zeros). Generalized to the matrix case, the spurious poles can happen when the determinant of the matrix denominator is zero (i.e. det (QN) 0).
However, in our case, the approximated function is (1 for |2| < 1, which only has one zero at 1 and does not have any poles. Therefore, the spurious pole does not exist in our approximation and there are no defects of our Padé approximants.
We can also prove this claim for the matrix case. Consider the denominator of our Padé approxi- mants:
N
QN=
- I-
Σ
In
(27)
n=1
Its determinant is calculated as:
N Σ n=1
A Σ x
det (QN)
II
(1
(28)
(=1
The coefficients On of our 55.5] Padé
approximant are pre-computed
as
[2.25, 1.75,, 0.54675, -0. -0.0500000 0.058559775, 0.0009765625).
Let 2; denotes (1
A
Then
E.
I; is in the range [0, 1]. and we have:
f(x)) = 1 - 2.25x, + 1.75x² 0.34675xz + 0.0585937522 0.0009763625x2
det (QN) = [I
(29)
The polynomial f(x;) does not have any zero in the range of x= 1). The minimal 0. 0.0108672 when r = 1. This implies that det (Qx) x A 0 always holds for any Q N and our Padé approximants do not have any pole. Accordingly, there will no spurious poles and defects. Hence, our MPA deemed stable. Throughout our experiments, we do not encounter any instability issue of our MPA.
Fig. depicts the function (1 2) and the corresponding approximation of Taylor polynomial and Padé approximants. As can seen, our approximating techniques do not suffer from any spurious poles or defeet regions, and the stability guaranteed. Moreover, when 2 is close to 1, the Taylor polynomial gets a relatively large approximation error but the Padé approximants still give a reasonable approximation. This confirms the superiority of our MPA.
17

---

Published as a conference paper at ICLR 2022
Function f(x) (1 -2) and Its Approximation
Original Function
0.9
Pade Approximants
Taylor Polynomial
0.8
0.7
0.6
0.4
0.3
0.2
0.1
0.2
0.4
0.6
0.8
Figure 4: The function (1 - 2) à in the range of | < 1 and its approximation including Taylor polynomial and Padé approximants. The two approximation techniques do not have any spurious poles or defect regions.
A.12 PADÉ COEFFICIENTS OTHER DEGREES THEIR STABILITY
The property of stability also generalizes to diagonal Padé approximants of other degrees there are no poles in the original function (1 2) for |a| 1. To better illustrate this point, here we attach the Padé coefficients from the degree [3, 3] to the degree [6, 6]. The numerator pm is:
P3 [1.75, 0.875, 0.109375).
P4 2.25 - 1.6875, 0.46875, 00.35515255..
(30)
P6 [2.75, 1..5,.22000 1181:122888155200222244257000100 0.11949977..00117718877777 [3.25.
4.0625, 2.4375, -0.7100.7..5 0.0888671875, -0.05,,000.000 019886771877,0000001782222111113
And the corresponding denominator is:
= [1.25, 0.335,
91 [1.75, -.9370 0.15625, -0.00990060
(31)
- 1.75, -15-5:55005500000000055100000000
96 (2.75, 22.125, 1.3125, -.27734375 0.20557715,. -.000444404444466221,
Similar to the deduction in Eqs.
29
can get the polynomial for deriving det (Qx)
f(x:) = 1.25r, 0
0.015625z
= [
0.003906252;
f(zx) = - 1.76, -1 0.000977552200
(32)
It can be easily verified that all of the polynomials monotonically decrease in the range of [0,
and have their minimal at =;=1 as:
min f(x:)gs 0.109375. 2.E|0.1]
min [xx4)s = 0.03515625. x;E(0.1]
(33)
min f(zz)g = 0.0108672.
x,E(0.1]
min f(x:)a0 = 0.003173828129
18

---

Published as a conference paper at ICLR 2022
As indicated above, all the polynomials are positive and we have det (QN) + 0 consistently holds for the degree 3 4, 5, 6.
19


Reviews:
[{"review_id": "nSYUq_Hu83", "reviewer": "Reviewer_J9af", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1.\tThis paper starts the point of matrix multiplication times in computation of matrix square root, and proposes some strategies to reduce number of matrix multiplications both forward and backward propagations, aiming to improve the efficiency. The motivation is clear and reasonable. \n2.\tThe idea of solving the Lyapunov equation based on matrix sign function and Newton-Schulz iteration seems interesting, which avoids huge memory consumption of the Kronecker product in the original closed-form solution and requirement of Schulz decomposition in Bartels-Stewart algorithm.\n3.\tThe experiments are conducted on two real-world applications, including decorrelated batch normalization and second-order vision transformer. The speed-up over Newton Schultz iteration is ok.\n4.\tThe paper is well-written and easy to read.\n\nWeaknesses:\n1.\tIn this paper, Matrix Taylor Polynomial or Matrix Pade Approximation are used for forward propagations, while approximate Lyapunov equation is used for backward propagation. Although Table 1 and Table 2 show MTP/MPA and Lya require less matrix multiplication than NS iteration, which one is most important for fast matrix square root? To verify it, MTP/MPA+NS based BP and NS baed FP + Lya are suggested to be compared in terms of accuracy and running time. \n2.\tFor previous SVD-based and NS-based methods, computation processes for forward and backward propagations are consistent. However, this work adopts Matrix Taylor Polynomial or Matrix Pade Approximation for forward propagations and uses approximate Lyapunov equation for backward propagation, leading variance in forward and backward propagations. The authors would better make some discussions about this issue.\n3.\tThe authors claimed BP of MPA is both time and memory-consuming. [r1] tries to respectively use SVD and MTP/MPA as forward and backward propagations, where the authors show BP of MPA is efficient (as shown in Table 6). The authors would better make some discussions about it.   \n\n[r1] Why Approximate Matrix Square Root Outperforms Accurate SVD in Global Covariance Pooling? ICCV, 2021.\n\nOther comments:\n1.\tIt is clear that MPA involves matrix inverse, which is very GPU-unfriendly.  As stated in the paper: \"Moreover, we note that the matrix inverse can be avoided, as Eq. (13) can be more efficiently and numerically stably computed by solving the linear system\". The authors would better provide more detailed computation and analysis.\n2.\tHow to compute the coefficients of $p_{m}$ and $q_{n}$ for the Matrix Pade Approximation in equation (12)? Do forward operations of MPA in Table 1 contain computation of coefficients $p_{m}$ and $q_{n}$?\n3.\tDoes Equation (2) lack a (·)_{sym} operation for $/frac{l}{U}$? \n4.\tIs equation (11) missing a term z^{k} in the left side?\n5.\tI am not sure why sign(B) in equation (21) can be calculated as identity matrix?\n6.\tP_{M} and Q_{N} are used to approximate the Taylor series. If I am not misunderstanding, does I- Q_{N}^{-1}P_{M} replace Q_{N}^{-1}P_{M} in Eqn. (13)? and do the terms I-X replace X in Eqn. (12)?\n", "summary_of_the_paper": "This paper aims to solve an important problem: how to efficiently compute matrix square root in both forward and backward propagations. To this end, the authors proposed to exploit Matrix Taylor Polynomial and Matrix Pade Approximation to compute matrix square root in forward propagations, while approximate Lyapunov equation based on Newton Schultz iteration is developed for backward propagation. From the perspective of computation complexity, the proposed method requires less matrix multiplications than one based on Newton Schultz iteration. Experimental results on the decorrelated batch normalization and second-order vision transformer show the proposed methods are faster than existing methods, including SVD based ones and NS based ones. ", "summary_of_the_review": "This paper proposes some strategies to reduce number of matrix multiplications both forward and backward propagations, aiming to improve the efficiency of computation of matrix square root. Particularly, the idea of approximate Lyapunov equation based on matrix sign function and Newton-Schulz iteration seems interesting. The more experimental comparisons and discussions could be strengthened. Overall, I like this work.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "f6Lb0UMDd2x", "reviewer": "Reviewer_JTDk", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n\n+ The proposed method is differentiable and can be applied to a variety of deep learning tasks.\n\n+ MTA and MPA tradeoff between speed and accuracy, which provide users flexible options. \n\n+ The writing is clear and easy to follow.\n\n\nWeaknesses and questions:\n\n- The novelty of the proposed method is limited and the theory behind it is well-known.\n\n- The speedup by the proposed method is not very significant compared to NS iteration for large matrix dimension, see Figure 2(b). \nMPA requires either matrix inverse or solving a linear system of equations, which should be more computationally expensive than matrix multiplication. This may explain that when increasing matrix dimension, in Figure 2(b) the runtime of MPA becomes closer to NS iteration. My question is that when further increasing the matrix dimension (>200), will MPA-Lya be slower than NS?\n\n- Most of the experiments are on matrices of small dimensions. I wonder what the maximal dimension the proposed method may handle (probably comparable to that of NS)? Of course it should also depend on the space complexity.\n", "summary_of_the_paper": "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.\n", "summary_of_the_review": "In summary, I believe that the proposed method can be useful for a variety of deep learning tasks, although clarification of scalability of the proposed method on large matrices is expected.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "_8T2_jDxhA_", "reviewer": "Reviewer_Jnyr", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I. Theoretical contributions\n\nI.I. Forward pass: MTP and MPA approximation approaches\n\nIn terms of math, the formulae are mostly correct except for minor issues. In eq (7), the notion $(1/2 \\quad k)^T$ should be further explained. In eq (11), $z^m$ is missing in the numerator and $z^n$ is missing in the denominator.\n\nThe approximation approaches were motivated by the slowness of SVD in GPU computing. We would like iterative approximations for speed. This is fair. However, while Jacobi-based SVD approaches abandon parallelism and a bit of speed to preserve stability, and orthogonality is crucial in it, the proposed approximation approaches, in my view, seem to trade stability for speed. \n\nThe proposed MTP based on Taylor expansion is fast. The authors thoroughly explained a stability issue MTP, which I find fair and honest. The explanation is used to motivate the invention of the second approach, MPA. Although it looks like MPA has addressed the stability issue in MTP with a bit of speed compromise, I think the authors have overlooked another stability issue with Pade approximants. Crucial to any Pade approximation approach is the idea that a function is approximated by the ratio of two low-degree polynomials. The problem though is that Pade approximants in scalar forms can have defects, i.e. small regions where the numerator and the denominator approaches zero. Defects can make the approximation unstable very quickly. See for example:\n\n[1] Baker, Defects and the Convergence of Padé Approximants, In Mathematics, 2000.\n\nOne can generalise the issue when scalar serieses are replaced by matrix serieses. In eq (13), it looks like solving a linear system is key to MPA. However, what if $det(Q_N) \\rightarrow 0$ and $det(P_N) \\rightarrow 0$ in some region in the matrix space? This case would be like the matrix version of scalar defects.\n\nI do not think the provided numeric tests are sufficient to spot the defects, because the defected regions can be small. Unfortunately, there is no treatment of Pade approximant defects in the paper. \n\nI.II. Backward pass: viewing matrix square root gradient as a continuous Lyapunov equation and using the matrix sign function to solve\n\nThere is some novelty in this section but I think the novelty is not much. The idea of viewing the matrix square root gradient as a Sylvester equation is interesting but it has been proposed before (i.e. Lin-Maji (2017) pointed out in the paper). There, Lin-Maji did not use any iterative approach but instead chose to solve directly the problem using the Bartels-Stewartz (1972) algorithm. Here, the authors proposed to solve the Lyapunov equation using the matrix sign function and Newton-Schulz iteration. However, if one looks at Benner et al (2006) section 2.3 then it is clear that the proposed solution here is the same as Benner et al (2006) section 2.3 with matrices A, B, C in the Sylvester equation being replaced by $A^{1/2}$, $A^{1/2}$ and $\\frac{\\partial l}{\\partial A^{1/2}}$.\n\nII. Literature Review\n\nThe review of relevant work is sufficient. However, grouping together the work of Huang et al (2019) with those of Lin-Maji (2017) and Li (2018) is, in my view, not appropriate. While the latter two papers use the Newton-Schulz iteration explained in this paper, which is Higham (2008)'s approach, Huang et al (2019) has a different kind of Newton-Schulz iteration, one which deals with the matrix square root inverse rather than the matrix square root, a direct application of [2] below. Their NS iteration is more closely related to Benner et al (2006) and the one proposed in the backward pass in here. I think they deserve a separate treatment.\n\n[2] Dario A. Bini, Nicholas J. Higham, and Beatrice Meini. Algorithms for the matrix pth root. Numerical Algorithms, 39(4):349–378, Aug 2005.\n\nIII. Experiments\n\nThe numeric tests in section 4.1 are sufficient for speed comparisons but are a bit lacking for accuracy comparisons. Fig 2c is the only test related to accuracy on general settings. While it reveals that on average, MPA is better than MTP and Higham (2008), there is no treatment about special cases where MTP or MPA may potentially fail. For example, a figure revealing when MTP fails can support the claimed \"hump phenomenon\" in sec 3.1.2, making a stronger point. It would be good to have some experiments on potential defects of MPA. It would be good to have some accuracy test for the backward pass as well, which is currently lacking.\n\nSection 4.3 is adequate, showing the proposed approaches inbetween two extrema, one with full SVD (i.e. SVD-Taylor) and the other with only the principle eigenpairs (PI). However, in section 4.2 I think a better candidate to compare against is not Higham (2008) but rather Bini et al (2000) ([2] above) that Huang et al (2019) have used. They directly targeted $A^{-1/2}$ which is appropriate for ZCA whitening. If there were space, I would recommend treating the inverse matrix square root problem as a separate, but related problem. However, if the scope of the paper is just matrix square root then presenting ZCA whitening is still good but it should be stated in the paper that you are handicapped against state-of-the-art by one extra linear system.\n", "summary_of_the_paper": "The paper addresses the problem of computing the matrix square root of a positive semidefinite matrix and computing the gradient of the matrix square root, so they can be used as a forward pass and a backward pass in a deep learning framework. \n\nThe paper has two separate contributions. In the forward pass, two approximation methods are introduced, Matrix Taylor Polynomial (MTP) and Matrix Pade Approximant (MPA). These are straight-forward generalisations of the scalar power series of, and the scalar Pade approximant to, the function $(1-z)^{1/2}$, to matrices.\n\nIn the backward pass, computing the gradient is viewed as a special case of the Sylvester equation, also known as the continuous Lyapunov equation, and is solved using a Newton-Schulz iteration first introduced by Benner (2006).\n\nNumerical tests were conducted to figure the speed and error of forward and backward passes as functions of batch size and matrix dimension.\n\nExperiments on ZCA whitening and covariance pooling in recent deep learning approaches were presented.", "summary_of_the_review": "The paper has two separate contributions, one in forward pass and one in backward pass in dealing with (positive-semidefinite) matrix square root in deep learning.\n\nIn the forward pass, two approximation approaches MTP and MPA both seem to trade stability for speed. MTP is the fastest but also the least accurate. MPA is slightly slower but on average is more accurate. However, the current paper is not complete in the sense that it lacks treatment of \"defects\" (regions where the numerator and the denominator approach zero) in Pade approximant approaches. Therefore, the stability and convergence of MPA is somewhat concerning.\n\nIn the backward pass, viewing the computation of the gradient of the matrix square root as solving a Sylvester equation is interesting. However, experiments in the generic settings lack accuracy tests to reveal how close these Newton-Schulz-based approaches are to the real gradient.\n\nThere is a bit of concern when in section 4.2 the proposed approaches are not compared against the closest possible candidate (Huang et al (2019)) for the ZCA whitening problem.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents some efficiency improvements over existing methods to compute matrix square root and its gradient.  Reviewers find that the novelty over existing methods is sufficient, and that the improvements are valuable.

I propose a poster despite the relatively high numerical scores, because the group of practitioners who will use the result is somewhat niche -- the reviewers are of course selected from this group and hence value the paper more highly.

In addition the real-world speedups are modest, but it is nevertheless important to document this approach.
Decision: Accept (Poster)