T-Router: Learning Thalamic Routing for Reasoning with Parameter-Efficient Reinforcement Learning

Liuxian Ma; Jiale Dai; Jiaqi Li; Lu Mi

cs.LGSep 30, 2026

T-Router: Learning Thalamic Routing for Reasoning with Parameter-Efficient Reinforcement Learning

Authors: Liuxian Ma, Jiale Dai, Jiaqi Li, Lu Mi

Organizations: College of Artificial Intelligence, Tsinghua University · State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · Beijing Institute for General Artificial Intelligence

Abstract

Parameter-efficient reinforcement learning aims to improve reasoning with a compact trainable interface to a pretrained model. We introduce the Thalamic Router (T-Router), which concentrates adaptation on the reuse of completed computations. A compressed, addressable bank preserves block changes; a depth-recurrent controller conditions their selection and relative-scale writeback. This coupling gives thalamic context-dependent routing a concrete computational form: learn which earlier contributions a receiving layer uses, and with what influence. Correctness rewards train the interface while preserving backbone parameters and layer order. On an 8.95B-parameter backbone, T-Router allocates 41.73M parameters (0.466% of the backbone) and achieves 83.64 +/- 1.16 MathAvg after GSM8K RL, compared with 73.79 +/- 1.83 for full-parameter GRPO across three evaluation rounds. At a comparable parameter budget and with matched retries, it exceeds LoRA's 77.28 +/- 1.95 MathAvg, improving all three task families and raising mean AIME accuracy from 48.33 to 60.56. Capacity-controlled comparisons favor addressable block changes and recurrent context; separate search training extends the interface to tool-mediated reasoning. These results establish controlled computation reuse as an effective route to parameter-efficient reasoning reinforcement learning.

Figures & tables

Figure 1 : From adapting transformations to learning their communication. Weight adaptation changes transformations (top). Thalamic circuits regulate cortical communication according to task context (middle). T-Router applies this routing principle across depth: select addressable block changes and regulate their influence before a later frozen layer executes (bottom).

Figure 2 : How T-Router reuses completed computation. Completed blocks supply compressed changes to the source bank. The depth controller updates its state

S

, then reads routing context

P

using the current hidden state. This context conditions which sources are mixed and how strongly the projected mixture is written back. The contribution is added before the receiving frozen layer; all layers execute in order. Teal and gold paths carry reusable content; dashed purple branches condition selection and strength. Equations 5 – 7 specify the operations.

Method	Params (M)	GSM8K	MATH-500	AIME mean	MathAvg
Frozen base	—	$88.73\,{\scriptstyle\pm 0.95}$	$75.60\,{\scriptstyle\pm 2.25}$	$31.67\,{\scriptstyle\pm 3.33}$	$65.33\,{\scriptstyle\pm 0.89}$
Full-parameter GRPO	8954.000	$93.86\,{\scriptstyle\pm 0.26}$	$83.07\,{\scriptstyle\pm 0.70}$	$44.44\,{\scriptstyle\pm 5.36}$	$73.79\,{\scriptstyle\pm 1.83}$
RFT	8954.000	$91.86\,{\scriptstyle\pm 1.16}$	$82.93\,{\scriptstyle\pm 2.47}$	$46.11\,{\scriptstyle\pm 2.55}$	$73.64\,{\scriptstyle\pm 1.33}$
LoRA-r2 + GRPO	5.410	$93.51\,{\scriptstyle\pm 0.23}$	$84.53\,{\scriptstyle\pm 0.81}$	$42.22\,{\scriptstyle\pm 12.73}$	$73.42\,{\scriptstyle\pm 4.28}$
LoRA-r16 + GRPO	43.278	$93.81\,{\scriptstyle\pm 0.77}$	$85.13\,{\scriptstyle\pm 2.53}$	$45.56\,{\scriptstyle\pm 6.31}$	$74.83\,{\scriptstyle\pm 1.41}$
LoRA-MoE + RO-GRPO	173.112	$95.30\,{\scriptstyle\pm 0.61}$	$87.47\,{\scriptstyle\pm 1.45}$	$44.44\,{\scriptstyle\pm 2.55}$	$75.74\,{\scriptstyle\pm 1.47}$

Table 1 : Mathematical reasoning after GSM8K RL. Accuracy (%), mean

\pm

SD across three evaluation seeds. AIME averages the two annual scores. LoRA-r16 (retries) matches T-Router 's retry rule. Parameters are in millions; a dash denotes no update. Bold/underlined means are best/second-best, including ties; fewer trainable parameters are preferred.

