{'fraction_non_alphanumeric': 0.04947079807777352, 'fraction_numerical': 0.01650992069341667, 'mean_word_length': 3.8257220088206005, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Low-Rank Updates of Matrix Square Roots

Models in which the covariance matrix has the structure of a sparse matrix plus a low rank perturbation are ubiquitous in machine learning applications. It is often desirable for learning algorithms to take advantage of such structures, avoiding costly matrix computations that often require cubic time and quadratic storage. This is often accomplished by performing operations that maintain such structures, e.g. matrix inversion via the Sherman-Morrison-Woodbury formula. In this paper we consider the matrix square root and inverse square root operations. Given a low rank perturbation to a matrix, we argue that a low-rank approximate correction to the (inverse) square root exists. We do so by establishing a geometric decay bound on the true correction's eigenvalues. We then proceed to frame the correction has the solution of an algebraic Ricatti equation, and discuss how a low-rank solution to that equation can be computed. We analyze the approximation error incurred when approximately solving the algebraic Ricatti equation, providing spectral and Frobenius norm forward and backward error bounds. Finally, we describe several applications of our algorithms, and demonstrate their utility in numerical experiments.


Introduction
In machine learning models or applications one often encounters matrices that are low-rank perturbations of another (perhaps simpler) matrix. For example, models in which the covariance matrix has the structure of a sparse matrix plus a low rank perturbation are common. In another example, it is common for algorithms to maintain a matrix that is iteratively updated by low-rank perturbations.
It is often desirable for algorithms to take advantage of such structures, avoiding costly matrix computations that often require cubic time and quadratic storage. An indispensable tool for utilizing low-rank perturbations is the famous Sherman-Morrison-Woodbury formula, which shows that the inverse of a lowrank perturbation can be obtained using a low-rank correction of the inverse. While the usefulness of the Sherman-Morrison-Woodbury formula cannot be understated, other matrix functions also frequently appear in applications. One naturally asks the following questions. Given a matrix function f , when does a low rank perturbation of a matrix A correspond to a (approximately) low-rank correction of f (A)? Can we find a high quality approximate correction efficiently, i.e. without computing the exact correction and truncating it using a SVD?
In this paper, we answer these questions affirmatively for two important and closely related functions: the matrix square root and inverse square root. The square root operation on matrices is a fundamental operation that frequently appears in mathematical analysis. Moreover, matrix square roots and their inverse arise frequently in machine learning applications, e.g. when sampling from an high dimensional multivariate Gaussian distribution, or whitening data. Computing the square root (or inverse square root) of low-rank perturbations of simple matrices (e.g. diagonal matrices) appear in quite a few machine learning applications, e.g. in the aforementioned applications when the covariance matrix follows a spiked population model. For more discussion on applications of low-rank perturbations of matrix square root, see Section 6.
In particular, we show that given a low rank perturbation D to a matrix A, we can approximate (A+D) 1 /2 well by a low rank correction to A 1 /2 . We do so by proving a geometric decay bound on the eigenvalues of (A + D) 1 /2 − A 1 /2 . We also provide a similar bound for the inverse square root. We then proceed to show that the exact update is a solution of an algebraic Ricatti equation, and discuss how a low-rank approximate solution to that equation can be computed. This allows us to propose concrete algorithms for updating and downdating the matrix square root and matrix inverse square root. Finally, we report experiments that corroborate our theoretical results.

Related Work
Most previous work on the matrix square root focused on computing A 1 /2 x and A − 1 /2 x for a given vector x using a Krylov method, possibly with preconditioning (Higham, 2008;Aune et al., 2013;Chow & Saad, 2014;Frommer et al., 2014). The motivation for most of the aforementioned works is sampling from a multivariate Gaussian distributions. Worth mentioning is recent work by Pleiss et al. (2020) which combines a Krylov subspace method with a rational approximation, and also allows preconditioning.
The problem of updating a function of a matrix after a low rank perturbation, i.e. computing f (A + D) given f (A) where D is low-rank, has been recently receiving attention. The Sherman-Morrison-Woodbury formula is a well known formula for updating the matrix inverse, i.e. f (x) = x −1 . Bernstein & Van Loan (2000) showed that that when f is a rational function of degree q, a rank one perturbation of A corresponds to a rank q perturbation of f (A). This paper also provides an explicit formula for the low rank perturbation. Higham (2008) showed that a rank k perturbation of A = αI corresponds to a rank k perturbation of f (A) for any f (see Theorem 1.35 therein). As for inexact corrections, Beckermann et al. (2018) proposed a Krylov method for computing a low-rank correction of f (A) that approximates f (A + D) well for any analytic f (however, approximation quality depends on properties of the function itself, e.g. how well it is approximated by a polynomial). In followup work, they proposed a rational Krylov method (Beckermann et al., 2020), citing the matrix square root as an example of a case in which their original method might have slow convergence.
The work most similar to ours is (Fasi et al., 2022). In that paper, the authors consider the problem of computing the square root of a matrix of the form αI + UV T , as a correction of √ αI. Due to (Higham, 2008, Theorem 1.35), the rank of the correction is the same as the rank of UV T . However, the formula in (Higham, 2008, Theorem 1.35) requires V T U to be non-singular, which is not required for the square root to be defined. The authors circumvent this issue by suggesting another formula for the square root, or by using a Newton iteration. In a way, (Fasi et al., 2022) is more general than our method since it allows non-symmetric updates. However, in others it is less general: the matrix to be perturbed must be a scaled identity matrix. Moreover, we also suggest a method for updating the inverse square root.

Notation and Basic Definitions
We denote scalars using Greek letters or using x, y, . . . . Vectors are denoted by x, y, . . . and matrices by A, B, . . . . The n × n identity matrix is denoted I n . We use the convention that vectors are column-vectors.
Given a symmetric positive-semidefinite matrix A ∈ R n×n , another matrix B ∈ R n×n is a square root of A if B 2 = A. There is a unique square root of A that is also positive semi-definite, which is called the principal square root and we denote it by A 1 /2 . Given two matrices F and G, the (F, G)-displacement rank of A is defined as the rank of FA − AG. Displacement structures and displacement rank are closely connected to the Sylvester equation.

Low-Rank Algebraic Ricatti Equation
Consider the following equation in X ∈ R n×n , where E ∈ R n×n is a symmetric full-rank matrix, G ∈ R n×k where k n, and α = ±1. Our algorithms are based on approximately solving Eq. (1) with a positive semidefinite low-rank X.
If α = +1, Eq. (1) is an instance of the algebraic Ricatti equation. For such equations, it is well known that when k n then indeed we can expect that the solution X to be positive semidefinite and to be well approximated by a low-rank matrix (Benner & Saak, 2013).
Recently, Mishra & Vandereycken (2014) proposed a meta-scheme for approximately solving a slightly more general version of Eq. (1) based on Riemannian optimization. When applied to Eq. (1) their scheme assumes the ability to take products of E by a vector, and to solve linear equations where the matrix is equal to E 2 + low-rank, which is easily achievable via the Sherman-Morrison-Woodbury formula if we have access to an oracle that multiplies E −1 by a vector. Under the assumption that each rank update in (Mishra & Vandereycken, 2014) requires O(1) trust-region iterations, and that the target maximum rank of X is r, the overall cost of the scheme in (Mishra & Vandereycken, 2014) is O((T E + T E −1 )r 2 + nr 4 ) where T E and T E −1 is the cost of multiplying E and E −1 by a vector (respectively). In most of our applications E is diagonal, so the cost reduces to O(nr 4 ).
The method described in (Mishra & Vandereycken, 2014) handles only the case of α = +1. However, it can be generalized to the case of α = −1. We give details in Appendix A.
In our algorithms, we denote the process of solving Eq. (1) via the notation where r is the target maximum rank, and U ∈ R n×r is a symmetric factor of the solution X (i.e., X = UU T ).
In the complexity analyses we assume that this process takes O((T E + T E −1 )r 2 + nr 4 ).

Problem Statement
Let A ∈ R n×n be a a symmetric positive semidefinite matrix A ∈ R n×n . Suppose we are given A 1 /2 , perhaps implicitly (i.e., as a function that maps a vector x to A 1 /2 x). Given a perturbation D ∈ R n×n of rank k n, our goal is to approximate (A + D) 1 /2 using a low-rank correction of A 1 /2 . That is, to find a∆ of rank r such that (A + D) 1 /2 ≈ A 1 /2 +∆. The rank r of∆ should be treated as a parameter, and should optimally be O(k). We show in Section 3 that we can expect to find a good approximation with r n. We make two additional assumptions on D. First, we assume that it is either positive semidefinite or negative semidefinite. Indefinite perturbations can be handled by splitting the update into two semidefinite perturbations and applying our algorithms sequentially. Secondly, we assume that D is given in a symmetric factorized form. We can combine the last two assumptions in a single assumption by assuming we are given a Z ∈ R n×k such that D = αZZ T where α = ±1. We refer to the case of α = +1 as updating the square root, and α = −1 as downdating the square root.
In the case of updating the square root, we are guaranteed that A + ZZ T is positive definite for any Z, but this does not necessarily holds for downdating. Thus, for downdates we further assume that A − ZZ T is positive definite. The following Lemma gives an easy way to test this condition in cases we also have access to the inverse of A 1 /2 or A.
Lemma 1. Suppose that A ∈ R n×n is symmetric positive definite, and Z ∈ R n×k for k ≤ n. are smaller or equal to 1. However, the non-zero eigenvalues of that matrix are equal to the eigenvalues of The correction∆ should be returned in factorized form. For reasons that will become apparent in our algorithm, the correction will be positive semidefinite when D is positive semidefinite, and negative semidefinite when D is negative semidefinite. Thus, we require our algorithm to return a U ∈ R n×r such that (A + αD) 1 /2 ≈ A 1 /2 + αUU T . The discussion so far was for updating or downdating the square root. We also aim at correcting the inverse of the square root, i.e. A − 1 /2 . Again, we will require D to be either positive semidefinite or negative semidefinite, and∆ to be factorized as well and definite. However, for updating the inverse of the square root, the definiteness of∆ will be opposite to the one of D.
We can capture the distinction between updating the square root and the inverse square root with an additional parameter β = ±1, where we wish to update A β /2 . Putting it all together, we arrive at the following problem:

Decay Bounds for Square Roots Corrections
Given a matrix A and a low-rank perturbation D, our goal is to find a low-rank correction∆ to A β/2 . However, one can ask whether such a correction even exists? Let ∆ denote the exact correction, i.e. ∆ : In this section we show that the eigenvalues of ∆ exhibit a geometric decay. Thus, by applying eigenvalue thresholding to ∆ we can obtain a low-rank approximate correction∆ (however, our algorithm uses a different method for finding∆).

Decay Bound for Square Root Corrections
We first consider the case that β = 1, so we are perturbing the square root.
Let us denote B := A + D. The following is a known identity: In our case, that is when B = A + D, we see that ∆ upholds the following Sylvester equation: Since the rank of D is k, we see that ∆ has a (A 1 /2 , −B 1 /2 )-displacement rank of k. Beckermann & Townsend (2017) recently developed singular value decay bounds for matrices with displacement structure. Using their results, we can prove the following bound. Theorem 1. Suppose that both A and B = A + D are symmetric positive definite matrices, and that D is of rank k. Let ∆ = B 1 /2 − A 1 /2 . Then for j ≥ 1, the singular values of ∆ satisfy the following bound Proof. Let µ = λ min (A) and δ = A 2 + D 2 . The matrix ∆ upholds the following Sylvester equation Thus, the results in (Beckermann & Townsend, 2017) show that we can bound where E is any set that contains the spectrum of A 1 /2 − (µ/2)I n , F is any set that contains the spectrum of −(B 1 /2 + (µ/2)I n ), and Z l (E, F ) is the Zolotarev number. Let a = µ/2. Since µ is the minimal eigenvalue of A 1 /2 , all the eigenvalues of A 1 /2 − (µ/2)I n are bigger than a or equal to it. Since B is by assumption positive definite, all the eigenvalues of −(B 1 /2 + (µ/2)I n ) are smaller than −a or equal to it. Let b = δ + µ/2. Obviously, all the eigenvalues of A 1 /2 − (µ/2)I n are smaller than b. Furthermore, since B 2 ≤ A 2 + D 2 , all the eigenvalues of −(B 1 /2 + (µ/2)I n ) are bigger than or equal to −b. Thus, we can take E = [a, b] and F = [−b, −a]. In (Beckermann & Townsend, 2017) it is also shown that plugging that into Eq. (3) gives the desired bound.
The theorem bounds the singular values. However, if D is positive semidefinite, then ∆ is also positive semidefinite, and the bound is actually on the eigenvalues. Figure 1 illustrates the bound vs. actual decay of the eigenvalues on two simple test cases. In both examples, A ∈ R 100×100 is a diagonal matrix. In the left graph , the diagonal entries are sampled uniformly from U (0, 1). In the right graph, diagonal entries are logarithmically spaced between 10 −3 and 10 3 . The perturbation is D = zz T , where z is a normalized Gaussian vector.
It has also been empirically observed that if the righthand side of a Sylvester equation is low rank, then the solution is well approximated using a low rank matrix (Benner & Kürschner, 2014).

Decay Bound for Inverse Square Root Corrections
Since D is rank k, so the matrix −A −1 DB −1 is of rank at most k. Thus, similarly to Theorem 1, we can prove the following bound on the singular values of ∆:

Equation for Square Roots Corrections and Error Analysis
Our algorithms are based on writing ∆ as a solution of an equation, and then finding a low-rank approximate solution∆. Seemingly, Eqs. (2) and (4) are the equations we need. However, these equations contain the unknown B β /2 so they are not useful for us algorithmically. We derive a different equation instead. In particular, we write ∆ as the solution of an algebraic Ricatti equation, i.e. in the form of Eq. (1).
We can combine Eqs.
(2) and (4) into a single equation: Recalling that B β /2 = A β /2 + ∆, and plugging it into the last equation we get In order for the equation to fit Eq. (1) we must write the right side as a positive semi-definite factorized matrix. The first step is finding a matrix V ∈ R n×k such that When β = 1, and recalling that in Problem 1 we have B − A = αZZ T , we can take V = Z. When β = −1, we obtain V using the Sherman-Morrison-Woodbury formula. Indeed, is positive definite and the inverse square root exists).
We now have the equation We do an additional change of variables to make the right-hand side positive definite even if α = β. Let C = αβ∆ (so B 1/2 = A 1/2 + αβC since α = ±1 and β = ±1). By multiplying the last equation on both sides by αβ we obtain the equation: Except C, all other quantities of the last equation are known, and solving Eq. (6) using a low-rank positive semidefinite C of the form C = UU T forms the basis of our algorithm (see next section). However, all our algorithms solve Eq. (6) approximately, since they output low-rank solutions, while the exact solution tends to be full-rank. We now analyze how errors in solving Eq. (6) translate to errors in approximating B β/2 . First, let us define the residual of an approximate solution: We start with a backward error bound, i.e. showing that if the residual has a small norm, then A β /2 + αβC is the square root of a matrix that is close to (A + αZZ T ) β .

