One encoder pass, calibrated probability for every option: open-source 149M WaterSheep

2026-10-02

A ModernBERT decision head gives calibrated per-option probabilities in one pass; 811k training decisions, 61.2% on held-out datasets, 1.2 ms per question, ~100x faster than a 4B LLM.

What problem this solves

Software keeps hitting small text decisions: which team gets this ticket, does it need a human, is this input a prompt injection, how angry is the customer, which required fields are missing. The two standard answers both hurt. A fine-tuned classifier per task means collecting labels and maintaining N sets of weights. An LLM call means one generation per question, hundreds of milliseconds, and a verdict whose confidence is rarely calibrated, so downstream threshold logic becomes guesswork. WaterSheep compresses these jobs into four typed decisions, yes/no, single choice, rating and multi-label, and answers them with one 149M-parameter encoder pass that emits a calibrated probability for every option.

Method

The trick is to score every option inside the same forward pass instead of asking one question at a time.

Training data comes from two places. First, 109 public datasets (MMLU, goemotions, RewardBench, prompt-injection detection and more) converted into a uniform decision format: 811k train, 25k validation, 25k test, plus 14k questions from datasets never seen in training. Second, synthetic data from Qwen3.5-4B across 57 task families, 40 industries, 12 writing styles and 3 difficulty levels. The same model then blind-verifies each example: it re-answers with options in shuffled orders, averaged to cancel position bias, and gets one yes/no question per option for multi-label; disagreement with the intended label kills the example. Of 48k generated examples, 30k survived (62.5%); 30% died on teacher disagreement.

Results

SettingAccuracyECE
In-distribution test (25k questions)77.8%0.026
Held-out datasets (14k, unseen in training)61.2%0.043
Mean over 16 external benchmarks80.2% (majority 51.8%)0.109

The LLM comparison covers two benchmarks (teacher Qwen3.5-4B, same GPU). Prompt injection: 91.4% vs 64.7%, ECE 0.079 vs 0.258. XSTest safety: 73.6% vs 88.4%, and the 4B model wins that one. The speed gap is stable: 1.2 to 5.2 ms per question against 202 to 450 ms, an 86 to 171x speedup, because any number of options costs one encoder pass and no generation.

Weak spots are visible. Goemotions, 27 emotion classes, lands at 22.4%. Inside LegalBench, Abercrombie terminology hits 21.1% and privacy-policy QA 58.9%. Fine-grained emotion and legal reasoning are too much for a 149M encoder.

Why it matters

High-frequency decision points in agent pipelines, routing, gating, safety checks, downgrade calls, need cheap, fast and honest confidence, not frontier reasoning. An LLM there is waste; a fleet of task-specific classifiers is maintenance drag. WaterSheep sits in between: one set of weights for all question types, probabilities that plug straight into threshold logic, an ONNX build that runs in the browser with a no-install JS runtime, Apache 2.0 throughout. A single author trained it on one consumer GPU and exported the model the same day (2026-09-28); the whole pipeline is scripted and re-runnable. As a template for small models taking over the long tail of decisions, this is a usable reference.

Limitations

From the author: English only; inputs truncated at 512 tokens; rating questions are the weakest type; calibration holds on data resembling the training distribution; not to be used alone for medical, legal, financial or hiring decisions. Additional concerns from this read: the held-out drop to 61.2% is a 16.6-point generalization cost, untested on production distributions; the generator, verifier and baseline of the synthetic data are the same Qwen3.5-4B, so the data tilts toward what a 4B model gets right; there is no comparison against nearer rivals such as SetFit or NLI-style zero-shot classification; and the paper is a single-author ResearchGate preprint that has not been reviewed (this writeup draws on the author's model card, full code and released result data, because the preprint PDF itself sits behind ResearchGate's bot wall). Days after release the model counts 28 downloads, so maturity is early.

Terms

Source

What people are saying

All paper explainers