A 5K-Parameter Router That Reads Your Intent Before Spending Your Money
Serban Mihai / 15 July 2026
~3 min read
I run an AI gateway between my coding agents and a pool of LLM providers. It handles health checks, failover, and model resolution. But it had one blind spot: it never looked at the content of what I was asking.
Every request was routed based on whatever model the agent picked. That works, but it wastes money. An exploration task β "search the codebase for User model references" β doesn't need a reasoning-heavy model at $0.56/M output tokens. A cheaper model at 20Γ less would do fine.
So I built a neural router that reads the first user message and classifies it into one of four task types: explore, plan, build, or quick. Based on the task type, the gateway picks the right model pool and reasoning effort. The classifier: a 5,120-parameter linear layer on top of Qwen-0.6B.
The Architecture
The router is a sidecar service. The gateway calls it over HTTP, falls back to the default model if it's unreachable, and opens a circuit breaker after the first failure. If the router dies, the gateway doesn't care β it just falls back to the safe default.
The head design is borrowed from TRINITY (Xu et al., ICLR 2026): a single weight matrix with no bias and no activation. At inference, the encoder runs the user message through Qwen3-0.6B, extracts a 1024-dimensional hidden state from the final layer (mean-pooled across all tokens, L2-normalized), and multiplies by the head. That's it.
The output maps directly into the gateway's routing config β task type selects the model pool, reasoning effort selects the reasoning depth.
Training Data Came From My Own Sessions
I didn't have to collect new data. opencode stores every session in a SQLite database β 18,000+ messages across 500 sessions. Each session has an agent field: the task type opencode already inferred. That was my label.
I mapped the existing agent labels to my four task types. The quick category was carved out heuristically from build sessions β anything containing "commit", "git push", "bash", "run" got labeled quick instead. 809 labeled samples total.
Training: From 47% to 79%
First pass was disappointing. Penultimate-token encoding + SGD got stuck at 47.5% task accuracy. The hidden states for different task types were almost indistinguishable β all coding prompts look similar under cosine similarity.
Two changes fixed it:
Mean pooling instead of penultimate token. The penultimate token only captures the last word's context. Mean-pooling across all tokens carries the full query semantics. The TRINITY paper uses penultimate for multi-turn transcripts where the last word is the role signal β but for single-message classification, mean pooling works better.
Adam with class-weighted loss. The build class dominated (383 samples vs. 53 explore). Without weights, the head leaned toward build. With inverse-frequency class weights and Adam at lr=0.001, convergence was smooth β 78.6% task accuracy and 91.2% reasoning accuracy at epoch 200.
Inference on Apple Silicon
Qwen3-0.6B runs on MPS. The entire stack β encoder + head β is a single Python process with FastAPI. Latency on an M3 Max: 96ms warm, 624ms cold start. The head weights file is 30KB.
The plan/build boundary is the hardest classification problem β "design system architecture" and "write a function" both look like structured tasks with similar vocabulary. The head leans toward build for anything code-related, which is the safe default: a coding task can't be ruined by too much capability.
What's Next
Deployment-level optimization. Right now the router picks a model pool. The next step is per-deployment scoring within a pool β the cheapest available provider that's healthy and close to the latency target gets the request. This is where CMA-ES (the TRINITY paper's training algorithm) comes in, optimizing for quality - Ξ» Γ cost.
Self-improving C-A-F loop. The Agent-as-a-Router paper (Zhou et al., 2026) describes a Context-Action-Feedback loop where the router learns from every outcome. My gateway already logs usage events to Postgres β token counts, latency, served deployment, cache hits. Adding a Verifier (execution success) and a Memory module (vector store of past routing decisions) turns the router from a static classifier into an online learner.
Avoiding routing collapse. The When Routing Collapses paper (Lai & Ye, 2026) shows a pervasive failure mode: routers default to the most expensive model as the cost budget increases, even when cheaper models suffice. Their fix β learning rankings instead of scalar scores β is something I'll adopt when moving to deployment-level routing.
Bottom Line
A 5,120-parameter linear layer trained on my own 809 coding sessions classifies requests with enough accuracy to route cheap models to exploration tasks. The entire router runs locally on Apple Silicon in under 100ms. If the router dies, the gateway doesn't care β it just falls back to the safe default.