Optimal Decay Spectra for Linear Recurrences

1. Problem Statement

Linear recurrent models (e.g., Mamba, RWKV, RetNet) offer $O(N)$ sequence processing time, making them efficient alternatives to Transformers. However, they suffer from sub-optimal long-range memory, particularly when extrapolating to context lengths longer than their training window.

The author identifies two fundamental failure modes in the decay spectrum (the collection of per-channel decay rates that control how quickly memory channels forget inputs):

1. Random Initialization Collapse: Standard random initialization causes the minimum spectral gap to collapse to $O(N^{-2})$. This leads to severe spectral redundancy (channels become indistinguishable) and sub-exponential error rates: $\exp(-\Omega(N/\log N))$.
2. Static Linear/Geometric Mismatch: Even well-structured initializations (like linear spacing in S4D or Mamba-2) suffer from a scale mismatch. A spectrum designed for a fixed context length $T$ wastes channels at early positions $t \ll T$ (where only short-term memory is needed) and fails to cover long-range dependencies during extrapolation ($t > T$). This degrades the approximation error to practically algebraic rates over long contexts.

2. Methodology: Position-Adaptive Spectral Tapering (PoST)

The paper proposes PoST, an architecture-agnostic framework that combines two mechanisms to enforce an optimal, scale-free memory structure without computational overhead.

A. Theoretical Foundation

The authors derive an optimal memory blueprint based on logarithmic information equipartition. Natural language data exhibits scale-free correlations (power-law decay), implying that predictive information is distributed uniformly across "octaves" of the dependency range.

• Optimal Structure: Memory channels should be geometrically spaced in log-decay space to minimize approximation error for power-law kernels.
• Optimal Scale: The spectrum must dynamically stretch to cover the actual dependency range $[1, t]$ at every position $t$, rather than a fixed design length $T$.

B. Core Components of PoST

1. Spectral Reparameterization (Structural Guarantee)

• Mechanism: Replaces independent parameterization of decay rates with a recursive structure:
  $$p_1 = \theta, \quad p_k = p_{k-1} + \text{Softplus}(\delta_{k-1})$$
  where $p_k$ are log-decay rates and $\delta$ are learnable gap parameters.
• Effect: This enforces a strict ordering $p_1 < p_2 < \dots < p_N$ by construction, preventing the minimum gap collapse.
• Optimality: Proven to achieve minimax-optimal exponential approximation rates for power-law kernels: $O(\exp(-cN/\log T))$.

2. Position-Adaptive Scaling (Dynamic Mechanism)

• Mechanism: Introduces a position-dependent scaling factor to the decay rates. The effective log-decay rate at position $t$ for channel $k$ is:
  $$p^{\text{eff}}_{t,k} = p_k \cdot t^{-\alpha_k}$$
  where the taper exponent is $\alpha_k = \frac{N-k}{N-1}$.
• Effect: This "stretches" the spectrum so that at any position $t$, the channels span the relevant dependency range $[1, t]$.
  • Slow channels ($\alpha \approx 1$) adapt to relative coordinates (scale-free).
  • Fast channels ($\alpha \approx 0$) adapt to absolute coordinates.
• Uniqueness: Proven to be the unique continuous mechanism that eliminates scale mismatch, sharpening the error bound to $O(\exp(-cN/\log t))$.
• Invariance: This scaling induces fractional invariance, making the impulse response scale-free and allowing channels to interpolate between relative and absolute temporal coordinates.

3. Key Contributions

1. Information-Theoretic Blueprint: Established that ideal memory channels must be geometrically spaced and dynamically scaled to match the logarithmic growth of predictive information in natural data.
2. Diagnosis of Failure Modes: Theoretically proved that random initialization leads to $O(N^{-2})$ gap collapse and that static linear spacing degrades performance over long contexts.
3. PoST Framework: Introduced a two-part solution (Reparameterization + Adaptive Scaling) that is architecture-agnostic and incurs zero computational overhead (only $O(N)$ pre-computation per layer).
4. Surjectivity & Expressiveness: Proved that PoST preserves the full representational power of the original architecture (it is a surjection onto the space of strictly ordered spectra) while preventing degenerate configurations.

4. Experimental Results

The authors instantiated PoST across five major architectures: Mamba-2, RWKV-7, Gated DeltaNet, Gated Linear Attention (GLA), and RetNet.

• Multi-Query Associative Recall (MQAR):

  • PoST significantly improved retrieval accuracy under length extrapolation (training at 512, testing up to 4096).
  • Mamba-2 saw the most dramatic gains (e.g., average accuracy jumping from 59.4% to 64.2% at 32K state size).
  • RWKV-7 and Gated DeltaNet also showed consistent improvements.

• Zero-Shot Language Modeling:

  • Pre-training at 180M and 440M scales on FineWeb-Edu showed consistent, albeit modest, improvements in perplexity and downstream benchmarks (LAMBADA, HellaSwag, etc.) across all architectures.
  • Mamba-2 440M with PoST achieved a lower perplexity (62.6 vs 77.3) and higher average accuracy (42.5 vs 42.2).

• Needle-In-A-Haystack (NIAH):

  • Mamba-2 exhibited massive gains in long-context retrieval. For example, at 4K context, Single-Needle retrieval jumped from 0.2% (baseline) to 2.0% (PoST), and Multi-Needle retrieval improved significantly.
  • RWKV-7, which already has a robust native spectrum, remained competitive, suggesting PoST provides the most benefit to architectures with poorly conditioned baseline spectra.

• Empirical Analysis:

  • Visualizations of learned timescales confirmed that standard Mamba-2 models suffer from "gap collapse" (heads clumping at fast timescales), whereas PoST models maintain a perfect geometric progression across all layers and heads.

5. Significance

• Spectral Hygiene: PoST acts as a "spectral hygiene" primitive, fixing fundamental structural flaws in linear recurrent models without requiring architectural changes or increased compute.
• Theoretical Rigor: It bridges the gap between information theory (scale-free memory) and practical deep learning, providing the first proof of minimax optimality for geometric spacing in this context.
• Scalability: The method is highly scalable, reducing to a simple two-line change in the forward pass of any diagonal linear recurrence, making it immediately applicable to existing large-scale models.
• Long-Context Future: By solving the scale mismatch problem, PoST enables linear recurrent models to effectively handle contexts far beyond their training length, a critical step toward efficient, long-context AI.

Code Availability: The implementation is open-sourced at https://github.com/SiLifen/PoST.