---
title: Token-level Adaptive Routing (TARo)
url: https://www.emergentmind.com/topics/token-level-adaptive-routing-taro
type: topic
---

# Token-level Adaptive Routing (TARo)

Searching arXiv for recent papers on Token-level Adaptive Routing and closely related usages of the term.
arxiv_search("Token-level Adaptive Routing TARo")

Token-level Adaptive Routing (TARo) designates a class of conditional-computation mechanisms in which the routing decision is made for individual tokens or generation steps rather than for an entire sequence, a fixed expert budget, or a uniform network path. In the recent literature, TARo appears in several distinct but structurally related forms: token-level dynamic routing to a linear-time State Space Model expert or a quadratic-time Transformer expert in clinical question answering, token-dependent variation of the number of active experts in Mixture-of-Experts layers, adaptive interpolation between frozen base-model and reward-model logits during test-time alignment, token-wise selection among multiple attention mechanisms, per-token recursion-depth assignment, per-token collaboration between small and large language models, and per-token switching between autoregressive and diffusion drafters or between discrete and latent reasoning modes [2601.01260][2603.01697][2603.18411][2605.28384][2502.01976][2606.07710][2606.05859].

## 1. Scope and problem setting

Across these works, TARo addresses a common limitation of static computation allocation: different tokens within the same input may have different computational requirements, different preferred inductive biases, or different risk profiles. In MoE settings, the motivating restriction is fixed Top-\(K\) routing, where exactly \(K\) experts are activated per token; DynaMoE explicitly relaxes this assumption and allows the number of active experts per token to vary based on input complexity [2603.01697]. In collaborative inference, CITER routes “non-critical tokens” to an SLM for efficiency and “critical tokens” to an LLM for generalization quality [2502.01976]. In test-time alignment, TARo replaces a fixed mixing coefficient \(\alpha\) between a base model and a reward model with a learnable token-level router because the “optimal \(\alpha\) varies across tasks, domains, and even decoding steps” [2603.18411].

A closely related line of work uses token routing even when the term TARo is not the sole organizing label. MEMatte inserts a router immediately before each global-attention layer and sends informative tokens to global attention while routing other tokens to a Lightweight Token Refinement Module, thereby reducing the quadratic burden of global self-attention on high-resolution image matting [2412.10702]. AdaMoE permits a variable number of “true” experts per token by augmenting the expert set with “null experts,” which consume zero FLOPs and fill otherwise fixed top-\(k\) slots [2406.13233]. MoLoRA makes a routing decision for every token over domain-specific LoRA adapters, while FusionRoute selects an expert at each decoding step and also adds a complementary logit from the router’s base LLM [2603.15965][2601.05106].

| Representative setting | Routed object | Representative decision |
|---|---|---|
| Clinical hybrid MoE | EMamba vs ET5 | hard top-1 token routing [2601.01260] |
| Dynamic MoE | number of active experts | percentile-threshold token routing [2603.01697] |
| Test-time alignment | base vs reward guidance | adaptive \(\alpha_t\) per token [2603.18411] |
| Efficient vision transformers | global attention vs LTRM | token-wise branch split [2412.10702] |
| Collaborative decoding | SLM vs LLM | token-level model choice [2502.01976] |
| Speculative decoding | AR vs diffusion drafter | per-step paradigm selection [2606.07710] |

This suggests that TARo is better understood as a routing principle than as a single architecture: the routed entities may be experts, model heads, adapters, attention mechanisms, drafters, or reasoning modes, but the granularity of control is consistently token-wise.

## 2. Routing functions and selection rules

A defining feature of TARo is that the routing function is evaluated from token-local or step-local state. In MambaFormer, the router input for token \(x_i\) is the fused feature
\[
R_i = [x_i; l_i; d] \in \mathbb{R}^{(d+2)},
\]
where \(x_i\) is the contextual embedding, \(l_i\in[0,1]\) is normalized sequence length, and \(d\in\{0,1\}\) is a binary domain flag. A 2-layer MLP computes
\[
g(x_i) = W_2\cdot \mathrm{ReLU}([x_i; l_i; d]\cdot W_1 + b_1) + b_2,
\qquad
S_i = \mathrm{softmax}(g(x_i)),
\]
followed at inference by the hard routing decision
\[
m_i^* = \arg\max_{k\in\{EMamba, ET5\}} S_{i,k}.
\]
The selected expert then processes the token, and the paper states explicitly that this is “equivalent to a ‘top-1’ selection per token” with “No soft mixture or \(k>1\) selection” [2601.01260].

