Recurrent Softmax Attention

• Recurrent softmax attention is a reformulation of standard softmax attention that integrates recurrence to aggregate and modulate information over time.
• It leverages Taylor series expansion and infinite-dimensional state representations to connect traditional softmax with recurrent neural networks and state-space models.
• Modern adaptations balance memory, computation, and accuracy via linear, gated, and higher-order variants to efficiently handle long sequences in large-scale architectures.

A recurrent form of softmax attention refers to the explicit or implicit recasting of the softmax-based attention mechanism in a recurrent or dynamical systems framework, enabling sequence models to aggregate, gate, or modulate information over time or positions using recurrences closely connected to those in classical RNNs. While softmax attention is canonically non-recurrent and quadratic in sequence length, recent works have sought to recast it in recurrent form to analyze expressiveness, efficiency, interpretability, and approximation by linear or gated variants. These perspectives are critical for understanding the trade-off between memory, computation, and accuracy in large-scale neural architectures.

1. Canonical Softmax Attention and Computational Limitations

Let $H \in \mathbb{R}^{n \times k}$ denote a matrix whose $t$-th row is the hidden state $\mathbf{h}_t$ of a processed sequence. Given a query vector $q \in \mathbb{R}^k$, classical softmax attention is computed as

$R(D, Q) = H^T \cdot \text{softmax}(H q).$

Here, $\text{softmax}(H q)$ maps the set of inner products between $q$ and all hidden states to a dense probability vector, which is used to reweight the states. This mechanism, while effective for context modeling and “soft” information retrieval, imposes significant computational and memory costs:

• Attention computation for one query scales as $O(nk^2)$.
• All $n$ hidden states must be retained in memory for weighting.

These properties restrict the use of softmax attention in very long sequences or applications with many queries per sequence, such as index-based retrieval or real-time large-scale search (Brébisson et al., 2016).

2. Recurrent Reformulations: Theory and Taylor Series Derivation

A key insight is that standard softmax attention can be rewritten as a recurrent process with an infinite-dimensional hidden state, as established through the Taylor expansion of the exponential kernel (Mongaras et al., 31 Jul 2025, Sieber et al., 24 May 2024).

Given a causal attention setup at timestep $t$: $$O_t = \frac{\sum_{s=1}^t \exp(Q_t \cdot K_s^T) V_s}{\sum_{n=1}^t \exp(Q_t \cdot K_n^T)}$$ Expanding $\exp(Q_t \cdot K_s^T)$ in a Taylor series: $$\exp(Q_t \cdot K_s^T) = \sum_{n=0}^\infty \frac{1}{n!} (Q_t \cdot K_s^T)^n = \sum_{n=0}^\infty \frac{1}{n!} (Q_t^{\otimes n}) (K_s^{\otimes n})^T$$ This allows reorganizing the attention computation as a sum over order-$n$ interactions, with each term in the sum corresponding to a hidden state $H_t^n = \sum_{s=1}^t (K_s^{\otimes n})^T V_s$: $$O_t = G_t \sum_{n=0}^\infty \frac{1}{n!} (Q_t^{\otimes n}) H_t^n$$ where $G_t$ is the inverse softmax normalization. This reveals that softmax attention is equivalent to an infinite ensemble of RNNs, each accumulating higher-order (Kronecker) powers of $K_s$, with memory cost scaling exponentially in $n$ (Mongaras et al., 31 Jul 2025). The first term ($n=1$) reproduces linear attention.

3. Dynamical Systems and State Expansion: Unification of Attention and Recurrence

The Dynamical Systems Framework (DSF) formalizes the equivalence between softmax attention, recurrent networks, and state-space models by expressing each as a linear time-varying system: $$h_k = \Lambda_k h_{k-1} + B_k u_k,\quad y_k = C_k h_k + D_k u_k$$ For standard softmax attention, the infinite expansion corresponds to an infinite-dimensional hidden state encoding all powers of the query and key vectors. In contrast, classical RNNs use $n=1$ (state of dimension $d$), and linear attention employs a $k\times k$ state (Sieber et al., 24 May 2024).

This framework reveals that increasing the state expansion parameter $n$ allows recurrent models to progressively approximate softmax attention. In the limit $n\to\infty$, one recovers the full expressiveness of the softmax kernel. Empirically, as $n$ increases, performance on sequence modeling tasks approaches that of softmax attention (Sieber et al., 24 May 2024).

