Backpropping through Cholesky

The Cholesky decomposition is the unique way to decompose a positive definite matrix $K$ into $K=LL^\top$ where $L$ is lower triangular and has a real, positive diagonal. To use Cholesky in a deep learning system, it is awkward to backpropagate through, yet possible, and we show a derivation here.

Backprop summary

Suppose we have,

training data $(X,\,y)$,
parameters $\theta = (\theta_1,\ldots, \theta_L)$,
a function approximator composed of $L$ layers ${f(X;\theta) = (g_L(\cdot; \theta_L) \circ \cdots \circ g_1(\cdot; \theta_1) ) (X)}$,
a loss function $\mathcal{L}$ that we use to optimize $f$. We evaluate the loss function as $\mathcal{L}(f(X;\theta), y)$.

We use gradients to optimize the parameters, and to obtain those gradients we use the chain rule,

$$ \theta_\ell \leftarrow \theta_\ell - \gamma \frac{\partial\mathcal{L}}{\partial g_L} \frac{\partial g_L}{\partial g_{L-1}}\cdots \frac{\partial g_\ell}{\partial \theta_\ell}. $$

When we actually compute this long chain, we do it recursively,

$$ \frac{\partial\mathcal{L}}{\partial g_k} = \frac{\partial\mathcal{L}}{\partial g_{k+1}}\frac{\partial g_{k+1}}{\partial g_{k}}. $$

We can view this as a sequence of VJPs (vector Jacobian products), since $\partial L / \partial g_{k+1}$ is an activation gradient (vector), and $\partial g_{k+1}/\partial g_{k}$ is a Jacobian.

Linear example

Linear layers are easy: $g(X; W) = XW$,

$$ \begin{align*} \frac{\partial \mathcal{L}}{\partial X} = \frac{\partial \mathcal{L}}{\partial g} W^\top,\\ \frac{\partial \mathcal{L}}{\partial W} = X^\top\frac{\partial \mathcal{L}}{\partial g}. \end{align*} $$

Cholesky?

It is not obvious how to do backprop with the Cholesky operation $\mathrm{chol}(K)$. We need to find $\partial{L}/\partial{K}$ in terms of $\partial \mathcal{L}/\partial L$. This is possible, but the solution is not very efficient (see [1]).

[1] provides an alternative. The idea is that we don't need to find $\frac{\partial{L}}{\partial{K}}$. If we can calculate $\frac{\partial{\mathcal{L}}}{\partial{K}}$, given gradients for $L$, $\frac{\partial{\mathcal{L}}}{\partial{L}}$, then we can backprop just fine.

Here is the method: by the chain rule,

$$ \begin{align*} \mathrm{d}\mathcal{L} &= \sum_{ij} \frac{\partial\mathcal{L}}{\partial L_{ij}} \mathrm{d}L_{ij} = \mathrm{Tr}\left(\left(\frac{\partial\mathcal{L}}{\partial L}\right) ^\top\mathrm{d}L \right). \end{align*} $$

If we can massage this into the form

$$ \begin{align*} \mathrm{d}\mathcal{L} &= \mathrm{Tr}\left(S \mathrm{d}K \right) \end{align*} $$

for some $S$, then we can identify

$$ \begin{align*} \frac{\partial\mathcal{L}}{\partial K} = S. \end{align*} $$

There is a complication - we want the gradient $\frac{\partial\mathcal{L}}{\partial K}$ to be symmetric because $K$ is symmetric. We are going to be careful later on in order to ensure this symmetricity.

The massaging begins. We have $\mathrm{d}K = L\mathrm{d}L^\top + \mathrm{d}L L^\top$, so

$$ \begin{align*} L^{-1}\mathrm{d}K L^{-T} = \underbrace{\mathrm{d}L^\top L^{-T}}_{\text{upper triangular}} + \underbrace{L^{-1} \mathrm{d}L}_\text{lower triangular}\\ \implies \text{tril}(L^{-1}\mathrm{d}K L^{-T}) - \tfrac{1}{2}\text{diag}(L^{-1}\mathrm{d}K L^{-T}) = L^{-1}\mathrm{d}L. \end{align*} $$

[1] calls this operation $\Phi(\cdot)$ for convenience,

$$ \begin{align*} \Phi(X) := \text{tril}(X) - \tfrac{1}{2}\text{diag}(X). \end{align*} $$

It turns out that $\Phi$ has the following nice property related to traces. Assuming that $A$ and $B$ are square, then since $\Phi(X)$ is always lower triangular,

