COBS: Cumulant Order Block Sparse Attention

Background Block-sparse attention reads only a fraction of the KV cache per decode step by scoring contiguous key/value blocks with a cacheable, query-independent summary and attending in full only to the top-k, which makes it the most hardware-friendly form of sparse attention since it yields contiguous memory access and needs a far smaller top-k than per-token selection. Yet major recent open-weight LLMs use everything except block sparsity: dense GQA, latent KV compression (MLA), local+global windows, or fine-grained token selection, leaving block sparsity conspicuously absent. DeepSeek's Native Sparse Attention (NSA) is the representative query-aware block selector: a three-branch design (compression, selection, sliding window) in which the selection branch alone decides which distant blocks are read at full resolution.

배경 블록 단위 sparse attention은 decode 스텝마다 KV 캐시 전체가 아니라 캐시 가능한(query-독립적) 요약으로 점수를 매긴 연속 블록 중 top-k만 골라 그 블록의 key/value만 읽는 방식이다. 연속 블록 접근은 메모리 접근이 연속적이고 요구되는 top-k도 개별 토큰 선택보다 훨씬 작아, 하드웨어 친화성 면에서 sparse attention 계열 중 가장 유리하다. 그런데도 최근 주요 오픈웨이트 LLM들은 dense GQA, latent KV 압축(MLA), 로컬+전역 윈도우, 세밀한 토큰 선택 등 다른 모든 방식을 채택하면서 블록 sparsity만 눈에 띄게 비어 있다. DeepSeek의 Native Sparse Attention(NSA)은 압축·선택·슬라이딩 윈도우 3-branch 구조를 가진 대표적인 query-aware 블록 선택 방식으로, 선택 branch 하나가 어느 먼 블록을 원해상도로 읽을지 전부 결정한다.

Problem The paper traces block sparsity's absence to the selection branch specifically: existing selectors such as NSA's learned MLP pooling score a block from a single cached, query-independent vector, so the score enters only linearly (to first order) in the query. This is like judging which of several sealed boxes holds the most valuable gem by comparing only their average weight: two boxes with an identical average can hold completely different value distributions, and the average alone cannot tell them apart. The paper shows this first-order ceiling is structural, not a training deficiency: a more expressive MLP pooling does not help, because whatever function of the keys gets cached still enters the score affinely in the query.

문제 이 논문은 하드웨어 친화성이 가장 높은데도 블록 sparsity가 주요 LLM에서 눈에 띄게 부재한 이유를 선택 branch 자체의 결함에서 찾는다. NSA의 학습된 MLP 풀링 같은 기존 선택기는 블록마다 캐시된 query-독립적 벡터 하나로 점수를 매기는데, 이 요약이 query에 대해 오직 1차(선형)로만 점수에 들어간다. 이는 여러 상자 중 어느 것에 가장 값진 다이아몬드가 들어 있는지를 상자의 평균 무게만 보고 판단하는 것과 같다. 평균이 같은 두 상자도 내용물의 분포는 전혀 다를 수 있는데, 평균만으로는 그 차이를 구분할 수 없다. 논문은 이 1차 한계가 학습 부족이 아니라 근본적 한계임을 보인다. 더 표현력 있는 MLP 풀링을 써도 소용없는데, 캐시된 것이 key의 어떤 함수든 여전히 q에 대해 아핀하게만 점수에 들어가기 때문이다.

Prior limitations NSA's block score is an inner product between the query and an MLP-pooled key summary, and DeepSeek's CSA adds a gated pooling with a ReLU indexer, but both still cache one query-independent vector and score affinely in the query, so both are first-order in the paper's taxonomy. Quest stores an axis-aligned min/max box per block and scores by the best possible inner product inside that box, which is piecewise-linear rather than strictly affine and is empirically the strongest prior baseline, but it is only a coarse proxy for how the keys are actually spread. None of these selectors captures the true softmax mass a block would receive if all of its keys were actually read.

기존 한계 NSA의 블록 점수는 query와 MLP로 풀링한 key 요약 벡터의 내적이고, DeepSeek의 CSA는 여기에 ReLU 인덱서를 더한 게이트 풀링을 쓰는데, 둘 다 여전히 query-독립적 벡터 하나를 캐시하고 q에 대해 아핀하게 점수를 매겨 논문의 분류상 1차에 속한다. Quest는 블록별로 축-정렬된 key의 min/max 박스를 저장하고 그 박스 안에서 가능한 최대 내적으로 점수를 매기는데, 이는 q에 대해 아핀이 아니라 조각별 선형이라 실증적으로 가장 강한 기존 baseline이지만, key가 실제로 어떻게 퍼져 있는지에 대한 거친 근사에 불과하다. 이 선택기들 중 어느 것도 블록의 모든 key를 실제로 읽었을 때 softmax가 부여할 진짜 질량(mass)을 포착하지 못한다.

Goal The paper wants to pin down exactly what a cacheable per-block summary should compute for block selection to recover dense-attention-level accuracy, then build a selector that computes it cheaply. It first asks whether the selection criterion itself, i.e. which blocks to rank highest, is the bottleneck, or whether estimating that criterion from a compressed summary is.

목표 논문은 캐시 가능한 블록별 요약이 정확히 무엇을 계산해야 블록 선택이 dense attention 수준의 정확도를 회복할 수 있는지를 규명하고, 그것을 싸게 계산하는 선택기를 만들고자 한다. 먼저 선택 기준 자체, 즉 어느 블록을 최우선으로 랭크할지가 병목인지, 그 기준을 압축된 요약으로부터 추정하는 과정이 병목인지를 묻는다.

Method The authors formalize block selection as reconstructing dense attention output from a top-k block subset and show that, under mild assumptions, the optimal criterion is to rank blocks by attention mass, the sum of a block's unnormalized softmax scores. An oracle that ranks by the exact mass (OSA, computed by reading every key) reaches 0.9010 on their 32k RULER suite, within 0.3 points of dense attention's 0.9040, confirming mass-ranking is essentially sufficient and the real difficulty is estimating mass cheaply. COBS (Cumulant Order Block Sparse Attention) gets there via a cumulant expansion: a block's log-mass equals ln L plus the cumulant-generating function of its key distribution evaluated at the query, whose Taylor series has the mean key as the first-order term (what mean-pooling already captures) and the within-block key covariance as the second-order term (what every existing selector discards). COBS stores, per block, the mean key plus a rank-r, FP4-quantized, query-subspace-projected version of that covariance, and scores blocks with the resulting quadratic-form estimate of log-mass.

방법 저자들은 블록 선택을 top-k 블록 부분집합으로부터 dense attention 출력을 재구성하는 문제로 정식화하고, 온건한 가정 아래 최적 기준이 attention mass(블록의 정규화 전 softmax 점수 합)로 블록을 랭크하는 것임을 보인다. 정확한 mass로 랭크하는 오라클(OSA, 모든 key를 읽어 계산)은 32k RULER 스위트에서 0.9010을 기록해 dense attention의 0.9040과 0.3점 차로 근접하며, 이는 mass 랭킹 자체는 사실상 충분하고 진짜 어려움은 mass를 싸게 추정하는 데 있음을 확인해준다. COBS(Cumulant Order Block Sparse Attention)는 cumulant 전개로 이를 달성한다. 블록의 log-mass는 정확히 ln L 더하기 그 key 분포의 cumulant-생성함수를 query에서 평가한 값이며, 이 함수의 Taylor 급수는 1차 항으로 평균 key(이미 mean-pooling이 포착하는 항)를, 2차 항으로 블록 내 key 공분산(기존 모든 선택기가 버리는 항)을 갖는다. COBS는 블록마다 평균 key와, 랭크-r로 압축하고 FP4로 양자화하고 query-부분공간에 투영한 공분산을 저장해, 그 결과로 얻는 log-mass의 이차형식 추정치로 블록 점수를 매긴다.

Key idea The core move is diagnosing exactly which term of an exact mathematical expansion every existing cacheable selector is missing, the within-block covariance, or curvature, rather than proposing a better heuristic. Two blocks can share an identical mean key while one has its key-spread aligned with the query direction (raising true softmax mass) and the other's spread is orthogonal to it (leaving mass unchanged); a mean-only score is blind to this by construction, like a compass that tells you a direction but not how far a ship has actually drifted off course along it.

핵심 아이디어 핵심은 더 나은 휴리스틱을 제안하는 것이 아니라, 정확한 수학적 전개 중 기존의 모든 캐시 가능한 선택기가 정확히 어느 항을 놓치고 있는지, 즉 블록 내 공분산인 곡률(curvature)을 진단하는 데 있다. 두 블록이 완전히 같은 평균 key를 가지면서도 한 블록은 key의 퍼짐이 query 방향과 정렬되어(진짜 softmax mass를 높이고) 다른 블록의 퍼짐은 query와 직교해(mass에 영향 없음) 있을 수 있는데, 1차(평균만 보는) 점수는 구조적으로 이 둘을 구별하지 못한다. 방향은 알려주지만 배가 그 방향을 따라 얼마나 항로를 벗어났는지는 알려주지 못하는 나침반과 같다.

Validation COBS is compared against a controlled NSA replication, a mean-pool ablation, and Quest's min/max scorer, all sharing the same base architecture, compression branch, and RULER-style supervised fine-tuning and differing only in the selection-branch summary, on an 11-task 32k-context RULER suite, seven short-context common-sense benchmarks (a parity check), and position-wise next-token perplexity on long natural text. Ablations sweep the stored covariance rank (r = 4 to 31), compare a fixed global query-subspace budget against an adaptive per-layer one, test FP4 versus bf16 quantization of the stored eigenvectors, and report two pre-registered negative results.

검증 COBS는 통제된 NSA 재현(NSA MLP), mean-pool ablation, Quest의 min/max 스코어러와, 모두 같은 기본 아키텍처·압축 branch·RULER 스타일 지도 미세조정을 공유하고 선택 branch 요약만 다른 채, 11-task 32k-context RULER 스위트, 7개 짧은-문맥 상식 벤치마크(패리티 체크), 긴 자연어 텍스트에 대한 위치별 다음-토큰 perplexity로 비교된다. Ablation은 저장 공분산 랭크(r=4~31)를 스윕하고, 고정 전역 query-부분공간 예산과 층별 적응형 예산을 비교하며, 저장된 고유벡터의 FP4 대 bf16 양자화를 시험하고, 사전 등록된 두 개의 부정적 결과를 보고한다.

Results COBS (subspace, rank-4, FP4) reaches a mean score of 0.8195 on 32k RULER, closing 86% of the gap between the NSA baseline (0.2999) and dense attention (0.9040), while using only 1.21x the KV-cache read traffic of that NSA baseline and 15.15x less than dense attention. Quest, the strongest prior selector in this controlled comparison, reaches only 0.5765 (46% of the gap), and NSA's learned MLP pooling, despite more pretraining tokens than every other variant, trails at 0.2999-0.2822. On held-out long natural text, COBS also attains the lowest position-wise perplexity of any tested variant (1.633 versus 1.727 for dense), with a flat long-context slope indicating genuine conditioning on distant tokens rather than reliance on the local window.

