스탠포드와 UC 버클리 연구원들이 발표했습니다 TokenGate긴 컨텍스트에 대한 프리필을 위한 모듈입니다. TokenGate는 토큰 프린싱을 학습하여 기존 LLM이 입력을 빠르게 처리할 수 있도록 합니다. 중요한 토큰만이 완전한 처리를 받습니다. 연구팀은 보고서를 발표했습니다. 1.4배에서 11.3배 빠른 첫 토큰 도달 시간(TTFT) 정확도 손실 없이. 피크 프리필 메모리가 1.5배에서 3.4배로 감소합니다. 코드와 체크포인트는 GitHub Apache-2.0에 따름.
TL;DR
- MLP 게이트: 작은 MLP 게이트와 LoRA(r=16)가 Qwen3-1.7B 또는 Qwen3.5-2B에 추가되었습니다. 128k 토큰까지의 맥락에서 테스트되었습니다.
- 실행 환경: 종이 기반 환경에서는 배치 크기 1에 1 A100이 사용됩니다. 훈련 시에는 4개의 A100이 사용되었습니다.
- 성능: 5개의 작업에서 6개의 기준값 대비 파레토 최적화를 달성했습니다.
- 최고: NIAH 16k는 11.3배 빠른 TTFT와 3.42배 낮은 메모리를 제공합니다. TTFT의 향상은 128k에서 13.0배에 달합니다.
- 최악: 일부 작업에서 속도가 1.4배로 감소합니다.
TokenGate는 무엇인가요?
pre-training은 일반적인 훈련과는 다른 이전 훈련 단계로, 사전 학습이라고 사용합니다. TokenGate는 사전 훈련된 LLM에 훈련 가능한 토큰 제거 기능을 추가하는 구성 요소입니다. Transformers는 모든 토큰을 모든 층을 통과시킵니다. 주의력 비용은 길이에 따라 4배로 증가합니다. 훈련 없이 하는 제거는 비용이 저렴하지만 높은 제거 비율에서는 정확도를 잃습니다. 구조적 변화는 종종 처음부터 다시 사전 훈련이 필요합니다.
TokenGate은 토큰 프린팅 작동 방식을 어떻게 배웠나요?
전진 패스는 3단계로 나뉩니다:
- 표준 전진: 게이트 계층 L 아래의 레이어는 일반적으로 N 개의 토큰을 처리합니다.
- 게이트 모듈: MLP는 레이어 L에서 포스트-앵커링 잔여값을 기반으로 각 토큰의 점수를 매깁니다. 상위-K 선택기로 가장 높은 점수를 받은 토큰을 유지합니다.
- pruned 전방 처리: L 이하의 레이어는 K 개의 토큰만을 처리하며, pruned된 값에 대한 KV 엔트리를 생성하지 않습니다.
프로ンプ스의 문서 부분만을 삭제할 수 있습니다. 질문은 절대 삭제되지 않습니다.
토큰 기반 접근 제한된 주의력이란 무엇인가?
Top-K 선택은 구별되지 않아 엔드투엔드 학습이 불가능합니다. TokenGate는 전문가 혼합 루터 기술을 사용합니다. 유지된 각 토큰의 점수는 게이트 위의 주의력 로지트에 편향으로 작용합니다. 그 편향은 그래디언트가 점수 계산기로 돌아가는 경로를 제공합니다.
학습 가능한 per-head 성과와 학습 가능한 온도가 이를 개선한다. 온도는 키가 적은 각 쿼리를 고려하여 수정된다. 훈련에서는 간단한 next-token 손실을 사용한다. 각 단계에서는 전체 컨텍스트 손실과 랜덤 예산 손실의 평균값을 사용한다. 따라서 1개의 훈련된 게이트는 모든 예산에서 작동한다.
TokenGate은 벤치마크에서 어떻게 작동하나요?
연구팀은 Qwen3-1.7B를 RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX 및 needle-in-a-haystack (NIAH)에서 테스트했습니다. 기준 모델로는 FastKV, SlimInfer, DASH, GemFilter, FlexPrefill 및 MInference가 있었으며, 이들 모두 동일한 LoRA 백엔드를 사용했습니다.
- Versus 전체 채우기: TTFT는 1.4배에서 11.3배 더 빠르고, 메모리는 1.5배에서 3.4배 더 적습니다.
- 다음으로 좋은 기준선에 비해: TTFT는 1.2배에서 6.4배 더 빠르고, 메모리 사용량은 2.1배 낮습니다.
- 2WikiMultihopQA: 전체 컨텍스트 정확도가 95%일 때, 프리필링 작업은 3.0배 더 빠르게 수행되며 메모리 사용량은 2.7배 줄어듭니다.
- 더 긴 컨텍스트 경우: 16k에서 훈련된 게이트는 32k에서는 12.8배, 64k와 128k에서는 13.0배의 성능을 보입니다. 메모리 절약은 약 3.43배입니다.
학습된 선택자는 동일한 계층에서 FastKV의 지능적 방법을 능가합니다. 심지어 무작위-K 선택도 TokenGate 백본에서 더 높은 점수를 얻습니다. 연구팀은 백본의 공동 적응에 기여했다고 평가합니다.
하이브리드-인지 모델에서도 작동하나요?
네, Qwen3.5-2B는 게이트드 델타넷과 전체 주의도 레이어를 혼합합니다. NIAH 32k에서 레이어 3의 토큰게이트는 전체 프리필 정확도와 일치했습니다. 이 모델은 전체 프리필보다 3.8배 빠르고 FastKV보다 2.1배 빠르게 작동합니다. 레이어 1 또는 2의 게이팅은 6%의 정확도로 7.8배의 속도를 달성했습니다. FastKV는 주의도 레이어만 삭제할 수 있습니다.
TokenGate vs FastKV, SlimInfer, GemFilter 및 DASH
| 기능 | TokenGate | FastKV | SlimInfer | GemFilter | DASH |
|---|---|---|---|---|---|
| 토큰 선택 | MLP 게이트를 배웠습니다 | 마지막 질문에서의 주의사항 | 주의를 끌기 위한 블록, 다단계 | 필터 레이어에서 마지막 쿼리에 주의 | 숨겨진 상태의 안정성 |
| 훈련이 필요함 | 예 (게이트 + LoRA r=16) | 아니요 | 아니요 | 아니요 | 아니요 |
| TokenGate의 Qwen3-1.7B 실행에서 가지치기 레이어 | 1 | 13 (기본값) | 8, 17, 25 | 12 | 11 |
| 게이티드 델타넷 레이어에서 프루인 | 예 | 아니요 (층면만 포함) | 공개되지 않았습니다 | 공개되지 않았습니다 | 공개되지 않았습니다 |
| 최적의 자가 보고된 프리필 이득 (자신의 설정) | 11.3x TTFT, 3.4x 메모리 | 1.82배 프리필 | 2.53x TTFT (LLaMA3.1-8B, RTX 4090) | 2.4배 빠른 속도, SOTA 대비 GPU 메모리 30% 감소 | 공개되지 않았습니다 |
| 코드 | GitHub | GitHub | GitHub | GitHub | 공개됨 |
자신이 보고한 성과는 서로 다른 설정으로 측정되기 때문에 직접 비교할 수 없습니다. 프린싱 레이어는 TokenGate의 부록 C.2에서 가져왔습니다.
TokenGate을 어떻게 배포하나요?
GitHub 저장소는 Python 3.10+와 CUDA PyTorch 2.9+가 필요합니다. 선택적으로 사용할 수 있는 FlashAttention-2 키-비아스 커널은 dense Qwen3를 더 빠르게 처리해 줍니다. 이 기능이 없을 경우 인출 과정은 FlexAttention으로 전환됩니다. 비용은 일반적인 호출의 1.0x에서 1.1x에 해당합니다.
{state} 플레이스홀더는 삭제 가능한 문서를 나타냅니다. generate는 max_kept_prefill_tokens 예산을 사용합니다. 체크포인트는 Qwen3-1.7B의 5개 작업 전체와 더불어 Qwen3.5-2B NIAH 게이트도 포함합니다. 예시: liana313/tokengate-niah-qwen3-1.7b.
훈련은 TRL/PEFT 레시피를 따릅니다. 각 단계는 일반적인 LoRA보다 1.3배 더 느립니다. 5개 작업 중 4개에서 훈련 과정이 더 빠르게 수렴되었습니다. 평균 훈련 시간은 A100 시간의 2.2에 비해 1.8 시간이었습니다.
주요 요약
- TokenGate은 작은 게이트와 LoRA를 통해 기존 LLM에 학습된 토큰 프린팅 기능을 추가합니다.
- 프로세스 속도가 1.4배에서 11.3배 더 빠르고, 최대 메모리 사용량은 1.5배에서 3.4배 더 낮습니다.
- 토큰 기반 접근 방식의 주의력은 MoE-러터 형태로 상위-K 선택을 훈련 가능하게 만든다.
- TTFT에서 다음으로 좋은 프린팅 기준선보다 최대 6.4배 더 높다.
- 이 기능은 하이브리드 Qwen3.5-2B에서 작동하며, FastKV는 DeltaNet 레이어를 삭제할 수 없습니다.
논문, GitHub 저장소 및 Hugging Face 논문 페이지을 확인해 보세요. 또한, Twitter에서 우리를 팔로우하거나, 150k+ML SubReddit에 가입하고 우리 뉴스레터을 구독하는 것도 좋아요. 잠깐! 당신은 Telegram에도 있나요? 이제 당신도 Telegram에서 우리와 함께할 수 있어요.
귀하의 GitHub 저장소, Hugging Face 페이지, 제품 출시, 웹 세미나 등을 홍보하기 위해 우리와 파트너 관계를 맺고 싶으신가요? 우리와 연락하세요
스탠포드와 UC 버클리의 연구자들이 TokenGate을 발표함: 토큰 기반의 애널리틱이 긴 컨텍스트 예비 작업을 11.3배 더 빠르게 수행한다라는 게시물은 MarkTechPost에 처음 등장했습니다.