PaperlayerView exact revision on arXiv

Source-linked AI summary

Dynamic Compression in Recurrent Networks

Jyothish Pari, Ryan Bahlous-Boldi, Pulkit Agrawal

arXiv:2608.17896v1cs.LG

TL;DR

Recurrent models must compress growing histories into fixed-size states before future information needs are known. This paper introduces dynamic compression, which selectively revisits past tokens to revise the state, substantially reducing the state required for accurate function reuse and scaling more favorably as stored functions increase.

  • Problem

    Recurrent models need fixed-size states to retain and reuse increasing amounts of information despite uncertainty about future task demands.

  • Method

    Dynamic compression selectively re-scans relevant past tokens and uses recurrent updates to revise their representation after later tasks reveal which information matters.

  • Results

    Selective re-scanning substantially reduces recurrent-state capacity needed for accurate function reuse and scales more favorably than single-pass compression as stored functions increase.

  • Takeaways & Limitations

    Dynamic compression trades additional computation for more effective use of fixed-size recurrent memory when reusing multiple functions.

  • Takeaways & Limitations

    The study uses a controlled synthetic setting, and applying dynamic compression to natural-data pretraining requires solutions for parameterizing re-scans and generating supervision.

Abstract

from arXiv · show

Recurrent models process long contexts efficiently by compressing their history into a fixed-size state, but modern architectures typically do so in a single causal pass over the sequence. Each input must therefore be compressed before the model knows how it will later be used, forcing a limited state to compromise across possible future demands. We introduce dynamic compression, which allows a recurrent model to selectively revisit past tokens and revise its fixed-size state through additional recurrent updates. The model need not preserve every part of the history at uniformly high fidelity in its recurrent state, because lower-fidelity information can be revisited from the retained raw sequence when it becomes relevant. We study this in a controlled setting where the model first learns multiple functions in-context and, later in the same sequence, encounters a series of few-shot tasks that each require it to identify and reuse one of those functions. A single-pass model must preserve every function at sufficient fidelity for any future task, whereas selective re-scanning allows the model to revisit and refine only the function currently needed. We find that dynamic compression substantially reduces the recurrent state required for accurate reuse and scales more favorably as the number of stored functions grows. These results demonstrate a computation--memory tradeoff in which recurrent models can spend more computation revisiting their history to make more effective use of a fixed-size state.

1 INTRODUCTION

Dynamic compression lets recurrent models revisit selected past tokens after later context reveals what information matters, revising a fixed-size state instead of relying on one task-agnostic compression. In a controlled function-reuse task, selective re-scanning reduces the recurrent-state capacity needed to identify and reuse stored functions.

  • Motivation: Recurrent models offer O(1) recurrent-state memory and O(1) recurrent computation per decoded token, unlike Transformers’ O(t) key–value memory and attention computation.This favorable scaling motivates studying recurrent models for long-context inference.
  • Problem: Single-pass recurrence must represent each token before knowing which details future tasks will require, forcing one fixed state to compromise across task demands.Once a token is processed, future use depends only on its contribution to the recurrent state.
  • Method: Dynamic compression revisits selected positions from the retained raw prefix and processes them again through the same recurrent update, revising the active state.The raw sequence uses O(t) storage, while the active recurrent state remains fixed-size; the model performs one selective re-scan before answering.
  • Evaluation: The controlled task presents several linear functions before few-shot tasks require identifying and reusing one, with identification needing coarse information but evaluation requiring higher fidelity.Single-pass models must preserve every function accurately enough for possible future use, whereas selective re-scanning can refine the relevant function.
  • Results: Selective re-scanning substantially reduces recurrent-state capacity requirements, evaluated first with oracle supervision and then with a self-supervised procedure that predicts re-scan targets.The controlled setting provides known function-token locations for the oracle evaluation; the self-supervised targets derive from write strengths on repeated context.

2 BACKGROUND

The paper uses Gated DeltaNet as a recurrent model that compresses history into a fixed-size, matrix-valued state updated online. Its gated delta rule separates global forgetting from key-specific update strength, enabling β_t to signal state revision during an additional context pass.

  • Recurrent state: Gated DeltaNet compresses history into a fixed-size state S_t, updates it online, and reads it to produce each output.The model is a linear-attention architecture with a matrix-valued state S_t ∈ R^(d_v×d_k).
  • Gated updates: The key-specific gate β_t controls how strongly the state updates toward v_t at k_t, while α_t controls global forgetting.The paper uses β_t as an update-strength signal during an additional pass after the relevant task is known.
  • State dimensions: The state uses key dimension d_k = d_head and value dimension d_v = 2d_head, with a fixed value-expansion factor of 2.The dimensions are defined per attention head.
  • Memory budget: The memory budget varies only d_head while fixing n_layer = 4, n_head = 6, and the value-expansion factor at 2.Total recurrent state size is determined by the number of layers, heads, and per-head dimensions.

