기술 요약: COBS (Cumulant Order Block Sparse Attention)

1. 문제 정의

긴 컨텍스트 길이를 가진 대규모 언어 모델(LLM)의 추론은 연산량(compute)보다는 메모리 대역폭(memory bandwidth)에 의해 병목 현상이 발생합니다. 구체적으로, 키-값(KV) 캐시 읽기 작업이 밀집 어텐션(dense attention)에서 요구하는 데코딩 속도를 제한하고 하드웨어 활용도를 낮춥니다. 블록 희소 어텐션(block sparse attention)은 연속적인 토큰 블록만을 읽음으로써 하드웨어 친화적인 해결책을 제공하지만, 주요 오픈 웨이트 모델들에서는 여전히 충분히 활용되지 못하고 있습니다. 기존 방식들은 밀집 어텐션, 미세한 토큰 선택 또는 KV 축출(eviction)에 의존하는 경우가 많습니다.

블록 희소 어텐션의 주요 과제는 **블록 선택(block selection)**입니다. 즉, 모든 키와 값의 전체 내용을 읽지 않고도 주어진 쿼리에 대해 가장 관련성이 높은 KV 캐시 블록을 식별하는 것입니다. DeepSeek의 Native Sparse Attention(NSA)과 같은 기존의 최첨단 방법들은 "1차(first-order)" 선택 전략을 사용합니다. 이러한 방법들은 블록을 쿼리와 독립적인 벡터(예: MLP 또는 평균 풀링을 통해)로 요약하고, 이를 쿼리에 대해 선형적으로 스코어링합니다. 본 논문은 이러한 1차 근사 방식이 어텐션 질량(attention mass)의 비선형적 곡률(non-linear curvature)을 포착하지 못하며, 이로 인해 밀집 어텐션과 상당한 성능 격차가 발생한다고 주장합니다.

2. 방법론

2.1 이론적 기초: 선택 오라클 (The Selection Oracle)

저자들은 블록 선택을 밀집 어텐션 출력을 일부 블록으로부터 재구성하는 문제로 정식화했습니다. 그들은 최적의 선택 기준이 블록 내의 어텐션 스코어 합인 **어텐션 질량($m_b$)**에 따라 블록의 순위를 매기는 것임을 도출했습니다:
$$m_b = \sum_{r \in b} e^{q^T k_r}$$
모든 키를 읽어야 하는 정확한 어텐션 질량에 기반하여 블록을 선택하는 "오라클 희소 어텐션(Oracle Sparse Attention, OSA)"은 밀집 어텐션 성능과 거의 일치함을 보여주며, 이는 선택 기준 자체가 타당함을 입증합니다. 여기서 병목 지점은 이 질량을 압축된 쿼리 독립적 요약으로부터 추정하는 것입니다.

2.2 누적 모멘트 전개 (The Cumulant Expansion)

본 논문은 누적 모멘트 전개(로그 모멘트 생성 함수의 테일러 전개와 동일)를 이용한 어텐션 질량 추정을 분석합니다. 로그 질량($\ln m_b$)은 다음과 같이 전개됩니다:
$$\ln m_b = \ln L + q^T \kappa_1 + \frac{1}{2} q^T \kappa_2 q + \dots$$
여기서:

• $\kappa_1$은 블록의 평균 키($\bar{k}_b$)입니다.
• $\kappa_2$는 블록 내 키의 공분산($\Sigma_b$)입니다.

기존의 1차 선택기(NSA의 MLP 또는 mean-pool 등)는 선형 항($q^T \bar{k}_b$)만을 포착하며, 2차 곡률 항($\frac{1}{2} q^T \Sigma_b q$)을 버립니다. 이러한 누락은 블로 내 키의 분포가 쿼리 방향에 대해 상당한 확산(spread)을 가질 때, 설령 평균이 유사하더라도 성능 저하를 야기합니다.

2.3 COBS 방법론

**COBS (Cumulant Order Block Sparse Attention)**는 선택 과정에 2차 항을 포함함으로써 이 문제를 해결합니다. 이 방법은 압축, 선택, 슬라이딩 윈도우 브랜치를 포함하는 NSA 아키텍처를 기반으로 하며, 다음과 같은 혁신을 도입합니다:

1. 2차 스코어링 (Second-Order Scoring): 단순히 평균 벡터만을 사용하는 대신, 블록 선택기는 블록 평균 키($\bar{k}_b$)와 블록 내 공분산 행렬($\Sigma_b$)의 압축된 표현을 저장합니다. 선택 스코어는 다음과 같이 계산됩니다:
   $$\text{score}_b \approx q^T \bar{k}_b + \frac{1}{2} q^T \Sigma_b q$$
2. 공분산 압축 (Covariance Compression): $D \times D$ 공분산 행렬 전체를 저장하는 것은 메모리 측면에서 불가능하므로, COBS는 세 가지 압축 기술을 사용합니다:
   • 저계수 인수 분해 (Low-Rank Factorization): $\Sigma_b$의 상위 $r$개 고유벡터와 고유값만을 저장합니다.
   • 서브스페이스 투영 (Subspace Projection): 공분산을 쿼리의 2차 모멘트(상위 고유벡터로 생성된 공간)로 투영하여, 저장 용량을 $rD$에서$rs$로 줄입니다.
   • 양자화 (Quantization): 고유벡터를 FP4 (E2M1) 형식으로 스케일링하여 저장하고, 블록 평균은 BF16으로 유지합니다.
