mirror of
https://github.com/cactus-compute/needle.git
synced 2026-10-29 07:23:21 +08:00
123 lines
3.8 KiB
Plaintext
123 lines
3.8 KiB
Plaintext
Cloud Routing: Confidence Head for Hybrid Inference
|
|
====================================================
|
|
|
|
Supports main.txt section 2.4.
|
|
The student knows when it can't reproduce the teacher.
|
|
|
|
|
|
## 1. Problem
|
|
|
|
A distilled model will fail on queries outside its training
|
|
distribution or beyond its capacity. The question is: how does it
|
|
know when to defer to the cloud (the teacher or a comparable model)?
|
|
|
|
|
|
## 2. Design Options
|
|
|
|
Option A - Special token (<cloud> instead of <tool_call>):
|
|
Asks the decoder to learn meta-cognition through the same
|
|
autoregressive mechanism it uses for generation. Two very
|
|
different skills in one pathway.
|
|
|
|
Option B - Synthetic routing tool (route_to_cloud as a fake tool):
|
|
Same problem plus pollutes the tool namespace.
|
|
|
|
Option C - Confidence head on encoder output (ours):
|
|
Separate head predicts confidence BEFORE the decoder runs.
|
|
Dedicated head, dedicated signal (perplexity). When confidence
|
|
is low, skip decoding entirely.
|
|
|
|
|
|
## 3. Architecture
|
|
|
|
Small MLP on the encoder output:
|
|
|
|
encoder_output (B, T, d_model)
|
|
-> mean_pool -> (B, d_model)
|
|
-> Dense(d_model // 4) -> ReLU -> Dense(1) -> sigmoid
|
|
-> confidence in [0, 1]
|
|
|
|
Same encoder output that feeds the decoder. When confidence < threshold,
|
|
skip the decoder entirely and route to cloud. Encoder compute is not
|
|
wasted - it was needed for the confidence prediction anyway.
|
|
|
|
MLP (not linear) because the encoder has no FFN - same argument as
|
|
the contrastive head (see tool-selection.txt section 3).
|
|
|
|
|
|
## 4. Training Signal
|
|
|
|
The model already tells you what it can't handle. High decoder
|
|
perplexity = struggling. Low perplexity = confident. Free,
|
|
automatic labeling - no manual "should route" annotations needed.
|
|
|
|
Sigmoid mapping calibrated to the model's own PPL distribution:
|
|
|
|
confidence = sigmoid(-k * (log(ppl) - log(mu)))
|
|
|
|
mu = median perplexity on training set
|
|
k = steepness (default: 3.0)
|
|
|
|
|
|
## 5. Post-Hoc Training
|
|
|
|
Trained AFTER the main model, not jointly:
|
|
|
|
1. Train full encoder-decoder as normal
|
|
2. Freeze encoder + decoder
|
|
3. Forward pass: compute per-example decoder PPL
|
|
4. Map PPL -> confidence labels via sigmoid calibration
|
|
5. Train only confidence head (2 Dense layers) with MSE loss
|
|
|
|
Why not end-to-end:
|
|
(a) Early-training PPL is noisy (reflects learning progress, not
|
|
true difficulty).
|
|
(b) Jointly-trained head would bias encoder gradients toward easy
|
|
examples.
|
|
(c) Simplicity: retrain head without touching main model.
|
|
|
|
|
|
## 6. Deployment
|
|
|
|
The threshold is a runtime knob:
|
|
threshold=0.4: aggressive routing, save device compute
|
|
threshold=0.8: prefer local, only route when very uncertain
|
|
threshold=0.0: never route (offline mode)
|
|
|
|
Continuous confidence lets the deployer choose their operating point
|
|
on the accuracy-cost curve.
|
|
|
|
When routing, the decoder never runs. Encoder is the cheaper half
|
|
(no causal masking, no autoregressive loop). For 30% routing rate,
|
|
~18% net compute reduction on-device.
|
|
|
|
|
|
## 7. Relationship to Contrastive Head
|
|
|
|
Both heads operate on mean-pooled encoder output:
|
|
|
|
Contrastive head: "which tools are relevant?" (filters WHAT)
|
|
Confidence head: "can I handle this locally?" (decides WHETHER)
|
|
|
|
Full pipeline:
|
|
1. Encode query + tools
|
|
2. Contrastive: retrieve top-k tools
|
|
3. Re-encode with focused context
|
|
4. Confidence: can I handle this?
|
|
5. If yes -> decode locally. If no -> route to cloud.
|
|
|
|
|
|
## 8. Ablations (maps to main.txt 4.5)
|
|
|
|
A. Calibration Quality [MUST-HAVE]
|
|
Reliability diagram, accuracy per PPL bucket.
|
|
|
|
B. Routing Accuracy Curve [MUST-HAVE]
|
|
Sweep threshold 0.1-0.9, measure % local and local accuracy.
|
|
|
|
C. Compute Savings [STRONG]
|
|
Wall-clock encoder-only vs full pass at each threshold.
|
|
|
|
D. vs Decoder-Based Routing [STRONG]
|
|
Compare to Option B (route_to_cloud as synthetic tool).
|