기술 요약: SPLA - 긴 문맥 모델링을 위한 블록 희소 플러스 선형 어텐션 (Block Sparse Plus Linear Attention)

1. 문제 정의

대규모 언어 모델(LLM)이 수백만 토큰의 문맥 창(context window)을 처리하도록 확장됨에 따라, 디코딩 단계는 메모리 제약으로 인해 결정적인 병목 현상이 됩니다. 전체 어텐션 메커리즘은 점점 더 커지는 키-값(KV) 캐시를 메모리에 로드해야 하며, 이는 메모리 바운드(memory-bound) 문제를 야기하여 추론 처리량(throughput)을 제한합니다.

기존의 블록 단위 희소 어텐션 방식(예: NSA, InfLLM)은 가장 관련성이 높은 KV 캐시 블록만을 로드함으로써 이를 완화하려고 시도합니다. 그러나 본 논문은 현재의 접근 방식에서 두 가지 근본적인 한계를 식별합니다:

1. 낮은 선택 충실도(Low Selection Fidelity): 기존의 블록 선택 방법은 고도로 압축된 표현(예: min/max 풀링, 단순 평균 풀링)에 의존하며, 이는 원래의 토큰 수준 어텐션 목적 함수와 수학적 연결성이 결여되어 있습니다. 이로 인해 높은 재현율 오류(recall error)가 발생하여, 의미적으로 중요한 블록을 놓치게 됩니다.
2. 누적 문맥 손실 (Long Tail Divergence): 기존 방식은 선택되지 않은 블록을 엄격하게 절단(기여도를 0으로 설정)합니다. 문맥 길이가 길어짐에 따라, 어텐션 분포의 "롱 테일(long tail)"에 존재하는 확률 질량(probability mass)은 무시할 수 없는 수준이 됩니다. 이 질량을 버리는 것은 희소 어텐션 출력을 덴스(dense) 베이스라인으로부터 크게 이탈시켜, 특히 긴 문맥 시나리오에서 생성 품질을 저하시킵니다.

2. 방법론: 희소 플러스 선형 어텐션 (SPLA)

SPLA는 어텐션 메커니즘을 고관련성 블록을 위한 **정확한 어텐션(Exact Attention)**과 나머지 "롱 테일"을 위한 **근사 선형 어텐션(Approximate Linear Attention)**이라는 두 가지 상보적인 구성 요소로 분해하여 이 문제를 해결합니다.

2.1. 2차 테일러 전개를 통한 원칙적인 블록 선택

선택 충실도를 높이기 위해, SPLA는 2차 테일러 전개를 사용하여 원래의 어텐션 목적 함수로부터 직접 블록 선택 지표를 도출합니다.

• 목적: 블록 관련성 점수는 블록 내 모든 토큰의 통합 확률 질량으로 정의됩니다.
• 근사: 직접적인 합산은 비용이 많이 듭니다. SPLA는 블록 평균($\bar{k}$) 주변에서의 지수 어텐션 점수의 기댓값을 근사합니다.
  • 이 유도는 1차 통계량(블록 평균)과 2차 통계량(블록 공분산/분산)을 모두 활용합니다.
  • 결과적인 점수는 다음과 같습니다: $\exp(q^\top \bar{k}) \left(1 + \frac{1}{2}q^\top \text{Cov}(k)q\right)$.
• 구현: $O(d^2)$ 저장 오버헤드를 피하기 위해 공분산 행렬을 대각 행렬로 근사합니다. 이를 통해 보조 학습 없이도 블록 통계량에 대한 요소별 연산만을 사용하여 효율적인 선택 지표 계산이 가능하며, 토큰 수준 어텐션과 블록 수준 검색 사이의 간극을 메웁니다.

2.2. "롱 테일"을 위한 잔차 선형 어텐션 (Residual Linear Attention, RLA)

선택되지 않은 블록을 단순히 버리는 대신, SPLA는 특수화된 선형 어텐션 메커니즘을 통해 이들을 압축된 순환 상태(recurrent state)로 응축합니다.

• 순환(Recurrence): RLA 상태는 선택되지 않은 블록으로부터만 정보를 축적합니다.
• IO 효율적 공식화: 단순한 구현은 선택되지 않은 블록을 위해 메모리를 다시 스캔해야 하므로 희소성의 목적을 저해할 수 있습니다. SPLA는 최적화된 **차감 기반 공식화(subtraction-based formulation)**를 도출합니다:
  $$o^{rla}_t = \bar{o}_t - \tilde{o}_t$$
  여기서:
  • $\bar{o}_t$는 모든 블록에 대해 표준 선형 어텐션을 적용하여 계산된 글로벌 문맥 추정치입니다.
  • $\tilde{o}_t$는 선택된 블록들에 대해서만 선형 어텐션을 통해 계산된 로컬 교정 항입니다.
• 효 efficiency: 선택된 블록들은 이미 정확한 희소 어텐션 계산을 위해 SRAM에 로드되어 있으므로, $\tilde{o}_t$는 동일한 퓨즈드 커널(fused kernel) 내에서 계산될 수 있습니다. 이는 선택되지 않은 블록이 희소 디코딩 과정 중 HBM으로부터 명시적으로 액세스되지 않도록 보장하면서도, 전체 글로벌 문맥을 보존하는 IO 프로파일을 유지합니다.

2.3. 아키텍처 및 통합

