PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models

William L. Tong, Aryo Lotfi, Emmanuel Abbe, Kostas Vaggelakos, Vishnu Banna, Etai Littwin, Josh Susskind, Cengiz Pehlevan, Eran Malach

arXiv:2609.16540v1cs.LGcs.AI

TL;DR

SSMs are computationally efficient but can underperform on in-context learning and retrieval. This paper combines theory and experiments to study gating, finding that it favors memorization during training while improving long-context generalization in certain settings.

  • Problem

    SSMs offer efficient sequence modeling but lag behind Transformers on tasks requiring in-context learning and precise retrieval.

  • Method

    The paper uses theory and experiments on simplified and practical SSMs to analyze how gating affects memorization, retrieval, and length generalization.

  • Results

    Gating can cause SSMs to prioritize in-weights memorization over in-context learning, while often improving generalization to longer sequences.

  • Takeaways & Limitations

    Gating is a crucial factor shaping SSM training dynamics and long-context behavior, but the paper does not immediately provide a solution for improving in-context learning.

  • Takeaways & Limitations

    A simple scalar gating mechanism may be insufficient for the complex memory management required by common long-context retrieval problems.

Abstract

from arXiv · show

State Space Models (SSMs) have emerged as a compelling alternative to Transformers, enabling sequence modeling with constant memory and linear compute. Although SSMs exhibit reasonable performance and favorable computational characteristics, they continue to lag behind Transformers on tasks that require in-context learning and precise retrieval, slowing their adoption for large-scale language modeling. In this work, we demonstrate that both the success and failure of SSMs in these domains can be explained by studying the role of the gating mechanism, a prevalent component in modern recurrent networks. Specifically, we show through theory and experiments that this gating mechanism causes SSMs to first learn an in-weights "memorization" solution, while delaying, or even preventing, convergence to a correct in-context learning solution. Importantly, this happens even in cases where there are no fundamental limitations due to the architecture or its memory capacity. On the other hand, we find that gating is often beneficial for improving generalization to long sequence lengths. Our results illuminate the crucial role of the gating mechanism in shaping both the training dynamics and generalization of SSMs, and provide a basis for understanding and improving linear-time models.

1 Introduction

Linear-time architectures offer efficient sequence modeling, yet SSMs can lag on in-context learning and retrieval. This work argues that gating shapes this gap by favoring memorization while also improving long-context generalization.

  • SSMs and related linear-time architectures provide linear computational complexity and constant memory, unlike Transformers’ quadratic computation and linearly scaling memory.
  • Gating biases SSMs toward in-weights memorization even when their memory capacity is sufficient for a correct in-context learning solution.The same mechanism can also improve generalization beyond the training sequence length.
  • Theory predicts that initial gating can make a simplified linear-time model converge first to memorization, although a better in-context solution remains reachable.
  • Synthetic retrieval and logical-rules experiments show that gated SSMs may miss the correct in-context solution, while changing initialization can accelerate convergence in some settings.
  • Tuning the gating parameter between pretraining and supervised fine-tuning changes hallucination rates on multi-tool in-context tool calling, and weaker-than-pretrained gating can improve performance.

2 Theory: Retrieval with memorization “shortcuts”

The paper analyzes how gating controls competition between retrieval and memorization in a tractable SSM setting. Strong gating can delay retrieval through a memorization plateau, whereas decay can limit long-context retrieval to an effective locality.

  • 2.1 Task: The synthetic task separates memorization from in-context learning by resampling context-label perturbations, making high accuracy require matching the query to its context occurrence.Memorizing the fixed underlying mapping can provide partial accuracy when perturbations are small.
  • 2.3 Short context retrieval: The model’s signed margin decomposes into memorization and retrieval margins, corresponding respectively to input weights and the label-coordinate readout.A correct ICL solution has zero input weights and a nonzero label readout.
  • 2.3 Short context retrieval: Weak gating favors rapid retrieval, while strong gating makes memorization dominate initialization and can delay recovery of the full retrieval solution.For strong gating, the escape time may grow exponentially with token distance and double exponentially with the gating initialization scale.
  • 2.3 Short context retrieval: The memorization plateau can eventually give way to retrieval, with escape time scaling as Θ((α log(1/α))^-1) under the analyzed reduced dynamics.
  • 2.4 Long context generalization: Decay reduces retained distractors but induces locality: the furthest reliably retrievable position scales as τ log(d/τ), so arbitrarily long contexts require relevant tokens to remain nearby.The analysis characterizes τ = e^-λ as the effective context length.
  • 2.4 Long context generalization: Gating therefore presents a tradeoff: weak gating helps retrieval in short contexts, whereas strong gating controls accumulated variance and supports long-context generalization within its effective context length.