결과 COBS(부분공간, 랭크-4, FP4)는 32k RULER에서 평균 0.8195를 기록해 NSA baseline(0.2999)과 dense attention(0.9040) 사이 격차의 86%를 좁히면서, NSA baseline KV 캐시 읽기 트래픽의 1.21배만 사용하고 dense attention보다 15.15배 적게 읽는다. 이 통제된 비교에서 가장 강한 기존 선택기인 Quest는 0.5765(격차의 46%)에 그치고, 다른 모든 변형보다 더 많은 사전학습 토큰을 쓴 NSA의 학습된 MLP 풀링조차 0.2999~0.2822에 머물러 있다. held-out 긴 자연어 텍스트에서도 COBS는 테스트된 모든 변형 중 가장 낮은 위치별 perplexity(1.633, dense는 1.727)를 달성하며, 평평한 장문맥 기울기는 로컬 윈도우에 기대지 않고 실제로 먼 토큰을 조건으로 삼고 있음을 보여준다.

Comparison Against NSA's MLP-pooling selector and DeepSeek-V4's CSA gated-pooling selector, COBS wins by construction, since both remain affine in the query regardless of how expressive their pooling function is; the controlled ablation confirms a fancier pooling network does not help, because the summary is still first-order. Against Quest's coarser min/max-box proxy for key spread, COBS's exact within-block covariance is the principled form of the same intuition, delivering roughly 40 points higher RULER accuracy at comparable or lower read traffic. The closest conceptual relative, CCQ, uses the identical cumulant-generating-function insight but applies it to query correction inside linear attention rather than to block selection in sparse attention: same mathematics, different premise and method.

비교 NSA의 MLP 풀링 선택기와 DeepSeek-V4의 CSA 게이트 풀링 선택기 대비, COBS는 구조적으로 이긴다. 풀링 함수가 아무리 표현력이 좋아도 둘 다 query에 대해 여전히 아핀이기 때문이며, 논문의 통제된 ablation은 더 화려한 풀링 네트워크가 도움이 되지 않음을 확인한다(요약이 여전히 1차이므로). key 퍼짐에 대한 더 거친 min/max-박스 근사인 Quest 대비, COBS의 정확한 블록 내 공분산은 같은 직관의 원리적 형태로, 비슷하거나 더 낮은 읽기 트래픽에서 약 40점 높은 RULER 정확도를 낸다. 가장 가까운 개념적 친척인 CCQ는 동일한 cumulant-생성함수 통찰을 사용하지만 sparse attention의 블록 선택이 아니라 linear attention 내부의 query 보정에 적용한다. 수학은 같지만 전제와 방법이 다르다.

Significance The paper reframes what should be cached for block-sparse attention: not a single learned or heuristic vector, but the specific next term of an exact statistical expansion, giving both a diagnostic lens (what cumulant order does a selector implicitly use?) and a concrete design axis for better selectors. Since block sparsity is the sparse-attention family most compatible with contiguous, hardware-friendly memory access, closing most of its accuracy gap to dense attention at a small traffic overhead is a step toward making it viable where it is currently absent (reviewer judgment).

의의 이 논문은 블록 sparse attention을 위해 무엇을 캐시해야 하는지를 재정의한다. 학습되거나 휴리스틱적인 벡터 하나가 아니라 정확한 통계적 전개의 바로 다음 항을 캐시하라는 것이며, 이는 진단 렌즈(선택기가 암묵적으로 몇 차 cumulant를 쓰는가)이자 더 나은 선택기를 위한 구체적 설계 축이 된다. 블록 sparsity가 연속적이고 하드웨어 친화적인 메모리 접근과 가장 잘 맞는 sparse attention 계열이라는 점에서, 작은 트래픽 오버헤드로 dense attention과의 정확도 격차 대부분을 좁히는 것은 현재 그것이 부재한 곳에서 상용화로 가는 한 걸음이다(리뷰어 판단).