Lemma 2. We have
In order to get a bound on the forward error in terms of the backward error, we need the following perturbation bound for the matrix square root: Lemma 3. (Schmitt (1992, Lemma 2.2)) Suppose that Re (A j ) µ 2 j I , µ j > 0, j = 1, 2. Then, both A 1 and A 2 have square roots satisfying Re A 1/2 j µ j I for j=1,2, and Lemma 4. Let W and H be two symmetric positive definite matrices. The following bounds hold: follows immediately from Lemma 3. The Frobenius norm bound follows from Wihler inequality the pth root of positive semidefinite matrices: for any two n×n positive semidefinite matrix X and Y and p > 1 we have X 1 /p −Y 1 /p p F ≤ n (p−1) /2 X−Y F (Wihler, 2009). We apply this inequality to X = W β and Y = H 2 .
2: Output: U ∈ R n×r such that Finally, we obtain the following bound: Corollary 1. Suppose that both A and A + αZZ T are positive definite where α = ±1. LetC be a positive semidefinite matrix. Assume that A β /2 + αβC is positive definite. The following bounds holds:

Algorithms
In this section we describe our algorithms for solving Problem 1. As alluded earlier, our algorithms are based on using the Ricatti low-rank solver of Mishra & Vandereycken (2014) as encapsulated by RicattiLRSolver. However, in some combinations of α and β there is a challenged: while the returnedC is guaranteed to be positive definite, there is no guarantee that A β /2 +αβC is positive definite as well. Such guarantee is necessary in order to apply Corollary 1. Thus, we split our algorithm to various cases based on the combination of α and β. Our proposed algorithms are summarized in pseudo-code form in Algorithm 1.

