Predicting whether a model will fail, before it answers

The short version

I took a recent paper called PIKA (Probe-Informed K-Aware Routing), reproduced it on my own benchmark pool, and found that its headline contribution does not hold up while its practical result holds up better than the authors claimed.

The idea is that you can predict whether a language model will get a question right before it writes a single token, by reading the internal activations of a completely different model as it processes the question. If that works, you can route each incoming question to the cheapest model likely to answer it correctly, and you get frontier-model accuracy at a fraction of the price.

The practical result reproduced. My router matches the best single model's accuracy at about 82% lower costper query, which actually beats the paper's reported 74.3% savings. But the paper's two representation-learning claims, the ones that are supposed to explain why it works, both failed to reproduce under paired bootstrap testing. I then spent four more experiment steps chasing the remaining performance gap and traced it to something neither the paper nor I had originally suspected.

The whole thing ran across a rented 80 GB GPU pod and my laptop, with 15 GB of cached activations moving between them.

Background: what the paper claims

Standard difficulty prediction uses a model's own internal states to predict its own failures. PIKA's more interesting move is encoder-target decoupling. You pick one open-weights model as an encoder, feed it the question, grab its hidden states, and train a small probe to predict whether some other model, possibly a closed API model you cannot see inside of, will answer correctly.

This matters because it means one encoder pass serves every target model at once. The activations depend only on the encoder and the question text, not on which model you are predicting for. Adding a new target to your pool costs you nothing but its labels.

The paper makes four claims worth testing:

1. Decoupled probes predict target-model correctness with high AUC.
2. Fisher separability (a cheap statistic, Fisher-J) picks a good layer without the expensive sweep of training a probe at every single layer.
3. SharedTrunkNet, a joint multi-output network trained across several target models at once, beats independent per-model probes by exploiting cross-model context. This is presented as the main contribution.
4. Joint training yields inherently calibrated probabilities, so you do not need post-hoc calibration.

And a practical claim on top: a cost-aware router built on these probes gets +10.9 accuracy points over the best single model, captures 45.6% of the oracle gap, and saves 74.3% on cost.

I tested all of it.

Setup

Encoder: Qwen/Qwen3.5-35B-A3B. This is a Mixture-of-Experts model with 256 experts, 8 routed per token, giving roughly 3B active parameters out of about 36B total. Important detail that bit me later: MoE reduces compute, not memory. All 36B parameters have to sit in VRAM regardless.

What gets extracted: hidden states from the prefill pass only. The question is rendered through the encoder's chat template, run forward once with output_hidden_states=True, and I take the last-token representation from a single chosen layer. No generation happens at any point. The encoder never produces output text. This is what makes the whole approach cheap enough to be worth doing, because the routing decision has to be cheaper than the thing it is routing.

The hidden dimension is 2048 across 40 layers, which is actually smaller per token than a dense 32B model at 5120 by 64. The cost here is entirely GPU time and VRAM, not disk.

Targets: eventually 11 models, from qwen3-235b-a22b-2507 at $0.00062 per query up to gemini-2.5-pro at $0.05548, roughly a 90x price spread.

Benchmarks: 14 of them, pooled. aime, arc-agi, arenahard and three of its subsets, gpqa, hle-test, livecodebench, livemathbench, mmlupro, simpleqa, swe-bench, tau2. Final pooled test set is n=2053, with a coding and agentic subset of n=274.

Phase 1: getting the data into shape

The labels came from LLMRouterBench/bench-release, a pile of JSON with per-model, per-question rollout records. I wrote scripts/bench_release_to_pika.py to convert it into PIKA's parquet layout, which also meant hand-building a vendor map from their directory names to HuggingFace style identifiers for every model in the pool.

Two extra artifacts fall out of that conversion and both turn out to matter a lot later:

• prices.json, per-model input and output rates per million tokens, fitted from the recorded costs rather than scraped from a pricing page. The rates are per benchmark, so 11 models by 14 benchmarks.
• median_out.json, median completion token counts per model per benchmark. You need this because a model's cost per query depends on how much it decides to write, and reasoning models write a lot more than chat models on the same question.

