pith. machine review for the scientific record. sign in

arxiv: 2604.12056 · v1 · submitted 2026-04-13 · 💻 cs.CL · cs.LG

Recognition: unknown

LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models

Haocheng Xi , Harman Singh , Yuezhou Hu , Coleman Hooper , Rishabh Tiwari , Aditya Tomar , Minjae Lee , Wonjun Kang
show 4 more authors
Michael Mahoney Chenfeng Xu Kurt Keutzer Amir Gholami
Authors on Pith no claims yet

Pith reviewed 2026-05-10 15:01 UTC · model grok-4.3

classification 💻 cs.CL cs.LG
keywords block-wise diffusion language modelssparse attentionKV cache reusedenoising stepslocality aware attentionattention efficiencynon-autoregressive generation
0
0 comments X

The pith

Block-wise diffusion language models reuse attention results for stable tokens to enable efficient sparse attention.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper tries to establish that block-wise diffusion language models can overcome their attention bottleneck in long contexts by exploiting the fact that most tokens remain nearly constant between denoising steps. It shows that only a small fraction of tokens change significantly, so cached attention results can be reused for them while sparse attention is applied only to the active minority. This avoids the KV inflation problem where different queries would otherwise pull in too many different key-value pages. A sympathetic reader would care because these models generate tokens in any order rather than left-to-right, and the method makes that advantage practical without dense attention costs. The outcome is higher accuracy at aggressive sparsity levels together with measurable GPU speedups.

Core claim

LOSA reuses cached prefix-attention results for stable tokens between consecutive denoising steps and applies sparse attention only to active tokens with significant hidden-state changes. This substantially shrinks the number of KV indices that must be loaded, yielding both higher speedup and higher accuracy than naive sparse attention across multiple block-wise DLMs and benchmarks.

What carries the argument

Locality-aware Sparse Attention (LOSA), the mechanism that separates active tokens from stable ones and reuses full cached attention computations for the stable majority.

If this is right

  • Preserves near-dense accuracy while maintaining up to 1.54 times lower attention density.
  • Delivers up to 4.14 times attention speedup on RTX A6000 GPUs.
  • Achieves up to 9 points higher average accuracy at aggressive sparsity levels compared with naive sparse methods.
  • Resolves the KV inflation problem that prevents uniform sparse attention from working in diffusion language models.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same stability pattern may appear in other iterative token-refinement processes, allowing similar caching outside diffusion models.
  • Pairing the method with hardware-aware KV management could extend usable context lengths further.
  • Direct measurements of active-token fractions on larger models would test whether the stability assumption holds at scale.

Load-bearing premise

Between consecutive denoising steps only a small fraction of tokens exhibit significant hidden-state changes while the majority remain nearly constant.

What would settle it

Count the actual fraction of tokens whose hidden states change by more than a small threshold between denoising steps on a long-context benchmark; if this fraction grows large, accuracy will drop when the reuse strategy is applied.

Figures

Figures reproduced from arXiv: 2604.12056 by Aditya Tomar, Amir Gholami, Chenfeng Xu, Coleman Hooper, Haocheng Xi, Harman Singh, Kurt Keutzer, Michael Mahoney, Minjae Lee, Rishabh Tiwari, Wonjun Kang, Yuezhou Hu.

