기술 요약: TokenButler

문제 정의

대규모 언어 모델 (LLM) 은 컨텍스트 길이가 64k~1M 토큰으로 확장됨에 따라 메모리 및 계산 병목 현상에 직면합니다. 효율적인 디코딩을 위해 토큰 히스토리를 저장하는 키 - 밸류 (KV) 캐시는 컨텍스트 길이에 비례하여 선형적으로 증가하며, 상당한 GPU 메모리와 대역폭을 소모합니다.

이 병목 현상을 완화하기 위한 기존 솔루션은 일반적으로 세 가지 범주로 분류되며, 각각 고유한 한계가 있습니다:

1. 정적 전략 (예: StreamingLLM): 캐시를 제한하기 위해 고정된 규칙 (예: 슬라이딩 윈도우) 을 사용합니다. 이러한 방법은 쿼리 인식 능력이 부족하여 동적인 입력 요구 사항에 적응할 수 없습니다.
2. 적응형 제거 (예: H2O, SnapKV): 어텐션 점수 등의 지표를 기반으로 중요도가 낮다고 판단된 토큰을 영구적으로 삭제합니다. 이는 중요도가 비단조적 (non-monotonic) 인 경우가 많아, 초기에 언급된 엔티티가 훨씬 나중에 참조되는 공동 참조 (co-referential) 컨텍스트에서 나중에 관련성이 생길 수 있는 중요한 토큰을 잃을 위험이 있습니다.
3. 적응형 동적 전략 (예: Quest): 전체 캐시를 유지하되 디코딩 시 토큰의 부분 집합 (페이지) 을 가져옵니다. 영구적 제거를 피하지만, coarse-grained(거친 입자) 페이지 수준 검색에 의존하기 때문에, 중요한 정보가 페이지 경계를 넘어 분산되거나 세밀한 선택이 필요한 밀집된 컨텍스트가 풍부한 작업에서는 실패합니다.

또한, 많은 기존 방법들은 토큰 중요도에 대한 부정확한 대리 지표 (예: 최근성 또는 집계 어텐션 점수) 에 의존하거나, 모든 토큰에 대한 전체 어텐션 점수를 계산할 때 높은 계산 오버헤드를 초래합니다.

방법론: TokenButler

저자들은 토큰을 영구적으로 제거하거나 거친 페이지링에 의존하지 않고 중요한 토큰을 식별하도록 설계된 고세분도 (high-granularity) 쿼리 인식 예측기인 TokenButler를 제안합니다.

핵심 아키텍처

TokenButler 는 고정된 깊이 간격 (producer layers) 에서 저차원 중요도 쿼리를 예측하고, 이를 후속 레이어 (consumer layers) 의 실제 KV 캐시 키에 대한 학습된 투영과 결합하여 작동합니다.