Then split_train_val_test.py does 70/15/15 stratified splits, and this script has one property I had to be careful about. The split is computed per question and applied identically across every model's parquet. If question 47 is in test for Claude, it must be in test for GPT-5 too, otherwise the [N, K] label matrix SharedTrunkNet needs later does not line up and the whole comparison is garbage. Stratification uses the per-question mean pass rate across all models, thresholded at 0.5, so splits get a balanced easy/hard mix without privileging any single target's labels.

Phase 2: the GPU side

This is the part that took the most real time and none of it was modeling.

Modifying PIKA

The upstream code loads the encoder with AutoModelForCausalLM. Qwen3.5-35B-A3B does not register against that class, because it registers as Qwen3_5MoeForConditionalGeneration and carries a pipeline_tag: image-text-to-text. It is technically a multimodal model. I use it text-only, but the loader still rejects it. The fix is a try/except that falls back to AutoModelForImageTextToText, plus a hard requirement of transformers 4.57 or later since the model's own config declares 4.57.0.dev0.

I made that edit on the laptop before the rsync, deliberately, so the GPU box never held a version of the file that differed from the one I was iterating on locally. Small discipline thing, but debugging a divergence between two copies of the same repo across an SSH boundary is miserable and I have done it before.

In total I modified five files upstream, about 267 lines added, mostly to thread encoder and target through as separate concepts where the original code assumed one model did both jobs.

GPU sizing

bf16 weights are about 72 GB, plus roughly 4 GB of per-forward activations at batch size 16 and sequence length 1024. That needs an 80 GB card, so A100 or H100. I documented FP8 (~36 GB, fits a 48 GB card) and GPTQ-Int4 (~18 GB, fits a 4090) fallbacks with an explicit note that Int4 hidden states are degraded enough that any AUC from them should be treated as a lower bound.

The daemon

This is the piece I am happiest with on the infrastructure side. Loading 72 GB into VRAM takes about ten minutes. The pod bills continuously whether or not it is doing anything. So paying that ten minutes once per job is pure waste, but so is keeping a model loaded that nothing needs.

scripts/encode_worker.py is a persistent daemon that solves both. It lazily loads: at startup it reads only the tokenizer, which takes seconds, and reports ready_no_model. The 72 GB encoder loads on the first genuine cache miss and then stays resident for the life of the process. Jobs whose activations are already cached get served with no model in memory at all.

Work arrives as JSON files dropped into data/encode_jobs/queue/. It writes a heartbeat to heartbeat.json with a status of loading_tokenizer, ready_no_model, loading_model, idle, or working.

encode_client.py checks liveness two ways, the worker PID via pgrep and heartbeat freshness with a 30 second staleness threshold. That redundancy is deliberate. A dropped SSH session or a pod migration should be detected loudly rather than silently triggering a second model load on top of the first, which on an 80 GB card means an OOM ten minutes into a job.

Critically, the daemon runs PIKA's own extraction path rather than a reimplementation, so the .pt caches it writes are byte-identical to what pika.main would produce and fully interoperable with the analysis scripts.

The cache key redesign

Upstream hashed the labels into the activation cache path. That is exactly wrong for decoupled mode. Activations are a function of the encoder and the rendered text and nothing else, so hashing labels in would write a separate 15 GB copy for every target model I added. I rewrote get_cache_path to key on (encoder, split, first-prompt hash, max_length, batch_size) and explicitly ignore labels, keeping the argument for back-compat.

This one change is what made everything downstream cheap. Because activations are shared, adding target model number 11 cost me a parquet of labels and about ninety seconds of numpy.

I did not just assume it worked. Step 3 hard-asserts that the train, val, and test activation tensors are byte-identical across targets, and separately asserts per-target text identity. If a shared encoder is your entire efficiency argument, that assertion is the thing standing between you and a silently invalid result.

