The Taylor Expansion of LogSumExp and Softmax

LogSumExp和Softmax的泰勒展开 苏剑林 · 2026-07-13 · Mathematics calculus approximation linear attention

I recently came across the paper The Key to Going Linear: Analysis-Driven Transformer Linearization, which linearizes attention by directly applying an approximate expansion to Softmax. I found it quite instructive, so I'm writing up some notes here.

Some Notation

First let's introduce the following notation:

\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}more

and let's agree that function operations on vectors are understood in the 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 should pay careful attention to the distinction between $\overline{e^{\boldsymbol{x}}}$ and $e^{\overline{\boldsymbol{x}}}$, and between $\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 — in general these are not equal.

Besides the definitions above, $\logsumexp$ and $\softmax$ can also be related to one another via gradients:

\begin{equation}\softmax(\boldsymbol{x})=\nabla_{\boldsymbol{x}} \logsumexp(\boldsymbol{x})\label{eq:grad-lse}\end{equation}

Basic Expansion

Next let's 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 uses $\log(1+t) = t - t^2/2 + t^3/3 - \cdots$, and if needed, one can continue expanding further. If the reader isn't familiar with this kind of manipulation, it can also be handed off to Kimi to work out.

Introducing the offset $\overline{\boldsymbol{x}}$ here is just a trick for simplifying the form; without it we can still get an equivalent result, except the terms of $\boldsymbol{y}$ would need to be expanded explicitly:

\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}

Taking a Gradient

For $\softmax(\boldsymbol{x})$, we can directly use the identity $\eqref{eq:grad-lse}$, differentiating both sides of equation $\eqref{eq:logsumexp-series}$ to obtain the expansion of $\softmax(\boldsymbol{x})$. To do this, 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}

to get

\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 to a finite number of terms, the right-hand side will still sum to 1 across all components (this can be seen directly from the gradient form of $\nabla_{\boldsymbol{x}} \overline{\boldsymbol{y}^m}$), but there is no guarantee that each component remains non-negative — something worth keeping in mind in practical applications (for instance, when taking logarithms to compute entropy).

Besides taking gradients, we can also expand using $\softmax(\boldsymbol{x}) = e^{\boldsymbol{x} - \logsumexp(\boldsymbol{x})} = e^{\boldsymbol{y} - \logsumexp(\boldsymbol{y})}$, which combined 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 these two expansions useful for? First, consider $\logsumexp(\boldsymbol{x})$, which relates to the block scoring used in Block Sparse Attention methods like MoBA. As we know, Full Attention assigns a score to every token via $e^{\boldsymbol{q}\cdot \boldsymbol{k}}$; for a given block $\mathcal{B}$, it's natural to define its score as the sum of the scores of the tokens 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$ — and this corresponds exactly to MoBA's empirical practice of using average pooling as the landmark vector within a block. If the first-order approximation feels too crude, we can consider a higher-order correction. 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}

so 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}$ turns out to be exactly the covariance matrix of $\boldsymbol{k}_t$ within the block. We could further consider a diagonal approximation to reduce computational cost. This idea of higher-order correction is essentially the same as SPLA. A related follow-up work is HiLS, which additionally learns a non-uniformly-weighted center vector to replace $\overline{\boldsymbol{k}}$.

Linear Attention

As for the application of the $\softmax(\boldsymbol{x})$ expansion, it's naturally linear attention. In Transformer Upgrade Path: 5. Linear Attention as an Infinite-Dimensional Approximation, we summarized three approaches to linearization, all of which approximate $e^{\boldsymbol{q}\cdot\boldsymbol{k}}$. However, after approximating $e^{\boldsymbol{q}\cdot\boldsymbol{k}}$ one still needs to normalize it — so it makes more sense to directly approximate the already-normalized $\softmax(\boldsymbol{x})$.

Writing $\boldsymbol{x}=(\boldsymbol{q}\cdot \boldsymbol{k}_1,\cdots,\boldsymbol{q}\cdot \boldsymbol{k}_n)$, we get 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 exactly the scheme discussed in the paper mentioned at the start, and it's quite intuitive on its own: the average vector $\overline{\boldsymbol{k}}$ provides a baseline for the relative magnitude of attention — tokens whose similarity ratio $\overline{\boldsymbol{k}}$ exceeds this baseline should receive more attention, and vice versa. To improve accuracy, Based expands $e^{\boldsymbol{q}\cdot\boldsymbol{k}}$ to second order (which happens to guarantee non-negativity — see here). But from our point of view, it's more principled to consider the second-order approximation of $\softmax(\boldsymbol{x})$ instead; 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 in Softmax Attention to second order also yields a form of linear attention.

Summary

In this post we derived the Taylor expansions of LogSumExp and Softmax, respectively, and discussed two potential applications of each.