스탠퍼드 대학과 캘리포니아 대학 버클리 캠퍼스의 연구진은 2026년 10월 25일, 대규모 언어 모델(LLM)의 긴 컨텍스트 사전 채우기 단계를 위한 학습 가능한 토큰 가지치기 기술인 TokenGate를 공개했다. 이 모듈은 기존 모델에 ‘플러그 앤 플레이’ 방식으로 통합되며, 원래의 작업 정확도는 유지한 채 첫 번째 토큰 생성 시간(TTFT, Time-to-First-Token)을 1.4배에서 최대 11.3배까지 향상시키고, 최대 메모리 사용량은 1.5배에서 최대 3.4배까지 절감한다.
TokenGate의 핵심 메커니즘은 Transformer의 특정 중간 계층(여기서는 L층으로 표기) 뒤에 소형 다층 퍼셉트론(MLP) 제어 모듈을 삽입하고, LoRA(r=16)를 결합하여 경량 미세조정을 수행하는 것이다. 이 모듈은 입력 문서의 각 토큰에 대해 중요도 점수를 산출한 뒤, Top-K 방식을 통해 핵심 토큰만 유지한다. 이후 상위 계층에서는 이러한 K개의 토큰만 처리하며, 해당 토큰들에 대해서는 KV 캐시를 생성하지 않는다. 한편, 질문 팁 부분은 강제로 유지되며 가지치기 대상이 아니다.
Top-K 선택이 미분 불가능으로 인해 엔드투엔드 학습이 불가능한 문제를 해결하기 위해, TokenGate는 MoE 라우터 메커니즘을 도입합니다. 즉, 보존 토큰의 점수를 주의력 로그릿스에 대한 학습 가능한 편향 항으로 변환하고, 각 헤드별 증폭( per-head gain)과 학습 가능한 온도 계수(temperature)를 함께 적용하여 그라디언트가 점수 산출 모듈로 되돌아갈 수 있도록 합니다. 학습에서는 표준 다음 단어 예측 손실을 사용하며, 전체 컨텍스트 손실과 무작위 예산 손실을 혼합하여, 한 번의 학습 과정에서 게이팅 모듈이 다양한 가지치기 예산에 적응할 수 있도록 보장합니다.
본 실험은 Qwen3-1.7B 및 Qwen3.5-2B 모델을 기반으로 하며, RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX 및 needle-in-a-haystack(NIAH) 등 5개의 벤치마크에서 테스트를 수행했으며, 컨텍스트 길이는 최대 128k 토큰에 달합니다. 결과는 다음과 같습니다: NIAH-16k 과제에서 최적 성능을 달성했으며(TTFT 가속 11.3배, 메모리 절감 3.42배), 128k 컨텍스트 환경에서는 TTFT 가속이 13.0배에 이르고 메모리 절감률은 약 3.43배로 안정적으로 유지되었습니다. 또한 2WikiMultihopQA 과제에서는 전체 컨텍스트 정확도 95%를 유지하는 대신 TTFT 가속 3.0배와 메모리 감소 2.7배를 달성했습니다.
TokenGate는 Qwen3.5-2B와 같은 혼합 주의력 모델(게이트형 델타넷과 전 주의력 층을 포함)에서도 안정적인 성능을 보입니다. NIAH-32k 작업에서 제3층에 게이트를 적용하면 전체 사전 채우기 정확도에 도달할 수 있으며, 속도는 전체 사전 채우기에 비해 3.8배, FastKV에 비해 2.1배 빠릅니다. 또한 제1층 또는 제2층에 게이트를 적용하면 가속 효과가 최대 7.8배에 달하지만 정확도 손실은 불과 6%에 그칩니다. 이에 반해 FastKV와 같은 기준 방법들은 순수 주의력 층에만 적용 가능하며, 델타넷과 같은 비표준 구조에는 대응하지 못합니다.
이 기술은 오픈소스로 구현되어 있으며, Python 3.10 이상과 CUDA PyTorch 2.9 이상을 지원합니다. 또한 FlashAttention-2의 키 바이어스 커널을 선택적으로 통합하여 추가적인 가속을 제공합니다. GitHub 리포지토리에서는 총 5가지 작업에 대한 Qwen3-1.7B 체크포인트와, NIAH 작업에 특화된 Qwen3.5-2B의 게이트 기반 모델(예: Hugging Face 모델 ID: liana313/tokengate-niah-qwen3-1.7b)을 제공합니다. 학습 소요 시간은 표준 LoRA 미세조정의 약 1.3배이지만, 5개 중 4개의 작업에서는 더 빠르게 수렴하며 평균 학습 시간은 A100 기준 1.8시간입니다(기준선 대비 2.2시간).