Downdating (α = −1) the inv. sqr. root (β = −1)
We first compute V = A −1 Z(I k − Z T A −1 Z) − 1 /2 . Along the way we can verify that I − Z T A −1 Z is positive definite, which is required for A − ZZ T to be positive semidefinite, and our algorithm to work. We now call U ← RicattiLRSolver(A − 1 /2 , V T , −1, r) and return U.
5.3 Downdating (α = −1) the square root (β = 1) Seemingly, we could simply call U ← RicattiLRSolver(A 1 /2 , Z T , −1, r) and return U. However, there is no guarantee that A 1 /2 − UU T is positive definite, and Corollary 1 no longer guarantees that we have an approximation to the principal square root.
If we want to approximate the principal square root, we can first solve for downdating the inverse square root (α = −1, β = −1), obtaining U 1 such that (A − ZZ T ) − 1 /2 ≈ A − 1 /2 + U 1 U T 1 . We now use the Sherman-Morrison-Woodbury formula to note that

Updating (α = 1) the inverse square root (β = −1)
Again, calling the Ricatti solver directly might return a corrected matrix which it not neccessarily positive definite, and it will not be a good approximation to the principal square root. To approximate the principal square root, we first solve the updating problem for the square root (A+ZZ T ) 1 /2 ≈ A 1 /2 +U 1 U T 1 (α = +1, β = +1), and then use the Sherman-Morrison-Woodbury formula to find a U such that (A 1 /2 + U 1 U T 1 ) −1 = A − 1 /2 −UU T . We omit the details and simply refer the reader to the pseudo code description in Algorithm 1.

