TokenGate: 읽은 문서 토큰을 중간층에서 골라 첫 답변을 빠르게

배경과 문제

긴 문맥 입력은 딥 리서치, 문서 처리, 코딩 에이전트 같은 최근 언어 모델 응용의 핵심이다. 하지만 트랜스포머는 입력의 모든 토큰을 각 계층에서 처리하므로, 주의(attention) 계산은 입력 길이의 제곱에 비례해 늘고 MLP 계산과 KV 캐시 메모리는 선형으로 늘어난다. 이 때문에 긴 문맥의 프리필(prefill)이 모델 서빙의 주요 병목이 된다. 프리필은 사용자가 프롬프트를 보낸 뒤 모델이 첫 토큰을 생성하기 전에 전체 입력을 처리하는 단계로, 이 단계의 지연 시간을 TTFT(time-to-first-token)라고 부른다.

전문가 독자는 긴 문서를 훑어보며 중요한 부분만 자세히 읽지만, 언어 모델은 정보가 중요하든 아니든 모든 토큰을 동일하게 처리한다. 이 연구는 '기존에 학습된 모델에 긴 문맥 프리필 효율을 개선하는 능력을 부여하면서 정확도는 유지할 수 있을까?'라는 질문에서 출발한다. 문서는 요청 시점에만 제공되므로 KV 캐시를 미리 계산해 둘 수 없고, 프리필 시간과 메모리 사용량을 동시에 줄여야 한다.

기존 접근의 한계

기존 연구는 크게 두 갈래다. 첫째는 학습 없는 토큰 가지치기로, 주의 점수 기반 휴리스틱으로 토큰을 평가해 상위 계층에는 점수가 높은 토큰만 전달한다. FastKV, SlimInfer, DASH, GemFilter 등이 여기 속한다. 이 방법들은 추가 학습이 필요 없어 배포 비용이 낮지만, 정확도가 떨어지기 전까지 얻을 수 있는 속도 향상이 제한적이다. 휴리스틱 점수 함수가 학습된 점수 네트워크보다 성능이 낮고, 모델 백본이 가지치기 함수와 함께 적응하지 못하기 때문이다.

둘째는 트랜스포머 구조 자체를 바꾸는 방법이다. 동적 토큰 라우팅이나 학습된 가지치기 결정을 사용하는 구조, 또는 서브쿼드러틱 주의 구조 등이 있다. 그러나 이런 방법들은 처음부터 사전 학습해야 하거나, 추론 시 토큰 예산을 고정하는 보조 손실로 비싼 미세조정을 해야 해서 예산이 바뀔 때마다 재학습이 필요하다. 또한 서브쿼드러틱 구조도 여전히 모든 토큰을 각 계층에서 처리한다. TokenGate는 기존 사전 학습 모델을 그대로 확장하면서 간단하고 효율적으로 학습할 수 있는 대안을 제시한다.

연구방법을 단계별로

TokenGate는 트랜스포머의 특정 계층 L에 게이트 모듈을 배치한다. 순전파는 세 단계로 나뉜다. 첫째, 게이트 아래 계층들은 모든 N개 토큰을 평소처럼 처리한다. 둘째, 게이트 모듈이 각 토큰에 점수를 매기고 상위 K개만 남긴다. 셋째, 게이트 위 계층들은 선택된 토큰만 처리하는데, 이때 토큰 게이트 주의(token-gated attention)라는 메커니즘으로 중요도 점수를 주의 계산에 반영한다. 가지치기된 토큰은 게이트 이후 계층을 통과하지 않고 KV 캐시 항목도 생성하지 않으므로 계산과 메모리가 줄어든다.

게이트 모듈은 점수 네트워크와 top-K 선택기로 구성된다. 점수 네트워크는 작은 MLP로, 가지치기 가능한 각 토큰의 정규화된 잔차 스트림을 읽어 스칼라 점수를 출력하고, 이를 RMSNorm으로 정규화해 중요도 점수를 만든다. 질문이나 생성 중인 토큰 같은 비가지치기 토큰의 점수는 0으로 고정한다. top-K 선택기는 점수가 가장 높은 K개의 가지치기 가능 토큰과 모든 비가지치기 토큰을 선택해 상위 계층으로 전달한다. MLP는 0을 출력하도록 초기화되어 학습 시작 시 모든 점수가 0이다.

