TokenGate는 긴 문맥 프리필에서 학습 토큰 선택으로 속도와 메모리를 줄인다

무엇인가

이 논문은 긴 문맥 입력에서 발생하는 프리필 병목을 다룬다. 트랜스포머는 입력의 모든 토큰을 각 레이어에 통과시키기 때문에 어텐션 연산은 입력 길이에 대해 제곱으로, MLP 연산과 KV 캐시 메모리는 선형으로 증가한다. 논문은 문서가 요청 시점에 주어져 KV 캐시를 미리 계산할 수 없는 상황을 전제로, 정확도는 유지하면서 TTFT와 프리필 메모리를 줄이는 방법을 묻는다. 기존의 학습 없는 토큰 프루닝은 배포 비용이 낮지만 정확도가 떨어지기 전까지 얻는 속도 향상이 제한적이고, 아키텍처를 바꾸는 접근은 처음부터 사전학습해야 해 비용이 크다는 문제를 지적한다.

어떻게 동작하나

TokenGate는 기존 학습된 모델에 붙이는 작은 드롭인 컴포넌트다. 특정 레이어 L에 게이트 모듈을 두고, 그 아래 레이어는 모든 토큰을 표준 방식으로 처리한다. 게이트 모듈은 프루닝 가능한 각 토큰의 정규화된 post-attention residual을 MLP로 읽어 스칼라 점수를 만들고, 시퀀스 단위 RMSNorm으로 중요도 점수를 정규화한 뒤 top-K로 남길 토큰을 고른다. 이때 답변 토큰, 디코딩 토큰, 사용자 질문처럼 지정된 비프루닝 토큰은 항상 유지된다. 게이트 위 레이어는 선택된 토큰만 처리하며, 프루닝된 토큰은 상위 레이어를 통과하지 않고 KV 캐시 항목도 생성하지 않는다.

무엇과 다른가

핵심 난점은 top-K 선택이 미분 불가능하다는 점이다. 저자들은 MoE에서 라우터가 전문가 출력을 가중하는 방식에서 착안한 token-gated attention을 제안한다. 게이트 중요도 점수 b_j를 softmax로 게이트 값으로 바꾸고, 상위 레이어의 어텐션 로짓에 softplus(γ)·b_j 항을 더해 선택된 토큰의 어텐션 가중치를 재조정한다. 레이어·헤드별 학습 가능한 게인 γ와, 프루닝 후 softmax가 날카로워지는 문제를 보정하는 온도 τ도 둔다. 온도는 유지된 토큰 순위 m_i와 원래 위치 p_i의 로그 비율로 스케일링해 길이 일반화를 노린다. 학습 시작 시 점수가 0이면 이 구조는 top-K 어텐션과 동등하다. 학습은 LoRA와 게이트 파라미터를 함께 표준 next-token cross-entropy로 진행하며, 전체 문맥 K_full과 로그균등 분포에서 뽑은 K_prune 두 예산의 손실을 평균한다. 전체 문맥 항은 정확도를 고정하고 모든 토큰에 그래디언트를 주며, 프루닝 항은 게이트가 다양한 예산에 노출되게 한다.

어떻게 쓰나

실험은 LongBench 계열 3개 데이터셋, 2WikiMultihopQA, RepoBench-P, MultiNews, BioDEX 약물 반응 추출, NIAH에서 수행됐다. 비교 대상은 FastKV, SlimInfer, DASH, GemFilter 같은 토큰 프루닝 방법과 FlexPrefill, MInference 같은 희소 어텐션 방법이다. 베이스 모델은 표준 트랜스포머 Qwen3-1.7B와 하이브리드 어텐션 모델 Qwen3.5-2B(실험 설정 일부에는 Qwen3-2B로 표기)이며, LoRA r=16으로 학습하고 4대의 A100을 사용했다. 평가 지표는 과제 정확도, TTFT, 프리필 피크 메모리다.

전제와 한계

결과적으로 TokenGate는 LoRA로 파인튜닝한 full-context 기준 모델과 비교해 정확도 저하 없이 TTFT를 1.4~11.3배, 프리필 피크 메모리를 1.5~3.4배 줄였다. 기존 최고 토큰 프루닝 방법과 비교하면 최대 6.4배 빠른 프리필과 최대 2배 낮은 피크 메모리를 보고한다. 본문 5.1절에서는 동일 정확도에서 차선 방법보다 최대 2.1배 낮은 피크 메모리라고도 적는다. 하이브리드 모델에서는 레이어 3에 적용했을 때 표준 프리필과 같은 정확도로 TTFT를 3.8배 줄였고, 같은 정확도에서 최고 FastKV보다 2.1배 낮은 TTFT를 보였다. 더 낮은 GatedDeltaNet 레이어 1·2에 적용하면 표준 프리필 대비 정확도 6% 이내에서 최대 7.8배 빠른 프리필이 가능했다. 절제 실험에서는 학습된 선택기가 같은 레이어의 랜덤-K와 FastKV 선택기보다 정확도가 높았고, 같은 랜덤-K를 쓰더라도 TokenGate로 학습한 모델이 파인튜닝 베이스보다 모든 예산에서 높은 정확도를 보여 백본과 게이트의 공동 적응 효과를 확인했다. 학습 오버헤드는 스텝당 1.3배 느리지만 5개 과제 중 4개에서 더 적은 스텝으로 수렴해 전체 학습 시간은 표준 파인튜닝과 비슷하거나 소폭 줄었다.

개발자 관점에서 TokenGate는 긴 코드 저장소, 다중 문서 QA·요약, 의료 추출, NIAH, RAG·에이전트처럼 프리필이 TTFT와 메모리를 지배하는 워크로드에 의미가 있다. 학습 없는 프루닝이 아니라 게이트와 LoRA를 파인튜닝해야 하며, 게이트 레이어 위치, 토큰 예산, 항상 유지할 비프루닝 토큰을 정해야 한다. 논문은 QA에서는 답 키워드, 요약에서는 구절 끝의 stop-word를 남기는 등 과제별로 다른 프루닝 패턴이 학습된다고 시각화해 보여준다. 또한 이 방법은 프리필을 줄이는 것이지 디코딩 비용을 줄이는 것이 아니므로, 디코드 단계의 KV 캐시 제거·병합 기법과 상보적으로 봐야 한다. 문서가 요청 시점에만 주어져 KV 캐시를 미리 계산할 수 없다는 전제도 실제 서빙 구조와 맞는지 확인해야 한다.

논문은 별도의 한계 절을 두지 않았지만, 전제와 범위는 분명하다. 요청 시점 문서 처리, 단일 게이트 레이어 중심의 설계, top-K 예산과 비프루닝 토큰 집합에 성능이 의존한다. 평가는 Qwen3-1.7B와 Qwen3.5-2B 계열 및 열거된 과제에 한정되며, 하이브리드 모델의 낮은 레이어처럼 공격적으로 프루닝하면 정확도가 6% 이내로 떨어지는 경우도 있다. 정확도 저하 없음은 원래 베이스 모델이 아니라 LoRA로 파인튜닝한 full-context 기준과의 비교라는 점도 실무 적용 전에 확인해야 한다.