MACRO: Markov Chain Routing of Transformer Layers

The paper introduces MACRO (Markov Chain Routing of Transformer Layers), a framework that learns task-specific routes over frozen LLM architectures without modifying underlying weights. As stated in the abstract:

"Standard Large Language Models (LLMs) execute layers sequentially. Dynamic layer routing, i.e. search for a different execution path through layers involving layer repetitions, skips and other moves, can improve performance. Existing routing approaches often require updating model weights, running expensive search loops per test instance, or demand ground-truth labels during inference. In this work, we propose Markov Chain Routing of Transformer Layers (MACRO), a framework that learns task-specific routes over LLM architectures without modifying underlying parameters."

The key design choice is modeling routing as a Markov chain:

"MACRO models layer routing as a context-dependent Markov policy conditioned on layer indices, computation budget phases, directional displacements, and operator context, supporting skip, repeat, and residual hidden-state addition operations."

The main results reported in the abstract are:

"MACRO achieves a +5.0% average accuracy improvement over the unrouted baselines, with largest gains on small models. We outperform the best dynamic routing approach Dr. LLM by +7.2%, while reducing route-search time 9.4× (from 14.8 to 1.6 hours)."

The authors motivate the work by noting that standard LLMs operate on a rigid computational schedule, executing transformer blocks sequentially from the first layer to the last, ignoring the potential for dynamic computation pathways. Prior work on adaptive depth focused on faster inference via early exits or structured layer dropping, but actively restructuring the execution route of a frozen model to improve reasoning accuracy on complex tasks remains a highly challenging and relatively unexplored frontier.

A compelling example is given in Figure 2: on GSM8K, rewinding Qwen3-1.7B from layer 7 back to layer 3 and replaying five blocks raises test accuracy from 43.4% to 69.5%, without changing a single weight. The route runs L0–L7, rewinds to L3, and then continues to L27, so L3–L7 execute twice and 33 block calls are made in total. The weights are frozen and only the execution order changes.

Existing approaches have bottlenecks: "Dr.LLM bypasses the need for test-time labels by training lightweight routers, yet it necessitates an expensive offline Monte Carlo Tree Search (MCTS)-driven data generation and training pipeline. Furthermore, purely exhaustive approaches, such as Neuroanatomy, operate in a constrained search space and are computationally prohibitive for per-task adaptation."

The authors chose Markov chains for specific reasons:

"First, unlike deep reinforcement learning or auxiliary neural routers, a Markov transition matrix is lightweight and requires no parameter updates or backpropagation through the LLM. Second, it naturally mitigates overfitting on small validation datasets by structurally constraining the decision space. Most importantly, modeling the routing as a Markov chain allows us to leverage exact, polynomial-time sequence decoding (the Viterbi algorithm) instead of relying on stochastic sampling or computationally expensive tree searches at test time."

The paper lists three main contributions:

1. We propose MACRO, a parameter-efficient framework for adaptive layer routing in frozen and quantized LLMs, modeled as a context-dependent Markov policy.

2. "We introduce an exact top-k Viterbi decoding scheme equipped with structural feasibility masks, which eliminates the need for per-sample test-time search, auxiliary neural network training, or ground-truth inference oracles."

3. "We demonstrate empirically that MACRO consistently outperforms prior dynamic routing methods across diverse reasoning and knowledge benchmarks, yielding a +5.0% average accuracy improvement over the non-routed baseline, +7.2% over Dr. LLM, and +1.5% over Dr. LLM with an equivalent action space, while reducing the route-search computational cost by 9.4×."

The paper situates itself against prior work on adaptive depth (early exits, structured layer dropping, LayerSkip, Mixture-of-Depths), noting Our objective differs: we use layer routing to improve accuracy while keeping the pretrained model fixed. The closest works are Dr.LLM, CoLA (per-sample MCTS with ground-truth reward), and Neuroanatomy (exhaustive search in a limited space). A qualitative comparison in Table 1 contrasts the methods on dimensions including accuracy-improvement objective, no test labels, no per-example test-time search, no offline route-label generation, and frozen weights.

The paper also relates to recurrent and expert routing: Universal Transformers, looped transformers, and recurrent-depth models show that repeated computation over depth can support algorithmic or reasoning behavior. In contrast to Mixture-of-Experts models, our experts are the existing blocks of a single pretrained transformer.