top-K 선택은 순위를 정해 잘라내므로 그 선택 자체를 부드럽게 미분하기 어렵다. TokenGate는 Mixture-of-Experts(MoE)의 게이트 학습에서 아이디어를 얻어, 선택된 토큰의 중요도 점수를 상위층의 주의 계산에도 넣는다. 기본식은 주의 가중치에 게이트 값을 곱한 뒤 다시 정규화하는 형태다. 실제 구현식은 정규화 전 점수인 주의 로짓에 중요도에 따른 항을 더하고, 머리별로 학습하는 배율과 남은 토큰 수에 맞춘 온도를 적용한다. 이 경로를 통해 최종 답변의 손실이 게이트 점수를 학습시킨다. 단순히 로짓에 게이트 값을 곱하는 연산과는 다르며, 게이트와 LoRA로 조정되는 모델이 함께 적응하는 것이 핵심이다.

학습 시에는 각 예제에서 전체 문맥 길이 Kfull과 로그 균등 분포에서 표집한 작은 Kprune 두 가지 예산으로 손실을 계산해 평균한다. Kfull 항은 전체 문맥에서의 정확도를 유지하도록 학습하고 모든 토큰에 기울기를 주며, Kprune 항은 다양한 가지치기 비율에 노출시켜 재학습 없이 넓은 예산 범위에서 일반화하게 한다. 게이트 계층까지의 순전파는 전체 문맥에서 한 번만 실행되고, Kfull과 Kprune 이후의 두 순전파가 계산 그래프를 공유하므로 학습 오버헤드는 Kprune 순전파에서만 발생한다. Kprune이 Kfull보다 훨씬 작아 오버헤드가 최소화된다.

결과와 숫자의 의미

실험은 코드 완성(RepoBench-P), 다중 홉 질의응답(2WikiMultihopQA), 다중 문서 요약(MultiNews), 의료 추출(BioDEX), 바늘 찾기(NIAH)의 다섯 과제에서 수행했다. Qwen3-1.7B 모델을 LoRA로 미세조정한 기준과 비교해, TokenGate는 정확도 저하 없이 TTFT를 1.4~11.3배, 최대 프리필 메모리를 1.5~3.4배 줄였다. 학습 없는 가지치기 기준선과의 별도 비교에서는 최대 6.4배 빠른 프리필과 2배 낮은 메모리를 보고했다. 앞의 1.4~11.3배와 뒤의 6.4배는 비교 대상이 다르다. 속도 측정은 A100 GPU 한 장, 배치 크기 1, 그리디 디코딩 최대 1024토큰 조건에서 측정했다.

그림 1(d)는 2WikiMultihopQA에서 TokenGate가 정확도 대 TTFT와 정확도 대 최대 프리필 메모리의 파레토 경계를 개선함을 보여준다. 전체 문맥 정확도의 95%를 유지하는 조건에서 3.0배 빠른 프리필과 2.7배 낮은 메모리를 달성했는데, 이는 약간의 정확도 손실을 허용한 결과다. 무손실 조건은 아니다. 하이브리드 주의 구조인 Qwen3.5-2B에서는 NIAH 32k 과제에서 전체 문맥 기준과 동일한 정확도로 TTFT를 3.8배 줄였고, 같은 정확도에서 최고 성능 기준선보다 2.1배 낮은 TTFT를 보였다. 게이트를 더 낮은 GatedDeltaNet 계층(1, 2층)에 적용하면 정확도가 전체 문맥 대비 6% 이내로 떨어지는 조건에서 최대 7.8배 빠른 프리필을 달성했다.

학습 효율성 측면에서, TokenGate 학습 스텝은 기준선보다 1.3배 느리다. 각 스텝에서 두 가지 토큰 예산을 표집하기 때문이다. 다섯 과제 중 네 과제에서는 더 적은 스텝으로 수렴했지만, 총 학습 시간이 줄어든 과제는 세 개다. 표 2에 따르면 총 학습 시간은 RepoBench-P 1.7시간→1.1시간, 2WikiMultihopQA 3.8시간→4.0시간, MultiNews 1.3시간→1.8시간, BioDEX 1.2시간→1.0시간, NIAH 2.8시간→1.0시간으로, 평균 2.2시간에서 1.8시간으로 감소했다. 모든 과제에서 학습이 빨라진 것은 아니며, LoRA와 주의 커널 수정이 필요하다.

구체적인 예시

