By 苏剑林 | July 13, 2026
Recently, I came across the paper "The Key to Going Linear: Analysis-Driven Transformer Linearization", which linearizes Attention by directly applying approximate expansions to Softmax. I found it quite insightful and would like to record it here.
Some Notations
First, let's introduce the following notations:
\begin{gather}\newcommand{\logsumexp}{\mathop{\text{logsumexp}}}\newcommand{\softmax}{\mathop{\text{softmax}}}
\boldsymbol{x} = (x_1, x_2, \cdots, x_n) \in \mathbb{R}^n,\qquad\overline{\boldsymbol{x}} = \frac{1}{n}\sum_{i=1}^n x_i \\
\logsumexp(\boldsymbol{x}) = \log\sum_{i=1}^n e^{x_i} = \log n + \log \overline{e^{\boldsymbol{x}}} \\
\softmax(\boldsymbol{x}) = \frac{e^{\boldsymbol{x}}}{\sum_{i=1}^n e^{x_i}} = e^{\boldsymbol{x} - \logsumexp(\boldsymbol{x})} \\
\end{gather}
We assume that functional operations on vectors are element-wise (Hadamard sense), for example:
\begin{equation}e^{\boldsymbol{x}} = (e^{x_1},e^{x_2},\cdots,e^{x_n}),\qquad\boldsymbol{x}^2 = (x_1^2,x_2^2,\cdots,x_n^2)\end{equation}
In particular, one must carefully distinguish between $\overline{e^{\boldsymbol{x}}}$ and $e^{\overline{\boldsymbol{x}}}$, as well as $\overline{\boldsymbol{x}^2}$ and $\overline{\boldsymbol{x}}^2$: the former applies the function to each component first and then averages, while the latter averages first and then applies the function; the two are generally not equal.
In addition to the definitions above, $\logsumexp$ and $\softmax$ can also be linked via the gradient:
\begin{equation}\softmax(\boldsymbol{x})=\nabla_{\boldsymbol{x}} \logsumexp(\boldsymbol{x})\label{eq:grad-lse}\end{equation}
Basic Expansion
Next, we expand $\logsumexp(\boldsymbol{x})$. Introducing $\boldsymbol{y} = \boldsymbol{x} - \overline{\boldsymbol{x}}$, we have:
\begin{equation}\begin{aligned}
\logsumexp(\boldsymbol{x}) =&\, \log n + \overline{\boldsymbol{x}} + \log \overline{e^{\boldsymbol{y}}} \\
=&\, \log n + \overline{\boldsymbol{x}} + \log \left(\,\overline{1 + \boldsymbol{y} + \frac{\boldsymbol{y}^2}{2} + \frac{\boldsymbol{y}^3}{6} +\frac{\boldsymbol{y}^4}{24} + \cdots}\,\right) \\
=&\, \log n + \overline{\boldsymbol{x}} + \log \left(1 + \frac{\overline{\boldsymbol{y}^2}}{2} + \frac{\overline{\boldsymbol{y}^3}}{6} + \frac{\overline{\boldsymbol{y}^4}}{24} + \cdots\right) \\
=&\, \log n + \overline{\boldsymbol{x}} + \frac{\overline{\boldsymbol{y}^2}}{2} + \frac{\overline{\boldsymbol{y}^3}}{6} + \left(\,\frac{\overline{\boldsymbol{y}^4}}{24} - \frac{(\,\overline{\boldsymbol{y}^2}\,)^2}{8}\,\right) + \cdots
\end{aligned}\label{eq:logsumexp-series}\end{equation}
The last step utilizes $\log(1+t) = t - t^2/2 + t^3/3 - \cdots$. If necessary, one can continue the expansion further. If the reader is not familiar with these derivations, they can be completed by an AI like Kimi.
Introducing the bias $\overline{\boldsymbol{x}}$ is a small trick to simplify the form. Without it, an equivalent result can be obtained, but it would involve explicitly expanding each term of $\boldsymbol{y}$:
\begin{equation}\logsumexp(\boldsymbol{x})= \log n + \overline{\boldsymbol{x}} + \underbrace{\left(\frac{\overline{\boldsymbol{x}^2}}{2}-\frac{\overline{\boldsymbol{x}}^{\,2}}{2}\right)}_{\overline{\boldsymbol{y}^2}/2} + \underbrace{\left(\frac{\overline{\boldsymbol{x}^3}}{6}-\frac{\overline{\boldsymbol{x}}\,\overline{\boldsymbol{x}^2}}{2}+\frac{\overline{\boldsymbol{x}}^{\,3}}{3}\right)}_{\overline{\boldsymbol{y}^3}/6} + \cdots\end{equation}
Calculating the Gradient
For $\softmax(\boldsymbol{x})$, we can directly use the identity \eqref{eq:grad-lse} and take the derivative of both sides of equation \eqref{eq:logsumexp-series} to obtain the expansion of $\softmax(\boldsymbol{x})$. To this end, we use:
\begin{equation}\nabla_{\boldsymbol{x}} \overline{\boldsymbol{x}} = \frac{\boldsymbol{1}}{n},\qquad\nabla_{\boldsymbol{x}} \overline{\boldsymbol{y}^m} = \frac{m}{n} \left(\boldsymbol{y}^{m-1} - \overline{\boldsymbol{y}^{m-1}}\right) \end{equation}
resulting in:
\begin{equation}\softmax(\boldsymbol{x}) = \frac{1}{n}\left(1 + \boldsymbol{y} + \left(\frac{\boldsymbol{y}^2}{2} - \frac{\overline{\boldsymbol{y}^2}}{2}\right) + \left(\frac{\boldsymbol{y}^3}{6} - \frac{\overline{\boldsymbol{y}^3}}{6} - \frac{\overline{\boldsymbol{y}^2}\, \boldsymbol{y}}{2}\right) + \cdots\right)\end{equation}
Note that if we truncate at a finite number of terms, the right side still maintains the property that all components sum to 1 (which can be seen from the gradient form of $\nabla_{\boldsymbol{x}} \overline{\boldsymbol{y}^m}$), but it cannot guarantee the non-negativity of each component. This needs to be considered in practical applications (such as when taking logs to calculate entropy).
Besides taking the gradient, we can also expand using $\softmax(\boldsymbol{x}) = e^{\boldsymbol{x} - \logsumexp(\boldsymbol{x})} = e^{\boldsymbol{y} - \logsumexp(\boldsymbol{y})}$. Combining this with equation \eqref{eq:logsumexp-series} gives:
\begin{equation}\begin{aligned}
e^{\boldsymbol{y}} e^{-\logsumexp(\boldsymbol{y})} =&\, \left(1 + \boldsymbol{y} + \frac{\boldsymbol{y}^2}{2} + \frac{\boldsymbol{y}^3}{6} + \cdots\right)\exp\left(-\log n - \frac{\overline{\boldsymbol{y}^2}}{2} - \frac{\overline{\boldsymbol{y}^3}}{6} - \cdots\right) \\
=&\, \frac{1}{n}\left(1 + \boldsymbol{y} + \left(\frac{\boldsymbol{y}^2}{2} - \frac{\overline{\boldsymbol{y}^2}}{2}\right) + \left(\frac{\boldsymbol{y}^3}{6} - \frac{\overline{\boldsymbol{y}^3}}{6} - \frac{\overline{\boldsymbol{y}^2}\, \boldsymbol{y}}{2}\right) + \cdots\right)
\end{aligned}\end{equation}
Sparse Attention
So, what are the applications of these two expansions? First is the expansion of $\logsumexp(\boldsymbol{x})$, which is related to block scoring in Block Sparse Attention like MoBA. We know that Full Attention scores each token by $e^{\boldsymbol{q}\cdot \boldsymbol{k}}$. For a certain block $\mathcal{B}$, its score can be reasonably defined as the sum of scores of every token within it, which is equivalent to:
\begin{equation}s_\mathcal{B} = \log \sum_{t\in\mathcal{B}} e^{\boldsymbol{q}\cdot \boldsymbol{k}_t}\end{equation}
Its first-order approximation is $\log |\mathcal{B}| + \boldsymbol{q}\cdot\overline{\boldsymbol{k}}$, where $\overline{\boldsymbol{k}}=\frac{1}{|\mathcal{B}|}\sum_{t\in\mathcal{B}} \boldsymbol{k}_t$. This exactly corresponds to MoBA's empirical practice of using AvgPooling as the Landmark vector within a block. If the first-order approximation feels too coarse, higher-order corrections can be considered. According to equation \eqref{eq:logsumexp-series}, the second-order term is:
\begin{equation}\frac{1}{2|\mathcal{B}|}\sum_{t\in\mathcal{B}} (\boldsymbol{q}\cdot (\boldsymbol{k}_t - \overline{\boldsymbol{k}}))^2 = \frac{1}{2}\boldsymbol{q}^{\top}\underbrace{\left(\frac{1}{|\mathcal{B}|}\sum_{t\in\mathcal{B}}(\boldsymbol{k}_t - \overline{\boldsymbol{k}})(\boldsymbol{k}_t - \overline{\boldsymbol{k}})^{\top}\right)}_{\boldsymbol{\Sigma}}\boldsymbol{q}\end{equation}
That is, the second-order approximation is $\log |\mathcal{B}| + \boldsymbol{q}\cdot\overline{\boldsymbol{k}} + \boldsymbol{q}^{\top}\boldsymbol{\Sigma}\boldsymbol{q}/2$, where $\boldsymbol{\Sigma}$ happens to be the covariance matrix of $\boldsymbol{k}_t$ within the block. Furthermore, we can consider a diagonal approximation to reduce computational costs. This idea of higher-order correction is essentially the same as SPLA. Subsequent related work includes HiLS, which additionally learns a non-uniformly weighted average center vector to replace $\overline{\boldsymbol{k}}$.
Linear Attention
As for the application of the $\softmax(\boldsymbol{x})$ expansion, it is naturally the linearization of Attention. In "Path to Transformer Upgrade: 5. Linear Attention as Infinite Dimensions", we summarized three approaches to linearization, all of which approximate $e^{\boldsymbol{q}\cdot\boldsymbol{k}}$. However, since one still needs to normalize after approximating $e^{\boldsymbol{q}\cdot\boldsymbol{k}}$, it is better to directly approximate the normalized $\softmax(\boldsymbol{x})$.
Let $\boldsymbol{x}=(\boldsymbol{q}\cdot \boldsymbol{k}_1,\cdots,\boldsymbol{q}\cdot \boldsymbol{k}_n)$, then we have the first-order approximation:
\begin{equation}\softmax(\boldsymbol{x})_i \approx \frac{1}{n} + \frac{1}{n}\boldsymbol{q}\cdot (\boldsymbol{k}_i - \overline{\boldsymbol{k}})\end{equation}
This is precisely the scheme discussed in the paper mentioned at the beginning. It is quite intuitive: the average vector $\overline{\boldsymbol{k}}$ provides a baseline for the relative magnitude of attention. Tokens with similarity greater than $\overline{\boldsymbol{k}}$ should receive increased attention, and vice versa. To improve precision, Based expanded $e^{\boldsymbol{q}\cdot\boldsymbol{k}}$ to the second order (which just happens to guarantee non-negativity, refer to here), but from our perspective, considering the second-order approximation of $\softmax(\boldsymbol{x})$ is more scientific. The second-order term is:
\begin{equation}\frac{1}{2n}\left[(\boldsymbol{q}\cdot(\boldsymbol{k}_i-\overline{\boldsymbol{k}}))^2 - \frac{1}{n}\sum_{j=1}^n(\boldsymbol{q}\cdot(\boldsymbol{k}_j-\overline{\boldsymbol{k}}))^2\right]\end{equation}
Where $(\boldsymbol{q}\cdot(\boldsymbol{k}_i-\overline{\boldsymbol{k}}))^2$ can be turned into an inner product by introducing an outer product:
\begin{equation}(\boldsymbol{q}\cdot(\boldsymbol{k}_i-\overline{\boldsymbol{k}}))^2 = \boldsymbol{q}^{\top}(\boldsymbol{k}_i-\overline{\boldsymbol{k}})(\boldsymbol{k}_i-\overline{\boldsymbol{k}})^{\top}\boldsymbol{q} = \left\langle \boldsymbol{q}\boldsymbol{q}^{\top},\, (\boldsymbol{k}_i-\overline{\boldsymbol{k}})(\boldsymbol{k}_i-\overline{\boldsymbol{k}})^{\top} \right\rangle_F\end{equation}
So, similar to Based, truncating the Softmax of Softmax Attention to its second-order approximation also results in linear attention.
Summary
In this article, we derived the Taylor expansions for LogSumExp and Softmax, respectively, and discussed two of their potential applications.