Limitations The authors state four limitations: the study is done at mechanism-study scale (about 1.2B parameters, 4k pretraining length extended to 32k via YaRN), not a deployment-scale system; their NSA replication differs from the original NSA's hyperparameters and is not meant to reproduce NSA's published numbers; the sparse variants use a NoPE (no rotary position encoding) scheme in the compression/selection branches that the dense baseline does not use, so long-context comparisons partly reflect this positional-encoding difference rather than sparsity alone; and their long-context signal comes from RULER-style supervised fine-tuning, a nonstandard protocol whose absolute numbers may not transfer, though the identical protocol across all selectors still makes relative comparisons meaningful. They also note COBS costs more read traffic than plain mean-pooling before quantization, and that KV-traffic accounting alone does not guarantee end-to-end runtime gains, which depend on kernels, batching, and hardware.

한계 저자들은 네 가지 한계를 명시한다. 이 연구는 배포 규모 시스템이 아니라 메커니즘 연구 규모(약 1.2B 파라미터, 4k 사전학습 길이를 YaRN으로 32k까지 확장)에서 수행됐고, NSA 재현은 원 NSA의 하이퍼파라미터와 달라 NSA의 발표 수치를 재현하려는 것이 아니며, sparse 변형은 dense baseline이 쓰지 않는 NoPE(RoPE 없는) 방식을 압축·선택 branch에 사용해 장문맥 비교가 부분적으로 이 위치 인코딩 차이를 반영할 뿐 sparsity 단독 효과가 아니며, 장문맥 신호는 비표준적인 RULER 스타일 지도 미세조정에서 나와 절대 수치는 다른 설정으로 전이되지 않을 수 있지만 모든 선택기에 동일하게 관대한 프로토콜이므로 상대 비교는 여전히 유의미하다. 또한 COBS는 양자화 전에는 단순 mean-pooling보다 더 많은 읽기 트래픽을 쓰며, KV 트래픽 회계 자체는 커널·배칭·하드웨어에 좌우되는 end-to-end 런타임 이득을 보장하지 않는다고도 명시한다.

Future work The paper states no explicit future-work section; its two labeled negative results point toward open directions instead: a query-centered cumulant expansion regressed performance because the tilted moments must be shared across a GQA group and degrade for outlier queries far from the calibration point, and a cheap signed third-cumulant correction hurt at the low ranks COBS actually operates in, though it partially repaired an observed high-rank regression caused by unsigned-variance false positives. Extending the selector to fix this high-rank regression, or to the token-level DeepSeek Sparse Attention and long-span HCA mechanisms the authors call orthogonal, are directions the paper gestures at without formally proposing (reviewer judgment).

향후 과제 논문에는 명시적 future-work 절이 없다. 대신 두 개의 부정적 결과가 열린 방향을 시사한다. query-중심 cumulant 전개는 기울어진 모멘트를 GQA 그룹 전체가 공유해야 하고 보정 지점에서 먼 이상치 query에서 저하되어 성능이 떨어졌고, 저비용 부호 있는 3차 cumulant 보정은 COBS가 실제로 쓰는 저랭크 영역에서는 오히려 해를 끼쳤다(다만 부호 없는 분산이 유발한 고랭크 영역의 관측된 성능 저하는 부분적으로 복구했다). 이 고랭크 저하를 명시적으로 고치는 방향, 혹은 저자들이 직교적이라 칭한 토큰 단위 DeepSeek Sparse Attention·장거리 HCA 압축 메커니즘으로의 확장은 논문이 공식적으로 제안하지는 않지만 자연스럽게 가리키는 다음 단계다(리뷰어 판단).

Released resources No public code or checkpoint release could be verified for COBS; the paper lists no GitHub or project-page link in its text, references, or arXiv listing, and the authors are affiliated with MatX, an AI training-hardware startup, rather than a lab with an established public-release pattern.

자원 공개 COBS의 공개 코드나 체크포인트 배포는 확인되지 않았다. 논문 본문·참고문헌·arXiv 목록 어디에도 GitHub나 프로젝트 페이지 링크가 없으며, 저자들은 공개 배포 관행이 확립된 연구실이 아니라 AI 학습 하드웨어 스타트업 MatX 소속이다.