I need an efficient way to take square root of a matrix which is a sum of diagonal matrix and rank-1 matrix.

More specifically it's the following matrix


Where entries of $u$ are non-negative and $\sum_i^d u_i=1$. This is also known as the covariance of the multinomial distribution with mean $u$

My first approach was to use Eigenvalue decomposition, but it was about 50 times too slow for my application.

$$ A^{1/2} = US^{1/2}U'$$ The second approach was to to ignore $uu'$ term. This was fast enough, but incurred large error.

$$ A^{1/2} \approx D^{1/2}$$

My third approach was to do some algebraic rearrangement on the error term and come up with the following rank-1 approximation

$$A^{1/2} \approx D^{1/2}+\frac{\sqrt{1+u'D^{-1}u}-1}{u'D^{-1}u}D^{-\frac{1}{4}}uu'D^{-\frac{1}{4}}$$

This is fast enough, and had a small approximation error on a couple of random examples (notebook). However, I'm not sure if there are cases when this approximation breaks down.

At this point I feel like I may be reinventing the wheel -- are there known results for this problem? Specifically, a cheap way to compute it exactly, or approximately with a guaranteed small approximation error.

The values of $u$ are highly skewed in my case, they decay at least as fast as $O(1/k)$ and possibly exponentially.


Would a decomposition of the form $A = XX^T$ suffice? This would be enough, e.g., if the end goal is sampling from the Gaussian distribution with this given covariance.

If so, you can use the following formula, which is quite similar to your approximation: $$X = D^{1/2} + \frac{\sqrt{u^T D^{-1} u+1}-1}{u^T D^{-1} u} u u^T D^{-1/2}$$ This follows from whitening the diagonal term \begin{align} D + uu^T &= D^{1/2}\left(I + D^{-1/2}uu^TD^{-1/2}\right)D^{1/2}, \end{align} then finding the square root of the middle factor, which is a rank-1 update of the identity.

If you really need the square root, you may consider using rational approximations, $$(D + uu^T)^{1/2} \approx c_0 + c_1 (\sigma_1 I + D + uu^T)^{-1} + c_2 (\sigma_2 I + D + uu^T)^{-1} + \dots.$$ then inverting each of these terms with the Sherman-Morrison formula. The number of terms you need for a very good approximation grows as $$O(\log \kappa)$$ where $\kappa$ is the condition number of $D + uu^T$. So, instead of a diagonal plus a rank-1 matrix, this becomes a sum of a small number of diagonal plus rank one matrices.

For more details on the rational approximation, here is a wonderful paper:

Hale, Nicholas, Nicholas J. Higham, and Lloyd N. Trefethen. "Computing A^α,\log(A), and related matrix functions by contour integrals." SIAM Journal on Numerical Analysis 46.5 (2008): 2505-2523. http://eprints.maths.manchester.ac.uk/834/1/hale_higham_trefethen.pdf

See, in particular, equation (4.4) and method 3 in that paper.

Edit: Here I wrote some Python code that does the rational method:

# Adaptation of Method 3 from Hale, Higham, and Trefethen, Computing f(A)b by contour integrals. SIAM 2008
import numpy as np
import scipy.linalg as sla
from scipy.special import *

def hht_isqrt_weights_and_poles(min_eigenvalue_m, max_eigenvalue_M, number_of_rational_terms_N):
    # 1/sqrt(z) = w0/(z - p0) + w1/(z - p1) + ...
    m = min_eigenvalue_m
    M = max_eigenvalue_M
    N = number_of_rational_terms_N
    k2 = m/M
    Kp = ellipk(1-k2)
    t = 1j * np.arange(0.5, N) * Kp/N
    sn, cn, dn, ph = ellipj(t.imag,1-k2)
    cn = 1./cn
    dn = dn * cn
    sn = 1j * sn * cn
    w = np.sqrt(m) * sn
    dzdt = cn * dn

    poles = (w**2).real
    weights = (2 * Kp * np.sqrt(m) / (np.pi*N)) * dzdt
    rational_function = lambda zz: np.dot(1. / (zz.reshape((-1,1)) - poles), weights)
    return weights, poles, rational_function

def diagonal_smw(dd, u):
    # (diag(dd)+uu^T)^-1 = diag(dd) - vv^T
    v = (u / dd)/np.sqrt(1. + np.dot(u, u / dd))
    return v

def diagonal_plus_rank_one_inverse_sqrt(dd, u, num_terms):
    # inv(sqrtm(diag(dd) + u*u')) = diag(ss) - V*V' + small error
    lambda_min_bound = np.min(dd) # These bounds could be improved
    lambda_max_bound = np.max(dd) + np.linalg.norm(u)**2
    ww, pp, _ = hht_isqrt_weights_and_poles(lambda_min_bound, lambda_max_bound, num_terms)
    dd_shift = [dd - p for p in pp]
    vv0 = [diagonal_smw(d_shift, u) for d_shift in dd_shift]
    ss = np.sum([w*(1./d_shift) for w, d_shift in zip(ww, dd_shift)], axis=0)
    V = np.vstack([np.sqrt(w)*v0 for w, v0 in zip(ww, vv0)]).T
    return ss, V

def diagonal_woodburyish(ss, V):
    # (diag(ss) - VV^T)^-1 = diag(1./ss) + W*M*W^T
    r = V.shape[1]
    W = V / ss.reshape([-1, 1]) # inv(diag(dd))*V
    capacitance_mtx = np.eye(r)-np.dot(V.T, W)
    M = np.linalg.inv(capacitance_mtx) # inverse of very small matrix, e.g., 5x5. Could do Cholesky if desired..
    return M, W

def diagonal_plus_rank_one_sqrt(dd, u, num_terms):
    # inv(sqrtm(diag(dd) + u*u')) = diag(iss) + W*M*W^T + small error
    ss, V = diagonal_plus_rank_one_inverse_sqrt(dd, u, num_terms)
    M, W = diagonal_woodburyish(ss, V)
    iss = 1. / ss
    return iss, M, W

n = 500
kappa_ish = 1e3
dd = kappa_ish * np.random.rand(n)
u = np.random.randn(n)

A = np.diag(dd) + np.outer(u, u)
sqrtA = sla.sqrtm(A)

for num_terms in 1+np.arange(10):
    iss, M, W = diagonal_plus_rank_one_sqrt(dd, u, num_terms)
    sqrtA2 = np.diag(iss) + np.dot(W, np.dot(M, W.T))
    err = np.linalg.norm(sqrtA - sqrtA2) / np.linalg.norm(sqrtA)
    print('num_terms=', num_terms, ', err=', err)

It works quite well, here is the output error for 1 to 10 terms in the rational approximation:

num_terms= 1 , err= 0.28856754578036703
num_terms= 2 , err= 0.04084537953169016
num_terms= 3 , err= 0.005770763245002039
num_terms= 4 , err= 0.000624851647505238
num_terms= 5 , err= 7.234944285553648e-05
num_terms= 6 , err= 1.0059955401611735e-05
num_terms= 7 , err= 1.22949322111201e-06
num_terms= 8 , err= 1.3750242705729841e-07
num_terms= 9 , err= 1.7564866678109768e-08
num_terms= 10 , err= 2.289881883503377e-09
  • $\begingroup$ My matrix is symmetric, so AA' works too. I think our formulas came out of the same approach -- find B=square root of rank-1 correction, and then approximate sqrt(DB) as sqrt(D)sqrt(B). It's interesting that our formulas are sightly different. I tried them out on 1 million random u vectors, and got better result using my formula 89.7% of the time, and yours was better the rest of the time -- notebook $\endgroup$ Sep 23 '19 at 2:25
  • 1
    $\begingroup$ @YaroslavBulatov Err, the matrix $X$ is not intended to approximation the true square root, in the sense that $X^2 \neq A$. But $A=XX^T$ exactly (no approximation). For the square root I strongly recommend the rational approach with a few terms. $\endgroup$
    – Nick Alger
    Sep 23 '19 at 2:28
  • 1
    $\begingroup$ Ah, I see what you mean. You were right, my underlying problem was to do XX' decomposition which I assumed would be the same as symsqrt (does XX' decomposition have a name?). Thanks for the solution! $\endgroup$ Sep 23 '19 at 2:33
  • $\begingroup$ Dunno of a word for this. I guess it's like the Cholesky decomposition, but without the requirement that the matrices are triangular. $\endgroup$
    – Nick Alger
    Sep 23 '19 at 2:42
  • 1
    $\begingroup$ @YaroslavBulatov I wrote some code that does the rational method and put it in my answer. Seems to work well. Feel free to use the code as you please $\endgroup$
    – Nick Alger
    Sep 23 '19 at 3:51

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.