cumulant-order-block-sparse-attention

name	cumulant-order-block-sparse-attention
description	COBS - block sparse attention selector that stores a compressed SECOND-order (cumulant) statistic per block instead of only first-order (attention mass), closing the gap to dense attention at ~15x less KV-cache read traffic. Use when building/improving block-sparse attention (NSA-style) for long-context LLMs and existing selectors under-select.

COBS — Cumulant Order Block Sparse Attention (from arXiv:2607.09052)

Fixes a specific failure mode in block-sparse attention (e.g. DeepSeek NSA): selection uses only a first-order approximation of per-block attention mass, so it under-selects important blocks.

When to use

• You run / build a block-sparse attention LLM and long-context retrieval quality is weak.
• You want to reduce KV-cache read traffic vs dense attention while keeping quality near dense.
• You're analyzing or extending NSA, Slide, or any "select top-k blocks" attention.

Key insight (the theory)

• Block selection reduces to ranking blocks by attention mass = sum of a block's attention scores. If you retrieve the highest-mass blocks, block-sparse ≈ dense quality.
• Exact mass needs reading every key (defeats the purpose), so selection approximates mass from a compressed summary.
• Existing methods approximate mass at first order only (mean of per-key importance). The paper shows via a cumulant expansion that first-order is insufficient; you need the second-order (variance/cumulant) term to select correctly.

Method

• Extend the per-block compressed summary to store a second-order statistic (e.g. mean AND a variance/cumulant of within-block attention scores), not just the mean.
• Selector scores each block by a cumulant-aware estimate of total mass; retrieve top-k.
• Build on NSA's three-branch design (compress / select / value) — only the select branch changes.

Results (paper, 32k RULER)

• NSA baseline mean 0.2999 → COBS 0.8195 (dense 0.9040) — closes ~86% of the gap.
• Only 1.21× the NSA baseline's KV-cache read traffic; 15.15× less than dense.
• Preserves short-context behavior; lower position-wise NLL than dense in their comparison.

Implementation steps

1. Locate the block-selection summary in your sparse-attention impl (NSA: per-block compression produces a single score used for ranking).
2. Add a second compressed channel per block storing the within-block second moment / cumulant.
3. Replace the ranking score with score_b = μ_b + α·(cumulant_b) (tune α); keep top-k blocks.
4. Keep compression/value branches unchanged.
5. Profile KV-read traffic and RULER/long-context retrieval to confirm the gap closes.

Pitfalls

• The second-order term must stay compressed (one or two scalars per block) or you lose the read-traffic advantage.
• α calibration matters; too high over-weights outliers, too low reverts to first-order behavior.
• Still needs the compression branch to produce a faithful low-order summary; garbage compression → garbage cumulant.
• Test on both long- and short-context — verify no regression at short context.

Verification

• Long-context retrieval benchmark (RULER / needle) score vs NSA baseline and dense.
• KV-cache read-traffic ratio vs dense and vs baseline.
• Short-context perplexity parity check.