Researchers from Stanford and UC Berkeley Introduce TokenGate: Token-Gated Attention Makes Long-Context Prefill Up to 11.3x Faster
Stanford·UC Berkeley 연구진, 긴 문맥 프리필을 최대 11.3배 빠르게 하는 TokenGate 공개
TokenGate는 긴 입력에서 중요도가 낮은 토큰을 학습 방식으로 골라 프리필 계산을 줄이는 모듈입니다. 연구진은 5개 과제에서 TTFT를 1.4~11.3배 단축하고 메모리를 1.5~3.4배 줄였다고 보고했습니다. 다만 성능은 과제와 토큰 예산에 따라 다릅니다.
- 주제
에디터 노트
TokenGate의 속도 향상은 현재 프롬프트를 처리하는 프리필 구간에서 측정됐습니다. 문서 토큰을 제거하고 KV 항목도 만들지 않으므로 다음 대화에서 해당 문서 내용이 다시 필요할 때는 이전 상태를 그대로 활용하기 어렵다는 지적이 나옵니다. Reddit 댓글의 prefix caching 문제와 미래 프롬프트에서의 정보 손실은 논문 수치로 검증된 결과가 아니라 적용 범위에 관한 우려입니다. 따라서 11.3배라는 최대치는 전체 작업의 체감 속도나 장기 대화 성능을 뜻하지 않습니다. 연구 결과는 과제별 정확도와 토큰 예산을 함께 봐야 하며, 2WikiMultihopQA에서도 전체 정확도의 95%를 기준으로 속도와 메모리를 비교했습니다. 댓글이 제기한 10만 토큰 이상 회수율과 서빙 시스템 통합은 제공된 평가에 답이 없습니다.
AI 요약
TokenGate는 기존 LLM에 학습 가능한 토큰 가지치기를 더하는 모듈입니다. Stanford와 UC Berkeley 연구진은 Qwen3-1.7B와 Qwen3.5-2B를 대상으로 평가했습니다. 핵심은 프롬프트의 문서 부분에서 일부 토큰만 골라 이후 레이어에 전달하는 것입니다. 질문 부분은 가지치기하지 않습니다.
작동 방식
게이트 레이어 아래에서는 모든 토큰을 평소처럼 처리합니다. 해당 레이어의 attention 뒤 residual을 작은 MLP 게이트가 평가하고, Top-K 선택으로 남길 토큰을 정합니다. 게이트 위쪽 레이어는 선택된 토큰만 처리하며, 제거된 토큰의 KV 항목도 만들지 않습니다.
Top-K 선택은 미분할 수 없어 그대로는 end-to-end 학습이 어렵습니다. TokenGate는 Mixture-of-Experts 라우터 방식에서 착안해 남긴 토큰의 점수를 상위 레이어 attention logit에 bias로 더합니다. 이 경로로 게이트에 gradient를 전달합니다. 학습에는 일반 next-token loss를 쓰며, 매 단계에서 전체 문맥 손실과 무작위 토큰 예산 손실을 평균합니다. 이에 따라 게이트 하나로 여러 예산을 처리하도록 학습합니다. 게이트와 함께 LoRA r=16을 사용합니다.
성능과 평가 조건
Qwen3-1.7B를 RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX, needle-in-a-haystack(NIAH)에서 평가했습니다. 전체 프리필과 비교해 TTFT는 1.4~11.3배 빨라졌고 최대 프리필 메모리는 1.5~3.4배 줄었습니다. 다음으로 좋은 비교 기준보다 TTFT가 1.2~6.4배 빨랐고 메모리는 최대 2.1배 적었습니다. 2WikiMultihopQA에서는 전체 문맥 정확도의 95%를 유지하는 조건에서 프리필이 3배 빨라지고 메모리가 2.7배 감소했습니다.
가장 높은 11.3배 속도 향상은 NIAH 16k에서 나왔습니다. 128k에서는 TTFT 향상이 13배에 이르렀으며 메모리 절감은 약 3.43배로 보고됐습니다. 반면 일부 과제의 속도 향상은 1.4배입니다. 논문 타이밍은 A100 한 장, batch size 1에서 측정했고 학습에는 A100 네 장을 사용했습니다. 비교 기법들의 자체 보고 수치는 하드웨어와 설정이 달라 직접 비교할 수 없습니다.
Qwen3.5-2B처럼 Gated DeltaNet과 full-attention 레이어를 섞은 모델에서도 시험했습니다. NIAH 32k에서 3번 레이어에 게이트를 두면 전체 프리필과 같은 정확도를 유지하면서 3.8배 빨랐고 FastKV보다 2.1배 빨랐습니다. 1번 또는 2번 레이어에 게이트를 두면 속도는 7.8배 빨라졌지만 정확도는 최대 6% 낮았습니다. TokenGate는 DeltaNet 레이어에서도 가지치기하지만 FastKV는 attention 레이어에서만 가지치기한다고 설명합니다.
설치와 학습
코드는 Apache-2.0 라이선스로 공개했습니다. 실행 환경은 Python 3.10 이상과 CUDA PyTorch 2.9 이상입니다. 선택 사항인 FlashAttention-2 key-bias 커널은 일반 호출보다 1.0~1.1배의 비용이 들며, 사용하지 않으면 FlexAttention으로 대체합니다. 학습 단계는 일반 LoRA보다 1.3배 느리지만, 5개 과제 중 4개에서 더 빨리 수렴했고 평균 학습 시간은 A100 기준 1.8시간으로 LoRA의 2.2시간보다 짧았습니다.
Reddit 반응
- @Kodix — 다른 방식으로 포장한 speculative prefill 같습니다. 같은 문제가 있을 것 같습니다. prefix caching을 깨고, 지금은 중요하지 않은 토큰도 몇 턴 뒤에는 중요해질 수 있으니 없애고 싶지 않습니다. 장점은 지금 무료로 쓸 수 있고 Qwen 3.5 아키텍처를 지원한다는 점입니다.
- @Previous-Nature7326 — 정확도 저하가 없다고 하네요.
- @Kodix — 현재 프롬프트에서 중요하지 않은 토큰을 삭제하고도 정확도 저하가 없을 수는 있습니다. 하지만 삭제한 토큰은 다음 프롬프트에서는 사라집니다. 에이전트 환경에서는 나중에 중요해질 가능성이 큽니다. 주장은 전부 맞을 수 있지만 적용 범위를 이해해야 합니다.
- @tired_of_mondays — 기존 모델에 붙이는 방식이라 새 모델 출시를 기다릴 필요가 없습니다. 제가 보고 싶은 실제 시험은 10만 토큰 이상에서 needle-in-a-haystack 회수율입니다. 학습 기반 가지치기는 평균 점수는 좋아도 코드베이스 한가운데 묻힌 한 줄을 조용히 놓치는 경우가 있습니다. 검색 성능이 유지된다면 프리필을 기다리는 시간이 긴 Mac이나 단일 GPU 사용자에게 큰 도움이 되겠습니다.
- @FullOf_Bad_Ideas — vLLM이나 SGLang에 통합하면 유용하겠습니다. decision 모델에서도 작동할 것 같습니다. FlashPrefill과 비슷합니다.
- @MundanePercentage674 — 실제로 작동하는 거였으면 좋겠네요.
원문: Marktechpost / 번역·요약: Trawling