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

2025/10/22 by Rosser, J, García, José Luis Redondo, Penha, Gustavo +2
#68T40 #Artificial Intelligence (cs.AI) #Computation and Language (cs.CL) #FOS: Computer and information sciences #I.2.11

paper · doi:10.48550/arxiv.2510.19875

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/.

Citations

• Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities
• Thought Anchors: Which LLM Reasoning Steps Matter?
• Understanding Reasoning in Thinking Language Models via Steering Vectors
• Why do LLMs attend to the first token?
• Gemma 3 Technical Report
• A Survey on Transformer Context Extension: Approaches and Evaluation
• Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
• DeepSeek-R1 incentivizes reasoning in LLMs through reinforcement learning
• GPT-4o System Card
• Automatically Interpreting Millions of Features in Large Language Models
• Round and Round We Go! What makes Rotary Positional Encodings useful?
• Qwen2.5-Math Technical Report: Toward Mathematical Expert Model via Self-Improvement
• Attention Heads of Large Language Models: A Survey
• Retrieval Augmented Generation or Long-Context LLMs? A Comprehensive Study and Hybrid Approach
• Understanding the RoPE Extensions of Long-Context LLMs: An Attention Perspective
• Transformers need glasses! Information over-squashing in language tasks
• Iteration Head: A Mechanistic Study of Chain-of-Thought
• Retrieval Head Mechanistically Explains Long-Context Factuality
• RULER: What's the Real Context Size of Your Long-Context Language Models?
• AtP*: An efficient and scalable method for localizing LLM behaviour to components
• How to think step-by-step: A mechanistic understanding of chain-of-thought reasoning
• Information Flow Routes: Automatically Interpreting Language Models at Scale
• A Mechanistic Analysis of a Transformer Trained on a Symbolic Multi-Step Reasoning Task
• YaRN: Efficient Context Window Extension of Large Language Models
• Code Llama: Open Foundation Models for Code
• Lost in the Middle: How Language Models Use Long Contexts
• Extending Context Window of Large Language Models via Positional Interpolation
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Locating and Editing Factual Associations in GPT
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• Measuring Mathematical Problem Solving With the MATH Dataset
• Longformer: The Long-Document Transformer
• SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference
• OpenAI o1 System Card

Related

• SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs
• SampleAttention: Near-Lossless Acceleration of Long Context LLM Inference with Adaptive Structured Sparse Attention
• Ltri-LLM: Streaming Long Context Inference for LLMs with Training-Free Dynamic Triangular Attention Pattern
• Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
• SALE : Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling