Element-wise Taylor Attention Mechanisms

• Element-wise Taylor Attention is a method that approximates the exponential kernel using a truncated Taylor series to balance computational efficiency and fidelity.
• It restructures self-attention computations into element- or channel-wise expansions, enabling sub-quadratic, linear, or even constant complexity during training and inference.
• Empirical results, such as those from TaylorShift, demonstrate significant memory savings and performance improvements on long-sequence tasks compared to standard Transformer models.

Element-wise Taylor Attention refers to a family of self-attention mechanisms that approximate the exponential kernel in softmax attention using Taylor series expansions applied element-wise or feature-wise, transforming the complexity profile and computational characteristics of attention in Transformer models. By systematically trading off approximation fidelity and computational structure, these methods enable sub-quadratic—sometimes linear or even constant—complexity in both training and inference, with tunable accuracy and direct preservation of token-to-token interactions (Nauen et al., 2024, Feng, 10 Jan 2025, Mercat, 2020, Heinsen et al., 30 Jan 2026).

1. Formal Definition and Taylor Expansion

In standard self-attention, the attention weights arise from exponentiating the dot products of queries and keys, i.e., $\exp(q^\top k/c)$. Element-wise Taylor Attention approximates $\exp(x)$ by a truncated Taylor series at each relevant point—either on entries of the attention score matrix, on channel-wise products, or generalized to higher-order symmetric polynomials. For instance, truncating at order $t$ yields

$\exp(x) \approx \sum_{n=0}^{t-1} \frac{x^n}{n!}$

This approximation is substituted for $\exp(q^\top k)$ or, in element-wise schemes, for $\exp(q_{ic}k_{jc})$ for each feature channel $c$ separately (Nauen et al., 2024, Feng, 10 Jan 2025, Mercat, 2020, Heinsen et al., 30 Jan 2026).

In the TaylorShift formulation, the attention output becomes

$$Y = \mathrm{normalize}(1 + X + \tfrac12 X^{\odot 2})V$$

where $X = d^{-1/2} Q K^T$ and the normalization is performed row-wise (Nauen et al., 2024). Similarly, in channel-wise policies, attention is computed by applying Taylor expansions to per-channel similarities and aggregating outputs dimension-wise (Feng, 10 Jan 2025).

2. Algorithmic Structure and Linearization Techniques

The primary technical advance of element-wise Taylor attention is the restructuring of the attention computation to enable efficient (i.e., linear or constant-cost) algorithms. For second-order Taylor expansions, quadratic terms are lifted using tensor product unrolling operators, e.g., for TaylorShift: $$(X_{ij})^2 = (Q_i K_j^T)^2 = Q_i^{\otimes 2} \cdot (K_j^{\otimes 2})^T$$ These higher-order terms can be computed efficiently without explicit $N \times N$ matrices by associating them to higher-dimensional but sequence-length-independent features ("quadratic features") and using accumulators over these features (Nauen et al., 2024, Mercat, 2020, Heinsen et al., 30 Jan 2026).

In the EA-series (Feng, 10 Jan 2025), similarity is first measured by the negative squared Euclidean distance,

$O_{ijc} = -(q_{ic} - k_{jc})^2$

and the crucial exponential term $\exp(2 q_{ic} k_{jc})$ is Taylor-approximated, allowing the weight computation and final output to be organized as sums over precomputed statistics, with channel and polynomial order separations, enabling both batch and incremental (RNN-style) evaluation.

Further, (Heinsen et al., 30 Jan 2026) introduces symmetry-aware tensor decomposition, mapping monomials of $(q^\top k)^p$ into symmetric polynomial feature bases, minimizing redundancy and achieving constant per-token inference cost for any fixed Taylor order $P$.

3. Computational Complexity and Crossover Analysis

Element-wise Taylor attention mechanisms enable varying computational scaling, depending on the order of expansion and chosen decomposition. For example, in TaylorShift (Nauen et al., 2024), the direct quadratic implementation incurs

$\mathrm{FLOPs} = 4N^2 d + 6N^2$

while the efficient linear implementation is

$\mathrm{FLOPs} = N (4 d^3 + 10 d^2 + 8d + 3)$

yielding a critical sequence length (speed crossover) $N_0 = d^2 + d + \frac12$ and memory crossover point $N_1 \lesssim \frac12 d^2 + 2d + \frac12$. Empirical measurements show TaylorShift becomes more efficient in memory for $N\gtrsim800$ (when $d=32$), and faster for $N\gtrsim1700$ (Nauen et al., 2024).

For feature-wise element-wise expansion (EA-series), both training and inference costs are $O(tLD)$ (for sequence length $L$, feature size $D$, and order $t$), compared to $O(L^2D)$ for standard attention (Feng, 10 Jan 2025). Inference cost is constant per token—$O(tD)$—versus linear in context for standard models.

Symmetry-aware schemes, by packing Taylor monomials using symmetric polynomial kernels, reduce memory and per-token computation to constants independent of context length (Heinsen et al., 30 Jan 2026):