3 Multi-Token Retrieval

The multi-token retrieval task tests whether SSMs retrieve perturbed values from context rather than rely on a fixed key-value mapping. Weaker initial gating accelerates in-context learning, while stronger gating improves generalization beyond training lengths.

  • Setup: The task presents in-context key-value pairs followed by multi-token query keys whose corresponding values must be predicted.Each key and query contains m tokens, and the number of demonstrations is sampled uniformly up to N.
  • Setup: Best-first ordering emphasizes retrieval because relevant keys occur farther from their queries, amplifying the effect of gating through decay.Best-last ordering is used for evaluating length generalization beyond the training context.
  • Setup: Training uses a static key-value dictionary, but independently perturbed value tokens make perfect accuracy require matching queried keys in context.With N = 512, ϵ = 0.1, q = 0.1, and m = 8, memorization reaches sequence accuracy approximately 0.43, while ICL can achieve perfect accuracy.
  • Results: Smaller initial gating enables the model to learn the in-context learning solution, whereas increasing gating progressively lengthens escape from memorization.Escape time is predicted to scale approximately as T ∝ 1/α up to logarithmic factors, with α proportional to e^-offset.
  • Results: Stronger gating at initialization yields better accuracy when contexts extend beyond the 512 key-value pairs seen during training.This length-generalization result is obtained in the best-last setting.

4 Logical Rules Retrieval

The logical-rules task evaluates retrieval of randomly named clauses needed to connect a starting proposition to a goal. Weaker gating improves in-distribution retrieval and escape from memorization, while stronger gating improves length generalization.

  • Setup: Horn clauses provide a controlled abstraction of tool calls because they map proposition inputs to an output proposition.The analogy bridges synthetic retrieval with the paper’s tool-use experiments but does not capture all function-call nuances.
  • Setup: The task requires selecting randomly named Horn clauses that advance deductions from a starting proposition toward a goal proposition.Each turn outputs one clause, whose output becomes the next starting proposition.
  • Results: Weaker gating promotes in-context retrieval, whereas stronger gating benefits generalization to longer contexts.This overall pattern is summarized across the logical-rules experiments.
  • In-context retrieval: The no-oracle condition primarily rewards memorizing name-to-clause mappings, while resampled oracle names require genuine in-context retrieval.High accuracy with a resampled oracle name is possible only when the model retrieves the name from context.
  • Results: Stronger-gated models memorize longer, succeeding without an oracle but struggling when the oracle name is resampled; very weak gating can instead bias training toward retrieval.For e^λ0 < 1, models do not reach a full memorization solution or saturate no-oracle accuracy.
  • Long-context generalization: Long-context evaluation trains on pools up to k = 8 and tests up to K = 128 clauses per transition, with all names resampled to force retrieval.Clauses are ordered so the most relevant clause is closest to the query.

5 Tool-Calling SSMs with Tool-Retrieval

The BFCL extension tests whether Mamba can retrieve the correct tool from variable in-context tool pools. Weaker gating modestly improves retrieval-related accuracy and hallucination, while all Mamba variants retain stronger length generalization than the Transformer baseline.

  • Setup: BFCL requires choosing a tool and producing a fully correct, well-formed function call from a natural-language query and suggested tool pool.A deterministic checker validates the complete call.
  • Setup: The experiment extends BFCL by constructing a shared tool pool and presenting variable-size subsets ranked by BM25 relevance.This creates a retrieval setting in which the model must select among multiple potentially relevant tools.
  • Setup: Gating strength is controlled by scaling Mamba’s pretrained decay rates before SFT, with larger scaling corresponding to stronger gating.In log-rate form, the intervention is λ ← λ + log s.
  • Results: Weaker-gated Mamba models score moderately better and hallucinate less, indicating more reliable retrieval of the correct tool from context.Hallucination counts syntactically valid calls to functions absent from the provided pool.
  • Results: All Mamba models length-generalize better than the Transformer baseline, while gating strength does not affect length generalization in the tested range.The models are trained with variable pools up to k = 24 and evaluated on larger pools.

6 Discussion