3 TASK DESCRIPTION

This section introduces a synthetic two-phase continual-learning task that tests whether recurrent models can store multiple matrix functions and identify, retrieve, and apply the relevant one from few-shot evidence. The task exposes the cost of static compression and motivates spending additional computation to revisit past information.

  • Task design: The task presents a single continuous sequence with a basis phase followed by a query phase to isolate functional reuse.The model first internalizes functions from context, then reuses them on new inputs.
  • Task design: Gated DeltaNet learns a vector-output function A ∈ R^d×d from context and predicts Axqry for a held-out query.The function is a full square matrix, extending scalar-output in-context learning to vector outputs.
  • Basis phase: In the continual setting, the model stores K basis matrices in one sequence and must identify and apply any requested basis to later queries.Basis tokens include one-hot identifiers, while few-shot and query tokens require inference of the relevant basis from examples.
  • Query phase: Each query group uses f < d few-shot pairs as a search signal, requiring the model to match them against stored bases before predicting Aiqxqry.The few-shot pairs cannot recover Aiq independently, so the model must retrieve the corresponding stored function.
  • Scaling challenge: ∼3M elements are required at K = 3 versus ∼3k at K = 1 to match the loss, showing that static compression becomes prohibitively expensive as functions accumulate.This scaling motivates using additional computation to compress past information more effectively within a smaller state.

4 SELECTIVE RE-SCANNING FOR DYNAMIC STATE COMPRESSION

Dynamic compression lets recurrent models revisit selected past tokens and revise their state after later context reveals which information matters. The section develops oracle and self-supervised selective re-scanning methods that trade additional computation for lower recurrent-state requirements.

  • Dynamic compression: Re-scanning revises a token’s contribution and fidelity through new recurrent writes, unlike transformer retrieval, which only reads past tokens.The model spends computation to raise fidelity where later context indicates it matters most.
  • Oracle selective re-scanning: A selection head predicts the relevant basis after each few-shot block, then the model re-scans its tokens before the query reads the revised state.The method augments Gated DeltaNet with separate prediction and selection heads.
  • Oracle selective re-scanning: Approximately 3k state elements suffice for near-perfect basis identification, compared with approximately 3M needed to store and apply three bases simultaneously.This motivates identifying the basis with a small state and refining it through re-scanning rather than retaining all bases at full fidelity.
  • Self-supervised re-scanning: The self-supervised method uses write strength β from repeated-context passes to infer which past tokens become relevant for the queried basis.Final-layer β patterns vary systematically with the queried basis and are clustered into a re-scan codebook with C = 3.
  • Self-supervised re-scanning: At inference, codebook dynamic processes the prefix, predicts a re-scan code, revisits its corresponding tokens, and produces an answer without repeated context.Its validation MSE lies well below the single-pass baseline and partway toward the oracle and repeat upper bounds.

5 RESULTS

Dynamic re-scanning outperforms single-pass compression, especially under tight memory and as more functions must be retained. Learned re-scanning captures substantial benefits without oracle supervision, supporting more scalable continual reuse.

  • Oracle dynamic re-scanning consistently outperforms single-pass compression across recurrent state sizes, with the largest gains under tight memory budgets.At K = 3, a 111k-element dynamic state achieves lower error than a 3.1M-element single-pass state.
  • At fixed state size, single-pass error grows substantially faster than dynamic re-scanning error as the number of stored functions increases.Both methods degrade with K, although one dynamic run at K = 6 failed to optimize and matched the single-pass level.
  • The codebook dynamic model substantially improves over the single-pass baseline and closes part of the gap to oracle dynamic re-scanning.This shows that useful re-scan decisions can be learned without oracle re-scan supervision.
  • Selective revisiting lets recurrent models retain raw context and dynamically recompress the past around the current task instead of preserving all information uniformly.The results suggest a more scalable approach to accumulating and reusing increasing numbers of earlier skills in lifelong settings.

6 RELATED WORK

