Stream: Scaling up Mechanistic Interpretability to Long Context in LLMs via Sparse Attention

TL;DR

The paper tackles the challenge of mechanistic interpretability in long-context LLMs by introducing Sparse Tracing and its Stream instantiation, which leverage hierarchical, dynamic sparse attention to analyze contexts of up to $10^5$–$10^6$ tokens. Stream prunes attention patterns in a top-$k$ blockwise manner, achieving near-linear time $O(T \log T)$ and linear space $O(T)$ while preserving key behavioral signals such as thought anchors and retrieval paths. Through two case studies—Thought Anchors and Needle in a Haystack—the approach demonstrates substantial sparsity (≈97–99% in reasoning traces and 90–96% in needle retrieval) with preserved outputs, enabling long-context interpretability on consumer GPUs. The work situates itself within broader literature on chain-of-thought monitoring and long-context limitations, and outlines practical avenues for future work, including extending coverage to residuals/MLPs and providing theoretical guarantees for information-path preservation.

Abstract

As Large Language Models (LLMs) scale to million-token contexts, traditional Mechanistic Interpretability techniques for analyzing attention scale quadratically with context length, demanding terabytes of memory beyond 100,000 tokens. We introduce Sparse Tracing, a novel technique that leverages dynamic sparse attention to efficiently analyze long context attention patterns. We present Stream, a compilable hierarchical pruning algorithm that estimates per-head sparse attention masks in near-linear time $O(T \log T)$ and linear space $O(T)$, enabling one-pass interpretability at scale. Stream performs a binary-search-style refinement to retain only the top-$k$ key blocks per query while preserving the model's next-token behavior. We apply Stream to long chain-of-thought reasoning traces and identify thought anchors while pruning 97-99\% of token interactions. On the RULER benchmark, Stream preserves critical retrieval paths while discarding 90-96\% of interactions and exposes layer-wise routes from the needle to output. Our method offers a practical drop-in tool for analyzing attention patterns and tracing information flow without terabytes of caches. By making long context interpretability feasible on consumer GPUs, Sparse Tracing helps democratize chain-of-thought monitoring. Code is available at https://anonymous.4open.science/r/stream-03B8/.

Figures (14)

• Figure 1: Illustration of the Stream hierarchical attention pruning algorithm. Starting from the full causally masked attention scores ($QK^T * C$), the pattern is recursively divided into branches and refined over successive iterations. At each step, less relevant regions are discarded, converging on a sparse mask $M$ that preserves only the top-$k$ most relevant key-query interactions. The algorithm produces a final pruned attention score pattern ($QK^T*C*M$), achieving near-linear time complexity $O(T log T)$ and linear space complexity $O(T)$.
• Figure 2: This figure displays attention patterns for three different context lengths (3,000, 6000, and 10000 tokens) using three different approaches: Full Attention Pattern (first column), Mean Averaged bogdan2025thought with $b=32$ (second column), and Stream with $b_y=b_k=32$ (third and fourth column).
• Figure 3: In this figure we plot the vertical attention scores for each 32 token block for three different context lengths (1,000, 10000, and 20,000 tokens).
• Figure 4: Visualization of sparse attention masks generated by Stream on the RULER needle-in-a-haystack benchmark hsieh2024ruler at three context lengths (1k, 10k, and 20k tokens). The needle is marked using a red vertical line in the Attention Pattern plots.
• Figure 5: Figure showing needle retrieval success as a function of context length (x-axis), needle depth (y-axis) and effective sparsity $s$ (subplots-axis). The parameter $s\in(0,1)$ controls the effective pruning of the attention pattern via $k$ with larger $s$ corresponding to higher pruning.