Costs
The main cost of the algorithms is in solving the Ricatti equation. Even when the Sherman-Morrison-Woodbury formula is needed to ensure positive definiteness of the correction, its cost of O(nk 2 ) is subsumed by the cost of solving the Ricatti equation. Overall, under our assumptions on the cost of solving the Ricatti equation, the overall cost of the algorithms is O((T A 1 /2 + T A − 1 /2 )r 2 + nr 4 ) where T A 1 /2 and T A − 1 /2 are the costs of the taking products of A 1 /2 and A − 1 /2 (respectively) with a vector. In many of the applications we discuss in the next section A is diagonal, in which case the cost of the algorithms reduces to O(nr 4 ).

ZCA Whitening of High Dimensional Data
Whitening transformations are designed to transform a random vector (or samples of that random vector) with a known covariance matrix into a new random vector whose covariance is the identity matrix. Suppose that x ∈ R p is a random vector with covariance matrix Σ ∈ R p×p . Let W be any matrix such that W T W = Σ −1 ; such a matrix is called a whitening matrix. Then the covariance matrix of the random vector z = Wx is the identity, so the random variable has been whitened. There are several possible choices for W, leading to different whitening transformations. ZCA whitening is the whitening transformation defined by In practice, the ZCA transformation is learned from data. Given samples x 1 , . . . , x n an estimateΣ of Σ is formed, andŴ =Σ − 1 /2 is used for the ZCA whitening matrix. A common choice is to use the sample covariance matrix S n = n −1 X T c X c forΣ where X c is the data matrix whose rows are x 1 , . . . ,x n after centering (subtraction of the mean).
However, it is well appreciated in the statistical literature that when the random vectors are high dimensional, i.e. when p is of the same order as n (or much larger), then the sample covariance S n is a poor estimate of Σ. Indeed, one can easily see that if p > n then S n is not even invertible so the ZCA transformation is not even defined. In high dimensional settings it is common to adopt the spiked covariance model of Johnstone (2001). In the spiked covariance model it is assumed that the covariance matrix has the following form: for some Z ∈ R p×k (k is a parameter). Suppose that we have formed an estimateΣ of Σ with the same structure as in Eq. (7), and we want to transform the samples using ZCA whitening. Explicitly computing the square root ofΣ will require O(p 3 ), which is prohibitive when p is large, and additional O(p 2 n) is required for applying the transformation the the data. Instead, we can use our algorithm to find a U ∈ R n×r with r = O(k) such that We can then apply the ZCA transformation in O(npk).

