Is there a more efficient way to compute the log-density of a Multivariate normal with a Symmetric Woodbury Covariance?

Suppose that I have \mathbf{y}\sim N(\boldsymbol{\mu}, \mathbf{A} + \mathbf{UCU'}), where \mathbf{A} and \mathbf{C} are diagonal matrices. Fortunately, the covariance term follows the pattern of a sparse symmetric Woodbury matrix.

I coded up the following function to evaluate the log-density of a multivariate normal:

function lpdfMVN_Woodbury(y, mu, A, U, C)

    k = length(mu)
    Const = -k/2*log(2*pi)
    invA = inv(A)
    invAU = invA*U
    inner_term = inv(C) + U'*invAU
    log_determinant = logdet(inner_term) + logdet(C) + logdet(A)
    wood_inv = invA - invAU*inv(inner_term)*(invAU)'
    ymu = y - mu

    return[Const - log_determinant/2 - ymu'*wood_inv*ymu/2]

end

While this is much quicker than logpdf(MvNormal(mu, A + U*C*U'), y), it’s still a bit slower than I’d like it to be. Does anyone have any suggestions to speed up this evaluation?

Can you profile the code to see which lines take longer on a typical example?

If they aren’t already, it may help if A and C are of the Diagonal type. Their inverses will take less memory and operations with them will be faster (no need to check for diagonality for inv and logdet, no need to add off-diagonals for +).

While inv and logdet on diagonal matrices are very cheap, they have a real cost on non-diagonal matrices like inner_term. I would replace a few lines with the following:

    inner_term = cholesky(Hermitian(inv(C) + U'*invAU))
    log_determinant = logdet(inner_term) + logdet(C) + logdet(A)
    wood_inv = invA - invAU * (inner_term \ invAU') # maybe wrap in Hermitian(...)

logdet and inv are both expensive operations that become very easy once inner_term has been factorized, so calling cholesky ensures that this only has to be done once rather than twice. You may need to apply a pivoting scheme if you run into definiteness issues with cholesky, or consider a bunchkaufman factorization instead.

Stylistically, I might replace some of the other inv(X)*Y statements with X \ Y, but for diagonal X (and since you use the inverses in additions later) I don’t foresee this making a big difference. There are further games you could play with in-place functions for your matrix operations, but they make things a bit less readable and may not be that impactful. You can look into LinearAlgebra.mul! to that end.

3 Likes

Great idea. The two lines that take noticeably longer are the computations for log_determinant and wood_inv. However, these appear to be sped up using @mikmoore 's suggestion!

1 Like