Stem: Rethinking Causal Information Flow in Sparse Attention

• URL: http://arxiv.org/abs/2603.06274v1
• Date: Fri, 06 Mar 2026 13:33:29 GMT
• Title: Stem: Rethinking Causal Information Flow in Sparse Attention
• Abstract summary: We propose Stem, a novel, plug-and-play sparsity module aligned with information flow.<n>First, Stem employs the Token Position-Decay strategy, applying position-dependent top-k within each layer to retain initial tokens.<n>Second, to preserve information-rich tokens, Stem utilizes the Output-Aware Metric.
• Score: 28.46760179998361
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: The quadratic computational complexity of self-attention remains a fundamental bottleneck for scaling Large Language Models (LLMs) to long contexts, particularly during the pre-filling phase. In this paper, we rethink the causal attention mechanism from the perspective of information flow. Due to causal constraints, tokens at initial positions participate in the aggregation of every subsequent token. However, existing sparse methods typically apply a uniform top-k selection across all token positions within a layer, ignoring the cumulative dependency of token information inherent in causal architectures. To address this, we propose Stem, a novel, plug-and-play sparsity module aligned with information flow. First, Stem employs the Token Position-Decay strategy, applying position-dependent top-k within each layer to retain initial tokens for recursive dependencies. Second, to preserve information-rich tokens, Stem utilizes the Output-Aware Metric. It prioritizes high-impact tokens based on approximate output magnitude. Extensive evaluations demonstrate that Stem achieves superior accuracy with reduced computation and pre-filling latency.

Related papers

• Context by Distinct Information: An Auditable Dirichlet-Process Working Memory for Long, Redundant Context Streams [0.48342038441006796]
  We develop that allocate-on-novelty cache as a working-memory component.<n>We show that novelty-gated attention reaches full-attention performance while attending to about half the tokens.<n>Experiments are small-scale and use only public data.
  arXiv  Detail & Related papers  (2026-07-11T18:58:22Z)
• How Does Reasoning Flow? Tracing Attention-Induced Information Flow for Targeted RL in LLMs [52.712343092801326]
  FlowTracer is an RL framework that traces answer-targeted reasoning flow on an attention-weighted directed acyclic graph.<n>It scores tokens by flow throughput, revealing high-impact hubs and aggregation checkpoints that mediate long-range dependencies.<n>These derived importances are used to shape token-level rewards, enabling learning signals to focus precisely on the tokens that route information toward correct answers.
  arXiv  Detail & Related papers  (2026-06-09T09:56:51Z)
• SinkTrack: Attention Sink based Context Anchoring for Large Language Models [57.87121346197518]
  Large language models (LLMs) suffer from hallucination and context forgetting.<n>SinkTrack treats BOS> as an information anchor and injects key contextual features into its representation.<n>SinkTrack is training-free, plug-and-play, and introduces negligible inference overhead.
  arXiv  Detail & Related papers  (2026-04-11T04:49:11Z)
• InfoFlow KV: Information-Flow-Aware KV Recomputation for Long Context [23.355776408834473]
  We show that a simple attention-norm signal from the query reliably identifies tokens that are semantically relevant and structurally positioned to propagate information.<n> Experiments on LLM and VLM benchmarks demonstrate consistent gains over prior methods under comparable efficiency budgets.
  arXiv  Detail & Related papers  (2026-03-05T16:33:20Z)
• Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection [13.937483869660648]
  Token Sparse Attention is a dynamic token-level sparsification mechanism that compresses per-head $Q$, $K$, $V$ to a reduced token set during attention.<n>We show that Token Sparse Attention consistently improves accuracy-latency trade-off, achieving up to $times$3.23 attention speedup at 128K context with less than 1% accuracy degradation.
  arXiv  Detail & Related papers  (2026-02-03T07:31:14Z)
• LMK > CLS: Landmark Pooling for Dense Embeddings [18.49372789918725]
  We introduce Landmark (LMK) pooling, which partitions a sequence into chunks, inserts landmark tokens between chunks, and forms the final representation by mean-pooling the landmark token embeddings.<n>This simple mechanism improves long-context extrapolation without sacrificing local salient features, at the cost of introducing a small number of special tokens.
  arXiv  Detail & Related papers  (2026-01-29T10:40:37Z)
• Hierarchical Token Prepending: Enhancing Information Flow in Decoder-based LLM Embeddings [52.49524240846879]
  We propose Hierarchical Token Prepending to mitigate attention-level compression and readout-level over-squashing.<n>HTP partitions the input into blocks and prepends block-level summary tokens to subsequent blocks, creating pathways for backward information flow.<n>As a simple, architecture-agnostic method, HTP enhances both zero-shot and finetuned models, offering a scalable route to superior long-document embeddings.
  arXiv  Detail & Related papers  (2025-11-18T19:37:40Z)
• Attention Illuminates LLM Reasoning: The Preplan-and-Anchor Rhythm Enables Fine-Grained Policy Optimization [56.083511902353365]
  Reinforcement learning (RL) typically applies uniform credit across an entire generation of Large language models.<n>This work positions attention as a privileged substrate that renders the internal logic of LLMs as a mechanistic blueprint of reasoning itself.<n>We introduce three novel RL strategies that dynamically perform targeted credit assignment to critical nodes.
  arXiv  Detail & Related papers  (2025-10-15T13:49:51Z)
• Multipole Attention for Efficient Long Context Reasoning [64.94673641704289]
  Large Reasoning Models (LRMs) have shown promising accuracy improvements on complex problem-solving tasks.<n>LRMs need to generate long chain-of-thought reasoning in order to think before answering.<n>We introduce Multipole Attention, which accelerates autoregressive reasoning by only computing exact attention for the most important tokens.
  arXiv  Detail & Related papers  (2025-06-16T03:00:40Z)
• Core Context Aware Transformers for Long Context Language Modeling [50.774702091154204]
  We propose a plug-and-play Core Context Aware (CCA) Attention for efficient long-context modeling.<n>Our method automatically focuses and strengthens core context while diminishing redundancy during the learning process.<n>Our method is able to replace the self-attention module in existing Large Language Models with minimal fine-tuning cost.
  arXiv  Detail & Related papers  (2024-12-17T01:54:08Z)
• RefreshKV: Updating Small KV Cache During Long-form Generation [54.00118604124301]
  We propose a new inference method, RefreshKV, that flexibly alternates between full context attention and attention over a subset of input tokens during generation.<n>Applying our method to off-the-shelf LLMs achieves comparable speedup to eviction-based methods while improving performance for various long-form generation tasks.
  arXiv  Detail & Related papers  (2024-11-08T18:57:07Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.