[KongchangAI]
· 3 min read· 1,749 words

26M Small Model Beats Fine-Tuned LLMs: 4ms Grammar Correction on Mobile with 0.799 Accuracy

26M Small Model Beats Fine-Tuned LLMs: 4ms Grammar Correction on Mobile with 0.799 Accuracy

26M GECToR tagger beats fine-tuned small LLMs in mobile grammar correction with 0.799 F0.5 and 2.4% false-edit rate.

The Mindscend team compared fine-tuned small language models with a from-scratch edit tagger for mobile keyboard grammar correction. Their 26.2M-parameter GECToR-style tagger achieves F0.5 of 0.799 with only 2.4% false-edit rate, defeating SmolLM2 360M LoRA (0.663 F0.5, 20.1% false edits) and zero-shot Gemma 3 270M. Three inference-time guardrails achieve 0.0% protected-span loss. iOS Core ML and Android ONNX models run at ~4ms per sentence.

Mindscend 26M Parameter Grammar Correction Model

Should grammar correction in a mobile keyboard use a fine-tuned small language model, or an edit tagger trained from scratch? The Mindscend team answered this with a detailed set of comparison experiments: a GECToR-style tagger with just 26.2M parameters comprehensively beats LoRA fine-tuned versions of Gemma 3 270M and SmolLM2 360M in both accuracy and false-edit rate, while running inference in only about 4ms per sentence.

Why Not Use an LLM for Grammar Correction?

Building grammar correction into a mobile keyboard extension comes with far stricter constraints than one might imagine. The memory limit for iOS keyboard extensions is around 60-80MB — exceed it and the system kills the process. The keyboard doesn't need GPU acceleration, but it must always remain responsive. More critically, the entire inference process runs completely offline — the iOS side never requests Full Access, and the Android side doesn't even have the INTERNET permission.

Under these constraints, the team also faced a core metric trade-off: precision matters far more than recall for grammar correction. Incorrectly changing a correct sentence damages user trust far more than missing an error. So they used F0.5 as the primary metric while separately tracking two key figures: the false-edit rate (the proportion of correct sentences that get changed incorrectly) and the protected-span loss rate (the proportion of names, numbers, emails, and URLs that get altered).

F0.5 score is a variant of the F-score that gives precision a higher weight than recall (β=0.5). The general formula for Fβ is (1+β²)·P·R/(β²·P+R). When β=0.5, precision is weighted four times as heavily as recall. In scenarios like grammar correction, incorrectly changing a correct sentence harms user experience far more than missing an error, so the team chose F0.5 rather than the standard F1 to amplify the importance of precision in evaluation, ensuring the model doesn't sacrifice accuracy just to catch more errors.

Technical Approach of the 26M-Parameter Tagger

The team's core idea comes from GECToR (Omelianchuk et al., 2020) — instead of generating text, the model predicts an edit tag for each word. For example, He don't know would be tagged as KEEP / REPLACE_doesn't / KEEP. After the tags are applied, the model processes its own output, iterating up to 3 rounds to support dependency-chain corrections.

The architecture choices are quite restrained: an 8-layer encoder-only Transformer, d_model 384, 6 attention heads, FFN dimension 1536, max 160 tokens, totaling 26.2M parameters. A custom 16k byte-level BPE tokenizer with 15,000 edit tags. Training is done in two stages: first, masked language model pre-training on approximately 20 million sentences from FineWeb-Edu, then grammar correction fine-tuning on 12.02 million (noisy, clean) pairs — 30% of which are identity pairs, teaching the model to treat $KEEP as a genuine answer. The entire training was completed on a single RTX 3080 Ti laptop GPU, taking about 9 hours in total.

Byte-level BPE (Byte Pair Encoding) is a subword tokenization method that sits between character-level and word-level approaches. It starts from individual bytes and repeatedly merges the most frequent adjacent byte pairs in the corpus until reaching a preset vocabulary size. The key advantage of byte-level BPE is that it naturally supports any Unicode character, never has out-of-vocabulary (OOV) issues, and the vocabulary can be kept very compact. Unlike conventional character-based BPE, byte-level BPE operates directly on raw byte sequences, enabling the model to handle any language and symbol — particularly important for mobile keyboards that need to cover diverse input scenarios.

Measured Results: Small Tagger vs. Small LLM

The comparison is strikingly direct. SmolLM2 360M with LoRA fine-tuning achieves an F0.5 of 0.663 with a 20.1% false-edit rate, while the 26M tagger v1.5 reaches an F0.5 of 0.799 with only a 2.4% false-edit rate and 0.0% protected-span loss rate. Zero-shot Gemma 3 270M is essentially unusable — F0.5 of only 0.122, with 99.6% of correct sentences being altered.

The gap in false-edit rates stems from the mechanics of generative correction: the model regenerates every token, and every token is an opportunity to introduce an error. Fine-tuning can substantially reduce this tendency but cannot eliminate it.

The real improvements in v1.5 over v1.2 show up in specific error types: F0.5 for real spelling errors ("recieve") jumped from 0.882 to 0.963, and missing words ("I going home") surged from 0.078 to 0.698. However, tense errors ("Yesterday I go") remain a weak spot with an F0.5 of only 0.119, which the team says is the next target.

Inference-Time Guardrails: A Few Lines of Code Beat Another Training Run