DynaMoE generalizes the selection rule by replacing fixed Top-\(K\) with a percentile threshold. For token representation \(x_t\in\mathbb{R}^d\), the gate computes \(z_t=W_g x_t\), \(g_t=\mathrm{softmax}(z_t)\), defines a threshold \(\theta_t=\mathrm{percentile}_{(\tau)}(g_t)\), and selects
\[
S_t \coloneqq \{i : g_{t,i} > \theta_t\}, \qquad K_t \coloneqq |S_t|,
\]
with \(1 \le K_t \le K_{\max}=\lceil(1-\tau)\cdot N\rceil\). The layer output is then a soft aggregation over the selected experts,
\[
y_t = \sum_{i\in S_t} \mathrm{softmax}(g_{t,i}/T)\cdot E_i(x_t),
\]
and a minimum-activation step guarantees \(K_t\ge 1\) [2603.01697].

In LLM test-time alignment, TARo does not route a token to one expert or another; instead it routes the *amount of guidance*. The router consumes either full-logit concatenations or top-\(k\) logits with index embeddings and produces
\[
\hat\alpha_t = \sigma\!\bigl(W_2\tanh(W_1 h_t+b_1)+b_2\bigr)\in(0,1).
\]
This adaptive mixing weight defines
\[
z_{\mathrm{guided}} = (1-\alpha_t)z_{\mathrm{base}} + \alpha_t z_{\mathrm{reward}},
\]
or equivalently
\[
p_{\mathrm{TARo}}(y_t\mid x,y_{<t}) \propto p_{\mathrm{base}}(y_t\mid \cdot)^{(1-\alpha_t)} \cdot \exp[\alpha_t\cdot R(y_t\mid \cdot)].
\]
The route is therefore continuous rather than categorical, but it remains token-level and inference-time [2603.18411].

Other instantiations preserve the same logic with different routed objects. Meta-Attention forms a token feature \( \mathrm{feat}[b,t]=[x_{\mathrm{norm}}[b,t]; s[b,t]; pos[t]] \in \mathbb{R}^{(D+2)} \), uses a 2-layer MLP to parameterize a Dirichlet posterior over routing weights \(\alpha\), and can route each token to full softmax attention, linear attention, or sliding-window local attention [2605.28384]. MoLoRA uses a two-layer MLP \(g_\theta(x_i)=W_2\,\mathrm{GELU}(W_1x_i+b_1)+b_2\in\mathbb{R}^K\), followed by TopK and temperature-scaled softmax over selected adapters [2603.15965]. Informed Routing replaces execute-or-skip with execute-or-approximate: the router predicts whether a transformer unit \(\mathcal{U}^{\ell,k}\) should run in full or whether the token should go through a Lightweight Feature Forecaster \(\mathcal{F}^{\ell,k}\) [2510.13831].

A plausible implication is that TARo should not be identified with a single routing algebra. The literature contains hard top-1 selection, variable-cardinality selection, weighted soft aggregation, binary action routing, and Bayesian posterior routing, all at token granularity.

## 3. Optimization objectives and training regimes

The optimization of TARo varies substantially across applications. In MambaFormer, the router parameters \(W_1,b_1,W_2,b_2\) are trained while the experts are frozen. The loss is
\[
L = L_{CE} + \lambda\cdot L_{bal} + \mu\cdot L_{pen},
\]
with
\[
L_{CE} = -\sum_{c=1}^{C} y_c\log p_c,
\]
\[
L_{bal} = (1/L)\sum_{i=1}^L \sum_{k\in\{EMamba,ET5\}} S_{i,k}\cdot \log(1/|E|),
\]
and
\[
L_{pen} = \sum_{i=1}^L \max(0, S_{i,ET5} - T_{ui}),
\]
using \(\lambda=0.5\), \(\mu=0.5\), and \(T_{ui}=0.08\). The paper states that this enforces that only \(\lesssim 8\%\) of tokens choose ET5 on average, yielding the 3.8% observed in practice [2601.01260].

CITER formulates token routing as policy optimization in a finite-horizon MDP with actions \(a_S\) and \(a_L\). Under the simplifying choices \(\beta=1\) and \(\mu(a_S)=\mu(a_L)\), the router is trained by minimizing token-wise cross-entropy on binary preference labels,
\[
L(\theta) = - \sum_{s\in P} \bigl[p(s)\log\pi_\theta(a_S\mid s) + (1-p(s))\log\pi_\theta(a_L\mid s)\bigr].
\]
The paper also introduces a three-case shortcut for reward estimation, and reports that “~80–90% of tokens fall into Cases 1–2,” reducing the cost of preference collection by \(\sim 10\times\) [2502.01976].