Sampling from a Multivariate Normal Distribution with Perturbed Precision Matrix
Consider a random vector x ∈ R n following a multivariate normal distribution with precision matrix Q, i.e.
x ∼ N (µ, Q −1 ). Suppose we want to sample x. This can be accomplished by sampling a vector z from the standard multivariate normal distribution (i.e. z ∼ N (0, I n )) and then computing the sample x = µ+Q − 1 /2 z. In certain cases the matrix Q has the structure of a low-rank perturbation of a fixed precision matrix, i.e. Q = Q 0 + ZZ T . Assuming we already have computed the inverse square root of Q 0 , we can use our algorithms to compute a U such that Q − 1 /2 ≈ Q − 1 /2 0 − UU T . We can then sample efficiently from x. Such cases can occur in Gibbs Samplers for Bayesian inference on spatially structured data. An example is the image reconstruction task discussed in (Bardsley, 2012). The computational bottleneck in the algorithm proposed in (Bardsley, 2012) is sampling from a conditional Gaussian distribution whose precision matrix has the structure Q = γ prior L + γ obs ZZ T where L is a fixed discrete Laplace operator that encodes prior smoothness assumptions on the image, while Z T encodes how the high-resolution images are blurred and downsampled to yield low-resolution images.

Preconditioned Second-Order Optimization
Recently introduced by Gupta et al. (2018), Shampoo is a preconditioned second-order optimization method for solving problems in which the parameter space is naturally organized as a m × n matrix or higher order tensor. Here, we consider the matrix-shaped case. In this case, at the core, Shampoo performs update steps of the form where η is the learning rate, {W t } are the parameters at time t, {G t } are the gradients at time t, and If m nt then it is better to avoid holding L − 1 /4 t explicitly (which is m × m ), and simply hold an implicit representation of both L − 1 /4 t and L − 1 /2 t as diagonal plus low-rank matrices. In each iteration we can update both by applying our algorithm twice. Since n m, we can store R t explicitly, and compute R − 1 /4 t in each iteration. If n mt we can reverse roles, implicitly keeping R t and explicitly keeping L t . Even if both m and n are of comparable size, then in some cases G t is of low-rank, and again we can track L t and R t using our algorithm.

