PaperlayerView exact revision on arXiv

Source-linked AI summary

Stream Recursion Model (SRM)

Asael Sorensen, Charles Brock, David Chamberlain, Jennifer Minnich, Matthew Hoffman, Ramyaa Ramyaa

arXiv:2609.28809v1cs.LG

TL;DR

Mechanistic interpretability seeks to analyze LLM internals, but large models make that work difficult and expensive. SRM modifies HRM with learned communication among streams; it reports transformer-comparable language performance and distinct routing behaviors across streams.

  • Problem

    Mechanistic interpretability aims to analyze LLM internals, but developing parameter-efficient models suited to interpretability is difficult as LLMs grow larger.

    Source: §1 · p. 1

  • Method

    SRM modifies HRM with a learnable communication layer that determines when streams communicate and how information is translated between them.

    Source: §2 · p. 3, §8 · p. 7

  • Results

    SRM performs comparably to standard transformers on language datasets, with analyses finding distinct routing structures and stream behaviors.

    Source: §8 · p. 7

  • Takeaways & Limitations

    SRM generalizes HRM to multiple streams through learned communication, opening avenues for mechanistic interpretability analysis.

    Source: §8 · p. 7

    The authors avoided post-steps because it was unclear whether their apparent performance gains came from added parameters or structural benefits.

    Source: §4.2 · p. 4

Abstract

from arXiv · show

Mechanistic interpretability seeks to make verifiable statements about the internal behavior of large language models (LLMs). Many interpretability techniques struggle to scale with the increasing size and depth of architectures. Our solution to this is to introduce smaller models with structures that lend themselves to interpretability. In this work, we introduce the Stream Recursion Model (SRM), a modification of the Hierarchical Reasoning Model (HRM) designed to expose internal computational structure while remaining scalable. SRM organizes computation into multiple interacting latent streams that are updated through recursive refinement, enabling direct analysis of stream dynamics, causal contribution, and routing behavior. SRM achieves performance comparable to GPT-2 on a per-parameter basis. Our analysis reveals consistent and distinct behavior across streams, indicating structured specialization and interaction. These results suggest that SRM provides a practical architectural foundation for scalable mechanistic interpretability and opens up promising avenues for future research in both reasoning performance and interpretability.

1. Introduction

The paper targets parameter-efficient models whose structure supports mechanistic interpretability as LLMs grow larger. It introduces SRM, a modification of HRM, and reports transformer-comparable language performance alongside distinct stream behaviors.

1. Introduction

  • Mechanistic interpretability aims to analyze LLM internals to verify functionality, security, and stability, while growing model size makes this work difficult and expensive.Source: §1 · p. 1

1.1. Related Works

  • Prior work reports small-parameter reasoning models and differing findings about the optimal number of streams, motivating an architecture that supports interpretability and stream scaling.Source: §1.1 · pp. 1–2

1.2. Contributions

  • SRM modifies HRM to scale performance with more streams.Source: §1.2 · p. 2
  • SRM achieves on-par performance with vanilla transformers and outperforms HRM on language datasets.Source: §1.2 · p. 2
  • SRM analysis identifies unique routing patterns and specialized behavior across streams.Source: §1.2 · p. 2

2. Motivation

The authors identify limitations in HRM's stream communication for language tasks and in manually specifying communication as stream counts grow. SRM addresses these issues with learned inter-stream translation and gating.

2. Motivation

  • HRM's additive stream communication can underuse streams, and increasing their depth and size did not fix its language-training issues.Source: §2 · p. 2
  • The number of possible communication configurations and parameters grows exponentially with stream count, making learned configuration necessary for practical multi-stream use.Source: §2 · p. 3
  • SRM introduces a learnable communication layer that learns when streams communicate and how information is translated between them.Source: §2 · p. 3

2.1. The Network Approach

  • The Network Approach treats streams as nodes, using stream-specific query, key, value, and output weights to learn translations and attention-based gating that can be inspected through the attention matrix.Source: §2.1 · p. 3

3. SRM

SRM combines token embedding and output functions with a recursive layer built from stream communication, state updates, and transformer step-functions. Its design uses an intermediate embedding and stream-specific attention and blocks.

3. SRM

  • SRM consists of input, output, and layer functions, with the layer containing connection, update, and step-functions.Source: §3 · p. 3

3.1. Input and Output Functions

  • An intermediate embedding reduces the parameter cost of embedding tokens into multiple streams and decoding their outputs.Source: §3.1 · p. 3

3.2. The Layer Function

  • The layer function is SRM's core component and encompasses the connection, update, and step functions.Source: §3.2 · p. 3

3.2.1. THE CONNECTION FUNCTION

  • The connection function is multi-head attention with stream-specific key, query, value, and output weights, also called streamwise attention.Source: §3.2.1 · p. 3