TARPO instead uses pure RL. A lightweight action head parameterized by \(\mathbf{W}_r\in\mathbb{R}^{2\times d}\), \(\mathbf{b}_r\in\mathbb{R}^2\) computes
\[
\rho_\theta(d_t\mid h_t)=\mathrm{Softmax}(\mathbf{W}_r h_t + \mathbf{b}_r),
\]
where \(d_t\in\{\texttt{hard},\texttt{soft}\}\). The LLM backbone and router are jointly optimized end-to-end with a shared group-relative advantage signal \(\hat A_i=(r_i-\mu_r)/(\sigma_r+\varepsilon)\), and the minibatch objective combines token loss, action loss, and KL regularization [2606.05859].

In Bayesian TARo, Meta-Attention places a compute-aware Dirichlet prior over routing weights and trains the amortized posterior \(q(\alpha\mid x_t;\phi)=\mathrm{Dir}(\alpha;\beta+\delta_\phi(x_t))\) using an ELBO objective. The KL term against the prior supplies a principled alternative to ad hoc balancing losses and also yields a posterior-entropy signal \(U_t\) for soft-to-hard routing [2605.28384]. By contrast, DynaMoE states that “No auxiliary balancing losses are required, but large-scale deployment may need capacity factors” [2603.01697].

In execute-or-approximate routing, training is two-stage. Informed Routing first freezes the base LLM and fits each Lightweight Feature Forecaster using
\[
\mathcal{L}_{\mathrm{fit}}^{\ell,k}
=
\mathbb{E}\Bigl[\bigl(1-\cos(\mathcal{F}^{\ell,k}(\mathbf{x}),\mathbf{z})\bigr)+\|\mathcal{F}^{\ell,k}(\mathbf{x})-\mathbf{z}\|_1\Bigr],
\]
then freezes the LFFs and trains routers with
\[
\min_{\mathcal{R}} \mathcal{L}_{\mathrm{LM}} + \lambda_1|\rho - S_{\mathrm{target}}|_1,
\]
using Gumbel-Softmax to sample hard masks [2510.13831]. MoLoRA, in turn, trains LoRA and router parameters jointly with a standard language-modeling loss plus a Switch-style load-balancing auxiliary term [2603.15965].

These differences matter conceptually. Some TARo systems train only the router, some train router and backbone jointly, some freeze the routed experts, and some incorporate explicit compute terms in the loss while others rely on priors, schedules, or threshold calibration.

## 4. Architectural realizations

The architectural role of TARo depends on what part of the model is being conditionally activated. In hybrid MoE clinical assistance, MambaFormer places TARo over two experts: EMamba, “a linear-time State Space Model (SSM) expert,” and ET5, “a quadratic-time Transformer (T5-Large) expert.” The router uses contextual embeddings, normalized sequence length, and a domain-aware flag to decide which tokens should incur Transformer cost [2601.01260].

In efficient vision transformers, MEMatte inserts a Router plus a routing decision immediately before each global-attention layer. Tokens are split into a Global-Attention branch and a Lightweight Token Refinement Module branch, and after each block the two branches’ outputs are “simply re-concatenated (in the original token order)” before the next router. Training uses Batch-constrained Adaptive Token Routing (BATR), with a batch-level average routed ratio
\[
\gamma = \frac{1}{B\cdot M\cdot N}\sum_{b=1}^{B}\sum_{m=1}^{M}\sum_{i=1}^{N}\delta_i^{b,m},
\]
and a compression loss \(L_{\mathrm{compress}}=(\rho-\gamma)^2\) [2412.10702].

In attention-mechanism routing, Meta-Attention maintains \(K=3\) experts: full softmax attention, linear attention, and sliding-window local attention, with normalized costs \(c_1=1.0\), \(c_2=0.15\), and \(c_3=0.30\). The Bayesian Meta-Controller can run all experts under soft routing or perform uncertainty-gated hard routing based on posterior entropy [2605.28384]. In recursive transformers, MoR’s token-choice router assigns each token an up-front recursion depth \(i_t=\arg\max_j g_t^j\), so that only tokens with \(i_t\ge r\) remain active at recursion \(r\); attention at step \(r\) is restricted to this active subset, and the framework supports both recursion-wise KV caching and recursive KV sharing [2507.10524].