Prior work develops recurrent models that compress past context into fixed-size states and methods that selectively access relevant portions of long contexts. These approaches include richer recurrent memory updates, sparse attention, parameter updates, and retrieval augmentation.

  • Modern RNNs and their limitations: Modern RNNs formulate linear attention recurrently by compressing past context into a fixed-size state.DeltaNet replaces additive fast-weight updates with the delta rule, while later work develops increasingly efficient and expressive variants.
  • Selective access to past context: Related approaches also extend recurrent models with richer test-time memory updates or retrieval from past context.Titans and Atlas introduce richer memory updates, while Resona augments recurrent models with retrieval.
  • Selective access to past context: Several methods selectively allocate additional computation to parts of long contexts rather than processing all past tokens uniformly.Sparse-attention methods restrict attention to selected subsets, while Self-Guided Test-Time Training selects question-relevant spans for parameter updates.

7 DISCUSSION AND LIMITATIONS · A ADDITIONAL WRITE-STRENGTH (β) VISUALISATIONS

The discussion frames continual learning as fixed-state compression improved by selective re-scanning, while noting open challenges for natural-data pretraining. The appendix provides additional visualizations of write strength β for repeat and dynamic models.

  • 7 DISCUSSION AND LIMITATIONS: Selective re-scanning improves compression by allocating additional computation where higher fidelity is needed, reducing the state required to reuse multiple functions.Continual learning is characterized as representing a growing context within a fixed-size recurrent state.
  • 7 DISCUSSION AND LIMITATIONS: Dynamic compression forms a two-level memory hierarchy in which raw context preserves lossless access while recurrent state acts as revisable working memory.The model can revisit and rewrite past information once its relevance becomes known.
  • 7 DISCUSSION AND LIMITATIONS: The study uses a controlled synthetic setting, leaving extension to natural-data pretraining as an open problem.The discussion identifies natural-data pretraining as a direction requiring further investigation.
  • 7 DISCUSSION AND LIMITATIONS: Natural-data applications require determining how to parameterize possible re-scans and supervise where to revisit when tasks are not explicitly delineated.Post-training is proposed as a complementary setting because task structure is more controlled and reinforcement learning can be applied.
  • A ADDITIONAL WRITE-STRENGTH (β) VISUALISATIONS: Figure 7 visualizes write strength β for a random sample of repeat-model sequences with K = 3.The query basis span rule marks the repeated block of the matched basis.
  • A ADDITIONAL WRITE-STRENGTH (β) VISUALISATIONS: Figure 8 visualizes dynamic-model β for a random sample of sequences, using the definition from Figure 5.The caption identifies these as additional dynamic-model write-strength visualizations.

B EXPERIMENTAL DETAILS · B.1 STATE-SIZE FIGURE (FIGURE 6)

The experiment evaluates Gated DeltaNet models on an in-context matrix-regression task while varying recurrent state size and comparing single-basis, single-pass, and dynamic re-scanning conditions. Figure 6 reports each run’s minimum validation MSE and summarizes per-state-size medians.

  • B EXPERIMENTAL DETAILS: All models use Gated DeltaNet with 4 layers, 6 heads, embedding dimension 256, and value expansion 2.The recurrent state size is set by head dimension d_head as 4·6·d_head^2·2.
  • B EXPERIMENTAL DETAILS: Training uses AdamW with β1=0.9, β2=0.95, weight decay 10^-2, learning rate 10^-4, gradient clipping at 1.0, and no schedule.Validation is evaluated 100 times during training.
  • B EXPERIMENTAL DETAILS: The underlying task is the in-context matrix-regression task from the main text with token dimension d = 8.
  • B.1 STATE-SIZE FIGURE (FIGURE 6): Head dimensions d_head ∈ {8, 48, 256} produce state sizes 3072, 110,592, and 3,145,728.These correspond approximately to 3k, 111k, and 3.1M states.
  • B.1 STATE-SIZE FIGURE (FIGURE 6): Each sequence contains M = 16 demonstration pairs per basis, F = 4 few-shot pairs, and 8 query positions.Training runs for 150,000 steps with batch size 512 and 10,000 held-out validation sequences.
  • B.1 STATE-SIZE FIGURE (FIGURE 6): The comparison includes single basis (K = 1), single pass (K = 3), and dynamic re-scanning (K = 3), with three seeds per condition.The first two conditions use the plain prediction head and are scored by validation MSE.
  • B.1 STATE-SIZE FIGURE (FIGURE 6): For each run, Figure 6 reports the minimum validation MSE over training, plots every seed, and connects per-state-size medians with dashed lines.

B.2 SCALING WITH THE NUMBER OF BASES (FIGURE 6)

