TokenGate: Learned selection, reduced prefill

1. Teaching existing models to select tokens

The approach lets models choose what to process, much as people skim documents and read important passages closely. Existing models usually process every token at every layer. TokenGate scores each document token at a designated layer and sends only a set number to upper layers. Dropped tokens are excluded from subsequent layers' computation and their key-value caches, which store earlier results for reuse.

Existing token-selection methods filter tokens without training, using rules such as attention scores that indicate relationships between tokens. The paper says methods relying only on early-layer scores lost accuracy when many tokens were removed. TokenGate jointly fine-tunes the existing model and selector to learn which tokens to retain.

2. Jointly fine-tuning the selector and model

During training, token-selection scores also affect attention computations in upper layers, helping the model learn which choices support correct answers. Training incorporates both full and reduced contexts, so one trained selector can serve multiple token budgets.

The repository's inference example marks the document segment and specifies how many tokens to retain. The relevant excerpt from the README code is reproduced below.

for k in (None, 2000, 500, 100):          # None = keep every state token
    out = model.generate([ids], tok, max_kept_prefill_tokens=k)
    print(k, f"{out.prefill_s * 1e3:5.0f} ms", out.answers[0])

3. Greater latency reductions with longer contexts

The paper trained Qwen3-1.7B on a 16k needle-in-a-haystack task and evaluated it on longer contexts. This task requires finding a specific fact embedded in a long document and answering a question about it.

The table below shows results at matched accuracy against a fine-tuned model processing the full context. It compares time to first token (TTFT), the time to process input and produce the first token, and peak memory use.

Table: TokenGate savings by context length

Conditions: Qwen3-1.7B, bf16, trained on a 16k needle-in-a-haystack task, batch size 1, 1 A100, relative to full context

Unit: x

Speedups grew with context length, while memory reductions stayed similar at 3.42–3.43x. The paper attributes this to layers below the gate processing every token and affecting peak memory.

4. Repository tests and task-specific differences

The repository's separate benchmark, a standardized performance test, used held-out test data on one H100 with batch size 1. In the table below, Original is the model before training. Unpruned is the accuracy baseline with all tokens retained after the gate and model were fine-tuned for the task.

K=2000 retains 2000 document tokens. The code-completion score measures similarity to reference code, while the drug-reaction extraction score measures extraction completeness and accuracy.

Table: Test results published in the repository

Conditions: Based on Qwen3-1.7B, 1 H100, batch size 1, held-out test data, K=2000

Unit: TTFT speedup in x relative to the original model before training

Code completion scored lower than the unpruned baseline, while drug-reaction extraction scored slightly higher. In a separate repository needle-in-a-haystack demo, retaining 100 tokens from about 32k documents found the answer in all five cases, with an 11.18x TTFT speedup. These results used five documents, needle positions of 11–100%, and one GPU.

5. Speed and memory gains and their costs

In its main experiments, the paper reports 1.4–11.3x faster prefill and 1.5–3.4x lower peak memory at matched accuracy against full-context fine-tuned models. These results cover multiple tasks using Qwen3-1.7B, bf16, one A100, and batch size 1. In Section 5.1, speedups reached 6.4x and memory reductions reached 2.1x against the next-best selection method.

Lower inference costs come with the costs of training the selector and preparing the runtime environment. In the appendix's Qwen3-1.7B experiments on one A100, each training step was 1.3x slower than standard fine-tuning. However, total training time averaged across five tasks was 0.8x that of the standard approach.

Building CUDA kernels is recommended for fast execution in the default configuration, though inference can use a fallback path without kernels. Using the fallback path for training requires explicitly selecting it.

Accuracy losses also occurred depending on the task and token budget. In the paper's summarization example, events were summarized correctly, but names and locations were wrong.

TokenGate offers a practical option for reducing long-context prefill, but its gains depend on training data and task-specific token selection.