Method	BrowseComp F1	$\Delta$ vs GRPO	ASearch	$\Delta$ vs GRPO
Frozen base	$5.11\,{\scriptstyle\pm 0.30}$	$-28.12\,{\scriptstyle\pm 1.25}$	$59.76\,{\scriptstyle\pm 0.41}$	$-10.19\,{\scriptstyle\pm 0.56}$
Full-parameter GRPO	$33.23\,{\scriptstyle\pm 1.54}$	—	$69.94\,{\scriptstyle\pm 0.53}$	—
RFT	$31.20\,{\scriptstyle\pm 0.75}$	$-2.03\,{\scriptstyle\pm 2.07}$	$65.81\,{\scriptstyle\pm 0.92}$	$-4.14\,{\scriptstyle\pm 1.18}$
LoRA-r2 + GRPO	$32.98\,{\scriptstyle\pm 0.56}$	$-0.25\,{\scriptstyle\pm 1.25}$	$69.19\,{\scriptstyle\pm 0.59}$	$-0.75\,{\scriptstyle\pm 1.01}$
LoRA-r16 + GRPO	$36.06\,{\scriptstyle\pm 0.56}$	$+2.83\,{\scriptstyle\pm 1.80}$	$70.59\,{\scriptstyle\pm 0.18}$	$+0.64\,{\scriptstyle\pm 0.64}$
LoRA-MoE + RO-GRPO	$\mathbf{37.00}\,{\scriptstyle\pm 1.17}$	$\mathbf{+3.77}\,{\scriptstyle\pm 2.56}$	$71.46\,{\scriptstyle\pm 0.55}$	$+1.52\,{\scriptstyle\pm 0.65}$

Table 2 : Agentic-search comparison. Mean

\pm

SD across three evaluation seeds. BrowseComp Plus uses answer token-F1

\times 100

; ASearch Test uses a 0–100 score.

\Delta

is the paired difference from full-parameter GRPO; dashes mark the reference. Rankings follow Table 1 .

Figure 3 : Reasoning quality and parameter efficiency. (a) MathAvg mean

\pm

SD. (b) Task-family gains over matched-retry LoRA, computed from displayed means. (c) Paired gains over full-parameter GRPO in each search metric, mean

\pm

SD. (d) Trainable allocations on a logarithmic scale; LoRA-r16 variants share the same budget. Error bars summarize three evaluation rounds.

Configuration	GSM8K	MATH-500	MathAvg	$\Delta$
Full T-Router	$\mathbf{97.62}\,{\scriptstyle\pm 0.79}$	$\mathbf{92.73}\,{\scriptstyle\pm 0.70}$	$\mathbf{83.64}\,{\scriptstyle\pm 1.16}$	—
Without cross-layer retrieval	$96.26\,{\scriptstyle\pm 0.59}$	$79.40\,{\scriptstyle\pm 1.56}$	$69.48\,{\scriptstyle\pm 2.36}$	$-14.16\,{\scriptstyle\pm 1.78}$
Hidden-state memory	$\underline{96.79}\,{\scriptstyle\pm 0.53}$	$77.80\,{\scriptstyle\pm 1.31}$	$67.46\,{\scriptstyle\pm 2.19}$	$-16.18\,{\scriptstyle\pm 2.96}$
Learned static routing	$96.01\,{\scriptstyle\pm 0.68}$	$77.27\,{\scriptstyle\pm 0.81}$	$64.98\,{\scriptstyle\pm 1.71}$	$-18.66\,{\scriptstyle\pm 2.78}$
State replaced by matched MLP	$95.02\,{\scriptstyle\pm 0.75}$	$\underline{83.93}\,{\scriptstyle\pm 1.92}$	$\underline{72.80}\,{\scriptstyle\pm 4.10}$	$\underline{-10.84}\,{\scriptstyle\pm 3.09}$
Without layer/block identity	$94.29\,{\scriptstyle\pm 0.42}$	$66.13\,{\scriptstyle\pm 1.80}$	$58.10\,{\scriptstyle\pm 1.06}$	$-25.53\,{\scriptstyle\pm 0.56}$

Table 3 : Components of the reuse interface. Accuracy (%), mean

\pm

SD across evaluation seeds.

\Delta

is the paired MathAvg difference from full T-Router ; a dash marks the reference. The MLP retains comparable parameter capacity; hidden-state memory keeps the full allocation. Appendix C.2 includes both annual AIME scores.

Figure 4 : Dissecting controlled reuse. (a) Paired MathAvg changes from full T-Router , mean

\pm

SD. (b) Family-wise changes in means; AIME averages both years. (c) Memory and controller replacements at matched capacity. (d) MathAvg versus recorded GPU-hours for all configurations. Numerical summaries in (c) and error bars in (d) show MathAvg SD across evaluation rounds. The no-retry setting also changes the generation budget.

Appendix figures & tables35 assets

Supplementary material from the paper’s appendix.

Appendix

Question	Material	Page
How is the module executed?	Tensor shapes, exact parameter counts, layer schedule, state lifetimes, training settings, prompts, scoring, and analytical storage	A
What properties follow from the design?	Causal source availability, transport subspaces, slot history, calibration derivatives, gradient flow, and finite group sampling	B
What do the configurations reveal?	Full-test comparisons, component results and costs, writeback and initialization comparisons, routing interactions, auxiliary diagnostics, evaluation variance, and metric sensitivity	C
How do the components fit together?	Thalamic functional correspondence, bank/controller roles, intervention coordinates, and adaptation interfaces	D