• 하이브리드 전략: 최종 출력은 희소 어텐션(정확)과 RLA 출력(근사)의 게이트형 합(gated sum)이며, RMS Norm을 통해 정규화됩니다.
• 파라미터 효율성: SPLA는 사전 학습된 덴스 모델에 적응하도록 설계되었습니다. 두 브랜치 모두 공유된 Query, Key, Value 프로젝션을 사용합니다. 추가되는 유일한 파라미터는 RLA 출력에 적용되는 RMS Norm의 스케일 파라미터뿐입니다.
• 특징 맵(Feature Map): 파라미터가 없는 지수 특징 맵($\phi(x) \sim \exp(x)$)이 사용되어 새로운 프로젝션 가중치가 도입되지 않도록 합니다.

3. 주요 기여

1. 원칙적인 블록 선택: 2차 테일러 전개를 기반으로 블록 평균 및 공분산 통계량을 활용하여, 휴리스틱 풀링 방법보다 높은 선택 충실도를 달-성하는 블록 선택 전략을 도출했습니다.
2. SPLA 프레임워크: 문맥을 "정확(Exact)" 세트와 "근사(Approximate)" 세트로 엄격히 분할하는 새로운 아키텍처를 제시합니다. 잔차 선형 어텐션을 통해 선택되지 않은 블록을 처리함으로써 전체 글로벌 문맥을 보존하고, 덴스 어텐션과의 성능 격차를 해소합니다.
3. IO 효율적 구현: 선택되지 않은 블록에 대한 명시적 메모리 액세스를 피하는 최적화된 차감 기반 RLA 공식화를 통해, 표준 희소 디코딩과 유사한 추론 효율성을 보장합니다.
4. 적응성: 덴스 어텐션 모델이 지속적 사전 학습(continual pretraining) 단계에서 최소한의 추가 파라미터만으로 SPLA 아키텍처로 전환될 수 있음을 입증했습니다.

4. 실험 결과

저자들은 일반 지식, 긴 문맥, 추론의 세 가지 설정에 대해 14B 파라미터 모델을 통한 지속적 사전 학습(CPT)을 통해 SPLA를 평가했습니다.

• 일반 지식: SPLA는 표준 벤치마크(ARC, MMLU, GSM8K 등)에서 덴스 어텐션 모델과 경쟁력 있는 성능을 유지했으며, 종종 덴스 베이스라인과 일치하거나 이를 약간 상회하였고 다른 희소 베이스라인(NSA, InfLLM-v2)보다 우수한 성능을 보였습니다.
• 긴 문맥 (RULER 벤치마크): SPLA는 다른 희소 방식들을 크게 압도했습니다.
  • 256k 문맥 길이에서, SPLA는 72.3점을 기록하여 덴스 베이스라인(69.3), NSA(32.5), InfLLM-v2(42.6)를 크게 앞질렀습니다.
  • SPLA는 다른 희소 모델들이 롱 테일 발산으로 인해 심각한 성능 저하를 겪는 구간에서도 견고함을 보여주었습니다.
• 추론: SPLA는 복잡한 추론 작업(AIME, HMT, LiveCodeBench)에서 NSA 및 InfLLM-v2보다 뛰어난 성능을 보였으며, 이는 원칙적인 선택과 하이브리드 설계가 필수적인 미세 정보(fine-grained information)를 유지함을 입증합니다.
• 절제 연구(Ablation Studies):
  • SPLA vs. SPA (RLA가 없는 Sparse Plus Attention): SPLA는 SPA보다 일관되게 우수한 성능을 보였으며, 이는 롱 테일 발산을 방지하는 데 있어 잔차 선형 어텐션 모듈의 결정적인 역할을 확인시켜 줍니다.
  • SPA vs. InfLLM-v2: 테일러 기반 선택을 사용하지만 RLA는 없는 SPA가 일반적으로 InfLLM-v2보다 우수한 성능을 보였으며, 이는 1차 근사보다 2차 선택 지표가 우월함을 검증합니다.

5. 의의 및 주장

본 논문은 SPLA가 효율성과 문맥 보존 사이의 트레이드오프를 해결함으로써 차세대 파운데이션 모델을 위한 확장 가능한 경로를 제공한다고 주장합니다.

• 격차 해소: SPLA는 희소 어텐션과 덴스 어텐션 사이의 성능 격차를 효과적으로 좁히며, 긴 문맥 벤치마크(최대 256k 토큰)에서 덴스 모델을 능가하면서도 희소 디코딩의 효율성을 유지합니다.
• 적응 전략: 본 연구는 비용이 많이 드는 사전 학습 단계에서는 표준 어텐션 아키텍처를 유지하다가, 지속적 사전 학습 중에 최소한의 추가 파라미터만으로 효율적인 희소 변형 모델로 매끄럽게 전환할 수 있는 유망한 방향을 제시합니다.
• 견고성: "하드 절단(hard truncation)"을 피하고 원칙적인 선택 지표를 활용함으로써, SPLA는 문맥 길이가 확장됨에 따라 발생하는 기존 희소 방식들의 고질적인 문제인 누적 문맥 손실을 완화합니다.

저자들은 현재 학습 효율성이 매우 최적화된 덴스 커널에 비해 약간의 오버헤드가 있지만(특정 14B 구성의 블록 선택 및 작은 GQA 그룹 크기 때문), 이 접근 방식이 대규모 학습에 실행 가능하며, 그룹 크기가 커짐에 따라 효율성 이득이 개선될 것으로 기대된다고 언급했습니다.