# COBS: Cumulant Order Block Sparse Attention

> COBS stores compressed per-block key covariances to raise block sparse attention selection from first-order to second-order approximations, closing 86% of the gap to dense attention with minimal extra KV cache reads.

- **Source:** [arXiv](https://arxiv.org/abs/2607.09052)
- **Published:** 2026-08-22
- **Permalink:** https://picx.dev/p/M6UsdR
- **Whiteboard:** https://picx.dev/p/M6UsdR/image

## Summary

# COBS: Cumulant Order Block Sparse Attention

## Summary (Overview)

- **Core problem**: Block sparse attention is hardware-friendly for reducing KV cache read bottlenecks in LLMs, but is underused in leading open-weight models. The paper studies why via DeepSeek's Native Sparse Attention (NSA) as a representative method.
- **Key theoretical contribution**: Block selection reduces to ranking blocks by their **attention mass** (sum of a block's softmax scores). An oracle selector (OSA) using exact masses essentially matches dense attention (0.9010 vs 0.9040 mean RULER score), proving the criterion is sound — the challenge is estimating mass from a cacheable summary.
- **Main limitation of existing selectors**: Via a **cumulant expansion** of the block mass, the authors show existing selectors (NSA MLP, mean-pool, Quest) are confined to **first-order approximations** in the query, discarding the within-block key covariance that supplies the second-order curvature term $\frac{1}{2} q^\top \Sigma_b q$.
- **Proposed method (COBS)**: Stores a compressed second-order statistic per block — the within-block key covariance $\Sigma_b$ in low-rank, query-subspace-projected, FP4-quantized form — raising the cumulant order of the cached summary while respecting the cacheability constraint.
- **Empirical results**: COBS (subspace $s \approx 85$, $r=4$, FP4) raises the 11-task 32k RULER mean score from 0.2999 (NSA baseline) to 0.8195, closing ~86% of the gap to dense attention (0.9040), while using only 1.21× the NSA baseline's KV cache read traffic and 15.15× less than dense.

## Introduction and Theoretical Foundation

### Background

Transformer inference at long context is bottlenecked by **KV cache reads**: at each decode step, attention reads the keys and values of every past token, making decoding memory-bandwidth-bound rather than compute-bound. Sparse attention reduces this by reading only a fraction of past tokens.

Sparse attention methods differ in what they keep and how they choose it:
- **Fixed-pattern**: Longformer, BigBird (local windows + global tokens)
- **KV-eviction**: StreamingLLM, H2O, SnapKV (drop low-importance tokens)
- **Low-rank**: compress keys themselves
- **Query-aware block selectors**: Quest, NSA — summarize contiguous blocks offline, then run fine-grained attention over the highest-scoring blocks

Block sparsity is the most hardware-friendly form (contiguous memory accesses, smaller top-k), yet it is largely absent from leading open-weight LLMs, which instead use dense GQA, latent KV compression (MLA), local windows with global layers, or fine-grained token selection (DSA).

### Theoretical Foundation

The paper studies NSA, which has three branches: (i) a compression branch over pooled block representations, (ii) a **selection branch** that ranks blocks by a lightweight score and runs fine-grained attention over the top-k, and (iii) a sliding-window branch. The selection branch is the focus: it alone determines which blocks fine-grained attention sees, and its scoring summary must be cached.

A selector is **cacheable** if its per-block summary can be precomputed and stored independently of the decode query $q$ — e.g., any function of the block's keys alone. This constraint makes the cumulant order of the stored summary the binding limitation.

## Methodology

### The Selection Oracle (Section 3)

The oracle formalizes block selection as reconstructing dense per-head attention outputs from a top-k block subset. With block masses $m_b$ and value centroids $v_b^c$ defined as:

$$s_r = q^\top k_r, \qquad m_b = \sum_{r \in b} e^{s_r}, \qquad v_b^c = \frac{1}{m_b} \sum_{r \in b} e^{s_r} v_r$$

The dense head output is a $P_b$-weighted mixture of centroids ($P_b = m_b/Z$). The exact per-head reconstruction error from dropping blocks is:

$$o^\star - \hat{o}_S = \frac{1}{1-\tau} \sum_{b \in S^c} P_b \left(v_b^c - o^\star\right)$$

Under three assumptions (value-agnosticism, disregarding per-head constants, and linear relaxation of the per-head penalty), minimizing the GQA selection objective reduces to ranking blocks by:

$$\boxed{\text{score}_b^{(h)} = \sum_{g=1}^{G} \frac{m_b^{(g,h)}}{Z^{(g,h)}}, \qquad Z^{(g,h)} = \sum_{b'} m_{b'}^{(g,h)}}$$

This is the **mass criterion**. OSA (Oracle Sparse Attention) applies this with exact masses (reading all keys), serving as a diagnostic that essentially matches dense attention.

### The Cumulant Expansion (Section 4)

The block mass is $L$ times the moment generating function of the block's empirical key distribution:

$$m_b = L \cdot \mathbb{E}_r\left[e^{q^\top k_r}\right] = L M_X(q)$$

so $\ln m_b = \ln L + K_X(q)$ where $K_X$ is the cumulant generating function. The multivariate cumulant expansion is:

$$K_X(q) = q^\top \kappa_1 + \frac{1}{2} q^\top \kappa_2 q + \frac{1}{6} \sum_{ijk} (\kappa_3)_{ijk} q_i q_j q_k + \dots$$

with first two cumulants:

$$\kappa_1 = \bar{k}_b = \frac{1}{L} \sum_{r \in b} k_r, \quad \kappa_2 = \Sigma_b = \frac{1}{L} \sum_{r \in b} (k_r - \bar{k}_b)(k_r - \bar{k}_b)^\top$$

### COBS Method (Section 5)

**Second-order truncation** gives the core estimator:

$$\ln m_b \approx \ln L + q^\top \bar{k}_b + \frac{1}{2} q^\top \Sigma_b q$$

The first-order term $q^\top \bar{k}_b$ is mean-pooling over block keys; the second-order term $\frac{1}{2} q^\top \Sigma_b q$ is the curvature that mean-pooling omits.

**Covariance compression** (for $D > L$): low-rank spectral decomposition keeping the top $r$ eigendirections:

$$\Sigma_b \approx \sum_{i=1}^{r} \lambda_i u_i u_i^\top = \sum_{i=1}^{r} \xi_i^{\text{fs}}(\xi_i^{\text{fs}})^\top, \qquad \xi_i^{\text{fs}} = \sqrt{\lambda_i} u_i$$

**Subspace method**: project covariance into the $s$-dimensional query subspace ($U_Q$ holds top $s$ eigenvectors of $\mathbb{E}[qq^\top]$):

$$B_b = U_Q^\top \Sigma_b U_Q \in \mathbb{R}^{s \times s}$$

with the projected quadratic form $q^\top \Sigma_b q \approx \tilde{q}^\top B_b \tilde{q}$ where $\tilde{q} = U_Q^\top q$.

**Quantization**: FP4 (E2M1) with one fp32 scale per eigenvector; block mean kept at bf16.

**Per-decode-step scoring cost**: $O(rs)$ per block (plus one $O(sD)$ query projection shared across all blocks) vs. $O(D)$ for mean-pool.

**Gram trick** for eigenvector computation: avoids materializing the full $D \times D$ covariance by using the $L \times L$ Gram matrix $\frac{1}{L}\tilde{K}\tilde{K}^\top$.

## Empirical Validation / Results

### Setup

- ≈1.2B parameter decoder-only transformer, 16 layers, model dim 2048, 16 query heads, 4 KV heads (GQA), head dim $D = 128$, RoPE with base $\theta = 10^6$
- Pretraining: LongCrawl64 for ≈20B tokens at 4k sequence length, then YaRN extension to 32k with RULER-style SFT
- NSA branches: $L = 32$-token blocks, top-k = 16, 256-token sliding window

### Headline Results (32k RULER)

**Table 1: 11-task 32k RULER breakdown** (gap closed measured from NSA MLP RoPE 0.2999 to dense 0.9040):

| Method | Mean | Gap closed | S1 | S2 | S3 | MK1 | MK2 | MK3 | MQ | MV | CWE | FWE | VT |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| Dense (full attention) | 0.9040 | 100.0% | 1.00 | 1.00 | 1.00 | 0.98 | 0.95 | 0.92 | 0.92 | 0.93 | 0.44 | 0.90 | 0.92 |
| OSA (mass oracle) | 0.9010 | 99.5% | 1.00 | 1.00 | 1.00 | 0.99 | 0.94 | 0.91 | 0.96 | 0.96 | 0.22 | 0.95 | 0.98 |
| NSA MLP (RoPE) | 0.2999 | 0.0% | 1.00 | 0.11 | 0.05 | 0.09 | 0.00 | 0.00 | 0.09 | 0.10 | 0.10 | 0.93 | 0.83 |
| NSA mean-pool (NoPE) | 0.5554 | 42.3% | 1.00 | 0.78 | 0.82 | 0.53 | 0.01 | 0.00 | 0.47 | 0.47 | 0.15 | 0.91 | 0.97 |
| NSA Quest | 0.5765 | 45.8% | 1.00 | 0.86 | 0.87 | 0.39 | 0.03 | 0.00 | 0.55 | 0.55 | 0.19 | 0.93 | 0.97 |
| COBS full-space r=4 | 0.8238 | 86.7% | 1.00 | 1.00 | 1.00 | 0.98 | 0.78 | 0.34 | 0.89 | 0.92 | 0.24 | 0.94 | 0.97 |
| COBS full-space r=6 | 0.8493 | 90.9% | 1.00 | 1.00 | 1.00 | 1.00 | 0.91 | 0.48 | 0.94 | 0.94 | 0.15 | 0.94 | 0.99 |
| **COBS (subspace s≈85, r=4, FP4)** | **0.8195** | **86.0%** | 1.00 | 1.00 | 1.00 | 0.97 | 0.79 | 0.31 | 0.89 | 0.91 | 0.23 | 0.94 | 0.97 |

Key findings:
- Removing RoPE from selection summaries (NoPE) is an additive improvement (mean-pool 0.4186 → 0.5554)
- NSA Quest improves only modestly over mean-pool (0.5765)
- COBS's covariance summary accounts for the large remaining gain

### KV Cache Read Traffic

**Table 3: Per-decode-step, per-layer KV cache read traffic at 32k (KiB)**:

| Method | Summary keys | Summary values | Window | Fine-grained | Per layer | vs. dense (×less) | vs. NSA MLP (×more) |
|---|---|---|---|---|---|---|---|
| Dense | - | - | - | 65,536 | 65,536 | - | 18.29× |
| OSA | 33,792 | 1024 | 512 | 1024 | 36,352 | 1.80× | 10.14× |
| NSA MLP | 1024 | 1024 | 512 | 1024 | 3584 | 18.29× | - |
| NSA mean-pool | 1024 | 1024 | 512 | 1024 | 3584 | 18.29× | 1.00× |
| NSA Quest | 3072 | 1024 | 512 | 1024 | 5632 | 11.64× | 1.57× |
| COBS full-space r=4 (FP4) | 2112 | 1024 | 512 | 1024 | 4672 | 14.03× | 1.30× |
| **COBS (subspace s≈85, r=4, FP4)** | **1767** | 1024 | 512 | 1024 | **4327** | **15.15×** | **1.21×** |

### Short-Context Parity and Language Modeling

- **Short-context common-sense reasoning**: All sparse variants preserve dense-like performance (38.0–38.6 avg vs. dense 38.2), as expected when the selection budget covers the input.
- **Position-wise NLL**: COBS has the lowest average NLL (1.633), below dense (1.727), NSA MLP (1.683), and mean-pool (1.745). COBS's NLL slope stays flat at long positions (no upturn), indicating it conditions on distant tokens rather than leaning on the local window.

### Ablations

**Rank sweep (Figure 5)**: Selection peaks at $r = 8$ (0.8539), regresses at $r = 16$ (0.8006), and stays below peak at max rank $r = 31$ (0.8135). The collapse concentrates in multi-key needle subtasks (MK3 falls from 0.470 at $r=8$ to 0.054 at $r=16$), attributed to the unsigned variance term boosting blocks with many distractors.

**Subspace (Table 4)**: Adaptive per-layer subspaces outperform a single global budget at comparable dimension. At average $s \approx 85$, score is within ≈0.005 of the full low-rank form ($s = 128$).

**Quantization (Table 5)**: FP4 (E2M1) is essentially lossless — full-space $r=4$ changes by only +0.0013 (0.8238 → 0.8251) while shrinking the descriptor ≈3.8× (1024 → 272 bytes per block).

**Negative results**:
- **Query-centered expansion**: Expanding the CGF around a calibrated origin $q_0$ instead of $q = 0$ regresses selection (0.8238 → 0.8100), concentrated on multi-key/multi-value needles.
- **Cheap diagonal skew**: Adding signed third-cumulant scalars $g_i$ per eigenvector hurts at low rank ($r=4$: 0.8238 → 0.7754) but partially repairs the high-rank regression ($r=16$: 0.8006 → 0.8252).

## Theoretical and Practical Implications

### Theoretical Implications

1. **Cumulant order as a diagnostic lens**: The paper establishes that existing cacheable block selectors are all confined to first-order approximations in the query. Any affine score of the form $\text{score}_b^{\text{aff}}(q) = a_b + q^\top \phi_b$ (mean-pool, NSA's MLP pooling, CSA's gated pooling) cannot capture the quadratic term $\frac{1}{2} q^\top \Sigma_b q$, which is not affine in $q$.

2. **The GQA cross-head nonlinearity cannot lift the ceiling**: An informal argument (Equation 20) shows that sharing a block set across heads is a constraint relative to independent per-head selection, so the first-order family's quality is bounded by a curvature-blind ceiling.

3. **The mass criterion is the right target**: OSA's near-dense performance (99.5% gap closed) proves that ranking by exact attention mass is sufficient; the bottleneck is estimation, not the criterion.

### Practical Implications

1. **Block sparse attention can approach dense quality**: COBS closes ~86% of the gap between first-order selection and dense attention at a fraction of the KV cache read traffic (15.15× less than dense).

2. **Hardware-friendly design**: The compressed covariance descriptor (low-rank + subspace + FP4) keeps the per-block summary small (272 bytes for full-space $r=4$), making the approach practical for deployment.

3. **Cumulant order is a design axis**: Raising the cumulant order of cached summaries is a principled way to improve cacheable selectors, applicable beyond NSA to the broader family of block sparse methods.

## Conclusion

The paper traces the challenges of block sparse attention to the selection branch and shows that selection reduces to ranking blocks by attention mass. A cumulant expansion reveals that existing selectors are first-order in the query, discarding the within-block key covariance. COBS keeps the cacheability constraint while raising the cumulant order, storing a compressed second-order statistic per block.

Empirically, COBS lifts 32k RULER from 0.2999 (NSA baseline) to 0.8195, closing ~86% of the gap to dense attention (0.9040), while using only 1.21× the NSA baseline's KV cache read traffic and 15.15× less than dense. The same model preserves short-context performance and attains the lowest position-wise NLL (1.633 vs. dense's 1.727).

**Future directions**: The paper is a mechanism study at ≈1.2B scale; deployment-scale validation is needed. The RULER-style SFT protocol is nonstandard, and the NoPE confound means long-context comparisons partly reflect position encoding differences. The authors see cumulant order as a step toward making block sparse attention more prevalent in leading open-weight LLMs.

**Limitations**: Scale (≈1.2B backbone), controlled NSA comparison (different hyperparameters, non-overlapping blocks), NoPE confound, RULER-style SFT protocol, and the fact that KV read accounting numbers don't by themselves imply end-to-end runtime gains (which depend on kernels, batching, hardware, and decoding regime).

---

_Markdown view of https://picx.dev/p/M6UsdR, served by PicX — AI-generated visual whiteboard summaries of research papers._