In collaborative LLM systems, the routed object is usually the model invocation itself. CITER switches between SLM and LLM at each timestep while maintaining separate KV caches so that switching back does not require recomputing history [2502.01976]. The edge-device inference system based on the CITER router keeps the SLM on-device, deploys the token router on-device, and defers low-confidence tokens to a cloud LLM served under SGLang [2504.07878]. FusionRoute uses a small linear routing head on a base LLM hidden state to select among experts at token step \(t\), but it also adds the router LLM’s own logit vector to the selected expert’s logits,
\[
\log \pi_{\mathrm{final}}(y\mid x,y_{<t}) = z_t^{e_t^*}(y) + r_t(y),
\]
so the router acts simultaneously as selector and complementary generator [2601.05106].

TARo also appears in multimodal and diffusion settings. MoS routes hidden states from an understanding tower into a generation tower, producing denoising-timestep- and token-dependent interactions. For each generation block \(j\), it uses \(\epsilon\)-greedy Top-\(k\) selection over context states and constructs a fused state \(S_j^c=\sum_{i\in I_j}\bar w_{i,j}S_i^c\) before projection into the generation block [2511.12207]. WhiFlash routes each decoding step between an autoregressive drafter \(m_a\) and a diffusion drafter \(m_d\), using either an entropy threshold on the target model’s next-token distribution or a learned MLP regressor that predicts the difference in acceptance lengths [2606.07710].

A plausible implication is that TARo is orthogonal to the backbone family. The same routing granularity has been applied to transformers, SSM–Transformer hybrids, recursive transformers, multimodal diffusion models, and speculative decoding systems.

## 5. Empirical trade-offs and observed routing behavior

The central empirical claim of TARo papers is that token-wise conditional computation yields better speed–accuracy or cost–quality trade-offs than static allocation. On PubMedQA, MambaFormer reports BERTScore F1 \(=0.9180\) at latency \(=0.077\,\mathrm{s/seq}\), compared with BioBERT at \(0.9000 @ 1.900\,\mathrm{s}\), Mamba at \(0.8500 @ 0.100\,\mathrm{s}\), and hybrid static baselines at BERTScore \(0.8750\)–\(0.8900 @ 0.38\)–\(0.45\,\mathrm{s}\). The paper states that MambaFormer is “24.4× faster than T5-Large” and that the token assignment distribution is 96.2% to EMamba and 3.8% to ET5, with ET5 handling “short, complex queries” and EMamba handling “long contexts” [2601.01260].

DynaMoE reports that, for a Small model on MNIST, the Descending schedule reaches 92.68% accuracy, compared with 91.35 for Uniform and 89.42 for the MLP Baseline. Cross-dataset gains reported in the paper are 89.42\(\rightarrow\)92.68 on MNIST, 84.15\(\rightarrow\)88.34 on Fashion-MNIST, and 62.38\(\rightarrow\)67.85 on CIFAR-10; the paper also reports expert-usage entropy \(+0.34\) bits versus fixed Top-2 routing [2603.01697]. MEMatte reports approximately 88% memory reduction and about 50% latency reduction on Composition-1K while keeping SAD close to ViTMatte: for ViT-S, 6.20 GB and 186.0 ms for ViTMatte-S versus 0.71 GB and 84.99 ms for MEMatte-S; for ViT-B, 12.53 GB and 340.2 ms versus 1.49 GB and 178.9 ms [2412.10702].

In reasoning-time alignment, TARo improves MATH500 accuracy for Llama-3.1-8B plus a distilled reward model from 32.0% for the base model to 54.4%, compared with 49.2% for GenARM with fixed \(\alpha=0.5\). The prompt-level routing ablation reaches only 33.2% on MATH500, versus 49.6% for token-level routing, and the qualitative analysis reports that high-\(\alpha_t\) tokens are often “mathematical operators, ‘Step,’ ‘Compute,’ etc.” while low-\(\alpha_t\) tokens are ordinary context words [2603.18411]. TARPO reports improvements over GRPO across Qwen2.5 model sizes and on Llama-3.1-8B, with additional evidence that generated token count on OOD evaluation fell from \(\sim 412\) to \(\sim 338\) [2606.05859].