3. NoPE (No Position Encoding): 압축 및 선택 브랜치는 회전 위치 인코딩(RoPE)을 제거하여, 블록 요약이 위치 노이즈에 의해 오염되지 않고 오직 콘텐츠에만 의존하도록 보장합니다.

3. 주요 기여

• 선택의 정식화: 본 논문은 블록 선택을 어텐션 질량에 의한 순위 매기기로 환원하였으며, 정확한 질량을 사용하는 오라클이 밀집 어텐션 성능을 거의 모두 회복함을 입증했습니다.
• 누적 모멘트 차수 분석: 기존 1차 선택기의 한계를 이론적으로 설명하며, 버려진 블록 내 공분산이 성능 격차의 원인임을 규명했습니다.
• COBS 알고리 알고리즘: 블록당 압축된 2차 통계량(공분산)을 저장하여, 전체 키를 읽지 않고도 어텐션 질량을 더 정확하게 근사하는 실용적인 방법을 제안합니다.
• 실증적 검증: 2차 선택 방식이 긴 컨텍스트 벤치마크에서 희소 어텐션과 밀집 어텐션 사이의 간극을 유의미하게 좁힌다는 것을 보여줍니다.

4. 실험 결과

4.1 긴 컨텍스트 검색 (32k RULR)

11개 태스크로 구성된 32k RULER 벤치마크 평가 결과:

• NSA Baseline (MLP): 평균 점수 0.2999.
• NSA Mean-Pool (NoPE): 평균 점수 0.5554 (NoPE의 이점 확인).
• COBS (Subspace, $r=4$, FP4): 평균 점수 0.8195.
• Dense Attention: 평균 점수 0.9040.
• 간극 해소: COBS는 NSA 베이스라인과 밀집 어텐션 사이의 격차를 약 86% 메웠습니다.
• 트래픽 효율성: COBS는 NSA 베이스라인의 1.21배의 KV 캐시 읽기 트래픽만을 사용하며, 밀집 어텐션보다는 15.15배 적은 트래픽을 사용합니다.

4.2 짧은 컨텍스트 및 언어 모델링

• 짧은 컨텍스트 추론: COBS는 표준 벤치마크(OpenBookQA, PIQA 등)에서 밀집 어텐션과 대등한 성능을 달성하며 짧은 컨텍스트 상식 추론 능력을 보존합니다.
• Position-wise NLL: COBS는 평균 다음 토큰 음의 로그 가능도(NLL) 1.633을 달성하여(밀집 어텐션의 1.727 대비), 검색 능력의 향상이 긴 컨텍스트 언어 모델링 능력의 희생을 초래하지 않음을 보여줍니다.

4.3 절제 연구 (Ablations)

• Rank Sweep: 성능은 $r=8$에서 정점을 찍고, $r=16$ 이상에서는 방해 요소로부터 오는 부호 없는 분산 질량(unsigned variance mass)의 누적으로 인해 하락합니다. 저자들은 $r \le 8$을 권장합니다.
• Subspace vs. Global: 적응형 레이어별 서브스페이스 차원($s \approx 85$)은 저장 비용을 줄이면서도 글로벌 차원 128과 유사한 성능을 보입니다.
• 양자화: 고유벡터의 FP4 양자화는 사실상 손실이 없는 것으로 나타났습니다.

5. 의의 및 주장

본 논문은 **누적 모멘트 차수(cumulant order)**가 기존 희소 어텐션 방법론에 대한 진단적 렌즈이자 개선을 위한 실질적인 설계 축 역할을 한다고 주장합니다. 1차(평균)에서 2차(공분산) 요약으로 이동함으로써, COBS는 블록 희소 어텐션이 상당한 메모리 대역폭 절감을 유지하면서도 밀집 어텐션의 품질에 근접할 수 있음을 입증합니다.

저자들은 이것이 $\approx 1.2$B 파라미터 규모에서 특정 훈련 프로토콜(RULR 스타일 SFT)을 통해 수행된 메커니즘 연구임을 강조합니다. 모든 LLM에 적용 가능한 배포용 시스템을 주장하는 것이 아니라, 블록 선택의 이론적 격차를 좁힐 수 있다는 증거를 제시하는 것입니다. 이 연구는 블록 희소성이 (선택 메커니즘이 충분히 표현력이 있는 경우, 즉 더 높은 누적 모멘트 차수를 갖는 경우) 주요 모델에서도 실현 가능하다는 점을 시사합니다.

저자가 언급한 한계점:

• 본 연구는 특정 백본 규모와 사전 훈련 시퀀스 길이에 국한되어 있습니다.
• 비교 대상인 NSA 베이스라인은 원래의 NSA 보고와 다른 특정 하이퍼파라미터(예: 겹치지 않는 블록, 큰 MLP)를 사용한 제어된 환경입니다.
• 밀집 브랜치의 RoPE와 희소 브랜치의 NoPE 사용은 긴 컨텍스트 비교 시 혼란 변수로 작용할 수 있습니다.
• 결과는 RULER 스타일의 파인튜닝에 의존하며, 이는 모든 긴 컨텍스트 시나리오에 완벽하게 일반화되지 않을 수 있습니다.