Appendix

Table 4 : Appendix guide. Each part develops a distinct aspect of the method and study.

Object	Tensor shape	Role
Residual $H_{\ell}$	$B\times T\times d$	Input to a receiving decoder layer
Block anchor $A$	$B\times T\times d$	Post-intervention input at the start of a block
Block displacement $D_{b}$	$B\times T\times d$	Completed block output minus its anchor
Compressed memory $M_{b}$	$B\times T\times r$	Source-specific record of that displacement
Visible bank	$B\times T\times J_{\ell}\times r$	Completed records available to a receiver
Controller slots $S_{\ell}$	$B\times T\times K\times p$	Depth history used to condition retrieval

Appendix

Table 5: Runtime tensors and their axes.

J_{\ell}=|\mathcal{J}_{\ell}|

is the number of completed sources visible before layer

\ell

. No source or slot axis is a token-history axis.

Component	Shape or multiplicity	Parameters
Block compressors $C_{b}$	$8\times(r\times d)$	8,388,608
Writeback projections $U_{\ell}$	$28\times(d\times r)$	29,360,128
Layer and source embeddings	$32e+8e$	1,280
Initial controller slots	$K\times p$	2,048
Routing query $W_{q}$	$r\times(d+p+e)$	1,122,304
Routing key and value	$2\times[r\times(r+e)]$	147,456

Appendix

Table 6: Canonical module dimensions and exact parameter counts. Counts include the allocated terminal source. The final row counts parameters with a path to the training loss.

Component	Initialization
Source compressors, source/layer embeddings, initial slots	Gaussian with standard deviation $10^{-3}$
Routing, controller-input and slot projections	Xavier initialization
Controller MLP output layer	Zero weights and biases
Receiving projections $U_{\ell}$	Orthogonal initialization, gain 1
Shared dynamic gate vector	Zero
Receiving-layer gate bias	$\operatorname{arctanh}(0.02/0.05)$

Appendix

Table 7: Initialization of the canonical learned-RMS module. Initialization assigns content transport and intervention magnitude distinct roles.

Receiving layers	Sources visible before layer	Appended after	New source	Read count
0–3	None	Layer 3	$m_{0}$	0
4–7	$m_{0}$	Layer 7	$m_{1}$	1
8–11	$m_{0},m_{1}$	Layer 11	$m_{2}$	2
12–15	$m_{0},m_{1},m_{2}$	Layer 15	$m_{3}$	3
16–19	$m_{0},\ldots,m_{3}$	Layer 19	$m_{4}$	4
20–23	$m_{0},\ldots,m_{4}$	Layer 23	$m_{5}$	5

Appendix

Table 8: Source availability over 32 decoder layers. Layer indices are zero based. Each row covers four receiver decisions; the source appended at the end is available from the next row onward.

Object	Lifetime	Update rule
Adapter parameters	Shared across examples and forward calls	Optimizer update during training
Initial slot array	Learned parameter	Copied at forward initialization
Controller state	Current forward call and token position	Recurrent update at each layer
Source bank	Current forward call and token position	Append after a completed block
Block anchor	Current block	Replace at block start; clear at block end
Backbone causal cache	Generation context	Managed by the backbone’s decoder

Appendix

Table 9: Objects with different lifetimes. The same word “memory” can refer to trainable state, temporary depth context, or the language model’s causal cache.

Setting	Value
Training data	GSM8K train: 7,473 questions
Evaluation data	GSM8K: 1,319; MATH-500: 500; AIME24/25: 30 each
Evaluation generation	One generated answer per question in each evaluation round
Evaluation aggregate	Equal weight for GSM8K, MATH-500, and mean AIME
T-Router allocation	41,730,332 parameters; 40,681,724 loss connected
Reference policy	Frozen backbone with the adapter disabled

Appendix

Table 10: Final mathematics comparison: training and evaluation scope. The full-test suite is shared by the methods in Table 1 .

Control	Value	Control	Value
Optimizer	AdamW	Weight decay	0.01
Peak learning rate	$10^{-4}$	Warmup	10% of scheduled updates
Schedule	Linear warmup/decay	Accumulation	2 prompt groups
Prompt batch size	1	Group size	4 completions
Additional retry limit	3 groups	Temperature	0.7
Top- $p$	0.95	Completion ceiling	512 tokens

Appendix

Table 11: Optimization controls exposed by the implementation. Values are the implementation’s configurable mathematics recipe. Sampling and regularization are specified separately from the benchmark aggregation rule.

Task	Questions	Scoring object
GSM8K	1,319	Normalized final answer
MATH-500	500	Mathematical equivalence of the final answer
AIME 2024	30	Mathematical equivalence of the final answer
AIME 2025	30	Mathematical equivalence of the final answer