Three lightweight inference-time techniques contributed substantial gains. KEEP bias — adding a constant to the $KEEP logit, where bias values of 1/2/3 directly correspond to user-visible strictness settings ("catch more"/"balanced"/"only change certainties"). Protected spans — upon detecting a name, number, email, or URL, only $KEEP is allowed, resulting in a 0.0% protected-span loss rate across the entire test set. Plausibility guard — $REPLACE is only accepted if it belongs to a known correction type or the edit distance is small enough; contextual rules also preserve prepositions when the next token starts with a capital letter, contains a digit, or is a symbol like $ £ € # @.

"Known correction types" in the plausibility guard refers to common categories of edit operations the model learned from training data, such as spelling corrections (recieve → receive) and subject-verb agreement (He don't → He doesn't). A $REPLACE tag is only accepted when the corresponding modification belongs to these predefined valid types, or when the edit distance (the number of character differences between the original and replacement word) is below a certain threshold. This prevents the model from making overly aggressive changes when uncertain, further reducing the false-edit rate.

A Counterintuitive Training Bug: The Longer You Train, the Worse It Gets

This is the most instructive part of the entire write-up. Versions v1.3 and v1.4 saw false-edit rates on correct text spike to 7.4% and 9.1% respectively, compared to only 3.1% for v1.2.

The root cause was unequal training lengths. The team's fine-tuning script converted a time budget into step counts based on the speed measured at step 300. v1.2 happened to get a slower speed measurement due to VRAM overflow (WSL shared memory fallback), which accidentally resulted in a shorter training schedule. After fixing the overflow, v1.3 and v1.4 ran at full speed, with actual training steps reaching 2-3x that of v1.2 (2.6-2.9 epochs vs. 1.4). The comparison itself was unfair, and the longer-trained versions happened to be the worse ones.

The extra epochs mainly amplified model priors — the surplus false edits were concentrated on preposition substitutions near names (e.g., "booked table 27 at Kyiv" changed to "in Kyiv"). Both synthetic noise and real GEC data contain preposition correction samples, and the longer the training, the more likely the model is to over-trigger. More insidiously, checkpoint selection used edit F0.5 on a synthetic validation set, which kept rising (0.9655 → 0.9739), but the false-edit rate on real correct text was rising in tandem.

The team's eventual fix was a controlled re-run: fixed step counts, changing only one variable at a time, training with 2-3 seeds to measure variance. The results showed that the new data itself was fine — the regression came purely from excess epochs. v1.5 achieved the best balance with a 30,000-step configuration: F0.5 of 0.799 and a 1.9% false-edit rate on real messages.

On-Device Deployment: 26MB Core ML, 4ms per Sentence

On iOS, the model is exported to Core ML with 8-bit weights, a model size of 26MB, and 99.96% label agreement with fp32. To adapt to the Neural Engine, the team used additive float attention masks, explicit attention, and cumsum positional encoding, with inputs bucket-padded to 16/32/64/96/160 tokens. In the keyboard, computation actually runs on the CPU unit, reading the fp16 output buffer in place and respecting strides, at about 4ms per sentence in release builds. The team specifically noted that early reports of "the model is slow" were actually due to Debug builds plus a 24ms MLShapedArray conversion overhead. On Android, they use ONNX opset 17 with int8 dynamic quantization, 27MB, 99.1% agreement with fp32, and likewise about 4ms per sentence on CPU. Swift and Python equivalence tests run on shared fixtures to ensure the phone and training code produce exactly the same tags.

Core ML is Apple's machine learning framework that allows developers to deploy trained models to iOS, macOS, and other devices, leveraging CPU, GPU, and the Neural Engine for hardware acceleration. 8-bit weights means quantizing model parameters from 32-bit floats to 8-bit integers, which can dramatically reduce model size and memory footprint while maintaining high inference accuracy. The Neural Engine is an accelerator in Apple chips designed specifically for neural network computation, but it requires models to satisfy specific operator constraints — hence the team used techniques like additive float attention masks to adapt.

Industry Takeaways

This work has several generalizable lessons for on-device AI deployment. For error-correction (rather than rewriting) tasks, small sequence taggers outperform small generative LLMs in precision, size, and speed. False-edit rates must be measured separately — F-scores mask the failure modes users care about most. Always fix step counts rather than time budgets during training, or comparability between runs is silently lost. Finally, a few lines of inference-time guardrail code may be worth more than another round of training.

Model weights are open-sourced on Hugging Face under the Apache-2.0 license, including the tokenizer, tags, Core ML and ONNX exports, and complete metrics.

One-line Summary

A 26M-parameter GECToR tagger comprehensively beats fine-tuned small LLMs in mobile keyboard grammar correction, achieving F0.5 of 0.799 with only a 2.4% false-edit rate.

Paragraph Summary

The Mindscend team compared fine-tuned small language models with a from-scratch edit tagger for mobile keyboard grammar correction. Under iOS keyboard extension constraints of 60-80MB memory and fully offline operation, precision matters more than recall, so they adopted F0.5 and separately tracked false-edit rates. Their 26.2M-parameter GECToR-style tagger (8-layer Transformer, d_model 384) achieves F0.5 of 0.799 with only 2.4% false-edit rate, while SmolLM2 360M LoRA gets 0.663 F0.5 with 20.1% false-edit rate, and zero-shot Gemma 3 270M is nearly unusable. Additionally, three inference-time guardrails (KEEP bias, protected spans, plausibility guard) achieve 0.0% protected-span loss. A training bug from converting time budgets to step counts caused extra epochs to amplify model priors, doubling the false-edit rate. For deployment, the iOS Core ML model is 26MB, the Android ONNX model is 27MB, and single-sentence inference takes about 4ms.

#grammar correction#GECToR#on-device AI#mobile keyboard#small language model#model quantization
Share:

Related articles