3.2.2. THE UPDATE FUNCTION

  • The update function connects step-functions, the connection output, and embedded input; its layer-normalization ordering affects performance and gradient flow.Source: §3.2.2 · p. 3

3.2.3. THE STEP-FUNCTION

  • Each stream has its own post-norm transformer block, computing attention and feed-forward functions in parallel; the layer hardcodes the order of step-functions.Source: §3.2.3 · p. 3

4. Hyper-Parameter selection

Hyperparameter experiments were limited to about 0.5B tokens, constraining the conclusions while identifying sensitivities and a relatively performant region. The section reports findings on computation depth, pre- and post-steps, layer-steps, and stream size.

4. Hyper-Parameter selection

  • The hyperparameter study was limited to about 0.5B tokens, restricting conclusion accuracy while still identifying SRM sensitivities and a relatively performant region.The experiments lasted about 1,000 training steps; the study varied layer counts, step-functions, stream dimensions and counts, and step depth and ordering.Source: §4 · pp. 3–4

4.1. Computation Depth

  • With depth fixed at 8 layers, returns diminished sharply when backpropagation covered more than 4 layers; this was a reasonable, not proven compute-optimal, choice.The authors recommend further study of depth and backpropagation-layer sensitivity across other hyperparameters.Source: §4.1 · p. 4

4.2. Pre- and Post-Steps

  • Pre-steps worsened performance, while post-steps produced similar or higher performance; uncertainty about whether post-step gains came from extra parameters or structure led the authors to omit them.Pre-steps may reduce the uniformity of token representations entering the input function and destabilize gradients.Source: §4.2 · p. 4

4.3. Layer-Steps

  • Step-function ordering had no noticeable effect, and matching the number of layer-steps to the number of step-functions gave good performance with the least compute.Experiments used three step-functions, varying their ordering and the number of layer-steps.Source: §4.3 · p. 5

4.4. Stream Size

  • A large performance jump appeared between 4 and 8 streams when stream width changed; the authors attribute it tentatively to stream dimensions exceeding the intermediate-embedding dimension.Experiments varied stream count, stream dimension, step-function width, and attention heads under different controls.Source: §4.4 · p. 5

5. Results

On OpenWebText, the authors compared three SRM models with GPT-2 trained under the same configuration. SRM-base matched GPT-2 with roughly half the parameters, SRM-med outperformed it at a similar parameter count, and longer-run outcomes were extrapolated from short training.

5. Results

  • The three SRM models were trained on OpenWebText; the reported sweep experiments used 1,000 steps (about 500M tokens) and were not representative of full training runs.The results are presented in Table 1 and a training graph.Source: §5 · p. 5, Figure 3 · p. 5
  • Roughly half the parameters let SRM-base match GPT-2, but required 7 times the compute; SRM-med outperformed GPT-2 at a similar parameter count, requiring 6 times the compute.SRM-large demonstrates scaling with increased parameters.Source: §5 · p. 5, Table 1 · p. 5
  • After about 5B tokens of training, extrapolation to 300B tokens projected SRM-base to remain similar to GPT-2, SRM-med to stall, and SRM-large to reach a loss of around 2.6-2.7.These are projections based on training graphs, not reported full-run results.Source: §5 · p. 5

6. Analysis

The analysis examines stream routing and alignment behavior across SRM configurations. It finds hub-like routing patterns, nearly orthogonal connection outputs and inputs, and distinct step-function alignment behavior, especially in SRM-large.

6.1. Routing Behaviors

  • SRM-large shows hub-like streams, sinks that rarely output to others, and routing behavior that changes across depth; adding streams makes strong routing sparser under the threshold while increasing seemingly specialized connections and streams.The routing analysis treats streamwise-attention edges above 0.10 after softmax as strongly significant.Source: §6.1 · p. 6

6.2. Connection Update Behavior Alignment

  • Across all three models, connection output ct is nearly orthogonal to token input xt, with a small positive cosine-similarity bias, suggesting directionally distinct information.The study measures cosine similarity between ct and xt in the update function.Source: §6.2 · p. 6

6.3. Step-Function Update Behavior

  • In SRM-base and SRM-med, most streams maintain cosine similarity above 0.7 and 0.5, respectively, between RMSNorm(ct + x) and step-function output.A few SRM-base streams are between 0.0 and 0.3.Source: §6.3 · p. 6
  • SRM-large has streams with cosine similarity near 0.0 or from −0.1 to −0.7 that shift at the layer’s last step, suggesting functionally different roles.Most other SRM-large streams remain above 0.5; the near-zero stream rises above 0.8 at the last step, while two negative-similarity streams reach around 0.3.Source: §6.3 · p. 6

