💬 NLP

MACRO: Markov Chain Routing of Transformer Layers

The paper introduces MACRO, a framework that enhances Large Language Models by learning task-specific, dynamic layer routing policies via Markov chains without modifying model weights, achieving significant accuracy improvements and faster search times compared to existing methods.

Original authors: Paweł Batorski, Abtin Pourhadi, Akylgali Aitaza, Przemysław Spurek, Paul Swoboda

Published 2026-08-07

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Paweł Batorski, Abtin Pourhadi, Akylgali Aitaza, Przemysław Spurek, Paul Swoboda

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart robot that reads books and answers questions. To do its job, the robot has a long assembly line of 28 or 32 tiny workers (called "layers"). In a standard setup, every single piece of information must walk past every single worker, from the very first to the very last, in a strict, unchangeable line. It's like a school bus that must stop at every single house on the street, even if the passenger only needs to get off three blocks down. This is how most AI models work today: they follow a rigid schedule, no matter how easy or hard the question is.

But what if the robot could be a little more flexible? What if, instead of marching in a straight line, it could decide to skip a worker who isn't needed, or even go back and ask a previous worker to double-check their math? This idea is called "dynamic routing." Think of it like a GPS that doesn't just follow a pre-printed map but can instantly reroute around traffic or take a shortcut. Scientists have been trying to teach AI to do this, but most attempts are like trying to teach a bus driver to drive by rewiring the engine every time they hit a new road, or by hiring a separate, expensive guide to shout directions for every single passenger. It's slow, costly, and often breaks the robot's brain.