Formally, for a pretrained transformer with frozen weights f1,...,fL, "A route r relaxes this schedule into a program over the blocks. At each step, it chooses which block to apply, so blocks may be skipped, repeated, or revisited by going back to earlier ones. The route may also add an earlier hidden state to the running hidden state before applying the next block." The objective is:

r⋆ = arg max Acc(r), with Acc(r) = (1/N) Σ 1 Mr(xi) = yi.

The dataset is split into Dtrain, Dval, Dtest. The training set is further partitioned into I disjoint subsets, and At search iteration i, candidates are evaluated on the current subset together with a small replay buffer of examples from previous rounds, with the buffer updated by adding a fraction λrep of the current subset each iteration.

The state factorizes as st = (l, ϕ, δ, o):

1. The current layer l ∈ 0,..., L − 1 is the most recently produced layer in the partial program.

2. The budget phase ϕ ∈ early, mid, late buckets the fraction of remaining route-length budget into three bins, biasing the policy toward Rejoin as the route nears its maximum length.

3. The incoming displacement δ = l−lprev ∈ −r,..., r, together with a dedicated start symbol, records how the current layer was reached.

4. The binary op-context o ∈ 0, 1 indicates whether the current hidden state ht was produced by a plain layer or by an add operation.

This gives a state space of size S = L · 3 · (2r+2) · 2.

Three action families are defined:

1. Local layer move: apply layer l + ∆, where ∆ ∈ −r,..., r and the destination layer is valid. This allows local skips, repeats, and backward revisits while biasing the search toward small perturbations of the standard forward pass.

2. Add-and-apply: "add a previous hidden state ht−b to the current hidden state ht, using ht + γht−b with γ > 0, and then apply a local destination layer l + ∆."

3. Rejoin: terminate the learned part of the route and execute the remaining standard suffix of transformer layers.

A feasibility mask removes invalid actions (outside layer bounds, exceeding budget, adding unavailable hidden states), and the policy is renormalized over the remaining feasible actions, so sampling and Viterbi decoding always operate over valid routes.

Crucially for efficiency: a revisited block fl keeps a separate key/value history for each application, so a route of T applied steps has the cache and decoding cost of a T-block network with tied weights.

The initial policy is a forward-biased prior, not uniform: We make the baseline route most probable because accuracy-improving programs are expected to be slight modifications of the standard forward pass rather than entirely random layer programs. The most likely route under πθ0 is the unmodified network, ⟨0, 1,..., L−1⟩. Unvisited states keep the initialization exactly; visited states move away from it only in proportion to the evidence collected for that state.

The policy is updated with an estimation-of-distribution procedure (which the appendix connects to the Cross-Entropy Method). At each iteration, K candidate programs are sampled from the current policy, evaluated on the training subset, and scores are converted into normalized update weights with an optional elite set. The update uses weighted maximum-likelihood estimation:

Ni(s, a) = α πθi(a s) + Σj wj Σt 1 sj,t = s, aj,t = a

and the new policy is the normalized weighted MLE. An optional exponential moving average smoothing is also available: This update increases the probability of high-scoring route programs while retaining probability mass on nearby alternatives through smoothing and feasibility-masked sampling.

After training, exploration halts. The authors decode the most likely programs under the final policy using a top-k Viterbi procedure with the recursion:

Vk(s) = TopKa∈A(s) [log πθ(a s) + Vk(T(s, a))]

The resulting candidate set is evaluated on the validation set, and the best route is selected: r̂ = arg max AccDval(r),... breaking ties by higher training accuracy and then by shorter route length. The test set is used only once, to report the final accuracy of the selected route r̂.

The appendix emphasizes the exactness of the decoding: The five final candidates are the exact k highest-probability programs under the learned policy, not a beam or sampling approximation — this is achieved by running the dynamic program over the augmented state, which is lossless and cheap because the incoming delta feature enlarges the state space by at most a factor of 2r+1.

The paper compares against: the standard Baseline (sequential execution), Dr.LLM (lightweight routers trained from MCTS-derived labels), Dr.LLM+Ext. (same Dr.LLM training but with the extended MACRO action space), MACRO (top-1) (rank-1 Viterbi candidate), MACRO (no Vit.) (highest training-accuracy route the search sampled), and full MACRO.

The main results table (Table 2) covers 13 benchmarks — GSM8K, MATH500, MedQA, OpenBookQA, SciQ, MMLU-Pro, SVAMP, ASDiv, MAWPS, GSM8K-Hard, GSM8K-Plus, multistep-arithmetic, and object counting — across six models: Qwen3-1.7B, Qwen3-4B, Qwen3-8B, Qwen3-14B (4-bit), Mixtral-8x7B (4-bit), and DeepSeek-R1-Distill-Llama-8B.

