A 26M-parameter tagger beat our fine-tuned small LLMs at grammar correction, and runs in ~4 ms on a phone

TL;DR: For grammar correction inside a phone keyboard (on-device only, no network), we compared LoRA fine-tunes of Gemma 3 270M and SmolLM2 360M with a GECToR-style edit tagger we trained from scratch. The 26M tagger (v1.5) wins on accuracy (F0.5 0.799 vs 0.663) and changes correct text 2.4% of the time vs 20.1%. It ships as a 26 MB Core ML model (iOS) and a 27 MB int8 ONNX model (Android), and takes ~4 ms per sentence on a phone CPU. Weights are on Hugging Face under Apache-2.0: Mindscend/mindtype-tagger-v1.5.

We also hit a nasty training bug: longer training made the model worse. That story, with the fix, is in the second half.

Constraints

• Keyboard extension: ~60–80 MB of memory on iOS before jetsam kills it, no GPU guarantees, and it has to stay responsive while you type.
• No network, enforced by the OS: the iOS keyboard never requests Full Access, and the Android app has no INTERNET permission. All inference is local.
• Precision over recall. A false correction costs more user trust than a missed one, so we score with F0.5 and separately track the false-positive rate (correct sentences we changed) and lost protected spans (names, numbers, emails, URLs we altered).

The eval suite

There are 4,356 held-out cases, kept out of every training set by a deterministic hash split.

There's also a separate set of 2,000 casual messages that's never used for training. We use it only to choose between runs.

Baselines: small LLMs

LoRA r32 on the same correction data, served with llama.cpp:

Generative correction regenerates every token, so every token is a chance to change something that was fine. Fine-tuning reduces that a lot but doesn't remove it.

The tagger

Following GECToR (Omelianchuk et al., 2020), the model predicts one edit label per word instead of generating text:

He       $KEEP
don't    $REPLACE_doesn't
know     $KEEP
when     $APPEND_the
meeting  $KEEP
starts   $KEEP
.        $KEEP

The labels are applied, then the model runs again on its own output (up to 3 passes), so dependent fixes can chain.

• Architecture: encoder-only transformer, 8 layers, d_model 384, 6 heads, FFN 1536, max 160 tokens, 26.2M parameters. Our own 16k byte-level BPE tokenizer. 15,000 edit labels.
• Stage 1, pretraining: masked LM on FineWeb-Edu (~20M sentences, then 6.8 h more on 61.7M). Computing MLM logits only at masked positions is what let a 16k-vocab head fit on a 16 GB laptop GPU.
• Stage 2, correction: 12.02M (noisy, clean) pairs: synthetic phone typos (QWERTY neighbours, apostrophes, their/there/they're, agreement, irregular verbs, articles, "should of"…), real misspellings, dropped words, tense errors, CoEdIT, and casual messages written by a local Qwen2.5-14B and filtered for correctness. 30% are identity pairs, so $KEEP is learned as a real answer.
• Compute: one RTX 3080 Ti laptop GPU, ~9 h in total.

Inference-time guards

Three cheap tricks do a lot of the work:

1. KEEP bias. A constant added to the $KEEP logit. Bias 1 / 2 / 3 maps directly to the user-facing strictness setting ("Catch more" / "Balanced" / "Only sure fixes").
2. Protected spans. Inside a detected name, number, email or URL, only $KEEP is allowed. Result: 0.0% lost protected spans across the suite, for every tagger version.
3. Plausibility guard. A $REPLACE is accepted only if it's a known kind of fix (confusion sets, agreement pairs, irregular forms, articles) or a small edit distance. A context rule also keeps a preposition when the next token is capitalized, contains a digit, or is a value symbol ($ £ € # @).

Results

On the 2,000 held-out real-style messages, v1.5 changes correct text 1.9% of the time (v1.2: 2.8%). v1.5's real gain shows up on error types v1.2 barely handled:

Tense is still weak, and it's next.

The bug: longer training made the model worse

v1.3 (new error types) and v1.4 (a longer-pretrained encoder) both lost to v1.2. On correct text, false positives went from 3.1% to 7.4% and 9.1%. Here's what we found:

1. Unequal training length. Our fine-tune script turned a time budget into a step count by measuring speed at step 300. v1.2 had measured its speed during a VRAM spill (WSL shared-memory fallback), so it got a short schedule by accident. After we fixed the spill, v1.3 and v1.4 ran at full speed and got 2–3× the steps (2.6–2.9 epochs vs 1.4). The comparison wasn't fair, and the long runs were the bad ones.
2. Overtraining amplifies priors. The extra false positives were mostly preposition swaps next to names ("booked table 27 at Kyiv" → "in Kyiv"). The synthetic noise and the real GEC data both contain preposition fixes, and more epochs made the model trigger-happy.
3. The checkpoint picker rewarded the wrong thing. We chose checkpoints by edit F0.5 on synthetic validation pairs. That number kept going up (0.9655 → 0.9739) while false positives on real correct text also went up.
4. One template dominated a test set. preserve_clean has 500 cases from only 102 templates, and one template alone is 16% of it. A single quirk moved the headline false-positive rate by points.

The fix was a controlled re-run: fixed --steps, one change per run, plus a second seed to measure spread.

The new data was fine all along. The regression came from the extra epochs. Our process changes:

• Always train with a fixed step count, never a time budget.
• Choose checkpoints with a score that includes false positives on real clean text, measured on a dev set rather than the test suite.
• Train 2–3 seeds and pick on held-out real messages. At 30k steps, seed-to-seed spread in suite false positives was 2.4–3.1%.
• Report false positives per template, not only per case.

Shipping it

• iOS: Core ML, 8-bit weights, 26 MB, 99.96% label agreement with fp32. For the Neural Engine export we used an additive float attention mask, explicit attention and cumsum positions, with inputs padded to 16/32/64/96/160-token buckets. In the keyboard we use CPU-only compute units and read the fp16 output buffer in place, respecting strides, for ~4 ms per sentence in a release build. (Our early "it's slow" reports were a Debug build plus a 24 ms MLShapedArray conversion.) Peak keyboard memory on an iPhone 16 Pro Max was 81 MB with the model loaded.
• Android: ONNX opset 17, int8 dynamic quantization, 27 MB, 99.1% label agreement with fp32, ~4 ms per sentence on CPU.
• Swift and Python parity tests run on shared fixtures, so the phone and the training code produce identical labels.

Takeaways

• For correction (as opposed to rewriting), a small sequence tagger beats a small generative LLM: more precise, much smaller and much faster.
• Measure false positives on their own. F-scores hide the failure users notice most.
• Train to a fixed number of steps. A time budget quietly makes runs incomparable.
• A few lines of inference-time guards can be worth more than another training run.

Links

• Model (weights, tokenizer, labels, Core ML, ONNX, metrics; Apache-2.0): huggingface.co/Mindscend/mindtype-tagger-v1.5
• Rewrite model (SmolLM2-360M LoRA, GGUF): huggingface.co/Mindscend/mindtype-rewrite-360m-v2
• Try it now in Telegram: @MindType_Mindscend_bot (same model, server-side demo, nothing stored)
• Original post: mindscendinc.com/blog

Questions welcome in the comments, especially from anyone shipping ML models inside keyboard extensions.