Figure 1
Figure 1. Figure 1: Illustration of the KV-union effect: in block diffusion, each query selects a small set of prefix KV positions, but the effec￾tive cost is determined by the union across the block, inflating KV access. Our method computes sparse attention only for selected (active) tokens during denoising and reuses cached results for the others (static tokens), substantially reducing the size of the KV union and the laten… view at source ↗
Figure 2
Figure 2. Figure 2: Visualizing representation change locality across denoising steps. Left and middle plots show the L2 norm of query vectors at timesteps t − 1 and t respectively. The right plot shows the per-token MSE between query vectors across the two timesteps. Only a small fraction of tokens exhibit large changes, motivating the reuse of cached-prefix attention for stable tokens. step, all tokens in the block must att… view at source ↗
Figure 3
Figure 3. Figure 3: KV-cache load reduction with locality-aware sparse prefix attention. X-axis: number of active query tokens in the block; Y-axis: percentage of the full KV cache that must be loaded. Evaluated on Trado-8B-Instruct on TriviaQA with block size 16 and 64K context length. LOSA only loads prefix KV positions in the union I = S i∈A Ii for the active tokens set A, rather than the union over all B queries. The dash… view at source ↗
Figure 4
Figure 4. Figure 4: Visualizing locality across denoising steps. We plot the distribution of per-token MSE changes in queries between steps t−1 and t. For each layer, we sort the tokens in decreasing order of MSE changes from top to bottom in this heatmap. We observe that only a small fraction of tokens exhibit large changes (we show a threshold line which corresponds to 50% of the total MSE change is exhibited by tokens abov… view at source ↗
Figure 5
Figure 5. Figure 5: Overview of the locality-aware sparse prefix attention workflow. We (i) measure per-token query changes and select the changed-token set A, (ii) run QUEST to obtain per-token prefix indices and take their union I, (iii) load prefix KV for I and compute updated prefix statistics for queries in A while reusing cached statistics for stable tokens, (iv) compute dense within-block attention, and (v) merge prefi… view at source ↗
Figure 6
Figure 6. Figure 6: Latency breakdown comparison on prefix attention using Trado-8B-Instruct on TriviaQA with RTX A6000. Left: 64K context length with block size 16. Right: 32K context length with block size 32. LOSA and QUEST (Adapted) achieve 4.14× and 3.26× speedup over Dense Attention, respectively, by reducing memory-bound attention operations through sparsity. LOSA further improves over QUEST through locality-aware reus… view at source ↗
Figure 7
Figure 7. Figure 7: Visualizing representation change locality across denoising steps for query, keys and values of the token in the block that is being decoded. (Top) L2 norm of vectors at step t−1. (Middle) L2 norm of vectors at step t. (Bottom) Per-token change in representations, measured using MSE between steps t−1 and t. Only a small fraction of tokens exhibit large changes, motivating reuse of cached-prefix attention f… view at source ↗
Figure 8
Figure 8. Figure 8: Combined QKV Normalized MSE Heatmap with Threshold Selection (τ = 0.5) for Data Sample 0. The red line indicates the boundary where cumulative normalized MSE reaches 50% of the layer total. Tokens above the line are selected; tokens below are dimmed [PITH_FULL_IMAGE:figures/full_fig_p015_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: Combined QKV Normalized MSE Heatmap with Threshold Selection (τ = 0.5) for Data Sample 100 [PITH_FULL_IMAGE:figures/full_fig_p015_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: Combined QKV Normalized MSE Heatmap with Threshold Selection (τ = 0.5) for Data Sample 200. 15 [PITH_FULL_IMAGE:figures/full_fig_p015_10.png] view at source ↗
Figure 11
Figure 11. Figure 11: Combined QKV Normalized MSE Heatmap with Threshold Selection (τ = 0.5) for Data Sample 300. 16 [PITH_FULL_IMAGE:figures/full_fig_p016_11.png] view at source ↗
read the original abstract

Block-wise diffusion language models (DLMs) generate multiple tokens in any order, offering a promising alternative to the autoregressive decoding pipeline. However, they still remain bottlenecked by memory-bound attention in long-context scenarios. Naive sparse attention fails on DLMs due to a KV Inflation problem, where different queries select different prefix positions, making the union of accessed KV pages large. To address this, we observe that between consecutive denoising steps, only a small fraction of active tokens exhibit significant hidden-state changes, while the majority of stable tokens remain nearly constant. Based on this insight, we propose LOSA (Locality-aware Sparse Attention), which reuses cached prefix-attention results for stable tokens and applies sparse attention only to active tokens. This substantially shrinks the number of KV indices that must be loaded, yielding both higher speedup and higher accuracy. Across multiple block-wise DLMs and benchmarks, LOSA preserves near-dense accuracy while significantly improving efficiency, achieving up to +9 points in average accuracy at aggressive sparsity levels while maintaining 1.54x lower attention density. It also achieves up to 4.14x attention speedup on RTX A6000 GPUs, demonstrating the effectiveness of the proposed method.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

3 major / 1 minor

Summary. The paper proposes LoSA for block-wise diffusion language models (DLMs) to address memory-bound attention in long contexts. It identifies a KV Inflation problem with naive sparse attention and observes that between consecutive denoising steps only a small fraction of tokens exhibit significant hidden-state changes (active tokens) while most remain nearly constant (stable tokens). LOSA reuses cached prefix-attention results for stable tokens and computes sparse attention only for active tokens, claiming this yields higher efficiency and accuracy than dense or naive sparse baselines.

Significance. If the locality assumption holds with bounded error, LoSA could enable practical long-context inference for non-autoregressive diffusion LMs by reducing attention density (claimed 1.54x lower) and delivering speedups (up to 4.14x on RTX A6000) while preserving or improving accuracy (up to +9 points). The approach is a targeted engineering insight rather than a new theoretical framework, but its impact would be high for deployment of block-wise DLMs if the empirical claims are reproducible.

major comments (3)
  1. [Abstract] Abstract: the central claims of +9 average accuracy points at aggressive sparsity, 1.54x lower attention density, and 4.14x speedup are stated without any implementation details on active-token detection, error bars, baseline comparisons, or ablation on the stability threshold, rendering the claims unverifiable from the provided text.
  2. [Method] Method description (core assumption): the claim that hidden-state changes for stable tokens are small enough to allow safe reuse of cached attention results without accuracy loss lacks quantitative bounds on per-step change magnitudes, analysis of error accumulation over multiple denoising iterations, or sensitivity tests to small key/value perturbations in long contexts.
  3. [Experiments] Experiments: no ablation studies or controls are described to isolate the contribution of the locality-aware reuse versus other factors, and the reported accuracy gains at low density are presented without statistical significance or comparison to standard sparse attention variants that might also mitigate KV inflation.
minor comments (1)
  1. [Method] Notation for active/stable token classification and the precise reuse rule for prefix attention should be formalized with pseudocode or equations to improve reproducibility.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the constructive and detailed feedback. We address each major comment point by point below, indicating planned revisions to improve clarity, rigor, and verifiability while preserving the core contributions of the work.

read point-by-point responses

  1. Referee: [Abstract] Abstract: the central claims of +9 average accuracy points at aggressive sparsity, 1.54x lower attention density, and 4.14x speedup are stated without any implementation details on active-token detection, error bars, baseline comparisons, or ablation on the stability threshold, rendering the claims unverifiable from the provided text.

    Authors: We agree that the abstract prioritizes brevity and therefore omits granular details. The full manuscript describes active-token detection via hidden-state change thresholds in Section 3.2, reports baseline comparisons against dense and naive sparse attention in Section 4, and presents ablations on the stability threshold in Section 4.3. Error bars appear in the main result figures. To address the concern, we will revise the abstract to briefly note the active-token mechanism and reference the experimental sections for supporting details and ablations. revision: partial

  2. Referee: [Method] Method description (core assumption): the claim that hidden-state changes for stable tokens are small enough to allow safe reuse of cached attention results without accuracy loss lacks quantitative bounds on per-step change magnitudes, analysis of error accumulation over multiple denoising iterations, or sensitivity tests to small key/value perturbations in long contexts.

    Authors: Section 3.1 provides empirical visualizations showing that stable-token hidden states remain nearly constant between denoising steps. We acknowledge the request for stronger quantitative support. In the revised manuscript we will add explicit bounds by reporting per-step L2-norm statistics of hidden-state differences for stable tokens, include an error-accumulation study measuring accuracy impact across increasing denoising steps when caching is enabled, and conduct sensitivity experiments that apply controlled perturbations to cached KV entries to quantify robustness. revision: yes

  3. Referee: [Experiments] Experiments: no ablation studies or controls are described to isolate the contribution of the locality-aware reuse versus other factors, and the reported accuracy gains at low density are presented without statistical significance or comparison to standard sparse attention variants that might also mitigate KV inflation.

    Authors: The manuscript already compares LoSA against dense attention and naive sparse attention to highlight the KV-inflation issue. We agree that further controls would strengthen the claims. We will add ablation experiments that disable the locality-aware caching component while keeping other factors fixed, report statistical significance via multiple random seeds with standard deviations and appropriate significance tests, and include additional baselines using standard sparse-attention techniques (e.g., importance-based or fixed-pattern sparsity) to isolate LoSA's specific benefit in the block-wise DLM setting. revision: yes

Circularity Check

0 steps flagged

No circularity: derivation rests on empirical observation without reduction to inputs or self-citations

full rationale

The paper's chain begins with an empirical observation about hidden-state stability between denoising steps in block-wise DLMs, which directly motivates the LOSA reuse of cached prefix attention for stable tokens and sparse computation for active ones. No equations, fitted parameters, or predictions are shown that reduce by construction to the inputs (e.g., no self-definitional reuse of the stability assumption as a derived result). The method applies this insight to shrink KV loading without invoking self-citations for uniqueness theorems, ansatzes smuggled via prior work, or renaming of known results. The central efficiency/accuracy claims follow from the practical application of the observation rather than tautological redefinition, making the derivation self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The central claim rests on one domain assumption about token dynamics during denoising; no free parameters or new entities are introduced in the abstract.

axioms (1)
  • domain assumption Between consecutive denoising steps, only a small fraction of active tokens exhibit significant hidden-state changes while the majority of stable tokens remain nearly constant.
    This observation directly determines which tokens receive new sparse attention versus cached results.

pith-pipeline@v0.9.0 · 5550 in / 1278 out tokens · 63638 ms · 2026-05-10T15:01:22.382968+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

5 extracted references · 4 canonical work pages · 3 internal anchors

  1. [1]

    Generating Long Sequences with Sparse Transformers

    1, 2, 6, 9 Child, R., Gray, S., Radford, A., and Sutskever, I. Gen- erating long sequences with sparse transformers.arXiv preprint arXiv:1904.10509, 2019. 9 Choromanski, K. M., Likhosherstov, V ., Dohan, D., Song, X., Gane, A., Sarlós, T., Hawkins, P., Davis, J. Q., Mo- hiuddin, A., Kaiser, L., Belanger, D. B., Colwell, L. J., and Weller, A. Rethinking at...

    work page internal anchor Pith review arXiv 1904

  2. [2]

    TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension

    9 Joshi, M., Choi, E., Weld, D. S., and Zettlemoyer, L. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension.arXiv preprint arXiv:1705.03551, 2017. 7 10 LoSA: Locality Aware Sparse Attention for Block-Wise Diffusion Language Models Kang, W., Galim, K., Oh, S., Lee, M., Zeng, Y ., Zhang, S., Hooper, C., Hu, Y ., Koo, H. ...

    work page internal anchor Pith review arXiv 2017

  3. [3]

    Reformer: The efficient transformer

    9 Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. In8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. 9 Koˇcisk`y, T., Schwarz, J., Blunsom, P., Dyer, C., Hermann, K. M., Melis, G., and Grefenstette, E. The narrativeqa reading comprehension cha...

    2020

  4. [4]

    Efficient Memory Management for Large Language Model Serving with PagedAttention

    7 Kwon, W., Li, Z., Zhuang, S., Sheng, Y ., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Ef- ficient memory management for large language model serving with pagedattention, 2023. URL https:// arxiv.org/abs/2309.06180. 9 Li, Y ., Huang, Y ., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., Cai, T., Lewis, P., and Chen, D. Snapkv: Llm kno...

    work page internal anchor Pith review arXiv 2023

  5. [5]

    Sparse-dllm: Accelerating diffusion llms with dynamic cache eviction, 2025

    7 Song, Y ., Liu, X., Li, R., Liu, Z., Huang, Z., Guo, Q., He, Z., and Qiu, X. Sparse-dllm: Accelerating diffusion llms with dynamic cache eviction, 2025. URL https: //arxiv.org/abs/2508.02558. 2, 9 Tang, J., Zhao, Y ., Zhu, K., Xiao, G., Kasikci, B., and Han, S. Quest: Query-aware sparsity for efficient long-context llm inference.arXiv preprint arXiv:240...

    work page arXiv 2025