Key findings:

"MACRO attains the best average accuracy on every model except the reasoning-distilled one, and that its margin over the unrouted baseline is widest exactly where that baseline is weakest, narrowing to a smaller but consistent gain once the model is already strong."

"Dr. LLM in its original skip-and-repeat action space frequently ends up below the unrouted baseline. Dr. LLM+Ext. keeps the same routers, supervision, and label budget and differs only in adopting the MACRO action space, yet is substantially stronger on every model we test, indicating that the action space itself is a decisive part of what makes layer routing help rather than hurt."

Regarding the components:

"Both stages appear to contribute, and each replaces a lucky draw with a more principled choice: Viterbi reconstructs candidates from the transition table rather than trusting whichever routes the sampler happened to visit, and validation then ranks those candidates on held-out data rather than on the training score the search has already optimized."

The paper reports significance testing using a one-sided Wilcoxon signed-rank test with benchmarks as the unit of analysis; Table 3 shows Every model improves significantly, degrades on no benchmark, and excludes zero in the confidence intervals.

Ablating state factors on Qwen3-1.7B shows: the largest drop coming from removing the budget-phase factor, indicating that routing decisions depend strongly on where the program sits in its computation budget. Full state obtains 63.03 avg, no delta 62.54, no op context 61.48, no phase 59.56.

Accuracy is essentially flat in β (selection sharpness), but The elite count matters more: m=1 commits too early, m=12 dilutes counts with noise, while Every setting still improves over the unrouted model, so MACRO does not depend on a narrow hyperparameter choice.

Using a logit lens, the authors find: "the standard forward pass can represent the answer internally and then suppress it near the output, while MACRO preserves this signal through the final readout. Because the model's own generated reasoning is held fixed in context, this suggests that routing improves how an already-computed answer is exposed rather than changing the reasoning itself."

Training a single shared route on a combined pool of eight in-domain benchmarks and evaluating on out-of-domain benchmarks yields positive transfer: for Qwen3-1.7B, in-domain gain +2.79, out-of-domain gain +9.02; for Qwen3-4B, +0.30 in-domain and +1.12 out-of-domain.

Cross-benchmark route transfer (Figure 6) shows: "The highest-accuracy cells are mostly on the diagonal, which is expected because each route is optimized for its home benchmark. Off-diagonal transfer still appears, especially among related mathematical-reasoning tasks, suggesting that some routes encode reusable model computations rather than only task-specific effects."

The paper concludes:

"We have introduced a lightweight dynamic routing method that outperforms the unrouted base model and Dr.LLM, the leading dynamic routing method across a plethora of benchmarks and LLM models. It is an interesting further research question why routing improves performance: Is it because routing improves capabilities or because an internally already present answer is better elicited? We argue that some evidence points towards the latter, including our limited mechanistic interpretability study. We argue that routing should be taken much further: we want models that profit from much more extensive input-depending dynamic routing, allowing for another dimension of flexible test-time scaling."

The appendix reports that across 13 benchmarks on Qwen3-1.7B, the mean overhead of the MACRO route is +16% latency, +4.6% peak memory, described as "modest and uneven: it is largest where the selected route is longest, negligible where the route stays close to baseline length, and occasionally negative, faster and lighter, when the route is shorter than the standard pass."

The appendix lists concrete discovered routes, e.g., for Qwen3-1.7B on GSM8K the route is L0-L7 L3-L9 RJ (28→33 layers, +26.08 accuracy), and on Obj-Count L0-L7 L7-L25 L24 L25 L24-L26 L26 L27 RJ (+65.00). Route composition statistics show Move actions are common across model families, while Add appears more selectively; every valid route eventually uses Rejoin.

Here are specific improvements to AI systems enabled by this paper, and what the improved system can do:

1. Test-time computation rewiring without weight updates.

The improved system can re-route computation through a frozen transformer: skipping layers, repeating earlier layers, rewinding, and adding residual hidden states — all without modifying weights. Concretely, on GSM8K it can rewind Qwen3-1.7B from layer 7 back to layer 3, replay blocks, and raise accuracy from 43.4% to 69.5%. Across 13 benchmarks and 6 models (including 4-bit quantized Qwen3-14B and Mixtral-8x7B), it achieves +5.0% average accuracy over the sequential baseline.

2. Polynomial-time optimal route discovery via Markov-chain policy + Viterbi decoding.