7. Future Works

The future-work section proposes using SRM’s recurrent structure to investigate interpretable routing and suggests language-based reasoning as another direction for study.

7.1. Mechanistic Interpretability

  • Dictionary learning could test whether interpretable features explain SRM’s routing switches, a study left for future work.The proposed techniques include Sparse Autoencoders and Cross Layer Transcoders.Source: §7.1 · p. 6

7.2. Reasoning Tasks

  • Further training on language-based reasoning tasks could explore SRM’s potential for logic and reasoning.The passage conditions this direction on further work and deep reinforcement learning.Source: §7.2 · p. 6

8. Conclusion

The conclusion presents SRM as a generalization of HRM with learned communication between streams, reports language-model performance comparable to standard transformers, and identifies distinct routing structures across model configurations.

8. Conclusion

  • SRM generalizes and improves upon HRM by adding learned communication between streams, enabling multiple-stream computation.The authors connect this design to expanded capability and new avenues for mechanistic-interpretability analysis.Source: §8 · p. 7
  • SRM performs comparably to standard transformer models such as GPT-2 on language datasets, whereas the original HRM was incapable of performing on them.Source: §8 · p. 7
  • Routing structures vary with model size and parameters, revealing distinct stream behaviors and motivating further study of individual streams.The passage suggests future work may identify critical roles and functionality for individual streams.Source: §8 · p. 7

A. Additional Tables and Results

The appendix describes experimental conditions and additional results for SRM, including baseline representational drift and convergence patterns.

A. Additional Tables and Results

  • Figure 6 examines how step-function ordering affects loss after 1,000 training steps, including a condition that randomly selects among three step-functions.The figure’s setup uses six layer steps per layer.Source: Figure 6 · p. 9

B. Model configuration and training

  • SRM-base uses 32 streams with 128-dimensional stream states, six layer steps, three step-functions, and one post-step; SRM-med and SRM-large use 8 and 16 streams, respectively.SRM-med uses three layer steps and three step-functions; SRM-large uses six of each.Source: B. Model configuration and training · p. 9

C. SRM-base: Experiments and Results

  • For the 128-dimensional baseline streams, KL divergence between consecutive layer-wise step-group contributions provides a stable measure of representational change.The appendix notes that this formulation becomes numerically unstable for the 512-dimensional streams in SRM-med and SRM-large.Source: §C.1 · p. 9

C.1. Representational Drift (KL Divergence)

  • Baseline streams vary in KL divergence: some change slowly, while others show sharp depth-specific spikes; convergence analysis distinguishes early, middle, and late convergers.Convergence is the earliest transition after which KL changes by less than 10% for the remainder of the forward pass.Source: §C.1 · pp. 9–10, Figure 7 · p. 10
  • A Layer-4 transition spike appears across all step groups, indicating a global structural transition in recursive computation rather than a group-specific effect.Source: §C.1 · p. 10

C.2. Stream Ablation

The appendix examines causal contributions, alignment, routing, and predictive behavior across SRM streams, finding both uneven stream roles and coordinated depth-dependent patterns.

C.2. Stream Ablation

  • The top nine streams account for approximately 50% of total causal impact, while the most impactful stream’s ablation KL is approximately 0.091 versus approximately 0.009 for the least impactful.The top sixteen account for roughly 75%, and the top nineteen approximately 80%; Figure 8 ranks streams by ablation impact.Source: §C.2 · p. 10, Figure 8 · p. 10

C.3.1. ATTENTION-TOKEN ALIGNMENT

  • Attention–token cosine averages ≈0.067 (range -0.070 to 0.175), with approximately 10.5% of stream-layer pairs negative; high-impact streams also show distinct alignments.Stream 16 has mean alignment -0.040, while Streams 26, 17, and 21 range from 0.024 to 0.068.Source: §C.3.1 · p. 11, Figure 9 · p. 11

C.3.2. UPDATE-STATE ALIGNMENT

  • The combined update aligns positively with the existing stream state across streams and layers, remaining stable across recursive steps despite bounded variation among streams.Figure 10 depicts the cosine similarity between the combined update and existing stream state.Source: §C.3.2 · p. 11, Figure 10 · p. 12

C.4. Routing Structure

  • Routing is concentrated in a subset of source streams that reach more targets, and this pattern remains stable across layers and heads; Figure 11 maps routing by layer and head.Under uniform attention across 32 streams, an individual stream would receive approximately 3.1% attention weight.Source: §C.4 · p. 11, Figure 11 · p. 12