The discussion identifies gating as a key factor in SSM training dynamics and length generalization. It also cautions that the paper does not establish a direct solution for improving retrieval and leaves richer memory mechanisms for future work.

  • Discussion: Gating can prioritize memorization over in-context learning while improving length generalization in certain settings.The paper characterizes this as a dual role of gating in SSMs.
  • Discussion: The study does not determine whether changing initialization or removing gating is a generally good solution, because gating often improves long-context performance.Its primary focus is scientific understanding rather than an immediate retrieval-improvement method.
  • Discussion: A scalar multiplicative gate may be insufficient for the complex memory management required by common long-context retrieval problems.The authors suggest higher-dimensional linear operators or nonlinear transitions as possible alternatives and leave them for future work.

A Related Work

Prior work shows that SSMs can perform in-context learning and short-sequence retrieval, but retrieval degrades with distance and pretrained models remain weaker on retrieval-intensive evaluations. The literature also links gating to both length generalization and recency bias, motivating this paper’s analysis of their trade-off.

  • SSMs achieve competitive in-context learning on many synthetic problems, but the gap between toy-task results and limited capabilities in pretrained models remains unresolved.
  • Retrieval studies report that SSMs can retrieve from relatively short sequences, while performance across arbitrarily long sequences is constrained by fixed-size memory.
  • Prior work associates recurrent gating with stable length extrapolation, recency bias, and local-pattern shortcuts that degrade long-context retrieval.
  • Mimetic initialization, which effectively disables gating initially, improves synthetic copying, while other work reports strong SSM length generalization on induction and tool-use tasks.
  • This work establishes a trade-off between performance on long sequences and in-context retrieval capabilities.
  • The paper develops formal statements, derivations, and proofs after introducing its theoretical findings informally.

B.1 Short context retrieval - extended version

The theory decomposes classification into memorization and retrieval margins and analyzes how gating determines which solution initially dominates. Strong gating favors memorization and can create a long-lived plateau, although retrieval eventually emerges with an escape time that can become enormous.

  • The signed margin decomposes as z = ηδm + δr, where δm = (1 + α)u measures memorization and δr = αv measures retrieval.The memorization term flips with label mutation, whereas the retrieval term remains consistent across mutations.
  • At initialization, the retrieval-to-memorization growth-rate ratio determines whether the model initially favors retrieval or memorization.For α > α∗ retrieval dominates; for α < α∗ memorization dominates, with α∗ having an approximate critical scale for large P and small α.
  • Small initial gating causes the model to become insensitive to the context label and converge first to a perfect memorization solution.With v fixed at zero, the expected loss has a unique minimizer on the memorization manifold.
  • Memorization is a plateau rather than a permanent sink: after v becomes positive, both v and α increase, allowing the model to transition toward retrieval.The transition is characterized by retrieval margin δr exceeding memorization margin δm.
  • As α0 approaches zero, escaping the memorization plateau takes a time proportional to 1/(α0 log(1/α0)).The escape time is defined when the retrieval margin reaches the memorization margin κ.
  • Because α = exp(−m e^λ), the required escape time may be exponential in token position or double exponential in the base initialization scale.For distant relevant tokens, the analysis says an initialization with very small λ, and thus α close to 1, is essential for reasonable training time.

B.2 Long context generalization - extended version

Long-context retrieval requires balancing decay: weaker decay preserves the queried token’s signal but increases distractor variance, while stronger decay limits retained context and induces locality. The effective context size and retrievable distance are controlled by the gating parameter, enabling long-context generalization only when relevant tokens remain sufficiently close.

  • Training dynamics: Gradient flow initialized with wx(0)=0 preserves wx(t)=0 and drives v(t)>0, yielding the pure retrieval trajectory rather than memorization through input weights.The expected-loss dynamics are rotationally invariant in wx, while the label coordinate supplies the retrieval signal.
  • Retrieval trade-off: Weaker decay increases the query signal weight α_∆ but also increases S2, so excessive retained context can drown the retrieval margin in variance.Successful retrieval therefore requires a moderate decay rather than uniformly favoring weaker decay.
  • Effective context size: τ controls the effective context size, with S2(ℓ, λ) scaling as Θ(τ) when ℓ/τ grows large.Here τ = e^-λ, and larger λ means stronger decay and a smaller effective context.
  • No-decay limit: Without decay, retrieval fails as distractors grow: when L/d →∞, the Bayes error approaches 1/2.The no-decay signal competes against all L = ℓ−1 distractors.
  • Locality and generalization: When τ ≪ d, the furthest retrievable distance scales as τ log(d/τ), so decay enables length generalization but imposes retrieval locality.Random retrieval becomes impossible as context length increases, whereas nearby relevant tokens can remain retrievable in arbitrarily long contexts.
  • Gate stability: Across Mamba heads, effective gating lengths span five orders of magnitude at initialization, remain ordered after 200k steps, and drift only modestly.This makes initialization a useful proxy for the post-training gating regime.