Appendix

Table 12: Full-test mathematics evaluation. Each evaluation round covers the listed questions. AIME’s two annual scores are averaged before the three-family MathAvg is computed.

Task	Dataset interface
GSM8K	openai/gsm8k
MATH-500	HuggingFaceH4/MATH-500
AIME 2024	HuggingFaceH4/aime_2024
AIME 2025	yentinglin/aime_2025
MuSR	TAUR-Lab/MuSR
GPQA-D	fingertap/GPQA-Diamond

Appendix

Table 13: Dataset interfaces in the evaluation implementation. These identifiers specify the task content accessed by the corresponding loaders.

Setting	ASearcher training/validation	BrowseComp Plus evaluation
Retrieval collection	Local wiki-18.jsonl	BrowseComp-Plus corpus
Retriever	e5-base-v2	Qwen3-Embedding-8B
Summary model	Qwen3-1.7B	Qwen3-8B ; three services
Retrieved documents	Top 10 per call	Top 10 per call
Initial excerpt	2,000 characters/document	10,000 characters/document
Summary ceiling	256 model tokens	256 model tokens

Appendix

Table 14: Default search-tool settings. Character limits apply to the initial document excerpts; summary limits use model tokens.

$BT$	Source bank (MiB)	Slots (MiB)	Anchor (MiB)	Sum (MiB)
1	0.0068	0.0078	0.0156	0.0303
128	0.875	1.000	2.000	3.875
512	3.500	4.000	8.000	15.500
2048	14.000	16.000	32.000	62.000
8192	56.000	64.000	128.000	248.000

Appendix

Table 15: Analytical storage of source bank, current slots, and anchor. Values use fp32 and

2^{20}

bytes per MiB.

BT

is the number of positions processed concurrently in one forward call.

Figure 5 : Analytical allocation of adaptation capacity. Left: area-proportional component allocation in the canonical module. Right: Equation 17 over source ranks and block sizes, with

L=32

d=4096

K=8

p=c=256

, and

e=32

. Cells give parameters in millions; the outlined cell is canonical. The last column counts source–receiver pairs.

Figure 6 : Analytical geometry of RMS calibration. Left: in two dimensions with

\|h\|=1

, positive multiples on the same ray map to one point on the target sphere; the double arrow marks a tangent. The local basis separates zero radial response from tangential gain

\kappa/q

. Right: the active-branch Jacobian eigenvalues from Equation 34 . The gate subsequently sets signed amplitude.

Region	Returned vector	Output norm	Derivative in $w$
$q>\epsilon$	$(\kappa/q)w$	$\sqrt{d}\kappa$	Tangential projector $\times\kappa/q$
$q<\epsilon$	$w$	$\sqrt{d}q$	$I$
$q=\epsilon$	$w$	$\sqrt{d}\epsilon$	Boundary; see one-sided limits

Appendix

Table 16: Calibration regimes for fixed

\kappa

. Values and derivatives refer to the piecewise mathematical map before the final dtype cast.

Per-completion correctness $p$	Mixed $q$	Retained mixed	Expected attempts	Expected responses
$0.1$	$0.3438$	$0.8146$	$2.3694$	$9.4774$
$0.5$	$0.8750$	$0.9998$	$1.1426$	$4.5703$
$0.9$	$0.3438$	$0.8146$	$2.3694$	$9.4774$

Appendix

Table 17: Analytical illustration of bounded retry. Four independent Bernoulli-correctness completions per attempt, all preference-valid, with four attempts allowed. Values are computed from

q=1-p^{4}-(1-p)^{4}

Method	AIME24	AIME25	AIME mean	MathAvg	$\Delta$ vs GRPO
Frozen base	$30.00\,{\scriptstyle\pm 5.77}$	$33.33\,{\scriptstyle\pm 3.33}$	$31.67\,{\scriptstyle\pm 3.33}$	$65.33\,{\scriptstyle\pm 0.89}$	$-8.46\,{\scriptstyle\pm 0.95}$
Full-parameter GRPO	$44.44\,{\scriptstyle\pm 6.94}$	$44.44\,{\scriptstyle\pm 3.85}$	$44.44\,{\scriptstyle\pm 5.36}$	$73.79\,{\scriptstyle\pm 1.83}$	$0.00\,{\scriptstyle\pm 0.00}$
RFT	$\underline{51.11}\,{\scriptstyle\pm 10.18}$	$41.11\,{\scriptstyle\pm 13.47}$	$46.11\,{\scriptstyle\pm 2.55}$	$73.64\,{\scriptstyle\pm 1.33}$	$-0.15\,{\scriptstyle\pm 2.54}$
LoRA-r2 + GRPO	$43.33\,{\scriptstyle\pm 17.32}$	$41.11\,{\scriptstyle\pm 10.72}$	$42.22\,{\scriptstyle\pm 12.73}$	$73.42\,{\scriptstyle\pm 4.28}$	$-0.37\,{\scriptstyle\pm 2.67}$
LoRA-r16 + GRPO	$50.00\,{\scriptstyle\pm 5.77}$	$41.11\,{\scriptstyle\pm 6.94}$	$45.56\,{\scriptstyle\pm 6.31}$	$74.83\,{\scriptstyle\pm 1.41}$	$+1.04\,{\scriptstyle\pm 2.69}$
LoRA-MoE + RO-GRPO	$41.11\,{\scriptstyle\pm 5.09}$	$47.78\,{\scriptstyle\pm 5.09}$	$44.44\,{\scriptstyle\pm 2.55}$	$75.74\,{\scriptstyle\pm 1.47}$	$+1.95\,{\scriptstyle\pm 0.41}$