$$ \begin{align*} \text{Tr}(A\Phi(B)) &= \sum_{ij} A_{ji}^\top \Phi(B)_{ji} \\ &= \sum_{j > i} A_{ji}^\top B_{ji} + \sum_{j} A_{jj}^\top (\tfrac{1}{2}B)_{jj}\\ &= \sum_{j > i} \Phi(A^\top)_{ji} B_{ji} + \sum_{j} \Phi(A^\top)_{jj} (B)_{jj}\\ &= \sum_{ij} \Phi(A^\top)_{ji} B_{ji} \\ &= \text{Tr}(\Phi(A^\top) B^\top). \end{align*} $$

Combining the above (using symmetry of $L^{-1}KL^{-T}$ and cyclic property of trace),

$$ \begin{align*} \mathrm{Tr}\left(\left(\frac{\partial\mathcal{L}}{\partial L}\right)^\top \mathrm{d}L \right) &= \mathrm{Tr}\left(\left(\frac{\partial\mathcal{L}}{\partial L}\right)^\top L \Phi(L^{-1}\mathrm{d}K L^{-T}) \right)\\ &=\mathrm{Tr}\left(\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right) L^{-1}\mathrm{d}K L^{-T} \right)\\ &=\mathrm{Tr}\left(L^{-T}\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right) L^{-1} \mathrm{d}K \right)\\ \end{align*} $$

We are almost done, however, we said earlier that we would be careful to ensure that the gradient of $K$ is symmetric. The implied gradient here, $S=L^{-T}\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right) L^{-1}$, is not symmetric. As a solution, we do some more manipulation. We enforce that $\mathrm{d}K = \mathrm{d}K^\top$, and use some trace properties,

$$ \begin{align*} \mathrm{Tr}\left(\left(\frac{\partial\mathcal{L}}{\partial L}\right)^\top \mathrm{d}L \right)=\mathrm{Tr}\left(L^{-T}\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right) L^{-1} \mathrm{d}K \right)\\ = \frac{1}{2}\mathrm{Tr}\left(L^{-T}\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right) L^{-1} \mathrm{d}K \right)+ \frac{1}{2}\mathrm{Tr}\left(\mathrm{d}K^\top L^{-T}\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right)^\top L^{-1} \right)\\ = \mathrm{Tr}\left(\frac{1}{2}L^{-T}\left[\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right) + \Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right)^\top\right] L^{-1} \mathrm{d}K \right) \end{align*} $$

Therefore the gradient of the loss w.r.t. $K$ is,

$$ \begin{align*} \frac{\partial\mathcal{L}}{\partial K} &= \frac{1}{2}L^{-T}\left[\Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right) + \Phi\left(L^\top\frac{\partial\mathcal{L}}{\partial L}\right)^\top\right] L^{-1}. \end{align*} $$

As we can see, backpropping through the Cholesky operation is slightly painful, because we need to do a matrix multiplication as well as two triangular solves. But it is do-able!

Pytorch implementation / sanity check

As a sanity check, we implement the above in pytorch, showing that it gives the same gradients as the pytorch Cholesky function. We also show that we get identical gradients when using a mathematically equivalent, but computationally different operation.

import torch

def Phi(x): return x.tril() - 0.5 * torch.diag_embed(x.diag())

class MyCholesky(torch.autograd.Function):
    @staticmethod
    def forward(ctx, K):
        L = torch.linalg.cholesky(K)
        ctx.save_for_backward(L)
        return L

    @staticmethod
    def backward(ctx, gradL):
        L, = ctx.saved_tensors

        S  = Phi(L.mT @ gradL)
        S = S + S.T
        S = torch.linalg.solve_triangular(L.mT, S, upper=True, left=True) # = L^{-T} @ S
        S = torch.linalg.solve_triangular(L, S, upper=False, left=False)  # = S @ L^{-1}
        return 0.5 * S

V = torch.randn(4, 4, dtype=torch.float64)
K = (V @ V.T).add_(torch.eye(4, dtype=torch.float64))

## we calculate f = trace(K^{-1}) in three different ways.
## then we show that calling f.backward() gives the same gradients on K
## in all three cases

K1 = K.detach().clone().requires_grad_(True)
K2 = K.detach().clone().requires_grad_(True)
K3 = K.detach().clone().requires_grad_(True)

K1inv = K1.inverse()  # A^{-1} = A^{-1}
cholK2 = K2.cholesky()
cholK3 = MyCholesky.apply(K3)

cholK2inv = cholK2.inverse()
K2inv = cholK2inv.T @ cholK2inv
cholK3inv = cholK3.inverse()
K3inv = cholK3inv.T @ cholK3inv

K1inv.trace().backward()
K2inv.trace().backward()
K3inv.trace().backward()

print('method 1\n', K1.grad)
print('method 2\n', K2.grad)
print('method 3\n', K3.grad)
## these 3 print statements should give the same gradients :)

References

[1] Iain Murray, Differentiation of the Cholesky decomposition