C.2 Results

With imperfect tool retrieval, moderate gating provides the strongest balance between retrieval and long-pool generalization. Low gating degrades as pool size grows, while strong gating weakens retrieval.

  • Weak tool retrieval: Moderate gating lets Mamba generalize to larger tool pools while retaining strong retrieval, benefiting most when the retriever is imperfect.Larger pools provide greater recall that can compensate for weaker retrieval quality.
  • Weak tool retrieval: Low gating prevents effective generalization to large pool sizes, causing performance to fall as the pool grows.The model misses the recall gains provided by larger pools.
  • Weak tool retrieval: Strong gating prevents the model from learning retrieval as effectively, despite improving locality for long contexts.The supplied result passage specifically identifies reduced retrieval effectiveness as the trade-off.

D.1 Theoretical validation experiments (Section 2)

The theoretical predictions were validated with simplified SSM experiments using stochastic gradient descent. Separate panels tested initialization bias, escape time, accuracy across effective context sizes, and the interaction between distance and gating.

  • Experimental setup: The validation experiments trained the simplified SSM with SGD at learning rate 0.1 and batch size 128.SGD was selected to match the gradient-flow analysis most directly.
  • Initialization bias: The initialization-bias panel used d=512, P=32, ℓ=8, ∆=4, and ϵ=0.1, measuring margins after 25 gradient steps.These settings isolate the early training behavior predicted by the theory.
  • Escape time: The escape-time panel used d=512, P=32, ℓ=4, ∆=1, and ϵ=0.1 with learning rate 0.3.The higher learning rate kept the escape time within a tractable range.
  • Long-context tests: The long-context accuracy panel used d=128 and ℓ=1024, while the heatmap used d=128 and ℓ=512; both trained for 10 thousand steps.The heatmap varied effective context size τ and distance ∆.

D.2 Multi-token retrieval task (Section 3)

The multi-token retrieval experiments used perturbed key-value demonstrations sampled from a static dictionary. Training varied the number of demonstrations, while evaluation included in-distribution maximal-context examples.

  • Task construction: A static dictionary D contains |D|=10000 key-value pairs, from which each example samples n uniformly between 1 and a fixed maximum N.The sampled demonstrations form the in-context retrieval task.
  • Perturbation procedure: Each sampled static value is independently token-perturbed with probability ϵ to produce the in-context demonstrations.This separates the static dictionary from the perturbed values presented during examples.
  • Evaluation: Evaluation uses the maximal N setting for the in-distribution experiments.The supplied passage specifies this evaluation choice but not its resulting accuracy.

D.2.1 Model Architecture and Initialization

The experiments use Mamba-2 and GPT-NeoX Transformer baselines with matched model dimension and tied embeddings, while varying Mamba’s gating initialization through additive decay offsets. Training uses AdamW and evaluates escape dynamics and length generalization across five seeds.

  • Model architecture: Mamba-2 uses 24 layers, dmodel = 768, 24 heads, head dimension 64, and SSM state dimension dstate = 128.Input and output embeddings are tied.
  • Initialization: Mamba’s default per-head decay parameters are sampled as λ ∼log(Uniform(1, 16)), and additive offset δ continuously varies gating strength while preserving head variation.δ = 0 recovers standard initialization; δ = −10 approximately recovers the no-decay regime.
  • Model architecture: GPT-NeoX provides the Transformer baseline with 12 layers, dmodel = 768, 12 attention heads, intermediate size 3072, and full RoPE.It uses Flash Attention 2, no dropout, and tied embeddings.
  • Training and evaluation: The study trains models with AdamW, linear warmup, cosine decay, gradient clipping, bfloat16 precision, and a global batch size of 128 for up to 100k steps.Experiments use five random seeds, with early stopping at 99% sequence accuracy.
  • Training and evaluation: Figure 2 measures median escape-time training curves, seed-level escape times versus initialization, and length generalization beyond 512 training key-value pairs.Escape time is the first step exceeding 0.95 evaluation sequence accuracy; length evaluations range from 512 to 2048.
Loading 2609.16540v1…