Approach	Training Complexity	Inference Complexity	Memory Usage
Standard Softmax Attention	$O(L^2 D)$	$O(L D)$ per token	$O(L^2)$
TaylorShift (efficient)	$O(N d^3)$	$O(d^3)$ per token	$O(N d^2)$
EA-series (element-wise)	$O(t L D)$	$O(t D)$ per token	$O(t D)$
Symmetry-aware Taylor (const.)	$O(M)$ per token	$O(M)$ per token	$O(M)$

where $M = \binom{d_K + P -1}{P-1}$ for Taylor order $P$.

4. Theoretical Properties and Approximation Fidelity

Approximation errors in element-wise Taylor attention are directly governed by the Taylor remainder. For second-order truncation,

$| e^x - (1 + x + x^2/2) | \leq (e^{|x|}/6)|x|^3$

provided attention logits are numerically bounded (e.g., via scaling or normalization), the per-element error is guaranteed to be $\lesssim10^{-2}$, which translates into tightly controlled per-row softmax errors (Mercat, 2020, Nauen et al., 2024). Higher-order expansions rapidly reduce this error, approaching machine precision for modest $t$ or $P$.

A crucial property is the preservation of "spikiness"—the sharp, selective weighting of large similarities that distinguishes softmax attention's representational power. By keeping enough terms (e.g., $t\geq6$), Taylor polynomial attention is empirically shown to retain such spikiness, outperforming other linear kernel approximations, which tend to oversmooth the attention weights (Feng, 10 Jan 2025). Furthermore, these schemes can avoid the information "compression" seen in linear RNNs or SSMs, as each channel and polynomial order maintains its own accumulator, mitigating information washout in long sequences (Feng, 10 Jan 2025).

5. Empirical Performance and Applications

TaylorShift (Nauen et al., 2024) achieves up to 65% reduction in peak VRAM for $N=2000$, matches or slightly outperforms the standard Transformer on long-sequence classification across five benchmarks (CIFAR10 Pixel, IMDB Byte, Long ListOps, ImageNet-Tiny, ImageNet-Small), and consistently exceeds all tested linear/efficient Transformer baselines (Linformer, Performer, Reformer, Nyströmformer) on 4/5 tasks. For example, TaylorShift averaged 62.7% accuracy versus 62.2% for standard Transformers.

The EA-series (Feng, 10 Jan 2025) matches or exceeds full self-attention on time-series and multivariate classification tasks when sufficient Taylor order is used (EA-6). In both causal and non-causal applications, it provides substantial boosts in memory efficiency (enabling 2–3× longer context for a given GPU budget), higher throughput at scale, and batch-size agnostic inference latency.

Symmetry-aware Taylor attention achieves constant memory and compute for unbounded contexts, with empirical validation on synthetic data showing negligible error for Taylor order $P=4$—errors on par with float16 rounding. These approaches create the possibility of Transformer architectures with arbitrarily long contexts at negligible incremental resource cost (Heinsen et al., 30 Jan 2026).

6. Implementation Considerations and Limitations

Stable implementation of element-wise Taylor attention requires careful normalization, including scaling $Q$, $K$ to unit length, introducing learnable temperatures, and managing the coefficients of the Taylor expansion to prevent exploding or vanishing intermediate values (Nauen et al., 2024). Layer normalization and input scaling are critical to ensure the logits remain in a regime where the Taylor approximation is accurate (Mercat, 2020).

A trade-off in Taylor-based methods is the increase in feature (or accumulator) dimensionality, scaling polynomially in the Taylor order and base feature size, which can be mitigated by favoring more attention heads with lower feature dimension per head (Heinsen et al., 30 Jan 2026). Current limitations include additional engineering required to fully exploit these advances at scale (such as custom GPU kernels for fused prefix scan operations), and further studies of optimization and training dynamics in real-world tasks (Heinsen et al., 30 Jan 2026).

7. Connections and Implications

Element-wise Taylor attention subsumes and generalizes several lines of research into fast and efficient Transformer architectures. It draws direct connections to kernel-based attention, linear transformers, and polynomial kernel learning, but with controlled approximation error and explicit representational fidelity to conventional softmax attention. The methods avoid the intrinsic performance degradations found in classical linear/RNN/SSM approaches due to information dilution or loss of attention "spikiness" (Feng, 10 Jan 2025).

A plausible implication is that Taylor-based schemes will enable both the scaling up of context length (for long-context or streaming tasks) and the reduction of inference/serving costs for large Transformer deployments (Nauen et al., 2024, Heinsen et al., 30 Jan 2026). The symmetry-aware features facilitate architectures with an increased number of smaller heads, as the per-head cost is fixed and independent of sequence length (Heinsen et al., 30 Jan 2026).

Together, the family of element-wise Taylor attention mechanisms constitute a robust, theoretically justified, and empirically validated toolkit for shifting the self-attention bottleneck from quadratic to linear or constant regimes, while retaining the principal accuracy and representational benefits of the softmax attention kernel (Nauen et al., 2024, Feng, 10 Jan 2025, Mercat, 2020, Heinsen et al., 30 Jan 2026).