YYash Rahurikar

Part 3 of 3 · Decision models in a real AI product

I taught a free, local model to make my AI's decisions as well as a paid API

On a laptop, for $0.25 of API calls, in two and a half hours.

TL;DR

  • My AI assistant makes a handful of small decisions on every message: which specialist should answer, does the user want an image, does this need a web search? A hosted "decision model" called Jev was the most accurate option I tested, but each call takes ~400 ms over the network.
  • So I used Jev as a teacher: it labeled 42,488 decisions on messages like my users', for $0.25. Then I fine-tuned a small open model, Laya (421M parameters), to copy it, on a laptop GPU in 2.5 hours.
  • The student matches the teacher on routing (89.1% vs 88.2%, no significant difference), edges ahead of it on tool-intent detection (not statistically significant), and answers in ~17 ms instead of ~415 ms on a GPU, for $0 per decision.
  • The catches: it's only that fast on a GPU (~570 ms on CPU), it changed a lot of its answers on tasks I didn't train it on, and it's weaker where I gave it little data.

The problem: a great decider that's too slow

At Nands, we've been building an internal multi-agent assistant. Before any agent writes a single word, the system makes several quick decisions about the message:

  • Routing: should the general assistant, the coding assistant, research, visual (images, figures, slides) or the website builder handle it?
  • Tool shortcuts: is this a request for an image? A web search? Something about the user's own uploaded documents?
  • Build type: if it's a website request, is it a single page, a multi-page site, or an interactive app?

Today those decisions are made by a mix of regular expressions, an embedding-based router and small LLM calls. In the first two parts of this series I benchmarked them against Jev, a new kind of model from TypeSafe AI that doesn't generate text at all. You give it a question and a fixed list of answers, and it returns the answer plus an honest probability. On these routing decisions, it was the most accurate option I tested:

  • it caught 92% of image requests, where my regex caught 31% ("draw me a cat" slipped past the regex because it never says "image");
  • it routed 97% of trick questions correctly, vs 53% for my router. "What's the best free software for 3D modeling?" is a question about 3D, not a request to make something; my router sent it to the visual agent anyway.

But each Jev call is ~400 ms over the network. For a decision that sits in front of every chat message, that's too slow: I'd set myself a 300 ms budget, and Jev missed it. Part 2 ended with a compromise — keep the fast router and ask Jev only about the messages the router is unsure of. This post is about the more ambitious option: a model with Jev's judgement that's fast enough to ask about everything.

I'd also tested Laya, an open-source model in the same family. It's small enough to run on my own hardware and answered in ~16 ms, but out of the box it was worse than my existing code on almost everything. Then I read its model card more carefully:

"This is a specialist… fine-tuned on four specific synthetic workflows [invoices, security incidents, customer service, agent traces]. Expect it to behave like the base checkpoint, or worse, on anything else."

And the detail that gave me the idea: that specialist had been trained on Jev's answers. Laya's authors had used Jev as a teacher for their tasks. Nothing stopped me doing the same for mine.

What "fine-tuning" and "distillation" actually mean

If you're new to this (I was), here's the short version.

A model like Laya is a very long list of numbers (weights), 421 million of them. Text goes in, gets multiplied through those numbers, and a score comes out for each possible answer. Everything the model "knows" lives in those numbers.

Fine-tuning means continuing to adjust those numbers using your examples. The loop is simple, even if it runs thousands of times:

1. show the model a small batch of examples            "which assistant should handle: 'build me a pomodoro timer'?"
2. let it answer                                       Laya: build 0.40 · coding 0.35 · general 0.25
3. compare with the answer I want                     Jev:  build 0.97 · coding 0.02 · general 0.01
4. measure how wrong it was (the "loss")               one number: big when far off, near 0 when it matches
5. nudge every weight slightly in the direction that would have made the loss smaller

Distillation is fine-tuning where the "answer you want" comes from a bigger, better model: the teacher. The student doesn't just learn the teacher's final pick; it learns the teacher's whole probability spread. "Build 0.55 / coding 0.45" tells the student this one is borderline in a way that just "build" never could, so the student learns both what to answer and how sure to be.

The analogy I kept coming back to: a smart new hire (the pre-trained model) shadowing your most experienced colleague (the teacher) on thousands of real tickets. They don't become a different person; they learn how this team decides.