What came back

15 GB of cached activations in data/activations/, covering all 14 benchmarks. Everything from Step 3 onward is marked in my notes as “pure cache analysis, no daemon, no model.” I load the 9 GB tensor once and run every subsequent experiment as numpy and scikit-learn on the laptop. Once the encoding was done I killed the daemon to stop the billing.

There was a real OOM detour here. Seven benchmarks failed to encode at batch size 16, including swe-bench, livecodebench, and the arenahard variants, which is awkward because those are precisely the hard ones. I added per-job batch_size support to the daemon and re-encoded them small: swe-bench at 2, arenahard_coding at 4, the rest at 8. Activations are batch-invariant, and the loader globs _bs*, so mixed batch sizes coexist in one cache. That is what closed the coverage gap and took me from 7 benchmarks at n=1265 to 14 at n=2053.

Step 2: layer selection, and a lesson about small samples

Which layer's activations do you probe? PIKA proposes Fisher-J, a cheap separability statistic, as a substitute for brute force sweeping every layer.

My first result said Fisher-J beatbrute force, and by a lot. On Claude, Fisher picked layer 40 and scored 0.8182 test AUC against brute force's 0.7386. On GPT-5, Fisher picked layer 20 and scored 0.9167 against 0.7083.

Those numbers were on n=19.

I did not publish that. I wrote step2_bootstrap.py and then step2_pooled_bootstrap.py, pooled to n=1265 across seven benchmarks, and ran 10,000 paired bootstrap resamples. The real answer:

Fisher-J is very slightly worse, and the confidence intervals span zero. The paper's claim holds: Fisher-J matches brute force within about 0.01 AUC. My exciting n=19 edge was noise, in the direction I wanted, which is the most dangerous kind.

This also confirmed claim 1. Pooled test AUC of 0.90 for Claude Sonnet 4 and 0.82 for GPT-5, matching the paper. Decoupling genuinely works. A Qwen model's prefill activations really do encode whether Claude is about to get this question wrong.

Every result after this point is reported with paired bootstrap CIs. That decision came directly from watching n=19 lie to me.

Step 3: SharedTrunkNet, and the first null

The paper's main contribution. Instead of training one probe per target, you train a joint multi-output network: concatenate each target's PCA-100 features at its own Fisher-J layer (Equation 1 in the paper), push through a shared trunk, and split into K heads. The claim is that each head benefits from cross-model context, from seeing what the other models' features say.

I built it faithfully. Linear(200, 256), ReLU, Dropout(0.3), K heads, BCEWithLogits, Adam, ensemble the top 5 of 10 runs, Platt scaling on 5-fold out-of-fold predictions. K=2 with Claude Sonnet 4 and GPT-5, both selecting layer 39.

It reproduced the paper's absolute number almost exactly. Mean AUC 0.8578 against the paper's 0.8560. Claude 0.8956, GPT-5 0.8201.

And it did nothing.

The comparison that matters is not against the paper's number, it is against independent probes on the same features and the same architecture, so the only difference is the concatenation and the joint training:

• Claude: dAUC +0.0010, CI [-0.0011, +0.0031]
• GPT-5: dAUC +0.0017, CI [-0.0010, +0.0044]

Both intervals span zero. Against independent linear probes it is actually slightly negative. The cross-model-context effect is null at K=2.

This is where reproducing the absolute number becomes important rather than incidental. If I had only gotten 0.79 I could not tell whether the missing lift was a broken implementation. Landing on 0.8578 against their 0.8560 means the implementation is right and the lift genuinely is not there.

Step 4: does the lift show up at larger K?

The obvious objection is that K=2 is minimal. Each trunk borrows from exactly one other model. The paper used a larger frontier pool. So maybe the effect scales.

Testing this would normally mean re-encoding everything for each new target. It did not, because of the shared-activation property, which I had already proven by assertion. New targets reuse the exact same Qwen activations and contribute only their labels. I loaded the 9 GB tensor once, no daemon, no model, and swept K from 2 to 6 across six targets.