Appendix

Table 18 : Annual mathematics results and aggregate contrasts. Accuracy and MathAvg are percentages; SD and

\Delta

are percentage points.

\Delta

is the paired MathAvg difference from full-parameter GRPO. GSM8K, MATH-500, and parameter allocations appear in Table 1 .

Method	BrowseComp F1	$\Delta$ vs GRPO	ASearch	$\Delta$ vs GRPO
Frozen base	$5.11\,{\scriptstyle\pm 0.30}$	$-28.12\,{\scriptstyle\pm 1.25}$	$59.76\,{\scriptstyle\pm 0.41}$	$-10.19\,{\scriptstyle\pm 0.56}$
Full-parameter GRPO	$33.23\,{\scriptstyle\pm 1.54}$	$0.00\,{\scriptstyle\pm 0.00}$	$69.94\,{\scriptstyle\pm 0.53}$	$0.00\,{\scriptstyle\pm 0.00}$
RFT	$31.20\,{\scriptstyle\pm 0.75}$	$-2.03\,{\scriptstyle\pm 2.07}$	$65.81\,{\scriptstyle\pm 0.92}$	$-4.14\,{\scriptstyle\pm 1.18}$
LoRA-r2 + GRPO	$32.98\,{\scriptstyle\pm 0.56}$	$-0.25\,{\scriptstyle\pm 1.25}$	$69.19\,{\scriptstyle\pm 0.59}$	$-0.75\,{\scriptstyle\pm 1.01}$
LoRA-r16 + GRPO	$36.06\,{\scriptstyle\pm 0.56}$	$+2.83\,{\scriptstyle\pm 1.80}$	$70.59\,{\scriptstyle\pm 0.18}$	$+0.64\,{\scriptstyle\pm 0.64}$
LoRA-MoE + RO-GRPO	$\mathbf{37.00}\,{\scriptstyle\pm 1.17}$	$\mathbf{+3.77}\,{\scriptstyle\pm 2.56}$	$71.46\,{\scriptstyle\pm 0.55}$	$+1.52\,{\scriptstyle\pm 0.65}$

Appendix

Table 19: Complete agentic search comparison. BrowseComp Plus uses token-F1 multiplied by 100; ASearch Test denotes the ASearcher held-out validation score (0–100). Appendix A.7 details the scorer and default penalty factors.

\Delta

is the paired difference from full-parameter GRPO; all entries are mean

\pm

SD in the corresponding metric's points.

Configuration	GSM8K	MATH	AIME24	AIME25	MathAvg	$\Delta$
Full T-Router	$\mathbf{97.62}\,{\scriptstyle\pm 0.79}$	$\mathbf{92.73}\,{\scriptstyle\pm 0.70}$	$\mathbf{63.33}\,{\scriptstyle\pm 6.67}$	$\mathbf{57.78}\,{\scriptstyle\pm 1.92}$	$\mathbf{83.64}\,{\scriptstyle\pm 1.16}$	$\mathbf{0.00}\,{\scriptstyle\pm 0.00}$
No cross-layer retrieval	$96.26\,{\scriptstyle\pm 0.59}$	$79.40\,{\scriptstyle\pm 1.56}$	$14.44\,{\scriptstyle\pm 3.85}$	$\underline{51.11}\,{\scriptstyle\pm 15.03}$	$69.48\,{\scriptstyle\pm 2.36}$	$-14.16\,{\scriptstyle\pm 1.78}$
Hidden-state memory	$\underline{96.79}\,{\scriptstyle\pm 0.53}$	$77.80\,{\scriptstyle\pm 1.31}$	$37.78\,{\scriptstyle\pm 10.72}$	$17.78\,{\scriptstyle\pm 5.09}$	$67.46\,{\scriptstyle\pm 2.19}$	$-16.18\,{\scriptstyle\pm 2.96}$
Learned static routing	$96.01\,{\scriptstyle\pm 0.68}$	$77.27\,{\scriptstyle\pm 0.81}$	$26.67\,{\scriptstyle\pm 5.77}$	$16.67\,{\scriptstyle\pm 6.67}$	$64.98\,{\scriptstyle\pm 1.71}$	$-18.66\,{\scriptstyle\pm 2.78}$
No state (matched MLP)	$95.02\,{\scriptstyle\pm 0.75}$	$\underline{83.93}\,{\scriptstyle\pm 1.92}$	$\underline{50.00}\,{\scriptstyle\pm 17.64}$	$28.89\,{\scriptstyle\pm 10.18}$	$\underline{72.80}\,{\scriptstyle\pm 4.10}$	$\underline{-10.84}\,{\scriptstyle\pm 3.09}$
No layer/block identity	$94.29\,{\scriptstyle\pm 0.42}$	$66.13\,{\scriptstyle\pm 1.80}$	$11.11\,{\scriptstyle\pm 10.72}$	$16.67\,{\scriptstyle\pm 3.33}$	$58.10\,{\scriptstyle\pm 1.06}$	$-25.53\,{\scriptstyle\pm 0.56}$

