기술 요약: Full Attention Strikes Back: Full Attention 을 수백 단계의 훈련으로 Sparse 로 전환하기

문제 제기

대규모 언어 모델 (LLM) 의 장문맥 추론은 현재 Full Attention 메커니즘과 관련된 2 차 계산 및 메모리 비용에 의해 병목 현상을 겪고 있습니다. 기존 효율적 대안들이 존재하지만, 이들은 종종 비용이 많이 드는 네이티브 Sparse 훈련이나 효율성, 훈련 비용, 정확도 간의 트레이드오프를 초래하는 휴리스틱 토큰 제거 전략에 의존합니다. 본 연구가 다루는 핵심 질문은, 이미 훈련된 Full Attention 모델이 내재된 Sparse 성을 보유하고 있어 최소한의 적응을 통해 고도로 Sparse 한 모델로 변환할 수 있는지, 그리고 이를 통해 네이티브 Sparse 사전 훈련의 높은 비용을 피할 수 있는지 여부입니다.

방법론: RTPurbo

저자들은 Full Attention LLM 을 최소한의 외과적 적응을 통해 효율적인 Sparse 추론 시스템으로 변환하도록 설계된 Head-wise Sparse Attention 프레임워크인 RTPurbo를 제안합니다. 이 방법은 세 가지 핵심 관찰과 이에 상응하는 기술적 구성 요소에 기반합니다:

1. Head 특화 및 오프라인 보정

저자들은 사전 훈련된 LLM 의 Attention Head 들이 동질적이지 않으며, 고유한 기능적 역할로 특화됨을 관찰했습니다. 구체적으로, 소수의 Head(검색 Head) 는 멀리 떨어진 의미상 관련 콘텐츠를 검색하는 책임을 지는 반면, 대다수 (로컬 Head) 는 주로 로컬 정보를 처리합니다.

• 구현: RTPurbo 는 시작과 끝에 "Needle"이 삽입된 경량 시퀀스를 사용하여 각 Head 가 후속 Needle 을 이전 Needle 로부터 검색하는 능력에 따라 점수를 매기는 오프라인 보정 단계를 수행합니다.
• 조치: 이러한 점수에 따라 Head 들은 **검색 집합 ($H_{ret}$)**과 **로컬 집합 ($H_{loc}$)**으로 분할됩니다. 추론 중 로컬 Head 는 Attention Sink 를 사용하는 슬라이딩 윈도우를 사용하며, 검색 Head 는 프리필 (prefill) 단계 동안 전체 KV 캐시를 유지하지만 디코딩 중에는 동적 Sparse 선택으로 전환합니다.

2. 검색을 위한 저차원 토큰 인덱싱

검색 Head 의 경우, 전체 차원 점수화 없이 관련 토큰을 효율적으로 식별하는 것이 과제입니다.

• 이론적 통찰: 저자들은 Rotary Position Embeddings(RoPE) 의 영향을 분석하여, 고주파수 성분은 거리에 따라 빠르게 변하며 장거리 친밀도를 저하시키는 반면, 저주파수 성분은 검색 신호를 유지함을 발견했습니다.
• 구현: RoPE 주입 이전에 쿼리와 키 특징에 저차원 프로젝터($r \ll d$, 구체적으로 16 차원) 가 적용됩니다. 이는 고차원 공간을 관련성을 추정하는 컴팩트한 부분 공간으로 투영합니다. 이 인덱서는 16 차원만 사용하여 90% 이상의 재현율 (recall) 로 효율적인 토큰 선택을 가능하게 합니다.

3. 동적 Top-p 선택 vs 정적 Top-k

저자들은 검색에 필요한 토큰 예산이 쿼리에 크게 의존한다고 주장합니다. 일부 쿼리는 광범위한 검색 (많은 토큰 필요) 을 유발하는 반면, 다른 쿼리는 매우 집중적 (적은 토큰 필요) 입니다.

• 구현: 고정된 Top-$k$ 예산 대신 RTPurbo 는 동적 Top-p 선택기를 사용합니다. 이는 누적 임계값 $p$(예: 0.9) 에 도달할 때까지 저차원 공간에서 Attention Mass 를 누적합니다. 이는 특정 쿼리 복잡도에 맞게 Sparse 수준을 조정하여 작은 $k$의 정보 손실이나 큰 $k$의 낭비 계산을 방지합니다.

4. 2 단계 훈련 파이프라인

Sparse 화 후 모델 성능을 복원하기 위해 RTPurbo 는 경량 2 단계 훈련 프로세스를 활용합니다:

• 1 단계 (프로젝션 훈련): 백본 LLM 은 고정됩니다. 식별된 검색 Head 에 대한 저차원 프로젝션 가중치만 원래 Full Attention 분포와 프로젝션된 저차원 분포 간의 KL 발산을 최소화하도록 훈련됩니다.
• 2 단계 (자기 증류): Sparse 모델은 밀집 Teacher 의 다음 토큰 예측을 맞추도록 학습하는 학생 역할을 합니다. 중요한 점은 이 단계가 복잡한 데이터 혼합 조정이 필요하지 않고 과적합을 방지하는 표준 지도 미세 조정이 아닌 자기 증류(Top-10 로짓 정렬) 를 사용한다는 것입니다. 이 단계는 수백 단계의 훈련 (약 1M 라벨 토큰) 만 필요합니다.

5. 하드웨어 인식 커널

저자들은 동적 Top-p 선택을 효율적으로 처리하기 위해 맞춤형 GPU 커널을 구현했습니다. 이러한 커널은 임계값 설정을 위한 정렬 없는 히스토그램 기반 접근 방식과 장문맥을 과도한 메모리 오버헤드 없이 관리하기 위한 대역폭 최적화 Sparse 디코딩을 활용합니다.

주요 결과

실험은 NVIDIA H20 GPU 를 사용하여 Qwen3-Coder-30B-A3B 및 Qwen3-30B-A3B-Think 모델로 수행되었습니다.

• 정확도 유지: RTPurbo 는 장문맥 벤치마크 (LongBench, RULER) 및 추론 작업 (AIME, MMLU-PRO) 에서 거의 손실 없는 정확도를 달성합니다.
  • 64K 문맥에서 RULER 기준 RTPurbo(Top-p) 는 평균 점수 **90.06%**를 기록하여 모든 Sparse 베이스라인 (예: Minference 65.61%, FlexPrefill 77.77%) 을 능가하고 Full Attention 베이스라인 (64K 에서 86.23%) 과도 일치합니다 (다만 일부 작업에서 더 긴 길이에서 베이스라인이 크게 하락하는 반면 RTPurbo 는 높은 정확도를 유지함).
  • AIME24/25 에서 RTPurbo 는 Full Attention 베이스라인과 정확히 일치합니다 (86.67).
• 효율성 향상:
  • 프리필 속도 향상: RTPurbo 는 1M 문맥 길이에서 FlashAttention-2 대비 최대 9.36 배의 속도 향상을 제공하며, 문맥 길이에 따라 속도 향상은 선형적으로 증가합니다 (32K 에서 2.83 배).
  • 디코드 속도 향상: RTPurbo 는 1M 문맥에서 디코딩 시 최대 2.01 배의 속도 향상을 달성합니다.
  • Sparse 성: 이 방법은 높은 Attention Mass (>0.93) 를 유지하면서 높은 Sparse 성 (64K 에서 최대 89.2%, 512K 에서 97% 초과) 을 달성합니다.
• 훈련 효율성: 전체 적응 과정은 수백 단계의 훈련 (약 1M 라벨 토큰) 만 필요하며, 이는 효과적인 Sparse 추론을 위해 네이티브 Sparse 사전 훈련이 필요하지 않음을 보여줍니다.

의의 및 주장

본 논문은 Full Attention 모델이 본질적으로 Sparse하며 최소한의 적응을 통해 고도로 효율적인 Sparse 추론 시스템으로 변환될 수 있다고 주장합니다. 이 연구의 의의는 다음과 같은 핵심 사항에 있습니다:

1. 효율성 달성 경로의 재정의: 저자들은 네이티브 Sparse 사전 훈련이 효율적인 장문맥 추론을 위한 유일한 경로라는 개념에 도전합니다. 그들은 표준 Full Attention 훈련이 여전히 매우 경쟁력 있는 선택지임을 입증하며, 결과적으로 생성된 모델은 매우 적은 추가 비용으로 사후적으로 "Sparse 화"될 수 있음을 보여줍니다.
2. 해석 가능성 및 최소 외과적 개입: 이 접근 방식은 Head 의 자연스러운 특화 (검색 대 로컬) 와 저차원 부분 공간에 의존하여 매우 해석 가능합니다. 모델이 학습한 Attention 패턴에 기반하여 Sparse 성을 근거로 삼음으로써 휴리스틱 토큰 제거의 "블랙박스" 성격을 피합니다.
3. 동적 적응: 고정된 Top-$k$ 제약에서 동적 Top-p 선택으로 전환함으로써, 이 방법은 기존 고정 예산 접근법의 한계로 지적된 다양한 쿼리의 복잡성을 효과적으로 처리합니다.
4. 실용성: 수백 단계의 훈련과 경량 하드웨어 커널만으로 거의 손실 없는 정확도를 달성할 수 있는 능력은, 이 패러다임이 처음부터 재훈련하는 데 prohibitive 한 비용 없이 기존 대규모 모델에 즉시 적용 가능함을 시사합니다.

요약하자면, RTPurbo 는 표준 Full Attention LLM 에 내재된 Sparse 추론 잠재력을 해제하는 실용적이고 효율적이며 이론적으로 근거를 둔 방법을 제공합니다.