4. Expressiveness, Universal Approximation, and Learning Dynamics

Recurrent softmax attention achieves much greater expressiveness than its linear or kernelized approximations. The infinite sum in the Kronecker expansion enables modeling of all orders of multiplicative interactions, which underpin the universal approximation property of even shallow self-attention networks (Hu et al., 22 Apr 2025). Specifically, two-layer softmax-attention-only architectures can approximate any continuous sequence-to-sequence function, subsuming classical ReLU networks in approximation power.

Learning dynamics in multi-head softmax attention under in-context learning settings show the emergence of canonical circuit patterns: homogeneous diagonal key-query weights and last-entry-only output-value weights (He et al., 17 Mar 2025). These learned recurrent structures implement debiased or preconditioned gradient descent predictors, supporting advanced in-context learning and adaptation capabilities not matched by linear attention.

Furthermore, the training-induced approximation by softmax of kernel interpolants (via near-argmax behavior as the temperature grows) provides a natural bridge between attention and classical statistical estimators. The emergent recurrence and normalization adapt to sequence length dynamically, supporting length generalization (He et al., 17 Mar 2025).

5. Efficiency Trade-offs: Linear, Gated, and Approximate Forms

To address the memory and computational bottlenecks of softmax attention in the recurrent form, several approximations have been proposed:

• Linear Attention: Dropping the softmax, defining $C = H^T H$ yields constant-time, fixed-size query lookups but reduced accuracy due to loss of nonlinearity (Brébisson et al., 2016).
• Higher-Order Linearization: Taylor-expanding the softmax exponential and retaining only low-order terms (typically $n=1$ or $n=2$), for efficiency that interpolates between linear and full softmax attention (Mercat, 2020, Mongaras et al., 31 Jul 2025).
• Gated and Forgetting Variants: Introducing learnable gates or data-dependent decay, as in gated linear attention or the Forgetting Transformer, augments recurrence with selective memory retention, further improving long-context performance over purely linear models while retaining greater efficiency than classic softmax (Brébisson et al., 2016, Lin et al., 3 Mar 2025).

These modifications typically trade off some expressiveness for reduced complexity, with the loss manifest in lower accuracy on complex contextual tasks relative to full recurrent softmax attention (Brébisson et al., 2016).

6. Gradient Analysis, Optimization, and Max-Margin Properties

The recurrent perspective on softmax attention clarifies its optimization landscape:

• The use of a regularized, convex mapping such as the softmax gradient (as in the smoothed max framework) yields differentiable mappings suitable for backpropagation and natural integration with RNN gates (Niculae et al., 2017).
• Under training, the softmax nonlinearity, when combined with gradient descent updates, implicitly enforces a max-margin property: attention parameters (notably the token selector or prompt vector) converge in direction to the solution of a maximum-margin SVM problem that selects the “most optimal” token from a sequence (Tarzanagh et al., 2023).
• Recurrent or greedy optimization of softmax units (via approximate Newton methods) benefits from strong convexity and Lipschitz properties in the loss, ensuring stable and efficient convergence (Deng et al., 2023).

These findings further explain the robust optimization dynamics and convergence properties of attention-based recurrent update schemes, with theoretical guarantees supporting the use of greedy or approximate algorithms in large-scale sequence models.

7. Practical Implications and Contemporary Extensions

The recurrent formulation of softmax attention guides the design of next-generation architectures and approximations:

• Efficient implementations leveraging finite-order expansions or agent-based intermediate aggregators (as in agent attention) scale well to high-resolution or long-context settings while retaining much of the power of softmax (Han et al., 2023).
• Efficient fixed-size context representations—enabled by linearized variants—support large-scale retrieval and indexing (Brébisson et al., 2016).
• Forgetting mechanisms or self-adjusting variants enhance stability, dynamic context management, and optimization, especially in very deep or adaptive transformers (Lin et al., 3 Mar 2025, Zheng et al., 25 Feb 2025).
• Universal approximation with interpolative recurrence—engineered solely in attention layers—eliminates the need for deep feedforward blocks, with implications for lightweight, flexible sequence models (Hu et al., 22 Apr 2025).

This line of work provides a rigorous platform for analyzing, approximating, and innovating upon softmax attention, as well as for integrating attention with the theory and practice of dynamical systems, recurrent neural architectures, and statistical learning in the context of very long or dynamically evolving sequences.