💬 NLP

LLM Router: Prefill is All You Need

This paper introduces LLM Router, a system that leverages internal prefill activations through an Encoder-Target Decoupling mechanism and the SharedTrunkNet architecture to achieve near-Oracle routing accuracy with significant cost savings by mathematically isolating optimal layer-wise signals for heterogeneous model pairing.

Original authors: Tanay Varshney, Annie Surla, Michelle Xu, Gomathy Venkata Krishnan, Maximilian Jeblick, David Austin, Neal Vaidya, Davide Onofrio

Published 2026-03-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tanay Varshney, Annie Surla, Michelle Xu, Gomathy Venkata Krishnan, Maximilian Jeblick, David Austin, Neal Vaidya, Davide Onofrio

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 are running a massive, high-end restaurant. You have a menu of 30 different chefs (AI models), ranging from a quick, cheap line cook to an expensive, world-famous Michelin-star chef.

The Problem:
Usually, when a customer orders a dish, the manager has to guess which chef will do the best job.

If they send a simple "grilled cheese" order to the Michelin chef, they waste a fortune.
If they send a complex "sushi platter" to the line cook, the food comes out terrible.

Current systems try to guess by reading the words of the order (e.g., "sushi" = send to expert). But this is like judging a book by its cover. Sometimes a simple-looking question is actually a trap that even experts struggle with, and sometimes a complex-sounding question is easy for a novice. The current "managers" (routers) are often wrong because they don't understand the actual difficulty of the task for a specific chef.

The Solution: "LLM Router"
This paper introduces a new kind of manager called LLM Router. Instead of just reading the order, it uses a special trick to peek inside the chefs' brains before they start cooking.

Here is how it works, broken down into simple concepts:

1. The "Prefill" Peek (The Crystal Ball)

When an AI starts thinking about a question, it goes through a "prefill" phase where it processes the input before generating an answer. During this split second, the AI's internal neurons fire in specific patterns.

The authors discovered that if you look at these internal "neuron sparks," you can tell before the AI even speaks whether it is going to get the answer right or wrong. It's like a chef glancing at an ingredient and instantly knowing, "I can't make this work," or "This is easy for me."

2. The "Foreign" Translator (Encoder-Target Decoupling)

Here is the magic trick: You don't need to ask the expensive chef to check their own brain. That's too slow and expensive.

Instead, the system uses a small, cheap, open-source AI (the "Encoder") to look at the question and mimic the thinking process of the expensive, closed-source AI (the "Target").

Analogy: Imagine you want to know if a famous, reclusive painter (the Target) can paint a specific landscape. Instead of asking the painter (who charges $10,000 just to look), you ask a talented art student (the Encoder) to look at the landscape and say, "If I were the famous painter, I'd struggle with this."
Result: Surprisingly, the art student is often better at predicting the famous painter's success than the painter's own internal signals! This allows the system to use cheap models to manage expensive ones.

3. Finding the "Sweet Spot" Layer

AI models have many layers of processing (like layers of an onion). The paper found that looking at the very last layer isn't always best.

They used a mathematical tool called Fisher Separability to find the exact "layer" where the difference between a "correct" thought and a "wrong" thought is the clearest.
Analogy: It's like tuning a radio. You don't just listen to the static; you find the exact frequency where the music is crystal clear. They found the perfect frequency for every AI model to predict its own success.

4. The "SharedTrunkNet" (The Smart Manager)

The system uses a central brain called SharedTrunkNet. It looks at the "neuron sparks" from the cheap translator and simultaneously predicts:

How likely is Chef A to get this right?
How likely is Chef B to get this right?
How much will it cost to hire Chef A vs. Chef B?

It then picks the best value: the cheapest chef who is still likely to get the job done.

The Results: Why This Matters

The researchers tested this on a huge pool of models (from small, cheap ones to massive, expensive ones).

The Oracle: Imagine a "God-mode" manager who knows the future and always picks the perfect chef. This is the theoretical best.
The Standalone: Just picking the single best chef for everything.
The Result: Their new router managed to close 45% of the gap between the "God-mode" manager and the "just pick the best chef" approach.
The Savings: They saved 74% of the money compared to just using the most expensive chef for every task, without sacrificing much accuracy.

In a Nutshell

This paper teaches us that we don't need to guess which AI is best by reading the question. Instead, we can use a cheap, smart "scout" to peek at the internal signals of expensive AIs before they answer. This allows us to route tasks to the right model instantly, saving massive amounts of money while keeping quality high.

The takeaway: Don't just ask the AI what it thinks; watch how it thinks before it speaks, and let a cheap scout do the watching for you.

1. Problem Statement

Current Large Language Model (LLM) routing systems primarily rely on semantic or intent-based signals (e.g., query embeddings, domain descriptions) to select the most appropriate model for a task. The authors identify two fundamental flaws in this paradigm:

Lack of Mechanistic Interpretation: Semantic routers ignore the intrinsic generative limits and internal states of the models they are routing to.
Semantic-Complexity Gap: Semantic similarity does not correlate with task difficulty. A query that looks simple semantically might be difficult for a specific model, and vice versa.

Furthermore, traditional evaluations often suffer from "agreement" bias (where all models succeed or fail together), masking the true value of routing. The goal is to create a router that can predict per-LLM confidence (probability of correctness) and cost prior to generation, enabling a selection policy that maximizes accuracy while minimizing cost.

2. Methodology

The paper proposes a shift from semantic analysis to mechanistic routing using internal model activations.

Core Concept: Encoder-Target Decoupling

Instead of using the target model's own internal states (which requires running the target model, defeating the purpose of routing), the authors introduce Encoder-Target Decoupling.

Encoder: A lightweight, open-source model (e.g., Qwen) that processes the input query.
Target: The model whose performance is being predicted (e.g., a closed-source API like Claude or GPT-5).
Mechanism: The router extracts hidden states from the Prefill phase of the Encoder. These activations serve as a predictive signal for the Target's likelihood of correctness.

Signal Extraction & Architecture

Layer Selection: The authors use Fisher Separability ( $J$ ) and Effective Dimensionality ( $d_{eff}$ ) as mathematical probes to identify the optimal transformer layers in the Encoder. They found that the most separable layers (where correct vs. incorrect predictions are most distinct) are not always the final layers.
SharedTrunkNet: A multi-output architecture that takes the concatenated, PCA-reduced hidden states from the Encoder (across multiple potential Targets) and outputs simultaneous correctness probabilities ( $P(\text{correct})$ $P (correct)$ ) for all candidate models.
- Advantage: Joint optimization provides cross-model context (evidence from one model's difficulty helps predict another's) and inherent calibration.
Cost Estimation: The router estimates cost based on input tokens and a median training-set output length, avoiding ground-truth leakage of output token counts.
Routing Policy: A scoring function balances accuracy and cost:
$\text{score}_{k,q} = \lambda \cdot P(\text{correct}_k) - (1 - \lambda) \cdot \text{cost}_{k,q}$
Where $\lambda$ controls the trade-off between accuracy and cost.

3. Key Contributions

Encoder-Target Decoupling: Demonstrated that "foreign" open-source encoders (e.g., Qwen-122B) can predict the performance of closed-source targets (e.g., Claude, GPT-5) more accurately than the targets' own internal states.
Fisher-Based Layer Selection: Moved beyond the assumption that the "final layer" is best, using Fisher Separability to mathematically identify the most separable intermediate layers for routing signals.
SharedTrunkNet Architecture: A novel multi-task network that leverages cross-model context to assess relative difficulty, outperforming independent probes.
New Evaluation Framework: Introduced P-AUCCC (Padded Area Under the Cost Coverage Curve) and Router Efficacy (distance to the theoretical Oracle) to measure routing gains against static model selection and theoretical limits.

4. Experimental Results

The authors evaluated their approach across three model pools: Frontier (high-end models), Small (7B–9B models), and Mixed (heterogeneous tiers).

Predictive Power: In the Frontier pool, the SharedTrunkNet using Qwen-122B as an encoder achieved the highest Area Under the Curve (AUC) for predicting correctness across all targets, significantly outperforming semantic backbones (like NV-Embed or DeBERTa).
Accuracy vs. Cost:
- Frontier Pool: SharedTrunkNet captured 45.58% of the accuracy gap between the best standalone model and the theoretical "Oracle" (perfect selector).
- Cost Savings: It achieved 74.31% cost savings relative to the highest-cost model in the pool.
- Global Efficacy: The router improved Router Efficacy by 53.62% and P-AUCCC by 14.67% over the static "Model-Only Pareto" baseline.
Robustness: The method held up across different pool sizes and model architectures, proving that internal activation geometry is a superior signal to semantic embeddings.

5. Significance

This work fundamentally changes the paradigm of LLM orchestration:

Zero-Cost Routing for Closed Models: It enables high-fidelity routing to proprietary, closed-source APIs using only the prefill activations of lightweight, local open-source models. This removes the latency and cost of running the target model just to decide whether to use it.
Theoretical Grounding: By linking routing performance to geometric properties of transformer activations (Fisher Separability, Effective Dimensionality), the paper provides a mechanistic explanation for why routing works, moving beyond heuristic "black box" approaches.
Economic Efficiency: The results demonstrate that significant cost reductions (up to 74%) can be achieved without sacrificing accuracy, making high-performance LLM usage more accessible and scalable.

In summary, "Prefill is All You Need" argues that the internal state of a small model during the prefill phase contains sufficient information to predict the performance of much larger, more expensive models, enabling a new class of efficient, mechanism-driven LLM routers.