This is where a new study called MACRO comes in. The researchers asked a simple question: Can we teach a frozen robot (one whose brain we aren't allowed to change) to find a better path through its own workers without breaking anything? They discovered that by treating the robot's path like a game of "choose your own adventure" based on a set of simple probability rules (called a Markov chain), they could find a much smarter route. Instead of retraining the robot or using expensive guides, they let the robot explore different paths on a few practice questions, learned a single, task-specific set of rules for navigating its workers, and then used a clever math trick to identify the absolute best route for that specific task. This best route is then used uniformly for all questions in that task, rather than calculating a new path for every single new question.

The results are surprisingly effective. On a variety of tricky math and science tests, MACRO helped the robot get significantly smarter. For example, on a difficult math benchmark called GSM8K, a small robot model jumped from getting 43.4% of the answers right to 69.5%—a massive leap—just by changing the order in which it visited its workers. The researchers found that this method was much faster to set up than previous methods. While other approaches took nearly 15 hours to figure out the best path for a task, MACRO did it in just 1.6 hours.

The key insight is that the robot already knew the answers; it just needed a better way to "listen" to its own internal thoughts. By allowing the robot to skip unnecessary steps or loop back to review its work, the researchers unlocked hidden potential without changing a single weight in the robot's brain. It's as if they didn't teach the robot new facts, but simply taught it how to think more efficiently. The study suggests that for many complex tasks, the answer was already inside the model; we just needed a better map to find it.

Technical Summary: MACRO – Markov Chain Routing of Transformer Layers

Problem Statement
Standard Large Language Models (LLMs) execute transformer blocks sequentially from the first layer to the last, adhering to a rigid computational schedule. This static approach ignores the potential for dynamic computation pathways that could improve reasoning accuracy on complex tasks without modifying the model's underlying weights. While prior work has explored adaptive depth primarily for inference acceleration (e.g., early exits, structured layer dropping), actively restructuring the execution route of a frozen model to enhance performance remains a challenging frontier. Existing dynamic routing methods face significant bottlenecks: some require expensive offline Monte Carlo Tree Search (MCTS) for data generation and training (e.g., Dr.LLM), others demand ground-truth labels during inference, and some rely on per-sample search loops that are computationally prohibitive.

Methodology
The authors propose MACRO (Markov Chain Routing of Transformer Layers), a framework that learns task-specific execution routes over frozen LLM architectures. MACRO models layer routing as a context-dependent Markov policy, $\pi_\theta(a_t | s_t)$ , which dictates the sequence of layer applications, skips, repeats, and hidden-state additions.

State and Action Space: The state $s_t$ is a compact summary comprising the current layer index ( $\ell$ ), the remaining computation budget phase (early/mid/late), the incoming layer displacement ( $\delta$ ), and the operator context (whether the current hidden state was produced by a plain layer or an addition). The action space includes:
1. Local Layer Move: Applying a nearby layer $\ell + \Delta$ (allowing skips, repeats, or backward revisits).
2. Add-and-Apply: Adding a previous hidden state $h_{t-b}$ to the current state $h_t$ (with a coefficient $\gamma > 0$ ) before applying a local layer.
3. Rejoin: Terminating the learned prefix and executing the remaining standard suffix of the transformer.
  Feasibility masks ensure all sampled programs remain within layer bounds and budget constraints.
Learning Procedure: MACRO employs an estimation-of-distribution approach. Starting from a forward-biased prior (where the standard sequential path is most probable), the framework iteratively samples candidate routing programs, evaluates them on a subset of training data, and updates the Markov transition matrix using weighted maximum likelihood. High-performing routes are weighted exponentially to refine the policy. A replay buffer ensures stable convergence by preserving coverage of examples from previous iterations.
Decoding and Selection: Unlike methods that rely on stochastic sampling or per-test-time search, MACRO uses a top-k Viterbi algorithm to deterministically decode the highest-probability valid routing programs from the learned policy. Structural feasibility masks are integrated into the Viterbi recursion. The final route is selected by evaluating the top-k candidates on a held-out validation set, choosing the one with the highest accuracy.

Key Contributions

Parameter-Efficient Framework: MACRO introduces a lightweight, context-dependent Markov policy for adaptive layer routing in frozen and quantized LLMs, requiring no updates to the model's weights.
Exact Decoding Scheme: The paper presents an exact top-k Viterbi decoding scheme equipped with structural feasibility masks. This eliminates the need for per-sample test-time search, auxiliary neural network training, or ground-truth inference oracles.
Empirical Superiority: The framework demonstrates consistent improvements over non-routed baselines and existing dynamic routing methods across diverse reasoning and knowledge benchmarks.

Results
Evaluated across multiple open-weight LLMs (including Qwen3 variants, Mixtral, and DeepSeek-R1-Distill) on benchmarks such as GSM8K, MATH500, MedQA, and MMLU-Pro:

Accuracy Gains: MACRO achieves an average accuracy improvement of +5.0% over unrouted baselines. The gains are most pronounced on smaller models (e.g., +12.68% on Qwen3-1.7B).
Comparison with SOTA: MACRO outperforms the leading dynamic routing method, Dr.LLM, by +7.2% in average accuracy.
Efficiency: The route-search process is significantly more efficient, reducing search time by 9.4× (from 14.8 hours to 1.6 hours per benchmark) compared to Dr.LLM.
Robustness: Statistical significance tests confirm that MACRO improves performance across all tested models without degrading any benchmark.
Mechanistic Insight: Logit-lens analysis suggests that routing improves performance not by changing the model's internal reasoning capabilities, but by better exposing answers that are already internally represented, preventing their suppression near the output.

Significance and Claims
The paper claims that MACRO represents a significant step forward in making dynamic computation practical for frozen LLMs. By formalizing routing as a Markov process, the authors overcome the computational costs and data requirements of previous methods (like MCTS-driven training). The work argues that routing serves as a mechanism to "elicit" internally present answers more effectively rather than fundamentally altering the model's learned representations. The authors posit that this approach should be extended further to enable input-dependent dynamic routing, offering a new dimension of flexible test-time scaling for LLMs. The method is particularly notable for its ability to improve performance on small models and its compatibility with quantized models, suggesting broad applicability in resource-constrained environments.

Technical Summary: MACRO – Markov Chain Routing of Transformer Layers

More like this