스탠퍼드대와 UC 버클리 연구진이 이미 학습된 대규모 언어 모델에 붙여 긴 입력의 프리필 단계를 빠르게 하는 모듈 TokenGate를 공개했다. 연구진은 전체 입력을 처리할 때보다 첫 토큰 시간(TTFT)이 1.4배에서 11.3배 빨라지고 프리필 최대 메모리는 1.5배에서 3.4배 줄었다고 보고했다. 정확도에 대해서는 일부 조건(NIAH 등)에서 전체 프리필과 같았지만, 2WikiMultihopQA에서는 전체 문맥 정확도의 95% 수준이었고 Qwen3.5-2B의 일부 설정에서는 정확도가 6% 이내로 차이 났다. 이 수치는 모두 연구진이 보고한 것이다.
코드와 체크포인트는 GitHub에 Apache-2.0 라이선스로 올라 있다. 실행에는 Python 3.10 이상과 CUDA용 PyTorch 2.9 이상이 필요하다. 논문의 시간 측정은 A100 한 장, 배치 크기 1 조건이고 학습에는 A100 네 장을 썼다.
작동 방식
트랜스포머는 모든 토큰을 모든 층에 통과시키고, 어텐션 비용은 길이의 제곱으로 늘어난다. 학습 없이 토큰을 가지치기하는 방식은 저렴하지만 가지치기 비율이 높으면 정확도를 잃는다. 구조를 바꾸는 방식은 처음부터 사전 학습해야 하는 경우가 많다.
TokenGate의 순전파는 세 단계다. 게이트 층 L 아래는 토큰 N개를 모두 처리한다. 층 L에서는 작은 MLP가 어텐션 이후 잔차값으로 토큰마다 점수를 매기고, Top-k 선택기가 점수 높은 토큰만 남긴다. L 위 층은 남은 K개만 처리하며 제거된 토큰의 KV 캐시는 만들지 않는다. 가지치기 대상은 프롬프트의 문서 부분뿐이고 질문은 그대로 둔다.
Top-k 선택은 미분이 되지 않아 끝까지 학습하기 어렵다. 연구진은 전문가 혼합의 라우터 방식을 가져와, 남은 토큰의 점수를 게이트 위쪽 어텐션 로짓의 편향으로 더했다. 이렇게 하면 기울기가 점수 모듈로 돌아간다. 헤드별 학습 이득과 온도도 학습하며, 온도는 질의가 보는 키가 줄어든 것을 보정한다. 학습 손실은 일반적인 다음 토큰 손실이고, 매 단계 전체 문맥 손실과 무작위 예산 손실을 평균한다. 그래서 게이트 하나로 여러 예산에 대응한다.
성능 결과
Qwen3-1.7B에 작은 MLP 게이트와 LoRA(r=16)를 붙여 RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX, 바늘 찾기(NIAH)에서 시험했다. 비교 대상은 같은 LoRA 백본의 FastKV, SlimInfer, DASH, GemFilter, FlexPrefill, MInference이며, 연구진은 다섯 과제에서 여섯 기법 대비 파레토 최적이라고 밝혔다. 컨텍스트는 최대 128k 토큰까지 시험했다.
- 전체 프리필 대비 TTFT 1.4~11.3배 단축, 메모리 1.5~3.4배 절감
- 차순위 기법 대비 TTFT 1.2~6.4배 단축, 메모리 최대 2.1배 절감
- NIAH 16k에서 TTFT 11.3배, 메모리 3.42배 개선. 일부 과제는 속도 향상이 1.4배에 그침
- 2WikiMultihopQA에서 전체 문맥 정확도의 95%를 유지하며 프리필 3.0배, 메모리 2.7배 절감
- 16k에서 학습한 게이트의 TTFT 이득은 32k에서 12.8배, 64k와 128k에서 13.0배였고 메모리 절감은 3.43배 안팎이다. 11.3배는 NIAH 16k 수치다.
같은 층에서 학습형 선택기가 FastKV의 휴리스틱 선택을 앞섰다. TokenGate 백본에서는 무작위 K 선택조차 비교 기법보다 점수가 높았는데, 연구진은 백본이 함께 적응한 결과로 풀이했다.
Gated DeltaNet과 전체 어텐션 층을 섞은 Qwen3.5-2B에서도 작동했다. NIAH 32k에서 3층에 게이트를 둔 결과는 전체 프리필과 정확도가 같았고, 속도는 전체 프리필의 3.8배, FastKV의 2.1배였다. 1층 또는 2층에 게이트를 두면 속도 이득이 7.8배였고 정확도는 6% 이내였다. FastKV는 어텐션 층만 가지치기할 수 있다.
| 항목 | TokenGate | FastKV |
|---|---|---|
| 토큰 선택 | 학습된 MLP 게이트 | 마지막 질의의 어텐션 |
| 학습 필요 | 필요(게이트와 LoRA r=16) | 불필요 |
| Qwen3-1.7B 가지치기 층 | 1층 | 13층(기본) |
| Gated DeltaNet 층 가지치기 | 가능 | 불가(어텐션 층만) |
| 자체 보고 최고 성능 | TTFT 11.3배, 메모리 3.4배 | 프리필 1.82배(자체 설정, TTFT와 다른 지표) |
사용법과 비용
프롬프트에서 {state} 자리표시자로 가지치기할 문서를 표시하고, generate에 max_kept_prefill_tokens 예산을 넘긴다. 체크포인트는 Qwen3-1.7B의 다섯 과제와 Qwen3.5-2B의 NIAH 게이트가 있으며, 예시는 liana313/tokengate-niah-qwen3-1.7b이다.
선택 사항인 FlashAttention-2 키 편향 커널은 밀집 Qwen3를 빠르게 하며 일반 호출의 1.0~1.1배 비용이 든다. 커널이 없으면 FlexAttention으로 대체한다. 학습은 TRL/PEFT 방식을 따르고 한 단계가 일반 LoRA보다 1.3배 느리지만, 다섯 과제 중 네 과제에서 더 빨리 수렴했다. 평균 학습 시간은 원문에 A100 기준 1.8시간과 2.2시간으로 제시됐으나, 어느 쪽이 TokenGate인지는 명시되지 않았다.
출처
이 기사는 아래 자료를 바탕으로 정리했습니다. 제목을 누르면 원문이 새 창으로 열립니다.