Updating/Downdating Polar Decomposition and ZCA Transformed Data
Given a matrix X ∈ R n×d where n ≥ d, a polar decomposition of it is where U has orthonormal columns and P is symmetric positive semidefinite. The matrix P is always unique, and is given by P = (X T X) 1 /2 . If X has full rank, then P is positive definite, and U = XP −1 . The polar decomposition can be computed using a reduced SVD, so the cost of computing a polar decomposition is O(nd 2 ). There are quite a few uses for the polar decomposition (Higham, 1986).
We now consider the following updating/downdating problem. Let us denote the rows of X by x 1 , . . . , x n ∈ R d , i.e. row j of X is x T j . Suppose we already have a polar decomposition X = UP of X. The downdating problem is: compute the polar decomposition of a matrix X − obtained by removing one row from X. The updating problem is: compute the polar decomposition of X + , a matrix obtained by adding a single row to X.
We describe an algorithm for downdating a polar decomposition. The updating algorithm is almost the same. Without loss of generality, assume we remove the last row from X: This is a rank-1 perturbation of X T X. The P-factor of X − , which we denote by P − , is obtained by computing the square root of X T − X − , and we already have the square root P for the unperturbed matrix X T X. So we can use the algorithm described in Section 5 to find a matrix Z ∈ R d×k for some small r (a parameter; e.g., k = 4) such that P − ≈P − := P − UU T Assume now that X − is full rank as well. We now have Using the Sherman-Morrison-Woodbury formula: Now notice that X − P −1 is just the first n − 1 rows of XP −1 = U so we do not need to recompute it. Multiplying X − P −1 by (I d + U(I r − U T P −1 U) −1 U T P −1 ) can be done, utilizing the low rank structure of that matrix, using O(ndr) operations. If r n this is a big save over O(nd 2 ). The polar decomposition is closely connected to ZCA whitening. Suppose that X is a data matrix whose rows are sampled from a zero mean random vector (or, alternatively, X has been centered). The U-factor is equal, up to scaling, to the ZCA transformed data, while the inverse of the P-factor is, up to scaling, the ZCA whitening matrix itself. Using the ability to update the inverse square root, the procedure for updating/downdating the polar decomposition can be adjusted to update/downdate ZCA. Updating/downdating ZCA can be useful if you want to transform data that arrives over time while the covariance matrix itself changes slowly. That is, data point x j is sampled with covariance Σ j . If we assume the covariance changes slowly, we can keep an approximate ZCA of the data by taking the sample covariance over a sliding window. To do so efficiently, we can use the proposed ZCA updating/downdating procedure to first remove outdated data (a downdate operation), and then add the newly arrived data (and update operation).

