Efficient Inversion of "Diagonal + Low-Rank" Triangular Matrices

“对角+低秩”三角阵的高效求逆方法 苏剑林 · 2025-07-01 · Mathematics computation matrix RNN attention

From the article A Brief History of Linear Attention: Imitation, Innovation, and Feedback we can see that DeltaNet and the linear attention models that followed it are essentially all connected to the inverse of the matrix $(\boldsymbol{I} + \boldsymbol{K}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-)^{-1}$. This post is devoted specifically to computing the inverse of this kind of triangular matrix with a "diagonal + low-rank" structure.

Basic Result

Let's define the problem in general terms as follows:

Given a matrix $\boldsymbol{Q},\boldsymbol{K}\in\mathbb{R}^{n\times d}$ and a diagonal matrix $\boldsymbol{\Lambda}\in\mathbb{R}^{n\times n}$ satisfying $n\gg d$, define
\begin{equation}\boldsymbol{T} = \boldsymbol{\Lambda} + \boldsymbol{Q}\boldsymbol{K}^{\top}\odot\boldsymbol{M}^-\end{equation}
where $\boldsymbol{M}^-=\boldsymbol{M} - \boldsymbol{I}$, and the matrix $\boldsymbol{M}$ is defined as
\begin{equation}M_{i,j} = \left\{\begin{aligned} &1, &i \geq j \\ &0, &i < j\end{aligned}\right.\end{equation}
We now want to find the inverse matrix $\boldsymbol{T}^{-1}$, and prove that its complexity is $\mathcal{O}(n^2)$.

First, if there were no lower-triangular constraint on $\odot\boldsymbol{M}^-$, the problem could be solved directly via the "Woodbury identity":

\begin{equation}(\boldsymbol{\Lambda} + \boldsymbol{Q}\boldsymbol{K}^{\top})^{-1} = \boldsymbol{\Lambda}^{-1} - \boldsymbol{\Lambda}^{-1} \boldsymbol{Q}(\boldsymbol{I} + \boldsymbol{K}^{\top}\boldsymbol{\Lambda}^{-1}\boldsymbol{Q})^{-1}\boldsymbol{K}^{\top}\boldsymbol{\Lambda}^{-1}\end{equation}

It's easy to verify that the computational complexity of the right-hand side is $\mathcal{O}(n^2)$. However, once we impose $\odot\boldsymbol{M}^-$, $\boldsymbol{T}$ itself no longer has the "diagonal + low-rank" structure, so the identity can't be applied directly anymore. To handle this triangular-matrix feature, a natural idea is recursion, since we have the block-matrix identity

\begin{equation}\begin{bmatrix}\boldsymbol{A} & \boldsymbol{0} \\ \boldsymbol{C} & \boldsymbol{B}\end{bmatrix}^{-1} = \begin{bmatrix}\boldsymbol{A}^{-1} & \boldsymbol{0} \\ -\boldsymbol{B}^{-1}\boldsymbol{C}\boldsymbol{A}^{-1} & \boldsymbol{B}^{-1}\end{bmatrix}\end{equation}

This allows us to turn $\boldsymbol{T}^{-1}$ into a recursive form (convention: in the absence of parentheses, slicing has the highest precedence)

\begin{equation}\boldsymbol{T}_{[:l+1,:l+1]}^{-1} = \begin{bmatrix}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{0} \\ -\boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\boldsymbol{T}_{[l:l+1,:l]}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\end{bmatrix}\end{equation}

The main computation here is $\boldsymbol{T}_{[l:l+1,:l]}\boldsymbol{T}_{[:l,:l]}^{-1}$, which is a product of a $1\times l$ matrix and a $l\times l$ matrix, with complexity $\mathcal{O}(\mathcal{l^2})$. That is, the complexity of each iteration step grows quadratically, so the total complexity is $\mathcal{O}(n^3)$.

Low-Rank Structure

Of course, this is because we haven't yet exploited the low-rank structure of $\boldsymbol{T}$ (i.e., before $\odot\boldsymbol{M}^-$). Let's now make use of it. This gives us $\boldsymbol{T}_{[l:l+1,:l]} = \boldsymbol{Q}_{[l:l+1]}\boldsymbol{K}_{[:l]}^{\top}$, and substituting into the equation above yields:

\begin{equation}\boldsymbol{T}_{[:l+1,:l+1]}^{-1} = \begin{bmatrix}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{0} \\ -\boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\boldsymbol{Q}_{[l:l+1]}\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\end{bmatrix}\end{equation}

Notice that $\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1}\in\mathbb{R}^{d\times l}$; if we can use it as the recursive variable, then the complexity of each iteration step becomes only $\mathcal{O}(l)$, and the total complexity can successfully be brought down to $\mathcal{O}(n^2)$. Following this idea, we obtain

\begin{equation}\begin{aligned} \boldsymbol{K}_{[:l+1]}^{\top}\boldsymbol{T}_{[:l+1,:l+1]}^{-1} =&\, \begin{bmatrix}\boldsymbol{K}_{[:l]}^{\top} & \boldsymbol{K}_{[l:l+1]}^{\top}\end{bmatrix}\begin{bmatrix}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{0} \\ -\boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\boldsymbol{Q}_{[l:l+1]}\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\end{bmatrix} \\[6pt] =&\, \begin{bmatrix}\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{0}\end{bmatrix} + \boldsymbol{K}_{[l:l+1]}^{\top}\underbrace{\begin{bmatrix}-\boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\boldsymbol{Q}_{[l:l+1]}\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\end{bmatrix}}_{\text{actually is}(\boldsymbol{T}^{-1})_{[l:l+1,:l+1]}}\end{aligned}\end{equation}

