GainRouter: Adaptive Inference Gate

• GainRouter is an inference-time gate that dynamically switches between fast LC-FT single-pass inference and slow explicit chain-of-thought generation based on latent strategy priors.
• It leverages codebook-derived thinking tokens and an MLP-based router to balance generation cost and accuracy through uncertainty and length difference signals.
• Empirical evaluations on math and programming benchmarks demonstrate that GainRouter maintains near chain-of-thought accuracy while significantly reducing token usage.

GainRouter is the inference-time gate in "Fast Thinking for LLMs" that adaptively switches between fast, codebook-guided single-pass inference and slow, explicit Chain-of-Thought generation. It is tightly coupled with Latent Codebooks for Fast Thinking (LC-FT), which learn discrete strategy priors distilled into continuous thinking vectors. At run time, GainRouter uses signals derived from those vectors, together with the input’s representation, to decide whether to remain in the fast LC-FT path or escalate to slow explicit reasoning. Its stated role is to suppress overthinking and reduce unnecessary token generation while preserving, and often matching, the accuracy of explicit reasoning (Zheng et al., 28 Sep 2025).

1. Functional role within LC-FT

GainRouter is designed for the setting in which explicit CoT improves accuracy on complex tasks but incurs latency and large token counts, and on simpler instances may cause overthinking and errors. LC-FT addresses part of this problem by providing strategy-level guidance through continuous thinking tokens without emitting textual rationales, enabling fast single-pass inference. GainRouter complements LC-FT by invoking explicit slow reasoning only when it is likely to be beneficial, thereby balancing accuracy and cost (Zheng et al., 28 Sep 2025).

The interface between GainRouter and LC-FT is explicit. LC-FT produces $K$ continuous thinking-token vectors by attending to a learned codebook of $M$ strategy prototypes. These vectors are injected into hidden states at a chosen layer $L$ and condition subsequent decoding in one pass. GainRouter consumes the pooled question representation and the codebook-derived thinking vectors to generate routing signals based on uncertainty, alignment, and predicted length difference. A positive routing decision triggers the slow explicit reasoning model, identified as Qwen3-4B-Thinking; otherwise the system stays in the fast LC-FT mode.

LC-FT’s codebook formalism provides the substrate on which GainRouter operates. The thinking vectors arise from learnable queries $\mathbf{Q}$ attending over a codebook $\mathbf{C}$ of $M$ prototypes in hidden size $H$:

$$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$

These tokens are injected at layer $L$ as

$$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$

and optionally refined by

$M$0

GainRouter’s signals directly inspect how these strategy tokens align with the question and how decisive the codebook composition is. This suggests that the router is not merely a confidence heuristic over output logits; it is a strategy-aware gate over latent reasoning priors.

2. Formal decision rule and routing signals

Given an input $M$1, the LC-FT non-thinking pass yields a pooled question vector $M$2 and $M$3 thinking-token vectors $M$4. After linear projection to a router space of dimension $M$5, these become $M$6 and $M$7. GainRouter aggregates token evidence by attention:

$M$8

$M$9

and forms the concatenated feature

$L$0

A small MLP maps this feature to a scalar routing score,

$L$1

Two auxiliary features are then computed. The first is cosine alignment between the projected question and the aggregated thinking evidence,

$L$2

which measures alignment between the question and strategy evidence. The second is normalized attention entropy,

$L$3

which captures uncertainty or dispersion over the codebook-derived strategies.

The router also predicts the normalized generation-length difference between slow and fast modes. The target is defined as

$L$4

where $L$5 is a fixed length scale, such as the mean non-thinking length on the router’s training set. A regression head predicts $L$6 from $L$7.

The final routing rule uses an adaptive threshold:

$L$8

$L$9

Here $\mathbf{Q}$0, $\mathbf{Q}$1, $\mathbf{Q}$2, and $\mathbf{Q}$3 are learned scalars. $\mathbf{Q}$4 means “escalate to slow explicit reasoning”; otherwise, the model stays in the fast LC-FT mode (Zheng et al., 28 Sep 2025).

The paper explicitly remarks that it does not define an explicit gain function $\mathbf{Q}$5 as an expected utility. Instead, the “gain” is implicit in supervised routing: the classifier learns when slow CoT helps accuracy, and the decision boundary is modulated by uncertainty and estimated extra cost. A plausible implication is that the name denotes operational benefit rather than a closed-form utility objective.