Appendix

Table 20 : Complete task-level component ablations. Entries are mean

\pm

SD across evaluation seeds.

\Delta

is each variant minus full T-Router in MathAvg percentage points.

Configuration	Parameters	Steps (est.)	GPU-hours
Full T-Router	41,730,332	3,737	121.23
No cross-layer retrieval	41,648,924	3,737	110.97
Hidden-state memory	41,730,332	3,737	122.42
Learned static routing	40,535,324	3,737	114.90
No state (matched MLP)	41,731,978	3,737	69.25
No layer/block identity	41,729,052	3,737	77.34

Appendix

Table 21: Resources for the component study. Parameter counts are exact allocations; steps retain their estimated status. GPU-hours retain their original training-cost records, separate from the evaluation-seed statistics.

Configuration	GSM8K	MATH-500	AIME24	AIME25	MathAvg
Standard gate †	$87.57\,{\scriptstyle\pm 0.87}$	$84.87\,{\scriptstyle\pm 3.29}$	$36.67\,{\scriptstyle\pm 0.00}$	$27.78\,{\scriptstyle\pm 10.72}$	$68.22\,{\scriptstyle\pm 2.83}$
Standard gate, no KL †	$89.03\,{\scriptstyle\pm 0.58}$	$71.67\,{\scriptstyle\pm 4.02}$	$13.33\,{\scriptstyle\pm 5.77}$	$4.44\,{\scriptstyle\pm 1.92}$	$56.53\,{\scriptstyle\pm 1.72}$
Local16 †	$90.75\,{\scriptstyle\pm 1.12}$	$82.00\,{\scriptstyle\pm 1.00}$	$33.33\,{\scriptstyle\pm 3.33}$	$35.56\,{\scriptstyle\pm 18.36}$	$69.07\,{\scriptstyle\pm 2.74}$
Local64 †	$89.54\,{\scriptstyle\pm 0.57}$	$71.07\,{\scriptstyle\pm 2.57}$	$5.56\,{\scriptstyle\pm 5.09}$	$11.11\,{\scriptstyle\pm 5.09}$	$56.31\,{\scriptstyle\pm 2.48}$
Wide512 †	$87.57\,{\scriptstyle\pm 1.18}$	$84.60\,{\scriptstyle\pm 1.93}$	$24.44\,{\scriptstyle\pm 10.72}$	$27.78\,{\scriptstyle\pm 5.09}$	$66.09\,{\scriptstyle\pm 2.85}$
Zero initialization †	$89.03\,{\scriptstyle\pm 0.35}$	$83.87\,{\scriptstyle\pm 1.81}$	$51.11\,{\scriptstyle\pm 10.72}$	$30.00\,{\scriptstyle\pm 6.67}$	$71.15\,{\scriptstyle\pm 2.29}$

Appendix

Table 22 : Complete historical task profiles and final reference. Generative entries are mean

\pm

SD;

\dagger

retains the historical IID-estimate attribute. The initialization and local/controller variants use the standard-gate branch. Full T-Router is the final configuration from Table 1 .

Figure 7 : Annual competition-mathematics profiles. The two annual means locate each configuration; whiskers show supplied SDs, and numerical summaries give annual means and their gaps. Pooled means average AIME24 and AIME25, and annual gaps are differences of those means. Each annual set contains 30 questions per evaluation seed.

Figure 8 : Where aggregate mean differences occur. Each family contributes one-third of its signed difference from the standard gate. Navy/dark red indicate positive/negative contributions; the net column sums those mean contrasts. These are analytical decompositions of the same reported evaluations.