We can see that this recursive process also doesn't involve any operation with $\mathcal{O}(l^2)$, so the approach is feasible — we just need to introduce a new variable to cache $\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1}$. If we replace $l+1$ with $l+c$, we get a chunk-wise recursive form.

Test code is as follows:

import numpy as np

n, d, c = 1000, 100, 200
Q = np.random.randn(n, d) / d**0.5
K = np.random.randn(n, d) / d**0.5
T = np.tril(Q @ K.T, -1) + np.eye(n)

Y, Z = np.zeros((n, n)), np.zeros((d, n))
for l in range(0, n, c):
    Y[l:l + c, l:l + c] = np.linalg.inv(T[l:l + c, l:l + c])
    Y[l:l + c, :l] = - Y[l:l + c, l:l + c] @ Q[l:l + c] @ Z[:, :l]
    Z[:, :l + c] += K[l:l + c].T @ Y[l:l + c, :l + c]

np.allclose(Y @ T, np.eye(n))

Computing the Product

Using the same idea, we can also prove the following:

For any matrix $\boldsymbol{V}\in\mathbb{R}^{n\times d}$, computing $\boldsymbol{T}^{-1}\boldsymbol{V}$ only requires complexity $\mathcal{O}(n)$.

The proof only requires a slight modification of the previous derivation. First, we have

\begin{equation}\begin{aligned} (\boldsymbol{T}^{-1}\boldsymbol{V})_{[:l+1]} =&\, \boldsymbol{T}_{[:l+1,:l+1]}^{-1}\boldsymbol{V}_{[:l+1]} \\[6pt] =&\, \begin{bmatrix}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{0} \\ -\boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\boldsymbol{Q}_{[l:l+1]}\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1} & \boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\end{bmatrix}\begin{bmatrix}\boldsymbol{V}_{[:l]} \\ \boldsymbol{V}_{[l:l+1]}\end{bmatrix} \\[6pt] =&\, \begin{bmatrix}\boldsymbol{T}_{[:l,:l]}^{-1}\boldsymbol{V}_{[:l]} \\ -\boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\boldsymbol{Q}_{[l:l+1]}\boldsymbol{K}_{[:l]}^{\top}\boldsymbol{T}_{[:l,:l]}^{-1}\boldsymbol{V}_{[:l]} + \boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}\boldsymbol{V}_{[l:l+1]}\end{bmatrix} \\[6pt] =&\, \begin{bmatrix}(\boldsymbol{T}^{-1}\boldsymbol{V})_{[:l]} \\ \boldsymbol{T}_{[l:l+1,l:l+1]}^{-1}(\boldsymbol{V}_{[l:l+1]} - \boldsymbol{Q}_{[l:l+1]}\boldsymbol{K}_{[:l]}^{\top}(\boldsymbol{T}^{-1}\boldsymbol{V})_{[:l]})\end{bmatrix} \end{aligned}\end{equation}

and then

\begin{equation}\begin{aligned} \boldsymbol{K}_{[:l+1]}^{\top}(\boldsymbol{T}^{-1}\boldsymbol{V})_{[:l+1]} =&\, \begin{bmatrix}\boldsymbol{K}_{[:l]}^{\top} & \boldsymbol{K}_{[l:l+1]}^{\top}\end{bmatrix}\begin{bmatrix}(\boldsymbol{T}^{-1}\boldsymbol{V})_{[:l]} \\ (\boldsymbol{T}^{-1}\boldsymbol{V})_{[l:l+1]} \end{bmatrix} \\[8pt] =&\,\boldsymbol{K}_{[:l]}^{\top}(\boldsymbol{T}^{-1}\boldsymbol{V})_{[:l]} + \boldsymbol{K}_{[l:l+1]}^{\top}(\boldsymbol{T}^{-1}\boldsymbol{V})_{[l:l+1]} \end{aligned}\end{equation}

Therefore, as long as we cache $\boldsymbol{K}_{[:l]}^{\top}(\boldsymbol{T}^{-1}\boldsymbol{V})_{[:l]}\in\mathbb{R}^{d\times d}$, the computational cost of each step becomes independent of $l$, so the total complexity is $\mathcal{O}(n)$. Likewise, replacing $l+1$ with $l+c$ gives us the chunk-wise form.

Test code is as follows:

import numpy as np

n, d, c = 1000, 100, 200
Q = np.random.randn(n, d) / d**0.5
K = np.random.randn(n, d) / d**0.5
V = np.random.randn(n, d) / d**0.5
T = np.tril(Q @ K.T, -1) + np.eye(n)

Y, Z = np.zeros((n, d)), np.zeros((d, d))
for l in range(0, n, c):
    X = np.linalg.inv(T[l:l + c, l:l + c])
    Y[l:l + c] = X @ (V[l:l + c] - Q[l:l + c] @ Z)
    Z += K[l:l + c].T @ Y[l:l + c]

np.allclose(T @ Y, V)

Summary

This post discussed the problem of inverting triangular matrices with a "diagonal + low-rank" structure, a type of matrix that commonly appears in modern linear attention models.