Faster Generalized Least Squares with a Spiked Weight Matrix
Let X ∈ R n×d and y ∈ R n . In Generalized Least Squares (GLS) we wish to find the minimizer where W ∈ R n×n is some symmetric positive definite weight matrix, and z W := √ z T Wz. In this section, we focus on the cases that W can be written as a diagonal plus a definite low rank perturbation W = D + αZZ T where D ∈ R m×m is diagonal, and Z ∈ R n×k .
A statistical motivation for this problem is generalized linear regression with a spiked covariance matrix. Assume that the rows of X correspond to data points x 1 , . . . , x n , the entries y 1 , . . . , y n of y are responses. We now assume that the responses follow the model y i = x T i w + i where the vector of noise elements 1 , . . . , m is distributed according to N (0, C) for some covariance matrix C = D + αZZ T . The optimal unbiased estimate of w is obtained by solving Eq. (9) with W = C −1 .
One can easily see that Eq. (9) is equivalent to Once we have efficiently computed W 1 /2 X and W 1 /2 y, we can leverage faster, sketching based, least squares algorithms (Woodruff, 2014;Drineas & Mahoney, 2016). Using the algorithms from Section 5 we can compute a U ∈ R n×r such that W 1 /2 = D − 1 /2 − αUU T with r = O(k). We can them compute W 1 /2 X and W 1 /2 y efficiently.

Experiments
We report experiments exploring the ability of our algorithm to find low-rank corrections to the square root or inverse square root of a perturbed matrix. In our experiments, we focus on the quality of the corrections found. We do not report running time since the code we used for the algebraic Ricatti solver (downloaded from the homepage of Mishra & Vandereycken (2014)) is not optimized to take advantage of the structures present in the input matrices for the specific algebraic Ricatti equations our algorithm solves.