Step 1: 4,015 messages that look like my users'

A student can only learn what it sees, so I needed a lot of realistic messages covering every decision. I had an LLM (gpt-oss:120b) write 4,015 messages across 16 kinds: coding questions, research requests, image requests without the word "image", website requests, questions about "the PDF I uploaded", news headlines for tagging, and deliberately messy ones (typos, Hinglish, one-word messages, three requests at once). Plus 300 trick questions that mention a topic without asking for the action.

Two rules I was strict about:

  1. No labels at this stage. The kinds only steered variety. If a "general" message was really a research question, that didn't matter, because the teacher labels everything next.
  2. No peeking at the exam. I had a frozen test set from earlier in the experiment. Any generated message too similar in meaning to a test item (cosine similarity ≥ 0.90 between embeddings) was thrown out, so the student could never be trained on near-copies of its own exam questions. One message was dropped.

Step 2: the teacher labels everything, for a quarter

Each message went to Jev once, with all its questions in a single call (routing, nine tool intents and build type, plus plan-first and news category where relevant). I'd measured earlier that asking Jev ten questions in one call gives the same answers as asking them one at a time, at about a third of the price, because the message is only billed once.

Jev calls 5,368
Labeled decisions 42,488
Input tokens 5.9 M
Cost $0.25

That's the entire API bill for the training data.

Step 3: training on an 8 GB laptop GPU without melting it

