Stream: 通过稀疏注意力扩展机械可解释性以分析LLM中的长上下文
Stream: Scaling up Mechanistic Interpretability to Long Context in LLMs via Sparse Attention
AI总结：
Stream通过稀疏注意力技术实现LLM长上下文的高效可解释性分析，显著减少内存消耗并提升推理轨迹追踪能力。
摘要：
随着大型语言模型（LLMs）扩展到百万token上下文，传统的机械可解释性技术分析注意力时会以二次方式随上下文长度增长，需要超过100000个token的terabytes内存。我们引入了稀疏追踪，一种利用动态稀疏注意力高效分析长上下文注意力模式的新技术。我们提出了Stream，一种可编译的层次剪枝算法，能够在近线性时间O(T log T)和线性空间O(T)内估计每个头的稀疏注意力掩码，使大规模可解释性成为可能。Stream通过二分查找式细化，仅保留每个查询的前k个键块，同时保持模型的下一个token行为。我们应用Stream到长链式推理轨迹，并在剪枝97-99%的token交互的同时识别思维锚点。在RULER基准测试中，Stream保留了关键检索路径，同时丢弃了90-96%的交互，并揭示了从针到输出的逐层路线。我们的方法提供了一种实用的工具，用于分析注意力模式和追踪信息流，而无需terabytes的缓存。通过使长上下文可解释性在消费级GPU上可行，稀疏追踪有助于普及链式推理监控。代码可在https://anonymous.4open.science/r/stream-03B8/上获得。
英文摘要：
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/.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
