09. Label-Efficient, Shift-Aware Model Routing for Reliable and Cost-Constrained Small Language Model Inference

2.1 One-Sentence Core Thesis

Replace data-hungry black-box prompt routers with a white-box prefill probe that reads the SLM's last-prompt-token hidden state and next-token entropy during the fast prefill phase (before generating a single output token!), calibrated to new domains with \(\le 100\) actively selected queries.

2.2 Research Motivation & Strategic Fit

• High Industry Relevance: Every AI startup wants to slash LLM API bills by \(60\%\) without losing accuracy on domain-specific queries.
• Strategic Focus: While closely related to Topic 01 (budget routing), this work focuses specifically on the cold-start label efficiency and prefill-vs-decode systems mechanics of generative SLMs, where evaluation requires objective exact-match / unit-test tasks (GSM8K, MBPP, MMLU-Pro, MedMCQA) to avoid subjective LLM-judge noise.

3.1 Three Flaws in Existing LLM Routers

1. 1
   Black-Box Prompt Routers Ignore SLM Internal State: RouteLLM (Ong et al., 2024) trains a separate classifier \(r_\phi(q)\) on prompt text alone. It cannot tell whether the SLM's internal representations are familiar or confused by the query.
2. 2
   Full-Generation Cascades Waste Decode Latency: FrugalGPT-style cascades first generate the entire \(300\)-token response from the SLM, score it, and then re-generate from the LLM if rejected—wasting the SLM's entire autoregressive decode latency!
3. 3
   Catastrophic Cold-Start Under Domain Shift: Supervised routers trained on general chat ( Arena / Alpaca ) drop to near-random routing efficiency when tested on specialized technical domains unless retrained on thousands of domain labels.

3.2 Explicit Research Gap

In LLM inference, the prefill phase (processing the input prompt \(q\) in parallel) takes \(< 5\%\) of total generation time, whereas autoregressive decoding takes \(> 95\%\). *No existing routing framework combines zero-decode-cost SLM prefill uncertainty probes (last-token hidden state density + first-step lookahead entropy) with active cold-start calibration (\(\le 100\) queries) to adapt LLM routers to unseen domains under a fixed dollar/token budget.*

4.1 Mathematical Formulation: Zero-Decode Prefill Probing + Active Calibration

Given query \(q = (t_1, \dots, t_L)\), we run only the parallel prefill pass of the local SLM \(f_{\text{SLM}}\) (generating \(0\) or \(k=3\) lookahead tokens). From this prefill pass, we extract a compact 4-dimensional white-box feature vector \(\boldsymbol{\psi}(q)\):

1. 1
   First-Step Lookahead Entropy: \(\bar{\mathcal{H}}_k(q) = \frac{1}{k}\sum_{j=1}^{k} \mathcal{H}\big(p_{\text{SLM}}(t_{L+j} \mid q)\big)\).
2. 2
   Middle-to-Final Layer Representation Drift: Mahalanobis distance \(d_{\text{Mahal}}\big(\mathbf{h}_{\text{SLM}}^{(L)}(q)\big)\) relative to the SLM's known-competent query cluster.
3. 3
   Linear Probe Competence Logit: \(\mathbf{w}_{\text{probe}}^\top \mathbf{h}_{\text{SLM}}^{(L)}(q)\), trained with \(L_2\)-regularized logistic regression.

When entering a new target domain \(\mathcal{D}_T\) with a labeling budget of only \(K_{\text{label}} \in \{25, 50, 100\}\) oracle queries, we select calibration queries via Expected Value of Information (Uncertainty + Cluster Diversity) to fit the low-dimensional router weights \(\mathbf{w}_{\text{probe}}\) without overfitting.