I reused the training recipe the Laya authors published (their notebook targets two free Kaggle GPUs). It combines two things:

  • Soft cross-entropy: "make your probability spread look like the teacher's." This is classic distillation and does most of the work.
  • A small reinforcement-learning term: the model tries slightly randomised versions of its answer, and each is scored by a proper scoring rule (a reward that's only maximised by reporting honest probabilities). This keeps its confidence well calibrated.

I held back 10% of the examples to calibrate the model's confidence afterwards, and trained for 4 passes over the data.

The first attempt ran out of memory. Training needs far more memory than just using a model: besides the weights, the optimizer (AdamW) keeps two extra numbers for every one of the 421 million weights. On an 8 GB laptop GPU that doesn't fit. The fix was an 8-bit optimizer, which stores that bookkeeping at a quarter of the size. Peak memory dropped to 4.8 GB.

The second concern was the laptop itself. Two-plus hours at full GPU load is a good way to freeze a desktop or cook a laptop, so I added guard rails. They cost about an hour, and I think it was worth it:

Risk What I did
Desktop becomes unusable Ran as a background service at the lowest CPU priority, capped at 4 of 24 cores, a 12 GB RAM ceiling, and a 30 ms pause after every batch
GPU overheats Checked the temperature regularly: pause at 83 °C, resume at 75 °C. (It paused 157 times and peaked at 90 °C)
Laptop sleeps mid-run Blocked sleep while the job ran
Crash or reboot loses hours of work Saved the complete training state (weights, optimizer, position in the data, random seeds) every 10 minutes, atomically

I didn't trust the crash recovery until I'd tested it: train, "crash" at step 7, resume with the same command, and compare against a run that was never interrupted. The final weights were bit-for-bit identical.

Then it ran: 150.8 minutes, no errors, no restarts. The distill loss (how far the student's answers are from the teacher's) fell pass by pass:

Pass 1 2 3 4
Average distill loss 0.249 0.182 0.142 0.121

Results: same exam, same rules

The fine-tuned model took exactly the same frozen test set as every other contestant in this experiment, under the same rules (question wording and confidence settings chosen on a practice split, scores reported on the held-out split, once).

Decision My existing code Jev (teacher) Laya, off the shelf Laya, fine-tuned
Routing accuracy 76.3% 88.2% 74.4% 89.1%
Routing, trick questions 52.9% 97.1% 41.2% 91.2%
Tool-intent detection (macro-F1) 0.675 0.711 0.350 0.812
… image requests / web-search requests (F1) 0.47 / 0.17 0.73 / 0.88 0.51 / 0.63 0.87 / 0.88
Website build type 94.5% 90.9% 78.2% 94.5%
News category (agreement with editors) – 85.7% 51.6% 76.9%
Time per decision 0–100 ms ~415 ms ~16 ms ~17 ms (GPU)
Cost per decision $0 ~$0.00002 $0 $0

On routing the student is 0.9 points behind the teacher, 95% CI [−2.8, +4.7]: statistically no difference. It's 12.8 points ahead of my existing router (p = 0.001) and 14.7 ahead of off-the-shelf Laya. Its confidence stayed honest too (calibration error 0.04).

It edged ahead of its teacher on tool intents (0.81 vs 0.71, though not statistically significant). That surprised me, but it's plausible: the student learns Jev's average judgement over tens of thousands of examples, which smooths out some of Jev's one-off mistakes. I'd still be careful here, because some intents have only a handful of positive examples in the test set, and the exact-match difference vs Jev isn't statistically significant.

It's weaker on news tagging (77% vs 86%), and that one is easy to explain: news headlines were only 302 of the 42,488 training decisions. The student is only as good as the examples it sees for each task.

The honest part

1. "Fast" means "fast on a GPU"

On the laptop GPU it answers in 17 ms. On a CPU with 4 threads, which is what most application servers have, it took ~570 ms, slower than calling Jev. So this isn't a drop-in production replacement yet. The options are a small GPU host, or making it faster on CPU (exporting to ONNX, 8-bit quantisation, or the smaller 322M-parameter Laya). I haven't tested those yet, and until I do, the speed claim comes with an asterisk.

2. It changed a lot of answers on things I didn't teach it

Before training I set four targets. One was that on questions I never trained on (a sample from a public benchmark), fewer than 20% of the model's answers should change. 42% changed. I count that as a miss, because that's the point of writing targets down in advance.

Look closer, though, and the changes weren't harmful: accuracy on those untrained questions went up (29% → 37%). My target measured change when it should have measured damage. I'd write that target differently next time ("accuracy on untrained tasks must not fall by more than X") and report the change rate only as a diagnostic.

3. The exam isn't real traffic (yet)

My test set is ~150 prompts I wrote plus ~150 synthetic ones, labeled by two independent LLM annotators with disagreements settled by hand. I haven't used real user messages, pending the right consent. A human spot-check of the labels is still pending. The comparison between contestants is fair, since everyone took the same exam, but the absolute numbers are provisional.

What I learned

  1. A decision model you can't afford to call on every message can still teach one you can. Twenty-five cents of teacher labels bought a model that makes the same routing decisions locally.
  2. Data volume per task is the lever. Where the student saw thousands of examples it matched or beat the teacher; where it saw 302, it didn't.
  3. Pre-register your targets, and choose them carefully. My drift target measured the wrong thing. I only know that because I wrote it down before seeing the result.
  4. Laptop training is practical with guard rails. An 8-bit optimizer made it fit, and throttling, a temperature guard and tested resumable checkpoints made it safe. It cost an hour of extra runtime, and the laptop stayed usable throughout.
  5. Hardware decides deployability. A 17 ms model on a GPU is a 570 ms model on a CPU. Plan where it will run before planning around its speed.

What's next

  • Make it fast on CPU (ONNX + quantisation), because that decides whether it can ship.
  • More training data for news tagging to close the gap with Jev.
  • A real-traffic test, once I can use real messages.
  • If those hold up: put it behind a feature flag as the router's second opinion on the messages my current router is unsure about.

Appendix: setup at a glance

Teacher Jev jev-1.13 (TypeSafe AI) via OpenRouter
Student Laya base (convaiinnovations/laya, 421M, ModernBERT-large encoder + decision head), pinned revision 55cf4c4
Training data 4,015 generated messages → 42,488 teacher-labeled decisions (routing, 9 intents, build type, plan-first, news category)
Recipe Soft cross-entropy + proper-scoring RL (the Laya authors' RLCD recipe), 4 epochs, LR 2.5e-5 encoder / 1e-4 head, effective batch 32, bf16
Deviations from the authors' recipe 8-bit AdamW (memory); micro-batch 4 × accumulation 8 instead of 8 × 4; bf16 instead of fp16
Hardware Laptop RTX 5070 (8 GB), 4.8 GB peak, 150.8 min with thermal throttling
Evaluation Frozen test set from earlier in the experiment; wording and temperature chosen on a practice split; McNemar mid-p and paired bootstrap for significance
Total API cost $0.25 for training labels (the whole experiment: $0.65)

Parts 1 and 2 of this series cover Jev in a RAG pipeline (reranking and filtering) and where decision models earn their place versus a plain regex.

More writing

All writing