부록 E의 시각화는 실제 평가 프롬프트에서 TokenGate가 어떤 토큰을 유지하는지 보여준다. 2WikiMultihopQA에서는 질문의 답에 해당하는 'Film Award for ... Director' 같은 구절과 인물 이름, 연도 같은 개체를 주로 유지했다. NIAH 16k에서는 15,843토큰의 건초 더미 속에서 78% 깊이에 있는 바늘 문장 24토큰 중 21토큰을 K=64 예산으로 유지해 정확히 답했다. 반면 MultiNews에서는 기능어를 주로 유지하고 이름과 장소를 적게 유지해, K=256에서 사건은 맞지만 이름과 직업, 장소를 틀리는 요약을 생성했다.

RepoBench-P에서는 코드 식별자보다 쉼표나 빈 인자 목록 같은 구문 토큰을 주로 유지했고, BioDEX에서는 진단명과 결과 용어('septic shock', 'failure')와 목록 연결어('and', 'or')를 유지했지만 숫자는 거의 유지하지 않았다. BioDEX의 한 예에서는 전체 문맥 답변이 두 반응을 제시한 반면, 가지치기 답변은 그중 하나만 남겼다. 원래 정답 목록에는 더 많은 반응이 있으므로 전체 문맥 답변도 완전하지 않았다. 이는 작은 K에서 정보 손실이 발생할 수 있음을 보여준다. 중간 계층 표현은 다른 토큰 정보를 이미 섞었기 때문에, 선택된 토큰이 문자 그대로의 중요도를 설명한다고 단정하기는 어렵다.

한계와 오해하기 쉬운 점

최대 11.3배 TTFT 감소는 16k 합성 NIAH 과제에서 나온 수치다. 모든 실제 문서나 코딩 과제에 같은 효과가 있는 것은 아니다. 부록 D.3의 표 3은 16k NIAH로 학습한 모델을 32k, 64k, 128k 합성 NIAH에 적용했을 때 TTFT 12.8배, 13배, 13배, 메모리 3.43배 감소를 보고하지만, 이는 합성 단일 바늘 길이 전이 실험이며 모든 128k 문서 추론의 일반화 증거는 아니다. 또한 그림 1의 3.0배 속도와 2.7배 메모리는 전체 문맥 정확도의 95% 보존 조건으로 약간의 손실을 허용한 결과다.

평가 지표와 샘플 수는 과제마다 다르다. RepoBench-P는 500건의 편집 거리 유사도, 2WikiMultihopQA는 200건의 F1, MultiNews는 200건의 ROUGE-L, BioDEX는 500건의 이상 반응 집합 F1, NIAH는 200건의 정답 활동 포함 여부다. 평균 문맥 길이는 각각 11k, 7.5k, 2.7k, 4.5k, 16k 토큰이다. 이 수치들을 섞어 비교하면 안 된다. 가지치기 기준선은 컴파일 친화적으로 재구현되었고, 학습 없는 기준선과 학습한 방법의 차이도 있다. 생성 중 새 토큰은 가지치기하지 않으며, 하위 계층 KV의 선택적 가지치기는 벤치마크에서 비활성화했다. 실서비스 배치, 하드웨어, 다중 턴 캐시, 출력 TPS, 배포 비용 절감은 별도로 검증되지 않았다.

실제로 어떻게 활용할까

TokenGate는 모델 내부의 게이트와 주의 연산을 수정해 기존 미세조정 파이프라인에 추가하는 구성 요소다. LoRA와 함께 표준 다음 토큰 예측 손실로 학습하므로, 추가 보조 손실 없이 기존 미세조정과 비슷한 비용으로 적용할 수 있다. 코드는 공개되어 있다. 긴 문서를 처리하는 의료 추출, 딥 리서치, 코딩 에이전트 같은 응용에서 프리필 지연과 메모리 사용량을 줄이는 데 유망하다.

그러나 과제별로 별도 학습이 필요하며, 한 체크포인트가 다섯 과제 모두에서 잘 작동한다는 실험은 없다. 배포 시에는 게이트 계층 선택, 토큰 예산 K, 하드웨어 환경을 고려해야 한다. 학습된 모델은 다양한 입력 길이와 토큰 예산에 일반화하지만, 이는 실험한 범위 내에서의 관찰이다. 실제 서비스에 적용하려면 배치 처리, 다중 턴 대화, 출력 속도 등 추가 검증이 필요하다.