3. Architecture, overhead, and optimization

Architecturally, GainRouter is a lightweight module composed of linear projections from $\mathbf{Q}$6 to $\mathbf{Q}$7 for the question vector and thinking-token vectors, an attention aggregator that forms $\mathbf{Q}$8, a small MLP gate producing $\mathbf{Q}$9, a regression head predicting $\mathbf{C}$0, and the scalar parameters $\mathbf{C}$1, $\mathbf{C}$2, $\mathbf{C}$3, and $\mathbf{C}$4. The scoring function is stated as

$\mathbf{C}$5

The decision margin is $\mathbf{C}$6, and this margin is passed to $\mathbf{C}$7 during training.

The reported overhead is negligible relative to the backbone. Projections, attention over $\mathbf{C}$8 tokens, and two tiny heads add minimal compute, and router features are computed in the same forward pass that produces thinking tokens, so no extra passes are required (Zheng et al., 28 Sep 2025).

The training labels are constructed by paired fast and slow runs. For each training input $\mathbf{C}$9, the fast LC-FT system and the slow CoT system are both executed. The binary target is

• $M$0 if only the slow model succeeds.
• $M$1 if the fast model succeeds, either alone or jointly.
• Both-fail cases are ignored.

This supervision directly encodes the observed advantage of slow versus fast reasoning.

The loss comprises three terms. The routing loss is class-balanced binary cross-entropy with logits:

$M$2

The length-regression auxiliary is

$M$3

and the regularizer is

$M$4

where $M$5 collects $M$6 and $M$7. The total objective is

$M$8

The router is trained for 150 epochs with learning rate $M$9 and batch size $H$0. The threshold parameters are learned end to end to produce an adaptive budget-aware gate that escalates more often when the fast path is likely to fail and when the extra generation length is warranted.

4. Inference procedure and execution semantics

Inference is described as a single cohesive runtime procedure. First, LC-FT computes the $H$1 thinking-token vectors via codebook attention, injects them at layer $H$2, and produces the pooled question vector $H$3 together with the token vectors $H$4. These are projected into the router space to obtain $H$5 and $H$6.

Second, GainRouter computes the attention weights $H$7, the aggregated thinking evidence $H$8, the feature vector $H$9, the raw routing logit, the alignment score $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$0, the normalized entropy $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$1, the predicted normalized length difference $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$2, and the adaptive threshold adjustment $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$3. The binary decision $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$4 is then applied.

Third, the system executes the selected reasoning mode. If $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$5, the model performs fast single-pass inference with LC-FT and no explicit textual CoT. If $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$6, the model escalates to slow explicit CoT and generates step-by-step tokens. The final answer is returned from the chosen path (Zheng et al., 28 Sep 2025).

Two operational characteristics are emphasized. First, the decision is made before generation; there is no described second-chance fallback after a fast failure. Second, because GainRouter uses the same hidden representations already produced for LC-FT, it introduces negligible runtime overhead. This pre-generation design distinguishes it from methods that perform halting checks or search during token generation.

5. Empirical behavior on reasoning and programming benchmarks

The reported empirical result is that GainRouter achieves near-CoT accuracy with substantially fewer tokens across mathematical reasoning and programming benchmarks. On the mathematical side, AIME and OlympiadBench are used; on the programming side, MBPP and HumanEval are used. The measurements compare the fast LC-FT path, the slow Qwen3-Thinking baseline, and the routed system (Zheng et al., 28 Sep 2025).

Benchmark	Slow CoT	Ours + Router
AIME	73.33%, 19,411 tokens	73.33%, 12,797 tokens
OlympiadBench	50.0%, 7,075 tokens	50.51%, 5,332 tokens
MBPP	80.39%, 3,573 tokens	78.20%, 1,348 tokens
HumanEval	91.93%, 4,061 tokens	92.55%, 1,401 tokens

The detailed comparisons to the fast LC-FT mode clarify the role of routing. On AIME, LC-FT fast reaches 55.0% accuracy with 7,690 average tokens, whereas the routed system matches the slow model’s 73.33% accuracy while reducing average tokens to 12,797, approximately 34% fewer than slow CoT. On OlympiadBench, LC-FT fast records 25.0% accuracy at 2,668 tokens; the routed system reaches 50.51% at 5,332 tokens, approximately 25% fewer than the slow model’s 7,075 tokens.

On programming tasks, the token savings are larger. For MBPP, LC-FT fast achieves 64.60% accuracy at 89 tokens, the slow model achieves 80.39% at 3,573 tokens, and the routed system attains 78.20% at 1,348 tokens, approximately 62% fewer than the slow model. For HumanEval, LC-FT fast achieves 87.58% at 98 tokens, the slow model 91.93% at 4,061 tokens, and the routed system 92.55% at 1,401 tokens, approximately 66% fewer than slow CoT and slightly higher accuracy than the slow baseline.

Routing analysis in Fig. 4b reports that, compared to a random router, GainRouter increases Fast→Correct and Slow→Correct counts while reducing Slow→Wrong. The stated interpretation is that selective escalation preserves accuracy while lowering token cost. This suggests that the router is effective not only at detecting difficult cases, but also at avoiding unnecessary slow reasoning on easier instances.

6. Comparisons, limitations, and deployment considerations

The paper positions GainRouter against several categories of alternatives. Relative to static prompting, LoRA with concise CoT, and SoftCoT, the fast LC-FT path already outperforms these non-thinking baselines in the reported accuracy–efficiency trade-off; GainRouter then recovers most or all of slow CoT’s accuracy with far fewer tokens. Relative to entropy-only routing and confidence-token style gates, GainRouter uses strategy-aware signals—particularly alignment $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$7 and codebook attention entropy $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$8 over LC-FT prototypes—and explicitly models length cost through $$\mathbf{A} = \mathrm{softmax}\!\Big( \frac{ \mathbf{Q}\mathbf{W}_q (\mathbf{C}\mathbf{W}_k)^\top }{ \sqrt{H} } \Big), \qquad \mathbf{T} = \mathbf{A} (\mathbf{C}\mathbf{W}_v) \mathbf{W}_o \in \mathbb{R}^{K \times H}.$$9. Relative to self-consistency, tree search, and dynamic halting, the stated distinction is that those approaches still force multi-sample reasoning or halting checks during generation, whereas GainRouter makes a pre-generation decision using compact features (Zheng et al., 28 Sep 2025).

Several limitations and trade-offs are described. Misrouting remains the central failure mode. False negatives, in which the system stays fast when slow reasoning is needed, reduce accuracy; false positives, in which the system escalates unnecessarily, increase cost. The adaptive threshold $L$0 is intended to mitigate this by learning to raise or lower the boundary based on uncertainty and predicted length. Sensitivity to $L$1, $L$2, $L$3, and $L$4 is also noted: poor calibration can lead to over-escalation or under-escalation. The paper addresses this by learning these scalars end to end and regularizing $L$5.

Domain shift is another reported concern. A router trained on OlympiadBench and MBPP generalizes well to AIME and HumanEval, but shifts may degrade calibration of $L$6 and $L$7. Monitoring routing rates and re-tuning $L$8 on a small held-out set is suggested as a mitigation. The absence of a second-chance fallback means that failure of the fast path is not automatically repaired post hoc; the routing decision is preemptive.

The deployment procedure follows four steps: train LC-FT in two stages, collect router labels from paired fast and slow runs while ignoring both-fail cases, train GainRouter with $L$9, $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$0, and $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$1 regularization on $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$2, and then deploy by computing router features from the LC-FT pass and executing the chosen path. Reported settings include codebook size $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$3 as a good accuracy–cost point; $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$4 task-dependent, with examples of 16 for code and 48 for math; insertion layer $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$5 around 32 as the best trade-off; and LoRA rank $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$6, $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$7, dropout 0.05 on layers $$\mathbf{Z}^{(L)} = [ z_1, \dots, z_n; \mathbf{t}_1, \dots, \mathbf{t}_K ],$$8. Experimental constraints and reproducibility notes specify a single NVIDIA A800 GPU, bf16 precision, greedy decoding up to 30,720 tokens, AIME 2024/2025 and OlympiadBench for math, MBPP train and HumanEval test for programming, Qwen3-4B-Instruct-2507 as the fast backbone, Qwen3-4B-Thinking-2507 as the slow CoT model, Qwen-Max as the teacher for concise rationale construction, and an LC-FT distillation to Qwen2-1.5B for cross-scale analysis.

In this formulation, GainRouter is best understood as a compact pre-generation controller for hybrid reasoning: it exploits LC-FT’s latent strategy priors, predicts whether explicit CoT is worth its token cost, and operationalizes an explicit accuracy–efficiency trade-off through a learned adaptive threshold.