C.5. Stream-wise Logit Lens

  • No single stream is a reliable standalone next-token predictor; across streams, predictive alignment improves early, drops at mid-layers, then partially recovers later.The coordinated mid-layer degradation is visible across nearly all streams in Figure 13.Source: §C.5 · p. 13, Figure 13 · p. 14

D.1. Representational Drift (Cosine Distance)

  • For SRM-med, 512-dimensional stream updates make KL divergence numerically unstable, so cosine distance is used; Groups A and B converge late at a mean of 6.0 ± 0.0 layers, while Group C averages 3.5 ± 2.1 layers.Figure 14 shows median cosine distance across streams with interquartile ranges and a pronounced mid-depth transition.Source: §D.1 · p. 14, Figure 14 · p. 14

D.2. Stream Ablation

SRM-med shows differentiated stream contributions alongside structured routing, attention-token decoupling, and stream-wise output alignment. At larger scale, representational change persists across depth and late convergence dominates across step groups.

D.2. Stream Ablation

  • S6, S4, and S3 have KL divergences of 0.339, 0.306, and 0.278 bits, respectively, and together account for about 50% of SRM-med causal impact.All streams contribute, with a minimum KL of approximately 0.13.Source: §D.2 · pp. 14–15, Figure 15 · p. 15

D.3.1. ATTENTION-TOKEN ALIGNMENT

  • SRM-med attention-token cosine similarity remains near zero (mean 0.056, range -0.022 to 0.154), consistent with decoupling beyond the baseline configuration.The near-orthogonal pattern persists at larger scale.Source: Figure 16 · p. 15, §D.3.1 · p. 15

D.3.2. UPDATE-STATE ALIGNMENT

  • The combined update is strongly aligned with the existing stream state across streams, layers, and recursive steps, remaining stable with modest increases at later steps and deeper layers.No systematic negative alignment is observed in SRM-med.Source: §D.3.2 · pp. 15–16, Figure 17 · p. 16

D.4. Routing Structure

  • At a 10% attention threshold, SRM-med routing concentrates in dominant sources, with 2–4 streams per head often self-routing above 50%.Routing varies by head: H0 is hub-like, H2 is channelized, and H3 is more context-dependent.Source: §D.4 · p. 16, Figure 18 · p. 17

D.5. Stream-wise Logit Lens

  • By L6–L7, S6 has the strongest output-head alignment (mean log-probability ≈-8.5), followed by S4 and S1; early layers are uniformly weak.No individual stream alone is sufficient to recover the model’s prediction.Source: §D.5 · p. 16, Figure 19 · p. 17

E.1. Representational Drift (Cosine Distance)

  • In SRM-large, cosine drift rises from near zero early to moderate levels at mid and late depth, and no step group stabilizes early across most streams.Late-converging streams dominate every step group; KL divergence is numerically unstable for its 512-dimensional stream representations, so cosine distance is used.Source: §E.1 · pp. 16–17, Figure 21 · p. 18

E.2. Stream Ablation

SRM-large distributes ablation impact across more streams while retaining differentiated contributions. Its streams show near-zero attention-token alignment, selective routing, and broadly similar output-space trajectories after an early improvement.

E.2. Stream Ablation

  • SRM-large ablation KL spans 0.013 to 0.155 bits; mean per-stream KL drops approximately 65% from SRM-med to SRM-large (0.224 to 0.078 bits).Impact separates into four tiers; S10, S13, and S14 have minimal impact, with KL changes below 0.03.Source: Figure 22 · p. 18, §E.2 · p. 18

E.3.1. ATTENTION-TOKEN ALIGNMENT

  • SRM-large attention-token cosine similarity remains near zero, with mean 0.064 (range -0.015 to 0.171), despite increased stream count and deeper recursion.The SRM-med and SRM-large values closely correspond.Source: Figure 23 · p. 19, §E.3.1 · p. 19

E.3.2. UPDATE-STATE ALIGNMENT

  • Most SRM-large streams remain positively aligned with their hidden state, but a small subset shows persistent reduced or negative alignment at specific recursive steps.The pattern persists across depth and is described as systematic rather than noise.Source: Figure 24 · p. 19, §E.3.2 · p. 19

E.4. Routing Structure

  • SRM-large routing is sparse and concentrated in a few stream pairs; self-routing is weaker than in SRM-med, while focal high-frequency edges persist across layers.With 8 rather than 4 heads, routing patterns diversify and individual heads specialize over narrower stream-pair subsets.Source: Figure 25 · p. 20, §E.4 · p. 20

E.5. Stream-wise Logit Lens

  • All streams improve sharply at L1, then cluster tightly with similar trajectories; the curves show coordinated mid-depth degradation and partial recovery.The ordering of streams does not change substantially with depth.Source: §E.5 · pp. 20–21, Figure 26 · p. 21
Loading 2609.28809v1…