The lift does not emerge and does not scale:

K=2: mean dAUC = -0.0000
K=3: mean dAUC = +0.0011
K=4: mean dAUC = +0.0013
K=5: mean dAUC = +0.0015
K=6: mean dAUC = +0.0023

Flat, tiny, no trend. Of 20 (K, target) cells, 19 are non-significant. The one that clears is gemini at K=6, +0.0055 with CI [+0.0006, +0.0105], which is about one expected false positive across 20 tests. Meanwhile Claude at K=5 is negative at -0.0027.

There is one honest nuance I want to keep in: all 20 cells are positive in sign, and a sign test on that gives p of about 2e-6. So there is very likely a real effect. It is just about +0.002 AUC, which is not a contribution, it is a rounding error.

Step 4b closed the last caveat by re-encoding the seven benchmarks that had OOM'd, taking the eval to 14 benchmarks and n=2053. The null held, including on the coding and agentic subset where the effect was most plausible. On the harder eval the Fisher layers shifted from L39 down to L21 and L27, which is interesting on its own.

The coding subset showed a hint: mean lift +0.0066 at K=6, with Claude, GLM, and Kimi all around +0.011. But every cell is non-significant at n=274 and the direction is mixed. I recorded it as suggestive and underpowered, not confirmed, and noted that resolving it needs more coding questions, not more models.

Step 5: the router, and the first thing that clearly works

Four steps of null results and then something worked.

Route each question to argmax_m [ λ × P_m − (1−λ) × cost_norm_m ], where P is the independent probe's confidence that model m gets this question right, and cost is Equation 2 from the paper, (r_in × median_in + r_out × median_out) / 1e6 dollars per query, using the fitted per-benchmark rates and median completion lengths from the data conversion, with median input tokens computed by actually tokenizing the prompts. λ sweeps from 0 (pure cost) to 1 (pure accuracy). The oracle is the same sweep run on true labels.

Results over 6 models, n=2053:

Three headline numbers:

1. Router peak accuracy 0.6551 versus 0.6220 for the best single model, a gain of +3.31 points. The paper reports +10.9.
2. Cost to match gemini's accuracy: $0.01043 versus $0.05548, which is 81.2% cheaper. The paper reports 74.3%, so I beat them here.
3. Router efficacy, the fraction of the oracle gap captured: 30.4%. The paper reports 45.6%.

So the practical claim reproduces, and on the cost axis it over-reproduces. The interesting part is that independent probes are enough. SharedTrunkNet contributed nothing, and yet the router built on plain per-model probes works fine. The paper's main architectural contribution is not load-bearing for its own headline application.

That left a real gap on the accuracy axis, +3.31 against +10.9. I had a hypothesis, and the next three steps are me testing it and being wrong.

Step 6: calibration, hypothesis refuted

My hypothesis was miscalibration. I was using raw sigmoid outputs and comparing them across models. If one model's probe is systematically overconfident relative to another's, the argmax is meaningless even when each probe is individually well ranked. Fixing that seemed like the obvious lever.

So I fit per-model Platt scaling and isotonic regression on validation predictions. As a sanity check, the retrained test AUCs reproduced Step 4 exactly, which is a nice property to have when you are retraining things.

Calibration did not help. It hurt:

The reason, which I should have checked first: the probes were already calibrated. Mean test ECE was 0.0331 raw, and mean predicted probability matched actual accuracy for all six models, Claude at 0.399 predicted against 0.408 actual, GPT-5 at 0.601 against 0.596, and so on down the list. Platt moved mean ECE to 0.0313 and isotonic to 0.0309. There was nothing to fix, so fitting a correction on validation data just added noise.

This also killed the paper's claim 4 from the other direction. The claim is that jointtraining produces inherent calibration. But SharedTrunkNet's mean ECE is 0.0365, which is worse than independent probes at 0.0331. Joint training gave no calibration benefit at all here. Independent BCE-trained probes were the best-calibrated thing in the experiment.

Lesson I wrote down for myself: check ECE and mean-P-versus-accuracy before assuming raw sigmoids need calibrating.

Conclusion: the gap to the paper is structural, not calibration.

Step 7: pool diversity, hypothesis refuted again

Second hypothesis. Maybe the paper's bigger accuracy gain comes from a more diverse model pool, giving the router more room to be clever.

I expanded from 6 models to 11, adding gemini-2.5-flash, deepseek-r1-0528, deepseek-v3.1-terminus, deepseek-v3-0324, and intern-s1. Five new probes trained, Fisher layers computed fresh, and the six existing probes reproduced exactly.

Diversity raises the oracle ceiling from 0.799 to 0.822, so in principle there is more to capture. But the router cannot capture it. Efficacy went down, from 30.4% to 28.0%, and peak accuracy gain stayed flat at about +3.3 points across every pool I tried.

So the bottleneck is not pool size. It is the probe's ability to predict per-model correctness on hard questions.

The consolation prize was a genuinely useful product finding. Sorting the 11 models by accuracy per dollar:

• qwen3-235b-a22b-2507 is the value king. 0.550 accuracy at $0.00062 per query, and it absorbs roughly 41% of routed volume at reasonable operating points.
• gpt-5 and gemini-2.5-pro are the escalation targets, used only when the probe flags a hard question.
• Eight models are dead weight, each taking 1 to 3% of traffic while adding nothing the others do not cover.

Operating points: 0.550 accuracy at $0.0006 per query, 0.602 at $0.0035 (the value sweet spot), 0.622 at $0.0096 (matching the best single model at 82.8% less), and 0.6551 peak at $0.0248, which both exceeds the best single model and costs 55% less than it.

Step 8: diagnosing the misroutes

Instead of guessing at a third lever, I went and looked at what the router actually got wrong. Pure numpy over saved predictions, runs in seconds.

Of 708 total routing errors at argmax confidence, 365 are unwinnable, meaning no model in the pool solves that question, and 343 are recoverable, meaning a capable model existed and the router picked a failing one. So 48% of errors are genuinely the router's fault and 16.7% of all questions.

Where they concentrate:

simpleqa alone is 44% of all recoverable errors.

I checked the obvious explanation, that the router over-trusts low-AUC models, and it is wrong. Mean probe AUC of the chosen model is 0.842 on correct routes versus 0.837 on recoverable errors. Essentially identical. The high-pick frontier models, gpt-5 and gemini, have routing precision around 0.60 and are the biggest single error sources purely because they get picked most.

The actual pattern only showed up when I bucketed questions by how many models in the pool can solve them at all:

That is the whole story. The router fails on needle questions, where only one or two of eleven models can answer and it has to pick that exact one. When six or more models can solve something, the router essentially never misses. Per-benchmark correlation between mean probe AUC and the recoverable gap is -0.58.

This reframes the gap to the paper. It is not architecture, not calibration, not pool size. It is probe discrimination on low-redundancy questions. And it explains the benchmark concentration: simpleqa is factual recall, which is a question about which model happens to know this fact, not about how hard the question is. That may be close to irreducible from question activations alone, because the difficulty is not in the question.

swe-bench is the exception and it is fixable. I encode only the first 1024 tokens of tasks that run to about 47,000 characters, so the probe is mostly blind on it. That is a truncation artifact, not a limit of the method.

Pruning. The last finding is my favorite because it is counterintuitive. Dropping 8 of the 11 models makes the router better. PRUNED3, just {qwen3-235b, gpt-5, gemini-2.5-pro}, peaks at 0.6605 against the full pool's 0.6551, with better AUCCC. Fewer weak models means fewer wrong things to route to. Adding deepseek-r1 back as PRUNED4 smooths the cheap-to-expensive escalation at the same AUCCC.

Settled product config: qwen3-235b as the cheap workhorse, gpt-5 and gemini-2.5-pro as probe-gated escalation. Best single model accuracy at roughly one sixth the cost.