Figure 6 evaluates single-pass and dynamic re-scanning models as the number of bases increases from K = 3 to 7, using a fixed state size and controlled training setup. Performance is compared using each run’s minimum validation MSE over training.

  • Experimental setup: The experiment compares single pass with dynamic re-scanning across K ∈ {3, 4, 5, 6, 7} bases, using three seeds per setting.Every seed is plotted individually.
  • Experimental setup: The recurrent state is fixed at size 6912, with head dimension d_head = 12 for every point.This keeps state capacity constant while varying the number of bases.
  • Experimental setup: Each sequence contains M = 16 demonstration pairs per basis, F = 4 few-shot pairs, and 8 query positions.Training uses 150,000 steps at batch size 512 with 10,000 validation sequences.
  • Evaluation: For each run, the reported metric is the minimum validation MSE over training.The comparison therefore uses the best validation MSE achieved during each run.

B.3 OFFLINE DYNAMIC-MODEL COMPARISON (TABLE 1) · B.4 SELECTION-HEAD RETRIEVAL ACCURACY (FIGURE 2) · B.5 CONTEXT SCALING (FIGURE 1)

The paper evaluates four offline compression families, selection-head retrieval across state sizes, and context scaling under controlled synthetic-task settings. These experiments vary the re-scan strategy, supervision, head dimension, and demonstration count.

  • B.3 OFFLINE DYNAMIC-MODEL COMPARISON (TABLE 1): The comparison fixes head dimension 16, state size 12,288, K = 3, M = 16, F = 4, and one query group per sequence.All families train for 50,000 steps at batch size 1024 and use 50,000 validation sequences, with three seeds per family.
  • B.3 OFFLINE DYNAMIC-MODEL COMPARISON (TABLE 1): Four model families are compared: single-pass, repeat, oracle dynamic, and codebook dynamic.Single-pass reads the prefix once; repeat re-reads the entire prefix; oracle dynamic re-scans the ground-truth basis block; codebook dynamic re-scans codebook-selected tokens without ground-truth labels.
  • B.3 OFFLINE DYNAMIC-MODEL COMPARISON (TABLE 1): The codebook dynamic model trains on 5 × 107 fixed, independently generated sequences containing basis, few-shot, re-scan, and query blocks.Each re-scan block is selected by the β-derived codebook rather than ground-truth basis labels.
  • B.3 OFFLINE DYNAMIC-MODEL COMPARISON (TABLE 1): Baseline and repeat models use validation MSE, whereas oracle dynamic and codebook dynamic use autoregressive re-scan MSE.Reported values are each run’s minimum validation MSE over training, summarized as mean ± sample standard deviation across seeds.
  • B.4 SELECTION-HEAD RETRIEVAL ACCURACY (FIGURE 2): The selection head predicts one of K = 3 query-basis indices using cross-entropy, with chance accuracy 1/3.Head dimension sweeps over {4, 8, 16, 24, 32, 48, 64, 96, 192, 256}, corresponding to state sizes from 768 to ≈3.1M.
  • B.4 SELECTION-HEAD RETRIEVAL ACCURACY (FIGURE 2): Selection-head experiments use M = 16, F = 4, and 8 query positions, training for 20,000 steps at batch size 512 with 128 validation sequences.One seed is used for each head dimension.
  • B.5 CONTEXT SCALING (FIGURE 1): Context scaling uses a single-basis task with K = 1, head dimension 16, no few-shot block, and M swept over {0, 1, 2, 4, 8, 16}.Each sequence contains M demonstration pairs followed by 8 query positions; validation-loss curves are plotted directly against training step on a log-y axis.

B.6 WRITE-STRENGTH (β) ANALYSIS AND CODEBOOK (FIGURES 4, 5)

The analysis records delta-rule write strengths β during forward passes and uses final-layer patterns to construct a three-codebook, equal-length re-scan scheme. The codebook dynamic model is then evaluated along its inference path of reading, predicting a code, re-scanning, and answering.

  • β analysis: Write strength β is recorded per layer, token, and head, then averaged across heads for the repeat and codebook dynamic models.Both models use head dimension 16 and are trained for 50,000 steps with batch size 1024.
  • β analysis: Repeat-format sequences expose β over the repeated region, with the final layer used because its pattern most clearly aligns with the queried basis.The study does not explore aggregating write strengths across layers, such as summing β.
  • Codebook: Each repeated-region β pattern is a 52-dimensional vector, clustered by k-means with the elbow selecting C = 3 codes.The codebook is derived from final-layer β values.
  • Codebook: Each code selects n = 21 prefix positions above βmax/2, producing equal-length re-scan blocks while keeping every selected position above threshold.The common length is the smallest number of above-threshold positions across centroids, and each code takes its top-n centroid positions.
  • Dynamic-model β: The trained dynamic model records β along inference as it reads the prefix, predicts a code, re-scans its tokens, and answers.The resulting sequence has the form basis | few-shot | re-scan | query.
Loading 2608.17896v1…