flowchart LR
    Q["User Query q"] --> PRE["SLM Parallel Prefill Pass Only<br/>(< 15 ms, No Full Decode!)"]
    PRE --> PROBE["Extract White-Box Probe ψ(q):<br/>Lookahead Entropy + Hidden Density"]
    PROBE --> ROUTER{"Shift-Aware Router<br/>P(SLM Correct | ψ(q)) ≥ τ_cost"}
    ROUTER -- "Competent" --> DEC_S["Complete SLM Autoregressive Decode<br/>(Low Cost)"]
    ROUTER -- "High Uncertainty / Domain Drift" --> LLM["Escalate Immediately to Frontier LLM<br/>(Zero Wasted SLM Decode Tokens!)"]

• Models: Qwen2.5-1.5B-Instruct / Llama-3.2-3B-Instruct (local SLM) routing to Qwen2.5-32B / GPT-4o (via RouterBench + local vLLM).
• Domains (Cold-Start Shift Protocol): Train base router on MMLU-General \(\to\) test zero-shot and \(\{25, 50, 100\}\)-shot active adaptation on GSM8K (Math), MBPP (Code), MedMCQA (Clinical), and FinQA (Finance).
• Baselines: RouteLLM (Matrix Factorization & BERT router, Ong et al., 2024), Full-Decode Confidence Cascade, KNN Prompt Router, and Random Routing.
• Metrics: Area Under the Performance-Cost Curve (\(\text{AUPCC}\)), Call-to-Performance Ratio (cost required to reach \(90\%\) of Large-LLM accuracy), and Wasted Decode Token Overhead.

1. 1
   Prefill-Only White-Box Routing: Eliminates the \(200\text{–}500\text{ ms}\) wasted SLM decoding latency of standard cascade deferral while outperforming black-box prompt routers.
2. 2
   Cold-Start Active Router Calibration: Demonstrates that \(50\) actively chosen prefill-probed queries match the routing accuracy of \(5,000\) randomly labeled prompt-router pairs on new domains.

• Feasibility: Running prefill passes on a \(1.5\text{B}\text{–}3\text{B}\) model for \(10,000\) prompts takes \(< 10\) minutes on a single 24GB GPU!
• Main Risk & Mitigation: On Chain-of-Thought (CoT) reasoning tasks, the very first token ("Let's", "To") has low entropy even when the SLM will fail at step 6 of the math proof. Mitigation: Rely on the penultimate layer hidden-state probe \(\mathbf{w}^\top \mathbf{h}_{\text{SLM}}^{(L)}(q)\) (which encodes global query representation before "Let's think step by step") combined with entropy over the first \(k=8\) tokens.

• [1]
  I. Ong et al., "RouteLLM: Learning to Route LLMs with Preference Data," arXiv:2406.18665, 2024.
• [2]
  D. Ding et al., "Hybrid LLM: Cost-Efficient and Quality-Aware Query Routing," ICLR, 2024.
• [3]
  Q. Hu et al., "RouterBench: A Benchmark for Multi-LLM Routing System," arXiv:2403.12031, 2024.
• [4]
  K. Kadavath et al., "Language Models (Mostly) Know What They Know," arXiv:2207.05221, 2022.

• Identifies the latency waste of full-decode cascades vs. prefill-only routingVerified
• Uses objective ground-truth benchmarks (GSM8K, MedMCQA, MBPP) to avoid noisy LLM judgesVerified
• Fast single-GPU execution on 1.5B–3B modelsVerified
• Extract prefill hidden states from Qwen2.5-1.5B on GSM8K and train a 50-shot linear competence probePending

1. 1
   Run a prefill pass of Qwen2.5-1.5B-Instruct on \(1,000\) MMLU-Pro and MedMCQA questions and save the last-token hidden states at layers \(\{14, 21, 28\}\).
2. 2
   Plot the sample-efficiency curve (\(N = 10, 25, 50, 100, 500\) labels) of our white-box prefill probe vs. a BERT prompt router.

• Which layer of a 28-layer SLM encodes the strongest linear separability for "I know how to solve this query"—the final layer \(L=28\) or the middle-late layer \(L \approx 20\)?