• 쿼리 예측: 경량 멀티 레이어 퍼셉트론 (MLP) 이 "프로듀서"레이어에서 은닉 상태를 입력받아 어텐션 헤드당 $G$개의 슬롯별 쿼리 벡터를 생성합니다. 이러한 벡터는 저차원 ($d' \ll E$) 입니다.
• 키 - 캐시 투영: 각 소비자 레이어에 대해 실제 고차원 KV 캐시 키는 학습된 행렬을 사용하여 동일한 저차원 공간으로 투영됩니다.
• 점수 매기기: 예측된 쿼리를 투영된 키와 점곱하여 토큰 중요도 점수를 저렴하게 생성함으로써, 고정된 예산 하에서 동적인 토큰별 선택을 가능하게 합니다.

학습

• 증류 (Distillation): TokenButler 는 고정된 베이스 LLM 의 마스킹된 인과적 어텐션 분포를 증류하여 학습됩니다. 목표는 교사 (Teacher) 의 어텐션 로짓과 학생 (TokenButler) 의 예측 로짓 간의 교차 엔트로피 손실을 최소화하는 것입니다.
• 효율성: 학습 중 긴 컨텍스트를 처리하기 위해 손실은 쿼리 위치의 부분 표본 집합에서 계산되며, 후기 컨텍스트 영역과 최종 토큰에 초점을 맞춰 복잡도를 $O(L^2)$에서 $O(RL)$로 줄입니다.
• 오버헤드: 예측기는 기본 모델에 1% 미만의 파라미터 오버헤드만 추가합니다.

추론 전략

TokenButler 는 메모리 접근을 최적화하기 위해 KV 캐시를 세 개의 버퍼로 분할합니다:

1. 싱크 버퍼 (Sink Buffer): 처음 $S$개의 토큰 (항상 유지됨).
2. 로컬 윈도우 버퍼 (Local Window Buffer): 가장 최근 $N$개의 토큰으로 구성된 순환 버퍼 (항상 유지됨).
3. 중요 버퍼 (Important Buffer): 예측된 점수에 따라 TokenButler 에 의해 동적으로 채워짐.

최적화:

• 배칭: 키의 투영은 매 디코딩 단계가 아니라 $N$개의 토큰 단위로 배치되어 수행되어 시간적 지역성을 활용합니다.
• 예측 간격 및 이웃 가져오기: 예측기의 비용을 분산시키기 위해 시스템은 모든 단계에서 예측기를 실행하지 않습니다. 대신 $N$단계마다 (예측 간격) 실행합니다. 간격 사이에서 중요도가 변하는 토큰을 놓칠 위험을 완화하기 위해 시스템은 "이웃 가져오기 (neighbor fetching)"를 사용하여 선택된 토큰에 공간적 이웃을 포함하도록 확장함으로써 멀티 토큰 엔티티 또는 추론 체인의 커버리지를 보장합니다.

주요 기여

1. 경량 예측기: 어텐션 증류를 통해 학습된 토큰 중요도 지표로, 최소한의 파라미터 오버헤드 (<1%) 로 세밀한 추정을 제공합니다.
2. 합성 벤치마크: 영구 토큰 손실 또는 페이지 미스로 인해 짧은 컨텍스트 길이 (300 토큰) 에서도 기존 제거 및 페이지 기반 방법이 실패함을 보여주는 새로운 공동 참조 검색 작업. TokenButler 는 이 작업에서 오라클에 가까운 정확도를 달성합니다.
3. 긴 컨텍스트 성능: 최첨단 희소성 및 압축 방법과 비교하여 RULER 및 LongBench 벤치마크 (64K 컨텍스트) 에서 경쟁력 있거나 우수한 결과를 달성합니다.
4. 효율성 향상:
   • 예측 간격을 사용하여 Dense Attention 대비 최대 1.6 배의 온 - GPU 속도 향상.
   • CPU 오프로딩이 포함된 Dense Attention 대비 최대 7.6 배의 지연 시간 감소로 메모리 제한 시스템에서 효율적인 디코딩 가능.
   • 예측기 호출 횟수를 16 배 줄여도 단계별 기준선 대비 정확도를 약 1.1% 이내로 유지.

실험 결과

• 합성 공동 참조 작업: TokenButler 는 Llama-3.1-8B 및 Qwen2.5-7B 에서 약 80-95% 의 정확도와 커버리지를 달성하여, 50% 미만의 정확도로 떨어진 제거 기반 방법 (H2O, SnapKV) 과 부분 페이지 미스로 고통받는 페이지 기반 방법 (Quest) 을 크게 능가했습니다.
• RULER 벤치마크 (64K 컨텍스트): TokenButler 는 Llama-3.1-8B 및 Qwen2.5-7B 에서 SnapKV, xKV, Quest 와 같은 베이스라인과 동등하거나 더 나은 성능을 보였으며, Qwen2.5-7B 에서 평균 점수가 가장 높았습니다.
• LongBench: TokenButler 는 다양한 데이터셋에서 기존 KV 희소성 및 압축 방법보다 더 나은 정확도를 보여주었습니다.
• 추론 모델 (AIME24): DeepSeek-R1-Distill-Llama-8B 에서 TokenButler 는 30.0% 의 정확도를 달성하여 밀집 디코딩 기준선 (33.3%) 에 근접했고, 동일한 8K 토큰 예산 하에서 다른 희소성 방법들을 능가했습니다.
• 지연 시간: CPU 오프로딩 시나리오 (256K+ 컨텍스트) 에서 TokenButler 는 토큰당 약 3.2 초 (Dense) 의 지연 시간을 약 0.6 초 (TokenButler) 로 줄였습니다.

의의

이 논문은 토큰 중요도가 매우 역동적이며 쿼리에 의존적이므로, 정적 휴리스틱과 거친 입자 검색이 복잡하고 긴 컨텍스트 작업에 충분하지 않다고 주장합니다. TokenButler 는 학습된 세밀한 선택 메커니즘이 이러한 휴리스틱을 효과적으로 대체할 수 있음을 보여줍니다. 전체 KV 캐시를 유지하면서 가장 관련성 높은 토큰만 선택적으로 접근함으로써 TokenButler 는 영구적 제거와 관련된 품질 저하 없이 메모리 대역폭 병목 현상을 완화합니다. 제안된 예측 간격 및 이웃 가져오기 전략은 또한 그러한 학습된 예측기의 계산 비용을 효과적으로 분산할 수 있음을 보여주어, 제한된 하드웨어에서의 실시간 긴 컨텍스트 추론을 위한 실용적인 솔루션임을 입증합니다.