Synthetic Experiments
We first test our algorithm on randomly generated matrices, and compare them to the approximation obtained using the algorithm of Beckermann et al. (2018), and the optimal correction obtained by zeroing out the smallest eigenvalues of the exact correction. The matrix A ∈ R 100×100 is a diagonal matrix, whose diagonal is either sampled uniformly from U (0, 1), or whose entries are logarithmically spaced between 10 −3 and 10 3 . The perturbation is D = zz T , where z is a normalized Gaussian vector. We consider both updates (α = +1) and downdates (α = −1). In case of downdates, we multiply z by 0.1 to ensure positive definiteness after the downdate. We consider both the square root (β = +1) and inverse square root (β = −1). We plot the relative error (A + αzz T ) β /2 − (A β /2 + αβUU T ) F / (A + αzz T ) β /2 F as a function of the rank of the update (#columns in U). Figures 2 and 3 show the result for all different combinations. Our algorithm is clearly able to find much better approximations than the Krylov method of Beckermann et al. (2018).

Matrices Arising from Second-Order Optimization
Our next set of experiments simulates the use of our algorithm to track L − 1 /4 t or R − 1 /4 t in Shampoo (see Section 6.3). We obtain and preprocess the data in a similar way to (Fasi et al., 2022), but the experiment  Beckermann et al. (2018) and to the optimal approximation, on randomly generated matrices of the form D + zz T where D is diagonal with uniformly sampled entries from N (0, 1).   Beckermann et al. (2018) and to the optimal approximation, on randomly generated matrices of the form D + zz T where D is diagonal with logarithmically spaced entries.
itself is different. We downloaded two test matrices available from the Lingvo framework for TensorFlow (Shen et al., 2019) and are available on GitHub 1 . These matrices are obtained by accumulating updates with α = 0. Unfortunately, the provided test matrices are only the final accumulated matrix, and do not contain the discrete updates themselves, so we need to extract updates that accumulate to the final matrix. We do so in a similar fashion to the one used by Fasi et al. (2022): we compute an eigendecomposition, and keep dominant factors that are bigger than 0.1. This yields a rank 82 approximation to the first matrix, and a rank 221 approximation to the second matrix.
For the experiment, we split the low rank approximation into discrete updates of rank 5. So we now have a set of {G s }, each G s having 5 columns. Our goal is to efficiently track L − 1 /4 t where L t = αI + t s=1 G s G T s where we set α = 0.001. We use our algorithm to form two sets of updates, {U s } and {W s }, each with 5 columns, such that L − 1 /2 t ≈ α − 1 /2 I + t s=1 U s U T s and L − 1 /4 t ≈ α − 1 /4 I + t s=1 W s W T s . We plot in Figure 4 the distance between L − 1 /4 t and its approximation, as it evolves over time. We do so for three different tolerances in the internal Ricatti low-rank solver. We see that our algorithm is able to track L − 1 /4 t well over time, without the errors blowing-up. (1) is an instance of the algebraic Ricatti equation, for which Mishra & Vandereycken (2014) proposed an algorithm for finding an approximate low rank solution. That algorithm can be adjusted to the α = −1, by making several small adjustments to the various Euclidean components (the Riemannian ones are obtained by converting the Euclidean components to Riemannian ones). We frame the expressions with an α = ±1 to cover both cases concurrently.
• Optimization problem: the new optimization problem is: • Gradient expression: The method in (Mishra & Vandereycken, 2014) keeps X in factorized low-rank form X = YY T . Let • Hessian expression: using a similar technique as (Mishra & Vandereycken, 2014), we calculate directional derivative in direction w by computing lim 1 (∇F (Y + w) − ∇F (Y)). In the limit, S(Y) and S(Y + w) include α and so are slightly different from the ones used in (Mishra & Vandereycken, 2014). Nevertheless, the expression for the Euclidean Hessian is: