PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Length-Extrapolatable Recurrent Models

Hanwen Jiang

arXiv:2609.09157v1cs.LGcs.CL

TL;DR

The paper asks why recurrent models trained with BPTT often fail beyond their training horizon, focusing on whether future losses provide useful state credit to earlier states. It proposes CST, which locally rescales state credit without rotating it or changing forward computation, with regime-specific variants for controlled and real data. CST improves extrapolation in both settings, with gains reported up to 128 times the training length.

  • Problem

    Recurrent models can fail to learn long-horizon behavior with BPTT, and gradient decay alone does not explain when learning fails because distant-loss credit to earlier states is task dependent.

  • Method

    CST directly rescales selected state-credit signals during backward propagation without rotating them or changing the forward trajectory, using task-specific stabilization for controlled and real-data regimes.

  • Results

    CST improves length extrapolation in both controlled and real-data regimes, including a 6.58-point average gain over BPTT across 16 long-context Meta-FSA evaluations and improvements across 14 real-data evaluations up to 128k.

  • Takeaways & Limitations

    State credit is a direct intervention point for learning recurrent models that extrapolate beyond their training horizon in the studied settings.

  • Takeaways & Limitations

    The empirical scope is limited by one 123M-parameter real-data model, one training run per method, and unresolved cross-domain evidence for distant-context use on GovReport.

Abstract

from arXiv · show

Recurrent models provide a natural path to long-context modeling, yet models trained with backpropagation through time (BPTT) often fail beyond their training horizon. Classical analyses emphasize gradients that vanish or explode along temporal paths. However, dense per-token losses can still train a shared recurrent rule despite severe decay, showing that decay alone does not determine whether learning fails. We instead study state credit: the signal through which future losses reach earlier recurrent states before contributing to parameter updates. Accordingly, we intervene directly on state credit and propose Credit Stabilization through Time (CST). During backward propagation, CST locally rescales the state-credit signal to stabilize its norm without rotating the component being corrected, while leaving the forward computation unchanged. Because controlled synthetic tasks and real data exhibit different credit dynamics, we specialize CST to each regime. In both settings, CST improves performance beyond the training horizon, with gains observed at up to 128x the training length.

1 Introduction

Long-context recurrent modeling is limited not only by temporal gradient decay but by whether future losses provide useful credit to earlier states. The paper studies state credit directly and proposes CST, which stabilizes that signal while preserving forward computation, improving extrapolation beyond training lengths.

  • Motivation: Recurrent models offer a natural route to long-context deployment, but standard BPTT can make future losses uninformative for earlier computation through products of state-transition Jacobians.This creates a learning problem beyond simply running the recurrence for longer.
  • Problem diagnosis: The paper analyzes state credit, ∇h_tL, the upstream signal through which future losses reach earlier recurrent states before parameter updates.This is presented as an upstream alternative to diagnostics based on aggregated parameter gradients.
  • Problem diagnosis: Gradient decay alone does not determine learning failure; failure depends on whether the task requires distant losses to credit earlier recurrent states.Dense local supervision can still train shared recurrent rules despite severe decay.
  • Method: CST locally rescales selected state-credit components without rotating them, stabilizing signals passed to earlier states while leaving the forward trajectory and training objective unchanged.The intervention acts during backward propagation rather than modifying the recurrent computation.
  • Method: CST is specialized to dependency structure: Event-CST targets substantial, repeatable contractions in controlled tasks, while real-data CST stabilizes credit independently and symmetrically across recurrent heads.EMA replay tracks slowly evolving contraction locations in the controlled setting.
  • Results: CST improves extrapolation in both regimes, including all 16 Meta-FSA aggregate evaluations from 16k to 128k, with a 6.58-point average gain over BPTT.Real-data experiments also reduce full-token negative log-likelihood in all 14 dataset–length evaluations up to 128k.

2 A Suite of Controlled Synthetic Tasks

The controlled task suite separates contextual retention, repeated state transformation, and their interaction. By varying memory load, input horizon, and transition depth independently, it isolates when local supervision suffices and when learning depends on distant consequences.

  • Task motivation: Synthetic tasks separate retaining earlier information, repeatedly transforming a state, and combining both demands, with lengths and computational requirements varied independently.This decomposition is intended to clarify why recurrent models fail as sequence length increases.
  • Task format: Each example contains a context of M addressable items followed by N task inputs, with targets supplied at every task-input position.Inputs may update a task state, retrieve a context item, or apply a context-defined transition.
  • State tracking: State-tracking tasks repeatedly apply a fixed rule from an initial state, with D = N dependent transitions extending computation without changing the rule or state space.Finite-state automata serve as the representative state-tracking task.
  • Retrieval: Memory-and-retrieval tasks use M addressable context items and later independent queries, with D = 0 and N controlling how long stored information remains accessible.MQAR represents this task family.
  • Hybrid tasks: Hybrid tasks supply an instance-specific transition rule in context and apply it across K independently initialized state-tracking segments while the recurrent state remains continuous.Meta-FSA uses a newly sampled transition table for each example, which must remain available across segments.
  • Evaluation scope: The study holds M fixed and evaluates temporal extrapolation in N and D rather than claiming memory-capacity extrapolation.Memory load is treated as bounded by model capacity and dependent on model size.
  • Learning demands: The suite distinguishes local supervision from distant consequences: nearby per-token losses can train a fixed transition, whereas context-specific behavior requires later predictions to credit earlier stored information.This distinction determines when BPTT must carry useful credit across long horizons.

3 BPTT and Temporal-Credit Diagnosis

The diagnosis separates gradient-level dependence on distant paths from state-credit transmission through recurrent states. Distant learning requires both boundary-crossing credit and a usable signal reaching earlier states.

  • BPTT transmits future-loss effects through products of state-transition Jacobians before they reach shared recurrent parameters.
  • The recurrent-parameter gradient combines local contributions with boundary-crossing components removed by state detachment at chunk boundaries.
  • At sequence length 1024 and H = 128, pfar is 0.021, 0.089, and 0.259 for Fixed-FSA, Meta-FSA, and MQAR, respectively.These values indicate weak, moderate, and strong dependence on boundary-crossing paths across the three tasks.
  • Removing the far component leaves MQAR accuracy at 0.6%, whereas full BPTT reaches 96.3%.This establishes that local supervision cannot replace the boundary-crossing update in MQAR.
  • State-credit transmission and parameter-gradient dependence need not share the same task ordering.Across approximately 895 positions, Fixed-FSA retains effectively none of late-loss credit yet depends little on distant parameter-gradient paths, while Meta-FSA and MQAR retain ratios of approximately 1.05 and 0.82.
  • Gradient decay alone does not determine learning: the task must require distant losses to credit earlier recurrent states.

4 Method

CST intervenes on state credit during backward propagation by locally rescaling its norm without changing its direction or the forward computation. Its controller and reference scale are specialized to the data regime.

  • CST modifies state credit before it is mapped and aggregated into the recurrent-parameter gradient.
  • CST locally rescales selected state-credit components, changing their strength without rotating them at the intervention point.
  • At each exposed internal boundary, CST rescales incoming state credit by a positive scalar before it crosses the preceding chunk.
  • The controller compares credit norm with a reference scale, amplifying contraction, optionally attenuating expansion, and leaving near-reference credit unchanged.
  • CST leaves the forward recurrence, loss, parameterization, and full-sequence connectivity unchanged.
  • CST retains cross-chunk credit with γk > 0, unlike exact BPTT’s γk = 1 and detachment’s γk = 0.
  • Controlled tasks use a one-sided event controller, whereas real text uses a symmetric head-wise controller because their credit dynamics differ.Both controllers instantiate the same intervention while choosing different components, references, and trigger rules.

5 Controlled-Task Specialization and Verification

Controlled-task evidence shows that Event-CST selectively stabilizes state credit and improves long-horizon extrapolation, especially when tasks require retaining context for later use. Its benefits depend on task regime, while selective corrections recover much of the throughput cost.

  • One-Sided CST: Event-CST restores substantial accumulated state-credit contraction through sparse, data-dependent corrections rather than correcting every boundary fluctuation.The controller triggers when the current norm falls below a fraction of a held reference, while still evaluating all candidate boundaries.
  • Meta-FSA Verification: 6.58 percentage points is Event-CST’s average gain across all 16 Meta-FSA depth–length cells from 16k to 128k.At the training setting it trails BPTT by 3.55 points, but every long-context aggregate cell improves.
  • CST Variants: 25.1% versus 22.1% and 18.5% is the long-context accuracy of Event-CST, Adjacent-CST, and BPTT, respectively.Accumulated event detection adds 3.0 points beyond one-sided adjacent stabilization.
  • Efficiency: 0.808M tokens/s is the projected EMA-replay throughput, reaching 84.6% of BPTT while retaining most Event-CST accuracy.EMA replay reduces exposed boundaries from 31 to 6.91 per step; its long-context accuracy is 23.2% versus 25.1% for Event-CST.
  • Task Families: Event-CST improves MQAR from 25.56% to 41.66% at 128k, while leaving Fixed-FSA’s nearly length-flat behavior largely unchanged.MQAR keeps the number of stored bindings fixed, so the gain reflects longer retention and reuse rather than increased memory load.
  • Scaling Across Regimes: Task regime determines depth transfer: Event-CST hurts unseen-depth transfer on S32/A4 but improves every evaluated depth on the harder S64/A8 task.The two comparisons jointly change state count, action count, training length, and controller hyperparameters, limiting direct attribution.

6 Real-Data Specialization and Language Modeling

Real-text recurrent language modeling requires a two-sided, head-dependent CST controller because state-credit dynamics vary across boundaries. CST improves extrapolation consistently, with modest all-token gains and larger gains on long-context-sensitive key tokens, while remaining limited in scope.

  • Real-data controller: The controller acts on unnormalized state credit at each exposed boundary, correcting its norm without altering the forward computation or training objective.Diagnostic normalization compensates for the number of remaining targets but is not used by the controller.
  • Real-data controller: Real-text state credit contracts or expands across recurrent heads, motivating symmetric stabilization rather than one-sided event correction.Adjacent-boundary changes have median multiplicative magnitude 1.50 and 90th-percentile magnitude 8.96.
  • Real-data controller: Two-sided corrections dominate observed real-text behavior, with 37.89% of gains amplifying and 60.52% attenuating the modified signal.Only 2.90% and 3.24% of gains reach the upper and lower caps; aggressive correction with α = 0.30 performs worse on average.
  • Language-modeling results: CST improves all-token NLL in all 14 dataset–length evaluations, by 0.00724 NLL on average.At the longest evaluated horizons, all-token gains range from 0.00595 to 0.00973 NLL.
  • Language-modeling results: Key-token improvements exceed full-sequence gains at long contexts, reaching 0.02393 NLL on LongData at 32K, or 3.50× its all-token gain.Corresponding gains are 0.01477 on Books3 at 128K and 0.04418 on GovReport at 32K.
  • Scope and limitations: The evidence supports improved distant-context use on held-out LongData, but the unresolved GovReport control limits cross-domain generalization.The broader empirical scope is also limited by one 123M-parameter model and one training run per method.

A.3 Causal Far-Gradient Ablation

The far-gradient ablation compares full BPTT with boundary-detached training to isolate the contribution of boundary-crossing paths. In MQAR, including the far gradient produces a large causal-comparison gain.

  • Causal far-gradient ablation: At length 1024, MQAR reaches 96.30% with α = 1 versus 0.60% with α = 0, while random guessing is 0.098%.Both conditions use matched initialization and data streams; the intervention depends on boundary spacing H.

B.1 Language-Model Implementation

The language-model experiments use a roughly 123M-parameter recurrent model trained on 4K sequences and compare it with a similarly sized Transformer. Evaluation extends beyond training context using fixed data, annotations, and specified extrapolation protocols.

  • Model and data: The recurrent language model has 123.30M parameters, width 768, one local-attention block, and 12 Gated DeltaNet blocks with 16 recurrent heads.GPT-2 embeddings and readout weights are tied over a 50,304-token vocabulary.
  • Model and data: LongData training uses packed 4,096-token sequences from 458 training shards, with AdamW, cosine decay, and parameter-gradient clipping at 1.0.Sixteen whole shards are reserved for validation.
  • Comparison protocol: Figure 5 compares CST with a similarly sized Transformer trained on the same 4K data and 7.86B-token budget.The Transformer uses native RoPE at 4K and dynamic YaRN beyond its training horizon.
  • Configuration: The selected CST configuration is (α, β, c) = (0.10, 0.90, 1.10), chosen using LongData-32K and Books3-128K NLL.Confidence intervals are conditional on this selection and do not account for configuration-search or training-seed variation.

B.4 Paired Confidence Intervals

Paired document bootstraps quantify uncertainty for language-model gains and context-use comparisons. LongData supports improved use of context beyond 4K, whereas GovReport does not resolve that specific effect.

  • Uncertainty protocol: Paired document bootstraps use 20,000 resamples, with documents retained as the paired unit when they contribute multiple windows.Intervals measure variation across fixed documents, not independent training runs, and are conditional on the selected CST configuration.
  • Uncertainty limits: The reported confidence intervals do not capture variation across independent training runs, and out-of-domain Books3 results are descriptive without paired document-level intervals.These limitations constrain how broadly the uncertainty estimates should be interpreted.
  • LongData: LongData context-use gains exclude zero at both 16K and 32K, supporting improved use of context beyond 4K.At 16K–32K, the key-token gain is 0.023537 with interval [0.011159, 0.037012], while context-use gain is 0.013666 with interval [0.008616, 0.020309].
  • GovReport: GovReport supports improved prediction of annotated key tokens but not a resolved increase in their use of distant context.At 32K, the context-use gain is 0.002684 with interval [−0.002305, 0.007471].

C.3 EMA Replay Implementation

EMA replay estimates boundary corrections from periodically refreshed profiles, replaying only predicted correction positions on ordinary steps. This reduces boundary-hook overhead while retaining empirical dependence on training stage and task.

  • EMA replay planning: EMA replay scans an exponentially averaged credit profile to select earlier boundaries when predicted cumulative drawdown reaches log ρ.Predicted peaks reset the anchor before scanning continues.
  • EMA replay planning: The planner initializes with a dense identity-gain probe, refreshes every P optimizer steps, and rejects invalid profiles with another dense probe.The first valid profile seeds the EMA; later profiles update it.
  • Overhead: Dense Event-CST costs O(B size(h)) boundary overhead, whereas EMA replay reduces amortized overhead to O((R + B/P) size(h)).Realized throughput also depends on segmented recurrence, sequence length, task, and training stage.

C.4 Event-CST Controller Behavior

Event-CST corrections are sparse, concentrated toward later sequence positions, and accumulate as credit propagates backward. EMA replay approximates this behavior with substantially fewer boundary hooks while preserving the late-boundary concentration.

  • Correction behavior: 1.79 corrections per state-component trajectory occur on average, with 5.8% of boundary–state-component decisions triggering and 403-token effective spacing.The mean gain over all candidates is 1.053×, and 4.3% reach the cap.
  • Correction behavior: Effective spacing averages distances between consecutive latest, triggered, and earliest candidate boundaries, excluding sequence endpoints 0 and 1024.Intervals are averaged over components and steps within runs, then across seeds.
  • Correction behavior: A geometric mean cumulative controller multiplier of approximately 3.7× appears at the earliest boundaries as local gains accumulate backward.This multiplier is the product of applied γk values, not the realized CST-to-BPTT credit-norm ratio.
  • Correction behavior: Corrections concentrate at several positions in the latter half of the sequence, while trigger probability remains low across most early boundaries.The figure reports controller coefficients and their cumulative products, not realized gradient amplification.
  • EMA replay efficiency: EMA replay uses 6.91 boundary hooks per step versus 31 for dense Event-CST, a 77.7% reduction or 4.49× fewer hooks.It probes the dense profile once every 100 steps and exposes only predicted correction positions between probes.
  • EMA replay efficiency: Verified plans select 4, 8, and 8 boundaries for seeds 42, 43, and 44, with mean spacings of 160.0, 77.7, and 82.3 tokens.All selected boundaries lie between positions 416 and 992.

D Research Process and Unsuccessful Explorations

The research process combined author judgment with high-throughput agent exploration, then documented approaches that did not yield reliable improvements. Several alternatives failed to improve consistently or introduced difficult optimization and complexity trade-offs.

  • Research process: The author set direction, supplied constraints, inspected experiments, and retained responsibility for interpreting observations, while Codex implemented, debugged, and ran experiments.The workflow reserved human attention for judgment while using agents for exploration and implementation.
  • Unsuccessful explorations: Length curriculum variants did not consistently improve extrapolation across schedules, so the approach was not pursued further.The schedules varied both growth rate and the interval between length increases.
  • Unsuccessful explorations: State self-distillation lacked a suitable privileged-information source for useful target states without imposing an inappropriate representation.
  • Unsuccessful explorations: Jointly learning quantizers, codebooks, and recurrent dynamics produced a difficult coupled optimization problem and no stable improvement.
  • Unsuccessful explorations: Fully adaptive temporal credit assignment became sensitive to noisy boundary statistics without fixed triggers or gain caps, and improvements failed to persist across tasks and configurations.These variants also increased computation and controller complexity.
  • Interpretive boundary: These unsuccessful explorations do not establish that the directions are fundamentally unworkable, because formulation, implementation, or experimental effort may have limited them.
Loading 2609.09157v1…