The interactive router

To make it something I could actually show someone, I built scripts/route_interactive.py. Five probes frozen to disk (qwen3-235b at L27, kimi-k2 at L27, deepseek-r1 at L21, gpt-5 at L21, gemini-pro at L27), each with its layer, PCA, scaler, and top-5 MLP ensemble, with AUCs reproducing Step 7. I added an encode_texts action to the daemon so it encodes a typed prompt through the same extraction path as training.

It is a REPL. Type a prompt, get P(correct), cost, and score for every model plus the routed winner. :lambda X changes the cost/accuracy tradeoff live, :sweep shows how the winner changes across λ.

Verified end to end. “capital of France” at λ=0.5 routes to Kimi (P=0.895, $0.0015) over Gemini (P=0.937 but $0.061). At λ=1 it goes to Gemini. That is exactly the behavior you want: on an easy question, do not pay for the frontier model.

And live testing surfaced a design flaw I would not have found in the aggregates. A flat λ is too crude. On a genuinely hard question, no model is confident, so all the P values are low and close together. The −(1−λ) × cost term then dominates those tiny differences and the router defaults to the cheapest model, which is precisely the one most likely to fail.

So a fixed λ systematically under-escalates on exactly the hardest inputs. This is the same failure as the Step 8 needle-question finding, showing up from a completely different direction, which is the kind of convergence that makes me believe it.

The fix, which I have specified but not built: adaptive λ. Make the weight a function of max_m P_m(correct). When the whole pool looks unconfident, push toward accuracy and spend the money. When some model is clearly confident, stay cheap.

What reproduced and what did not

Two of the paper's four claims held. The two that failed are the two that were supposed to be the novel representation-learning contribution.

Limitations I want stated plainly

• Accuracies are on the pooled difficulty-probe test set, not public benchmark leaderboard scores. They are not comparable to published numbers.
• AUCCC, MDP-AUCCC, and router efficacy are my implementationof the paper's Equations 7 through 10. I did not have the exact equations. Relative comparisons between router, single, and oracle are robust; absolute values may not be directly comparable to theirs.
• 1024-token input truncation blinds the probe on long tasks, swe-bench especially.
• Cost granularity is per (model, benchmark) median, not per question.
• One encoder. A stronger encoder could move probe AUC and I have not tested that.
• There is a known minor bug in the Step 3 Platt path: out-of-fold probabilities use a different procedure than the final ensemble, so the logit map does not transfer and ECE degrades from 0.034 to 0.13. AUC is Platt-invariant so no conclusion changes, and the raw ensemble is already well calibrated at ECE 0.035. Logged, not hidden.

What I would do next

1. Un-truncate swe-bench. Part of that 0.303 gap is a 1024-token blindspot rather than a real limit.
2. Needle-question discrimination, probably via hard-question-weighted probe training. Though simpleqa-style factual recall may be near-irreducible from question activations, since the difficulty lives in the model's knowledge and not in the question.
3. Adaptive λ, from the live testing insight.
4. More coding and agentic questions to resolve the underpowered coding hint from Step 4b. More questions, not more models.

What I actually learned

The most useful thing I did in this project was write step2_pooled_bootstrap.py after an n=19 result told me what I wanted to hear. Everything downstream of that decision is trustworthy specifically because I stopped believing a good number.

The second most useful thing was making the failures cheap. Because activations are shared and I proved it with an assertion instead of assuming it, Steps 4 through 8 each cost minutes of numpy rather than hours of GPU. Five hypotheses tested and four refuted, which is only a viable way to work if refuting one is cheap. The infrastructure decision at the start, the cache key redesign, is what bought the scientific thoroughness at the end.

And the honest headline is that most of my results are negative. SharedTrunkNet does not lift. Calibration does not help. Pool diversity does not close the gap. What survives is a router that works well for reasons the paper does not fully explain, and a diagnosis of why it does not work better that points somewhere nobody was looking. I would rather have that than a confirmation I did not stress.