In inference systems, WhiFlash reports, for Qwen3-8B overall average, acceptance length 5.12 for EAGLE-3, 5.29 for DFlash, 6.81 for Oracle-Token, 6.08 for WhiFlash-Entropy, and 6.26 for WhiFlash-Neural. Throughput speedups over AR decoding are 3.55× for DFlash, 3.82× for WhiFlash-Entropy, and 3.87× for WhiFlash-Neural; category-specific peak gains on Qwen3-8B reach +69.6% TPS over AR for Math and +37.3% TPS over diffusion for Chat [2606.07710]. The edge-device routing system reports a “60% performance gain on CommonsenseQA” using a 0.5B model on an M1 MacBook, with “under 7% of tokens generation uploaded to the large model in the cloud” [2504.07878].

Adapter and specialization routing also shows strong empirical effects. MoLoRA reports that Qwen3-1.7B + MoLoRA surpasses Qwen3-8B on GSM8K, MATH, BBH, and GPQA while being 4.7× smaller, and that on a mixed-modality workload with \(K=4\) adapters, per-sequence routing takes 5.88 ms whereas per-token TARo takes 1.43 ms, a 4.1× speedup, with 5.5× end-to-end improvement under full CUDA graph capture [2603.15965]. FusionRoute reports average cross-domain accuracy 0.566 on Llama-3, compared with 0.536 for fine-tuning, 0.502 for Collab, and 0.466 for sequence-level selection, while ablations show that removing complementary logits loses approximately 4–5 points of average accuracy [2601.05106].

These results do not establish a single Pareto frontier for all TARo systems, but they do show that token-wise routing repeatedly outperforms fixed per-sequence or fixed per-layer choices when the routed resource is expensive and heterogeneous.

## 6. Misconceptions, limitations, and open directions

A common misconception is that TARo is synonymous with hard expert selection in MoE. The literature is broader. MambaFormer uses hard top-1 routing [2601.01260], but DynaMoE uses a token-dependent \(K_t\) with soft aggregation [2603.01697], TARo alignment uses a continuous \(\alpha_t\) for logit interpolation [2603.18411], Meta-Attention uses Bayesian routing weights and uncertainty-gated hardening [2605.28384], and Informed Routing uses a binary execute-or-approximate policy rather than execute-or-skip [2510.13831].

Another misconception is that token routing is always equivalent to permanent token pruning. MEMatte explicitly avoids permanent token loss: routed tokens go either to global attention or to LTRM, and the outputs are re-concatenated in original token order [2412.10702]. MoR similarly routes tokens to different recursion depths rather than deleting them from computation entirely [2507.10524]. AdaMoE’s “null experts” make the number of *true* experts per token adaptive, but the top-\(k\) routing interface itself remains intact [2406.13233].

The limitations reported in the papers are also heterogeneous. TARo for test-time alignment notes that current reward models are trained only on math stepwise preferences and suggests richer, domain-specific rewards for stronger OOD performance [2603.18411]. WhiFlash depends on cache-management optimizations such as Lazy Catch-up and KV-only Prefill to keep high-frequency switching overhead below 7% of per-round latency [2606.07710]. The edge-device collaboration system reports that current ONNX-based SLM runtimes do not support incremental KV-cache injection, forcing full re-prefill on each deferred token [2504.07878]. Meta-Attention identifies tuning burdens for the prior strength \(\beta_0\), floor \(\epsilon\), and hard-routing threshold \(\eta\), and reports a +6.3% relative PPL overhead for the Bayesian controller in its Phase 1 Tiny LM benchmark [2605.28384].

Informed Routing states an upper bound on sparsity for simple forecasters: beyond \(\sim 40\)–50%, the LFF cannot capture complex transformations well, and at 70% sparsity both PPL and reasoning accuracy degrade sharply [2510.13831]. FusionRoute argues more fundamentally that pure expert-only routing is limited: unless strong global coverage assumptions hold, token-level routing based solely on fixed expert outputs cannot in general realize the optimal decoding policy, which motivates the addition of a trainable complementary generator [2601.05106].

Open directions are stated explicitly across several papers. TARo alignment suggests extending routing to more than two experts and incorporating backtracking or global rollout checks [2603.18411]. WhiFlash proposes extension to any heterogeneous pool of draft experts [2606.07710]. MoS points to dual-way routing between modalities and adaptation to audio, video, and early-fusion architectures [2511.12207]. Informed Routing suggests richer forecasters, adaptive forecaster capacity, and multimodal or long-context extensions [2510.13831]. Taken together, these proposals indicate that TARo is evolving from a narrow MoE routing tactic into a general mechanism for token-wise allocation of heterogeneous computation, model specialization, and reasoning control.

Source: https://www.emergentmind.com/topics/token-level-adaptive-routing-taro