Fractional State Space Transition for Long Sequence Modeling

AI summary

Overview

• Research area: Efficient long-sequence modeling with State Space Models (SSMs); specifically, replacing ODE-based exponential memory with fractional-calculus-derived power-law memory.
• Technical level: Advanced. The architecture is built on Caputo fractional derivatives, Mittag–Leffler functions, and diffusive representations of fractional kernels.
• Scope: The paper proposes Frac, a selective SSM layer whose memory follows fractional (heavy-tailed) dynamics, derives a finite-state sum-of-exponentials realization of that kernel, and validates it on synthetic probing tasks, 1.3B-parameter language models, and 7M-parameter DNA models.

What This Paper Is About

Most modern State Space Models compress sequence history into a bounded recurrent state whose decay law is exponential, meaning old inputs are forgotten at a fixed rate. That is a poor match for long sequences where distant events still matter. The paper's goal is to replace exponential forgetting with power-law (heavy-tailed) memory derived from fractional differential equations, while keeping the recurrent layer finite-state, parallelizable during training, and usable for autoregressive decoding.

Key Contributions

1. Frac, a selective SSM derived from fractional dynamics. Frac replaces the exponential memory law of ODE-based SSMs with the Mittag–Leffler relaxation and power-law tail induced by the Caputo fractional relaxation equation.

2. A finite-state realization of fractional memory. Theoretically, the paper shows the fractional kernel admits a nonnegative diffusive representation (Theorem 1), then approximates it on a bounded horizon with a geometrically (log-spaced) bank of exponential modes (Theorem 2), and shows this is equivalent to a bank of M first-order ODEs whose readout approximates the fractional solution uniformly (Proposition 3).

3. A selective, hardware-efficient recurrence. The mode bank is discretized exactly with zero-order hold and made token-dependent through learned Δ_t, α_t, λ_t, plus mode-wise read and write weights with a log-timescale prior. The result is an affine recurrence supporting chunked parallel scan for training and prefill, and one-step recurrent decoding with cached states, implemented in a custom Triton kernel.

4. Empirical validation at scale. Frac is tested on synthetic heavy-tail and MADLab benchmarks, 1.3B-parameter language models trained on 100B tokens from deduplicated FineWeb-Edu, LongBench, NIAH, LM Harness, recall-retrieval tasks, and 7M-parameter DNA language models on HG38.

Main Findings

• Heavy-tail synthetic extrapolation: In a single-layer, 200K-parameter setting trained only on length-512 sequences and evaluated up to 128K tokens, Frac shows the smallest performance decay over distance. GDN and Mamba3 are the closest competitors, while vanilla self-attention drops to near-random performance starting at 8K.

• MADLab: In four-layer models of roughly 500K parameters, Frac attains the best average (75.4) versus Mamba3 (74.8), GDN (74.3), and Mamba2 (73.8). It scores 90.4 ± 0.2 on Compression and 11.5 ± 0.2 on Fuzzy ICR, with Mem. 99.9 ± 0.1 and SC 99.9 ± 0.1.

• NIAH length generalization: Evaluated far beyond the 4K training length up to 64K tokens, Frac matches baseline linear models at ≤4K context and degrades more slowly from 8K to 64K. The main exception is GDN on S-NIAH-1. Transformer performance is zero after 4K.

• LongBench: Frac reaches 17.9 average across 14 tasks, beating the strongest linear baseline GDN (16.0) by 1.9% and posting the best result on 8/14 tasks. Gains are smaller than on NIAH because more than half of LongBench tasks have average input lengths below 8K tokens, keeping the Transformer competitive.

• Short-context LM Harness: Frac averages 55.1, only 0.2% behind the Transformer (55.3) and Mamba3-MIMO (55.3), with comparable perplexity. These tasks involve roughly 32–256 tokens.

• Recall-Retrieval: Frac averages 36.4, ranking third overall: 2.1% behind the Transformer (38.5) and 0.4% behind Mamba3-MIMO (36.8). Evaluation used truncated 2K-token sequences.

• DNA modeling: On 7M-parameter causal models trained on human genome HG38 across sequence lengths from 1K to 64K, Frac achieves on-par or lower perplexity than Mamba3 and GDN at all lengths, with larger gains at longer contexts.

Methodology in Plain English

The authors start from the observation that standard SSMs behave like simple first-order differential equations, and such equations always produce exponential decay. They instead write down a fractional differential equation — one where the derivative is of non-integer order α between 0 and 1 — which produces memory that decays like a power law rather than an exponential.

The problem is that fractional equations are non-Markovian: the state at any moment depends on the entire past, so they cannot be used directly in a fixed-size recurrent layer. The authors solve this with two theoretical steps. First, they use a known result that the fractional kernel can be written as a continuous mixture of ordinary exponentials with different timescales (a "diffusive representation"). Second, they approximate that infinite mixture on a bounded time horizon with a finite weighted sum of exponential modes, placed on a logarithmic timescale grid because fractional memory spreads mass across many orders of magnitude.

Each of those modes is then a simple scalar linear system, so the whole thing becomes a bank of first-order ordinary differential equations. The authors discretize each mode exactly using the standard zero-order-hold method, yielding per-mode retention factors ρ and injection factors β. Self-similarity in λ lets them keep one shared geometric bank of base timescales and express λ through a token-dependent effective timescale.

Finally, they make the recurrence selective: token-dependent Δ_t, α_t (sigmoid-constrained to stay between 0 and 1) and λ_t control the transition, while softmax-normalized read weights use a log-timescale prior −α_t log τ_m with a learned residual, and write weights use the same form with an independent learned map. That makes the whole layer an affine recurrence that can be computed with a chunked parallel scan during training and prefill, and in one-step form during decoding.

Why This Matters

The paper argues that the memory law induced by a model's underlying dynamics is itself a powerful architectural axis, separate from refinements in information routing or selection. If long-context performance is limited by how an architecture forgets, then changing the forgetting law — not just adding capacity or better gating — is a distinct and apparently effective lever.

• Long-document question answering and summarization: The LongBench gains, especially on retrieval-heavy summarization tasks, point toward models that retain relevant information across spans far longer than their training context.
• Retrieval-augmented and recall-intensive systems: The recall-retrieval results suggest usefulness for tasks where a model must retrieve specific facts from long inputs.
• Genomic and scientific sequence modeling: Frac's gains on HG38 DNA language modeling at lengths up to 64K are relevant to biological sequence analysis, where long-range dependencies are common.
• Efficient deployment: Because the method keeps a bounded recurrent state with parallel training and prefill, it targets the same efficiency profile as existing linear-time models rather than quadratic attention.

Industry relevance: The work comes from Huawei Noah's Ark Lab, Montreal Research Center, and compares directly against Mamba2, Mamba3, GDN, and Transformer baselines at the 1.3B-parameter scale on 100B tokens — the scale at which architectural choices are actually deployed. The released code (https://github.com/anasiri/frac-ssm) and the Triton kernel indicate an emphasis on practical implementation.

Future Directions

• Combining fractional memory with delta-rule updates. The conclusion explicitly proposes merging fractional long-memory transitions with delta-rule style update mechanisms to get both stronger long-range inductive bias and more precise associative retrieval.
• Closing the gap on recall-intensive short-context tasks. Frac trails Mamba3-MIMO by 0.4% and the Transformer by 2.1% on the truncated 2K-token recall-retrieval evaluation, leaving room to understand whether the heavy-tail prior costs something at short range.
• Scaling the study. The language-modeling experiments use 1.3B parameters and 100B tokens; behavior at larger parameter and token budgets is not reported.
• Theory-to-practice gap in the SoE approximation. The error guarantees hold on a bounded horizon [0, T] with tolerance ε, while deployed models extrapolate well beyond their training context; how the approximation quality interacts with extrapolation beyond the training length is not resolved in the reported results.

Target Audience

Researchers and practitioners working on efficient sequence modeling, long-context language models, and state-space architectures, particularly those already familiar with Mamba-style selective SSMs and interested in alternatives to exponential-decay memory. The theoretical sections assume comfort with fractional calculus, Mittag–Leffler functions, and zero-order-hold discretization, so the paper is best suited to readers with a strong mathematical or systems background rather than beginners.

Authors’ abstract

State Space Models (SSMs) compress sequence history into a bounded recurrent state, making the resulting memory law a central architectural choice for long-context performance. Most modern SSMs rely on ODE-based dynamics that lead to exponential forgetting, limiting their ability to retain information over broad temporal ranges. We introduce FRAC, a selective SSM architecture derived from fractional dynamics that replaces this exponential decay with power-law long memory. To make fractional dynamics practical, FRAC approximates the heavy-tailed target kernel with a finite-state, log-spaced sum of exponential modes. This construction turns fractional memory into an efficient recurrent module with parallel training and prefill, while retaining bounded-state autoregressive decoding. Extensive experiments, including 1.3B-parameter language modeling, demonstrate that FRAC consistently improves long-context performance over state-of-the-art SSM baselines while staying competitive on short-context. These results show that fractional dynamics provide a practical and effective prior for long-context SSMs.