PAPER DEEP DIVE
TokenGate: Learnable Token Pruning for Fast, Memory-Efficient LLM Prefill
TokenGate learns a residual scorer at a middle layer, retains top-K document tokens plus request scaffolding, and uses token-gated attention to train through the discrete selection. It reports 1.4-11.3x faster TTFT and 1.5-3.4x lower peak prefill memory on Qwen3 and Qwen3.5 without accuracy loss; robotics performance remains untested.
TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning
Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia, Carlos Guestrin (Stanford University and UC Berkeley)
arXiv:2610.26754v1 [cs.CL], submitted 19 October 2026 · Paper · Working public repository: token-gate-llms/TokenGate (Apache-2.0)
Code status: the repository organization printed in the paper does not match the repository inspected here. The implementation and benchmark paths below were checked against token-gate-llms/TokenGate, default branch main, commit 6e51aac5.
One-sentence summary
TokenGate inserts a small learned gate at one Transformer layer, keeps the highest-scoring $K$ document tokens plus all non-prunable instruction tokens, and sends only that subset through the upper layers; token-gated attention restores a training path through the otherwise non-differentiable top-$K$ choice, yielding $1.4$–$11.3\times$ faster time-to-first-token and $1.5$–$3.4\times$ lower peak prefill memory on Qwen3-1.7B and Qwen3.5-2B without reported accuracy loss.
1. Why prefill is the bottleneck
Long-context inference spends a large fraction of its latency before the first output token exists. Every supplied token must traverse the model, and each layer must create keys and values that later decoding will reuse. Attention work grows roughly quadratically with sequence length, while MLP work and KV-cache storage grow linearly. Decode-time cache eviction helps later generation, but it cannot remove this initial full read. The paper therefore targets prefill directly, with two metrics tied to serving economics: time-to-first-token and peak GPU memory during prefill.
The motivating analogy is expert skimming: a reader does not give equal attention to every sentence in a long repository or document. The analogy does not license arbitrary deletion, but it suggests a useful architecture. A middle layer has already mixed enough context to expose which positions matter for the current request. If the model could keep a small set after that layer, later layers would avoid both MLP work and KV-cache growth for discarded positions.
Training-free pruning is attractive because it can be dropped into a served model without adaptation. Existing methods score tokens with attention statistics or other heuristics and retain a budget at an intermediate layer. Their weakness appears when the budget becomes small: early-layer scores may not identify query-relevant content reliably, and the upper backbone has not learned to operate on a suddenly shorter context. TokenGate treats these as one problem rather than trying to make a better fixed heuristic. The selector and the receiver of its selection are trained together.
The alternative of designing a new sparse architecture can lower asymptotic cost more aggressively, but it may require pretraining from scratch and a different kernel stack. TokenGate is deliberately incremental. It preserves the pretrained backbone, adds a cloned scoring branch and a gate module, and adapts the model with LoRA. The method's practical claim is thus narrower and more testable: can an ordinary fine-tuning pass turn an existing model into a budget-aware prefill system without sacrificing full-context accuracy?
No robotics experiment appears in the paper. For embodied systems, the relevance is infrastructural. A vision-language-action model may prefill long camera histories, retrieved procedures, logs, or multiple sensor descriptions before producing a short action chunk. Learned prefill pruning could matter when most observations are irrelevant to the current instruction, but that remains a hypothesis. Language-benchmark TTFT reductions do not establish safer or faster physical control.
2. Problem definition and Transformer baseline
The input is split into a request and request-time documents: $x=[x_R,x_D]$. The request $x_R=(x_1,\ldots,x_r)$ contains the question, instruction, and answer scaffold. The document portion $x_D=(x_{r+1},\ldots,x_N)$ may come from an upload, a retriever, or an agent tool and can be thousands of tokens long. Because documents arrive only when the request arrives, their KV cache cannot be computed ahead of time. Prefill is where both the latency and the transient memory spike occur.
A standard layer first mixes positions through attention and then applies a position-wise MLP. With normalized residual inputs, one layer can be written as
$$a_i^{\ell}=h_i^{\ell}+\operatorname{Attn}^{\ell}\!\left(\operatorname{LN}_1(h^{\ell})\right)_i,\qquad h_i^{\ell+1}=a_i^{\ell}+\operatorname{MLP}^{\ell}\!\left(\tilde a_i^{\ell}\right),\qquad \tilde a_i^{\ell}=\operatorname{LN}_2(a_i^{\ell}) \tag{1}$$
Here $a_i^\ell$ is the post-attention residual. TokenGate uses this state at its chosen gate layer because it has already accumulated contextual information from lower layers, while the current layer's MLP remains available to process only retained positions. This placement is not arbitrary; the ablation later shows that gating at residual depth 0 is materially worse than gating at layers 1 or 2.
Causal self-attention itself is
$$\operatorname{Attn}_i=\sum_{j\le i}\alpha_{ij}v_j,\qquad \alpha_{ij}=\operatorname{softmax}_{j\le i}(z_{ij}),\qquad z_{ij}=\frac{q_i^{\top}k_j}{\sqrt{d_{\mathrm{head}}}} \tag{2}$$
Prefill evaluates this expression for all positions and stores per-layer keys and values. TokenGate changes the token set at the gate rather than merely masking a few attention edges. Rejected positions stop at the gate, do not enter upper-layer MLPs, and are absent from upper-layer KV caches. This creates the efficiency opportunity, but it also explains why savings depend on gate depth: layers below the gate still execute the full dense prefill.
3. The TokenGate mechanism
3.1 Scoring and top-K selection
Let $\mathcal P$ denote the prunable document positions and $\mathcal N$ the positions that must survive, such as the question, response template, and optionally target tokens during training. A small MLP maps each normalized post-attention residual to a scalar, followed by sequence-axis RMS normalization:
$$\tilde b_j=\operatorname{MLP}_{\theta}(\tilde a_j^L),\qquad b_j=\frac{\tilde b_j}{\sqrt{\frac{1}{|\mathcal P|}\sum_{j'\in\mathcal P}\tilde b_{j'}^2+\epsilon}},\qquad b_j:=0\quad(j\in\mathcal N) \tag{3}$$
Normalizing over positions gives scores a request-relative scale rather than allowing unusually large hidden norms to dominate the top-$K$ ordering. The score MLP starts at zero, so the gate initially has no preferred document position. The retained set is
$$S=\operatorname{topK}_{j\in\mathcal P}(b_j)\cup\mathcal N \tag{4}$$
The top-$K$ operator is discrete. Gradients cannot flow through a token merely because it crossed or failed to cross the selection boundary. TokenGate does not pretend otherwise. Instead, it gives selected scores a second role inside upper-layer attention, where smooth derivatives exist and the task loss can still teach the scoring network what matters.
3.2 Token-gated attention
The basic form weights each attention edge by the score of its key:
$$\alpha_{ij}=\frac{\exp(z_{ij})g_j}{\sum_{j'\in S,j'\le i}\exp(z_{ij'})g_{j'}},\qquad g_j=\operatorname{softmax}_{j\in S}(b_j) \tag{5}$$
This resembles a mixture-of-experts router in structure, but the sparsity axis is input length rather than parameter activation. Once $K>1$, changing a retained token's score changes its share of attention mass and therefore the task loss almost everywhere. A token that contributes useful information can receive a larger score; a distractor can be suppressed even after it survives selection.
One global score is too rigid for multiple heads and layers, so the paper introduces a nonnegative learned gain $\gamma^{\ell h}$:
$$g_j^{\ell h}=\operatorname{softmax}_{j\in S}\!\left(\operatorname{softplus}(\gamma^{\ell h})b_j\right) \tag{6}$$
Softplus prevents a learned negative gain from inverting the intended ordering. The gain is initialized at the standard deviation of that layer's attention logits, keeping the new term at a comparable numerical scale. Per-head gains also let different heads express different preferences without forcing one shared interpretation of importance.
Pruning changes softmax sharpness. A query that originally sees $p_i$ positions may now see only $m_i$ retained ones, so the same logits can produce a more concentrated distribution. The paper calibrates a learnable temperature with the ratio of these counts:
$$\tau_i^{\ell h}=\operatorname{softplus}(\sigma^{\ell h})\frac{\log m_i}{\log\max(2,p_i)} \tag{7}$$
Combining gain and temperature gives the operative attention rule:
$$\alpha_{ij}^{\ell h}=\operatorname{softmax}_{j\in S,j\le i}\left[\tau_i^{\ell h}\left(\frac{q_i^{\top}k_j}{\sqrt{d_{\mathrm{head}}}}+\operatorname{softplus}(\gamma^{\ell h})b_j\right)\right] \tag{8}$$
The causal mask remains in force, the gate score enters as a per-key bias, and the temperature depends on the query position and current pruning ratio. A modified fused attention kernel adds the per-key, per-head bias without abandoning online softmax. The appendix reports that this biased call costs about $1.0$–$1.1\times$ a plain fused call, so the mechanism preserves most of the hardware efficiency that makes full prefill fast in the first place.
3.3 Training with two budgets
The score network should preserve full-context ability and learn aggressive pruning at the same time. Training only at a small budget can erode the original model; training only at full budget provides no pressure to identify a useful subset. TokenGate samples both and averages their losses:
$$K_{\mathrm{full}}=|\mathcal P|,\qquad K_{\mathrm{prune}}\sim\operatorname{LogUniform}\left(1,|\mathcal P|/2\right),\qquad \mathcal L=\frac{1}{2}\left(\mathcal L_{K_{\mathrm{full}}}+\mathcal L_{K_{\mathrm{prune}}}\right) \tag{9}$$
Each term is ordinary next-token cross-entropy under its token budget:
$$\mathcal L_K=-\sum_{t\in\mathrm{target}}\log p_{\phi}\left(x_t\mid x_{<t};K\right) \tag{10}$$
The full-budget branch anchors the model and gives every token a gradient from the task. The log-uniform branch exposes the selector to budgets spanning orders of magnitude, which lets one trained checkpoint serve several $K$ values at inference. Implementation packs the two targets into one variable-length call. Phase 1 is evaluated once; the two post-gate targets share much of the later computation, so a training step is not simply two complete model passes.
flowchart LR
X[request tokens plus document tokens] --> P1[full prefill below gate layer L]
P1 --> R[post-attention residual a_i^L]
R --> SCORE[MLP clone of gate-layer MLP
sequence RMSNorm scores b_j]
R --> KEEP[non-prunable N
instruction, question, target scaffold]
SCORE --> TOPK[top-K over prunable P]
TOPK --> SET[S = top-K(P) union N]
KEEP --> SET
SET --> SLICE[slice hidden states
recompute kept positions]
SLICE --> UP[upper layers with per-head gate bias
and pruning-ratio temperature]
SCORE -. task gradient through bias .-> UP
UP --> OUT[first token and decode]
3.4 How the released code implements the paper
The finalized class in tokengate/e2e_gate/gate_module/e2e_gate_module.py:204-239 records the intended invariants: sequence-axis RMS normalization, softplus gate scale, a learnable temperature, per-layer and per-head bias, uniform head aggregation for selection, and a scoring branch cloned from the gate layer's MLP. Lines :278-305 create the raw gain and temperature parameters; the forward paths convert them with softplus so positivity is guaranteed rather than learned by accident.
The actual prune path is compact. At tokengate/e2e_gate/gate_module/gate_forward.py:1155-1207, the code scores article tokens, clamps the budget, selects local indices, builds a full-sequence keep mask, slices the hidden states, recomputes positions for the shorter sequence, and installs the gate bias for Phase 2:
b = gate(h_art_cast)
b_scores = gate.aggregate(b)
topk_local = _select_indices(b_scores, K_clamped, gate)
keep_indices = _keep_indices(topk_local, article_start, article_end, seq_len, hs_device)
hs_pruned = hidden_states[:, keep_indices, :]
kept_pos = position_ids[:, keep_indices.to(position_ids.device)]For Qwen3.5, tokengate/e2e_gate/attn/qwen3_flex_attention.py:584-628 implements the Phase 2 attention path. The learned temperature is factored into the query scale, while the per-key gate bias is passed to the attention kernel. Comments in that block show the algebraic rewrite that keeps the scale outside a compiled score modifier:
attn_scale * (Q @ K^T / sqrt(d) + b)
= (attn_scale * Q) @ K^T / sqrt(d) + attn_scale * bThe benchmark harness also exposes an important measurement choice. benchmarks/run_benchmark.py:1-12 defines three arms: the full base model, the TokenGate model at full budget, and TokenGate at each pruned budget. Lines :117-123 run one sample at a time and record TTFT through prefill_s. Batch size 1 isolates per-request latency, but it does not establish serving throughput under concurrent load.
4. Evidence across tasks and architectures
The evaluation spans RepoBench-P for next-line code completion, 2WikiMultihopQA for multi-hop question answering, MultiNews for multi-document summarization, BioDEX for biomedical reaction extraction, and NIAH-16k for long-context retrieval. The main transformer is Qwen3-1.7B. A second model, Qwen3.5-2B, tests compatibility with an architecture that alternates full attention and GatedDeltaNet layers. Baselines include FastKV, SlimInfer, DASH, GemFilter, FlexPrefill, and MInference.
Against full-context prefill, the reported aggregate is $1.4$–$11.3\times$ lower TTFT and $1.5$–$3.4\times$ lower peak prefill memory with no accuracy degradation. Against the strongest prior method at matched accuracy, the headline is up to $6.4\times$ faster and about $2\times$ less memory. A concrete 2WikiMultihopQA point is easier to interpret: at $95\%$ of full-context accuracy, TokenGate prefills $3.0\times$ faster and uses $2.7\times$ less peak memory.
Figure 3 compares FastKV at its default layer and at layer 1, matching TokenGate's gate depth. This matters for attribution. If TokenGate won only because it pruned at a more favorable depth, end-to-end training would not be the explanation. The authors report that TokenGate still outperforms FastKV at the same layer, which supports their claim that a learned selector plus backbone adaptation adds value beyond gate placement.
Figure 5 separates the two mechanisms. With the same TokenGate backbone, learned top-$K$ beats random top-$K$, indicating that the score network carries task-relevant signal. With the same random selection, the TokenGate backbone beats the LoRA-only backbone, indicating that training changes representations or redistributes information so the model tolerates missing positions. This second result is easy to overlook: the method is not just a scoring head attached to a frozen model.
Figure 2 shows that pruning policy changes with the task. On 2WikiMultihopQA, entities, years, and answer phrases receive high scores. On MultiNews, the gate often retains function words and scattered event words. This diversity cuts against a universal saliency rule and helps explain why an attention heuristic tuned on one workload may not transfer cleanly to another.
4.1 Hybrid attention
Qwen3.5-2B uses full attention at layers 3, 7, 11, and so on, with GatedDeltaNet layers in between. FastKV relies on attention scores and therefore cannot natively prune the non-attention layers. TokenGate scores residuals and applies a bias through the attention kernel, so it can place the gate at either attention or recurrent layers. On NIAH-32k, a gate at layer 3 matches standard prefill with $3.8\times$ lower TTFT and is $2.1\times$ faster than the best FastKV run at the same accuracy. Layers 1 and 2 reach as much as $7.8\times$ faster prefill while staying within $6\%$ of full-context accuracy.
The hybrid result should not be reduced to “shallower is better.” Earlier gates save more work but have less contextual evidence for scoring. Layer 3 is the matched-accuracy operating point; layers 1 and 2 expose a faster but slightly less accurate frontier. That is a deployment control, not a contradiction. A scheduler could choose among them according to latency budget and task criticality, provided the application validates the worst case.
4.2 Length extrapolation and training cost
The model is trained on NIAH-16k and evaluated without retraining at 32k, 64k, and 128k. Table 1 reproduces the reported result. TTFT speedup rises from $11.3\times$ to $13.0\times$, while peak-memory reduction stays near $3.43\times$. The stable memory ratio is consistent with the architecture: the lower layers still execute a dense pass over the whole input, so their footprint bounds the peak even as post-gate work remains small.
| Test length | TTFT speedup | Peak-prefill memory reduction |
|---|---|---|
| 16k | $11.3\times$ | $3.42\times$ |
| 32k | $12.8\times$ | $3.43\times$ |
| 64k | $13.0\times$ | $3.43\times$ |
| 128k | $13.0\times$ | $3.43\times$ |
Training overhead is modest in aggregate but not free per step. On one A100, the mean SFT run takes 2.2 hours and TokenGate takes 1.8 hours, a ratio of $0.8\times$. The mean step is slower, 0.58 seconds versus 0.45 seconds, or $1.3\times$. Four of five tasks converge in fewer steps; MultiNews is the exception and takes 1.4 times longer overall. Reporting both values avoids the misleading shortcut of calling a method “faster to train” based only on a total that depends on convergence behavior.
| Dataset | SFT total | TokenGate total | Total ratio | SFT / step | TokenGate / step | Step ratio |
|---|---|---|---|---|---|---|
| RepoBench | 1.7h | 1.1h | $0.6\times$ | 0.50s | 0.64s | $1.3\times$ |
| 2Wiki | 3.8h | 4.0h | $1.1\times$ | 0.34s | 0.43s | $1.3\times$ |
| MultiNews | 1.3h | 1.8h | $1.4\times$ | 0.21s | 0.27s | $1.3\times$ |
| BioDEX | 1.2h | 1.0h | $0.8\times$ | 0.23s | 0.30s | $1.3\times$ |
| NIAH | 2.8h | 1.0h | $0.4\times$ | 0.99s | 1.24s | $1.3\times$ |
| Mean | 2.2h | 1.8h | $0.8\times$ | 0.45s | 0.58s | $1.3\times$ |
4.3 Ablations and qualitative behavior
The left panel of Figure 7 isolates the training recipe on BioDEX. Removing the full-context loss has the largest effect because the resulting model no longer preserves its full-budget reference accuracy. Removing temperature scaling or head-wise scaling also pushes the frontier inward. The right panel compares gates at layers 0, 1, and 2. Layer 0 is Pareto-suboptimal, while layers 1 and 2 are competitive. The likely cause is signal quality: layer-0 residuals have not yet combined enough context for the scoring MLP to distinguish request-relevant content.
The qualitative analysis also records failures, not only successful examples. On NIAH, 21 of the 24 tokens in a needle sentence at about 78% depth enter the $K=64$ set, and the answer is exact. On MultiNews, the gate under-retains names and places, so the generated summary can describe the right event with the wrong entity. On BioDEX, one example recovers only one of two reactions. RepoBench-P shows a different policy: many retained tokens are punctuation or empty argument syntax, which the authors interpret as structural anchors rather than copied content. Table 3 summarizes these patterns.
| Dataset | What is retained | Observed failure |
|---|---|---|
| RepoBench-P | code syntax rather than identifiers; commas and empty argument lists are prominent; selections are scattered | wrong generated line at small budgets |
| 2WikiMultihopQA | entities and dates, with relatively few function words | small budgets confuse awards; multi-hop answers need larger budgets |
| MultiNews | function words and event words; names and places are under-retained | right events but wrong name, occupation, or place |
| BioDEX | diagnosis and outcome terms plus list connectives; numbers rarely survive | some adverse reactions are missed |
| NIAH-16k | the needle sentence as a contiguous span, with a few rare subwords elsewhere | none in the displayed example |
5. Limitations
The most important architectural limit is that memory savings are bounded by the dense lower prefix. Every token still passes through all layers before the gate, so peak memory does not shrink proportionally with the number of post-gate tokens. This explains the nearly constant $3.43\times$ ratio across 16k to 128k in Table 1. Moving the gate earlier can lower the prefix cost, but the layer-0 ablation shows why that is not free: the selector has less contextual evidence and produces a worse accuracy-latency frontier.
The authors document semantic failure modes as well. Names and locations can be lost in summarization, and a medical extraction example misses one of two reactions. For a robot, losing a unit, joint limit, object identity, or safety condition could be more consequential than a small average F1 drop. A practical deployment would need a task-specific non-prunable set and worst-case validation, not only a mean score averaged over samples.
The measurements use batch size 1 on one A100. They report TTFT and peak prefill memory, not multi-tenant throughput, end-to-end latency, energy, or robotics closed-loop behavior. The authors intentionally leave lower-layer KV pruning disabled to isolate prefill savings, so the total benefit in a production decoder depends on complementary cache work. Dynamic budgets across concurrent requests may also change utilization and kernel shapes in ways this benchmark does not capture.
There is a reproducibility caveat around the repository link. The organization printed in the paper differs from the public repository used for this source audit. The checked repository contains the gate implementation, attention integration, and benchmark harness, but independent reproduction of all figures also depends on data preparation, checkpoints, model access, and the exact hardware/kernel stack. Public source availability should not be conflated with full numerical replication.
6. Conclusion
TokenGate's main contribution is the coupling of selection and computation. A residual-based gate chooses a subset, while token-gated attention gives selected scores a differentiable role in upper layers. The model can therefore learn what to keep through the same next-token objective that adapts the backbone. A single training run supports multiple budgets, and the reported training total is comparable to conventional parameter-efficient fine-tuning.
For embodied foundation models, the next experiments should test whether visual and sensor sequences contain enough removable redundancy, whether safety-critical state can be made non-prunable at acceptable cost, and whether dynamic selection can meet a hard control deadline under continuous camera input. Those questions remain open. The paper supplies a concrete mechanism and an honest set of measurements for testing them, rather than evidence that token pruning already improves robot behavior.
Golden quote
TokenGate does not ask the model to see less of the world; it teaches the model where deeper computation should begin.