Configuration	GSM8K	MATH	AIME mean	MathAvg	Final WB	Tail
Standard gate †	$87.57\,{\scriptstyle\pm 0.87}$	$84.87\,{\scriptstyle\pm 3.29}$	$32.22\,{\scriptstyle\pm 5.36}$	$68.22\,{\scriptstyle\pm 2.83}$	0.1254	—
Fixed RMS 0.5% †	$87.47\,{\scriptstyle\pm 0.99}$	$76.20\,{\scriptstyle\pm 1.59}$	$35.56\,{\scriptstyle\pm 3.47}$	$66.41\,{\scriptstyle\pm 1.69}$	0.5005	—
Fixed RMS 1% †	$89.31\,{\scriptstyle\pm 0.67}$	$79.67\,{\scriptstyle\pm 3.11}$	$40.00\,{\scriptstyle\pm 5.00}$	$69.66\,{\scriptstyle\pm 2.34}$	1.0010	—
Fixed RMS 2% †	$96.99\,{\scriptstyle\pm 0.56}$	$\underline{87.60}\,{\scriptstyle\pm 2.11}$	$47.22\,{\scriptstyle\pm 2.55}$	$77.27\,{\scriptstyle\pm 1.16}$	2.0020	—
Learned RMS, init 0.5% †	$96.99\,{\scriptstyle\pm 0.29}$	$84.67\,{\scriptstyle\pm 0.99}$	$\underline{55.00}\,{\scriptstyle\pm 8.33}$	$78.89\,{\scriptstyle\pm 2.39}$	3.4987	3.7596
Learned RMS, init 1% †	$\mathbf{97.83}\,{\scriptstyle\pm 0.46}$	$84.07\,{\scriptstyle\pm 1.67}$	$\underline{55.00}\,{\scriptstyle\pm 12.02}$	$\underline{78.96}\,{\scriptstyle\pm 4.40}$	3.6535	3.8762

Appendix

Table 23 : Complete writeback exploration and final-model reference. Scores are mean

\pm

SD;

\dagger

retains historical estimates. Final WB and tail median are original percentage-valued endpoint diagnostics. A dash denotes an unavailable trailing median.

Figure 9 : Writeback scale, task response, and structure. (a) MathAvg versus final relative writeback. (b) Mean family-score differences from the standard gate. F/L denote fixed/learned RMS, with coefficients or initializations in percent. (c) Original final and trailing-median diagnostics. (d) Structure-dependent mean changes under fixed 2% RMS. Whiskers in (a,d) show evaluation-seed SD. Tables 23 and 24 give the complete statistics.

Figure 10 : Writeback endpoint summaries. Initial coefficients are paired with Final WB (filled circles) and trailing-median WB (open circles). The table gives the original scalar diagnostics, their differences, and the largest-to-smallest initial and final ratios. Final WB is the latest routed layer's whole-tensor writeback-to-residual norm ratio.

Structure	Original writeback	Fixed RMS 2%	Paired change
Dense, standard gate †	$\underline{68.22}\,{\scriptstyle\pm 2.83}$	$\underline{77.27}\,{\scriptstyle\pm 1.16}$	$\underline{+9.05}\,{\scriptstyle\pm 2.26}$
Wide512 †	$66.09\,{\scriptstyle\pm 2.85}$	$\mathbf{78.69}\,{\scriptstyle\pm 0.66}$	$\mathbf{+12.60}\,{\scriptstyle\pm 3.41}$
Top2 composite †	$\mathbf{77.95}\,{\scriptstyle\pm 2.12}$	$63.98\,{\scriptstyle\pm 1.62}$	$-13.96\,{\scriptstyle\pm 1.95}$
Local64 †	$56.31\,{\scriptstyle\pm 2.48}$	$62.44\,{\scriptstyle\pm 0.90}$	$+6.13\,{\scriptstyle\pm 1.62}$

Appendix

Table 24: Structure–writeback combinations. MathAvg is in percent; SD and paired changes are in percentage points;

\dagger

retains the historical estimate attribute. Change is fixed RMS minus original writeback, paired by evaluation-seed index.

Figure 11 : Structure-dependent response to fixed RMS writeback. Original and fixed-RMS scores in (a) show means and evaluation SDs. The dashed reference adds the dense row's 9.05-point change. The paired changes use supplied SDs; differences between those changes are descriptive contrasts of means.

Configuration	MuSR	GPQA-D	IFEval strict@50	WikiText-2 PPL
Frozen base	—	—	—	—
Full-parameter GRPO	63.00 ‡	40.00 ‡	—	—
Standard gate	63.00 ‡	41.00 ‡	$31.33\,{\scriptstyle\pm 3.06}$	7.0391 ‡
Fixed RMS 2%	61.00 ‡	39.00 ‡	$\underline{36.00}\,{\scriptstyle\pm 4.00}$	7.0890 ‡
Learned RMS, init 1%	61.00 ‡	40.00 ‡	$34.00\,{\scriptstyle\pm 7.21}$	7.0758 ‡
Full T-Router	61.00 ‡	38.00 ‡	$\mathbf{38.67}\,{\scriptstyle\pm 4.16}$	7.1185 ‡

Appendix

Table 25: Complete auxiliary task profiles. IFEval is mean

\pm

sample SD over evaluation seeds.

\ddagger

marks original non-generative point references for MuSR, GPQA-D, and perplexity; no evaluation-seed dispersion is assigned to those entries. Lower perplexity is preferred; dashes denote unavailable entries.