Instead of expensive Monte Carlo tree search or per-sample beam search, the system learns a lightweight Markov transition matrix over routing decisions and decodes the exact top-k highest-probability routes with a Viterbi dynamic program. This eliminates per-instance test-time search, removes the need for ground-truth labels at inference, and cuts route-search cost from 14.8 hours to 1.6 hours (9.4× faster).

3. Task-adaptive computation with budget awareness.

The system tracks a budget phase (early/mid/late) in its routing state and learns when to rejoin the standard suffix of the network. This means the improved system can allocate computation adaptively: routes that need extra iterations get them, and the system stops early when additional routing would not help. Mean overhead is only +16% latency and +4.6% peak memory, and it is sometimes negative (faster and lighter than the unrouted model) when the learned route is shorter than the standard forward pass.

4. Safe, always-valid route generation via feasibility masking.

The system renormalizes its policy over only feasible actions, so every sampled or decoded route is guaranteed to respect layer bounds, the remaining computation budget, and the availability of hidden states for add operations. The improved system cannot generate invalid execution programs even during exploration.

5. Efficient repeated-layer execution with cache reuse.

When a route revisits a block, the system maintains a separate key/value history per application, so a route of T applied steps costs the same as a T-block tied-weight network. This makes rewinding practical for real inference rather than a toy experiment.

6. Better elicitation of already-computed answers.

Using logit-lens analysis, the improved system can detect when a standard model has computed the correct answer internally but suppresses it near the output layers. The routing mechanism preserves and exposes that latent signal through the final readout. This means the improved system can boost performance on tasks where the base model already knows the answer but fails to express it — a different mechanism than improving reasoning itself.

7. Cross-task route transfer.

The system can train a single shared route on a pool of benchmarks and deploy it on unseen tasks. For Qwen3-1.7B, one shared route yields +2.79 in-domain and +9.02 out-of-domain accuracy gain; for Qwen3-4B, +0.30 in-domain and +1.12 out-of-domain. Off-diagonal transfer is strongest among related mathematical-reasoning tasks, so the improved system can reuse computation patterns across tasks rather than overfitting to a single benchmark.

8. Stronger routing than prior methods due to a better action space.

The system's action space — local layer moves, add-and-apply with residual hidden states, and rejoin — is itself the decisive improvement. With identical routers and supervision, the system outperforms Dr.LLM by +7.2% overall, and outperforms Dr.LLM configured with the same action space by +1.5%. The improved system benefits most where the base model is weakest, degrading no benchmark in significance testing.

9. New dimension of flexible test-time scaling.

The improved system enables a new axis of test-time scaling: increasing effective depth by lengthening routes, repeating layers, or adding residual connections — per task, per budget, without any training or fine-tuning. This can be combined with existing scaling approaches.

10. Robust, hyperparameter-insensitive deployment.

Accuracy is essentially flat with respect to selection sharpness, and every tested hyperparameter configuration still beats the unrouted baseline. The system degrades gracefully with suboptimal settings, making it practical to deploy without extensive tuning.

11. Parameter-efficient adaptation for quantized models.

Because only a lightweight transition matrix is learned (no backpropagation through the LLM), the improved system can adapt 4-bit quantized models that cannot be easily fine-tuned — enabling routing-based improvement on models where traditional parameter updates are infeasible.

Standard Large Language Models (LLMs) execute layers sequentially. Dynamic layer routing, i.e. search for a different execution path through layers involving layer repetitions, skips and other moves, can improve performance. Existing routing approaches often require updating model weights, running expensive search loops per test instance, or demand ground-truth labels during inference. In this work, we propose Markov Chain Routing of Transformer Layers (MACRO), a framework that learns task-specific routes over LLM architectures without modifying underlying parameters. MACRO models layer routing as a context-dependent Markov policy conditioned on layer indices, computation budget phases, directional displacements, and operator context, supporting skip, repeat, and residual hidden-state addition operations. The Markov route distribution is updated via feedback on training data and decoded using a top-k Viterbi algorithm to isolate high-probability candidate programs. We evaluate MACRO across diverse reasoning and knowledge benchmarks on multiple open-weight LLMs. MACRO achieves a +5.0% average accuracy improvement over the unrouted baselines, with largest gains on small models. We outperform the best dynamic routing approach Dr. LLM by +7.2%, while reducing route-search time 9.4x (from 14.8 to 1.6 hours). Our code is publicly available at https://github.com/Batorskq/MACRO.