---
title: 'TaylorShift: Linear Self-Attention via Taylor Expansion'
url: https://www.emergentmind.com/topics/taylorshift
type: topic
---

# TaylorShift: Linear Self-Attention via Taylor Expansion

TaylorShift is a self-attention mechanism for Transformers that achieves linear time and memory complexity in sequence length while retaining full token-to-token interactions. Based on a low-order Taylor expansion of the exponential in Softmax, TaylorShift replaces standard Softmax-based attention with a polynomial-based normalization. This reformulation avoids the quadratic bottleneck of classic attention and provides precise crossover analyses, making it practically attractive for long-sequence tasks and high-resolution vision applications [2403.02920][2411.10231].

## 1. Mathematical Construction of TaylorShift

The canonical attention mechanism in Transformers operates via Softmax normalization:
\[
\text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
\]
TaylorShift replaces the exponential with its $k$-th order Taylor expansion and normalizes the sum:
\[
\operatorname{T\!-\!SM}^{(k)}(x) := \text{normalize}\left(\sum_{n=0}^k \frac{x^{\odot n}}{n!}\right)
\]
where $x^{\odot n}$ is the elementwise $n$-th power, and normalization is division by the $\ell_1$ norm. For $k=2$ (default), this yields:
\[
1 + x + \frac{1}{2}x^{\odot2}
\]
Thus, attention becomes:
\[
Y = \operatorname{T\!-\!SM}^{(k)}(d^{-\frac12} QK^\top) V
\]
where $Q,K,V \in \mathbb{R}^{N\times d}$, $N$ is sequence length, and $d$ is embedding dimension. The 2nd-order Taylor expansion provides guaranteed positivity for even $k$ and preserves the distributional properties post-normalization [2403.02920][2411.10231].

## 2. Efficient Algorithmic Formulation

A direct materialization of $\operatorname{T\!-\!SM}$ scales quadratically, $\mathcal{O}(N^2 d)$. However, TaylorShift leverages algebraic factorization to reduce this to linear in $N$. The core identity is:
\[
(QK^\top)^{\odot2} = (Q^{\boxtimes2})(K^{\boxtimes2})^\top
\]
where $Q^{\boxtimes2}\in\mathbb{R}^{N\times d^2}$ denotes the flattened outer products of rows of $Q$. This structure allows the computation of the attention numerator and denominator via:

- Constant: sum over $V$
- Linear: $Q(K^\top V)$
- Quadratic: $\frac12 (Q^{\boxtimes2})[(K^{\boxtimes2})^\top V]$

A high-level pseudocode of the efficient TaylorShift workflow is as follows [2403.02920]:
```python
# Inputs: Q, K, V ∈ ℝ^{N×d}; order k=2
# Pre-normalization
for each i: 
    q_i ← τ q_i / ||q_i||₂
    k_i ← k_i / ||k_i||₂  # τ is a learned temperature
Q, K ← d^{1/4} ⋅ Q, K
V ← (1/N)(sqrt(d/N) ⋅ 1ₙ || V)  # append 1s for denominator

# Quadratic (efficient) TaylorShift
A_mod ← (K⊠K)ᵀ V                # ⊠: pairwise outer product flattening
Yhat ← ½·(Q⊠Q) A_mod + d^{½}·Q(KᵀV) + d·sum_i V_{i:}
# Split Yhat into numerator and denominator, final output:
Y ← numerator ⊘ denominator
```
Each main operation scales as $\mathcal{O}(N d^{3})$ or less. The interface and normalization remain consistent with quadratic TaylorShift, allowing seamless switching.

## 3. Computational Complexity and Crossover Regimes

The TaylorShift mechanism rigorously characterizes its complexity trade-offs:

- Direct (quadratic) TaylorShift:
  - FLOPs: $4N^2d + 6N^2$
  - Memory: $2N^2 + dN$
- Efficient (linear) TaylorShift:
  - FLOPs: $N(4d^3 + 10d^2 + 8d + 3)$
  - Memory: $d^2(d+1) + 2dN + (d+1)N + d^2N$

Key crossover points:
- Compute $\text{ops}_{\text{triv}} = \text{ops}_{\text{eff}}$ yields $N_0 = d^2 + d + \frac12$
- Memory crossover: $N_1 = \frac14[d^2 + 2d + 1 + \sqrt{d^4 + 12d^3 + 14d^2 + 4d + 1}] \le \frac12 d^2 + 2d + \frac12$

Empirically, memory savings emerge at $N \approx 800$ tokens, speedup at $N \approx 1700$ tokens for typical embedding dimensions ($d=64 \Rightarrow N_0 \approx 4161$, $N_1 \approx 2174$) [2403.02920][2411.10231].

## 4. Empirical Evaluation and Benchmarks

TaylorShift has been validated in both natural language and vision domains:

- **Classification (Long-Sequence) Benchmarks**: On CIFAR-pixel, IMDB-byte, Long ListOps, ImageNet-Tiny, and ImageNet-Small, TaylorShift matches or slightly exceeds vanilla attention and surpasses prior linear attention algorithms on most tasks.

| Model              | CIFAR | IMDB | ListOps | ImageNet Tiny | ImageNet Small | Average |
|---------------------|-------|------|---------|--------------|----------------|---------|
| Linformer           | 29.2  | 58.1 | ---     | 64.3         | 76.3           | 57.0    |
| Performer*          | 34.2  | 65.6 | 35.4    | 62.0         | 67.1           | 52.9    |
| Reformer            | 44.8  | 63.9 | 47.6    | 73.6         | 76.2           | 61.2    |
| Nyströmformer       | 49.4  | 65.6 | 44.5    | 75.0         | 78.3           | 62.6    |
| Transformer         | 44.7  | 65.8 | 46.0    | 75.6         | 79.1           | 62.2    |
| TaylorShift (Ours)  | 47.6  | 66.0 | 45.6    | 75.0         | 79.3           | 62.7    |

- **Vision Super-Resolution (SR)**: TaylorShift enables pixel-level $1\times1$ attention windows in SwinIR, reducing VRAM by up to 60%. For $48\times48$ windows, VRAM is reduced from 78.50GB (SwinIR) to 49.10GB (TaylorSwinIR) [2411.10231]. Quality as measured by PSNR and SSIM remains at or above SOTA baselines:

| Method              | Set5 PSNR/SSIM | Urban100 PSNR/SSIM | VRAM (GB) | Reduction |
|---------------------|:--------------:|:------------------:|:----------:|:---------:|
| SwinIR              | 38.35/0.9620   | 33.49/0.9393       | 78.50     | 0%        |
| TaylorSwinIR (ours) | 38.46/0.9627   | 33.71/0.9417       | 49.10     | 37%       |

Across five SR datasets, TaylorSwinIR consistently matches or surpasses SwinIR on both PSNR and SSIM, demonstrating the feasibility of full-range attention at linear cost.

## 5. Integration into Transformer Architectures

TaylorShift is a drop-in replacement for Softmax-based attention in both encoder-only and encoder-decoder Transformers, with seamless sharing of normalization and interface. In SwinIR-like architectures, TaylorShift allows the use of 1×1 patch embedding—previously impractical due to quadratic scaling—for pixel-level attention, substantially increasing the contextual receptive field. It naturally integrates with existing windowing, local-global, and positional encoding schemes without modification [2403.02920][2411.10231].

The attention block with TaylorShift executes as:
- For $N\leq N_0$ (below crossover): use direct quadratic TaylorShift.
- For $N \gg N_0$: use efficient tensorized variant.

## 6. Practical Considerations and Stability

- **Taylor order**: $k=2$ is empirically optimal, balancing computational overhead with attention expressivity and regularization.
- **Input normalization**: Per-token $\ell_2$ normalization with a learnable temperature $\tau$ is essential for numerical stability; improper scaling leads to overflow or non-convergence.
- **Multi-head scaling**: Efficient TaylorShift scales with the number of heads $h$ as $\propto N[4(d_{\text{embed}}/h)^{3} + \cdots + 3h]$, permitting typical head configurations up to $h=32$.
- **Switching**: Selection between direct and efficient mode is automatic based on sequence length; both methods implement identical mathematical forms and output interfaces, ensuring compatibility.
- **Error bounds**: The polynomial expansion yields provable bounds on the approximation error of Softmax, and the attention weights remain strictly positive and normalized.

## 7. Significance and Limitations

TaylorShift advances the landscape of efficient Transformer attention by enabling full-range, dense interactions at linear sequence scaling. Unlike sparse or kernel-based methods, it does not compromise token-to-token connectivity or introduce stateful recurrence, preserving model expressiveness in long-sequence regimes. Its use in pixel-level SR demonstrates capability in high-resolution, memory-constrained tasks, setting new practical baselines for efficiency and accuracy [2403.02920][2411.10231].

A plausible implication is that, for sufficiently large $N$, TaylorShift offers the best-available trade-off between resource consumption and modeling power without loss of accuracy. Remaining limitations include the bottleneck shifting to embedding dimension $d$, especially in regimes where $d^3$ is large, and the necessity of careful normalization for stable training dynamics. For very short sequences, TaylorShift reverts to quadratic scaling, preserving compatibility with traditional Bottleneck attention settings.

**References**:  
Nauen, J., et al. "TaylorShift: Shifting the Complexity of Self-Attention from Squared to Linear (and Back) using Taylor-Softmax" [2403.02920].  
Chang, L., et al. "A Low-Resolution Image is Worth 1x1 Words: Enabling Fine Image Super-Resolution with Transformers and TaylorShift" [2411.10231].

Source: https://www.emergentmind.com/topics/taylorshift