Punctuation-aware Hybrid Trainable Sparse Attention for Large Language Models

TL;DR

The paper tackles the high computational cost of long-context processing in large language models by combining a trainable sparse attention approach with punctuation-aware semantics. It introduces Punctuation-aware Hybrid Sparse Attention (PHSA), which builds dual-branch block representations—global and punctuation-enhanced—fusing them via a gate and scoring blocks to preserve fine-grained semantics while reducing attention to a Top-K subset. The method includes a block-level pipeline with initial and local-window priority and a flexible extreme-sparsity training regime that stabilizes performance when only a small fraction of tokens are activated. Empirical results show PHSA outperforms dense attention and state-of-the-art sparse baselines on general benchmarks and long-context tasks, reduces information loss by about 10% under high sparsity, and demonstrates cross-lingual benefits when punctuation sets are extended; this supports efficient long-context processing and broader deployment on resource-constrained devices. $PHSA$ thus offers a practical, scalable path for extending long-context capabilities in LLMs without prohibitive compute or memory costs, while also providing theoretical and empirical insights into boundary-aware sparse attention.

Abstract

Attention serves as the fundamental mechanism for long-context modeling in large language models (LLMs), yet dense attention becomes structurally prohibitive for long sequences due to its quadratic complexity. Consequently, sparse attention has received increasing attention as a scalable alternative. However, existing sparse attention methods rely on coarse-grained semantic representations during block selection, which blur intra-block semantic boundaries and lead to the loss of critical information. To address this issue, we propose \textbf{P}unctuation-aware \textbf{H}ybrid \textbf{S}parse \textbf{A}ttention \textbf{(PHSA)}, a natively trainable sparse attention framework that leverages punctuation tokens as semantic boundary anchors. Specifically, (1) we design a dual-branch aggregation mechanism that fuses global semantic representations with punctuation-enhanced boundary features, preserving the core semantic structure while introducing almost no additional computational overhead; (2) we introduce an extreme-sparsity-adaptive training and inference strategy that stabilizes model behavior under very low token activation ratios; Extensive experiments on general benchmarks and long-context evaluations demonstrate that PHSA consistently outperforms dense attention and state-of-the-art sparse attention baselines, including InfLLM v2. Specifically, for the 0.6B-parameter model with 32k-token input sequences, PHSA can reduce the information loss by 10.8\% at a sparsity ratio of 97.3\%.

Figures (3)

• Figure 1: Comparison of the computational complexity of single-layer attention under dense (blue) and sparse (orange) settings for different sequence lengths.
• Figure 2: Pipeline of PHSA. By splitting key blocks, we first compute the dual-branch representative tokens for each block, then calculate the relevance score with the query. Finally, we select the Top-K blocks, which, together with the initial block and local window blocks, form the to-be-computed indices for sparse attention.
• Figure 3: Comparison of NIAH scores for PHSA and InfLLM v2 (Training vs. Training-free) at 32k sequence.