Figure 12 : Task-specific diagnostic scales. Each task retains its native units. IFEval means carry evaluation-seed SDs;

\ddagger

point references retain their original non-generative scoring status. The full-GRPO reference is available for the two multiple-choice tasks.

Figure 13 : Analytical sensitivity to benchmark-family weights. The simplex shows the largest weighted mean among 14 complete profiles. The one-dimensional slice gives equal weight to MATH-500 and AIME and varies GSM8K weight; the crossover is near

w_{\rm G}=0.9713

. Scores are analytical reaggregations of the supplied means using Equation 58 .

Figure 14 : Aggregation at full-test sizes. The same nine method mean profiles are summarized by equal-family MathAvg and equal-question reaggregation. The size/resolution panel reports per-evaluation answer increments and family weights. The alternative weighted means are calculated from reported mean scores.

Organizing principle	T-Router realization	Computational consequence
Distributed processing	Frozen decoder blocks	Every block produces its original nonlinear transformation on the received residual
Selective inter-area communication	Attention over completed block records	Later computations select from origin-indexed intermediate changes
Context-dependent coordination	Token-specific slots recurrent over depth	Retrieval uses an auxiliary history of preceding computational states
Modulation of influence	Signed RMS-calibrated writeback	Direction and relative magnitude are controlled separately
Coordinating system around existing processing	Full source–controller–writeback pathway	Adaptation capacity is concentrated in communication between frozen blocks

Appendix

Table 29: Functional correspondence and its computational realization. Each row specifies an operation in the design.

Property	Source bank	Depth controller
Index meaning	Completed block identity	Slot identity
Cardinality	Grows as blocks finish	Fixed at $K=8$
Content width	$r=256$	$p=256$ per slot
Write frequency	Once per completed block	Once per decoder layer
Write rule	Append compressed block change	Decay plus allocated shared proposal
Read function	Source attention produces $c_{\ell}$	Slot attention produces $P_{\ell}$

Appendix

Table 30: Source storage and controller state are complementary interfaces. Their different update rules organize the computation even when their widths coincide.

Coordinate	Governing objects	Quantity it controls
Source preference	Query, source keys, depth context	Relative attention across completed computations
Content transport	$C_{b}$ , source value projection, $U_{\ell}$	Source-to-receiver transformation within $\operatorname{col}(U_{\ell})$
Relative influence	Target-RMS calibration and signed gate	Norm relative to the receiving residual and sign along the direction

Appendix

Table 31: Coordinates of controlled reuse. Distinct parameter groups control the origin, direction, and relative strength of the intervention.

Explore similar work

PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning
May 29, 2026Daize Dong, Junlin Chen, Haolong Jia +9Large Language Model Reinforcement Learning Replay
♣Rutgers University · ♦AMD · ⋆MBZUAI
Random Recursive Models
Sep 30, 2026Jama Hussein Mohamud, Mirco RavanelliRecursive Models Deep Learning Architectures
Mila – Quebec AI Institute · Université de Montréal · Concordia University
EntroRouter: Learning Efficient Model Routing via Entropy Regulation
Jun 28, 2026Kaiyi Zhang, Xueliang Zhao, Zhuocheng Gong +2Large Language Model Routing Trust Region
Gaoling School of Artificial Intelligence, Renmin University of China · Ant International · The University of Hong Kong +1

May 29, 2026, cs.LGJ/K move · Enter open · S save

PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning

Daize Dong, Junlin Chen, Haolong Jia, Jiang Liu+8

♣Rutgers University · ♦AMD · ⋆MBZUAI

Large Language Model Reinforcement Learning Replay

Open paper AskPDF arXiv

Mixture of Experts (MoE) Large Language Models (LLMs) achieve strong performance at scale. However, reinforcement learning (RL) on MoE-based LLMs often suffers from training instability. A root cause is router drift, i.e., expert activations can change drastically across model updates and differ between disaggregated rollout and training phases, causing large rollout--training mismatch and unstable importance sampling weights in PPO-style RL algorithms. Routing replay mitigates this issue by freezing the replay route within each reasoning trajectory, but it ignores how the router evolves under off-policy updates and thus causes router staleness. To address this limitation, we propose Predictive Routing Replay (PR2), which augments each router with a lightweight evolution predictor that learns to anticipate short-horizon router evolution. During the rollout phase, we use the predictive routing distribution to apply top-

k

routing, enabling gradients to reach experts that are likely to become active after updates. During the training phase, we replay the resulting predicted route to retain consistency for stable importance estimation. Theoretical analysis and experiments support that PR2 reduces routing-induced mismatch, improves RL stability, and yields stronger performance across various reasoning benchmarks.

Abstract

Explore similar work

PR2: Predictive Routing Replay for MoE-Based LLM Reinforcement Learning

Random Recursive Models

EntroRouter: Learning Efficient Model Routing via Entropy Regulation