Model soups failed until we fixed one line: our 26M grammar tagger from 11% to 1.9% false fixes

#machinelearning #nlp #ios #pytorch

Our keyboard's grammar model is a 26.2M-parameter tagger we trained from scratch: one label per word (KEEP / DELETE / REPLACE_x / APPEND_x), no text generation, Core ML on the phone. Over two weeks it went from F0.5 0.693 with 11.1% false positives to 0.803 with 1.9%, and from fixing 22% to 77% of phone typos exactly. Here's what moved the numbers, including one bug that made model soups fail.

Scoreboard (4,356-case held-out suite, plausibility guard)

Version	Change	F0.5	False positives
v1	from-scratch baseline	0.693	11.1%
v1.1	plausibility guard	0.748	8.7%
v1.2	better (noisy, clean) pairs	0.800	3.1%
v1.3 / v1.4	2–3× more steps (rejected)	0.784 / 0.758	7.4% / 9.1%
v1.5	real misspellings, dropped words, tense; fixed `--steps`	0.80	2.4%
v1.6	past-tense errors	0.802	2.8%
v1.9	+ phone-typo slips, keep bias 4	0.784	2.2%
v2.0	40% identity pairs, more chat text, 3-run soup, keep bias 3	0.803	1.9%

Phone-typo set: F0.5 0.689 (v1.6) → 0.912 (v2.0); exact 22% → 77%. Real messages changed (5,556 held out): 1.64% → 1.26%.

1. Fixed step count, not a time budget

We used to train for N minutes. v1.2 happened to measure its speed during a VRAM spill and got a short schedule; v1.3/v1.4 ran at full speed and got 2–3× the steps. They were worse: v1.4, on exactly v1.2's data, changed correct text 9.1% of the time vs 3.1%. Typical damage: booked table 27 at Kyiv → in Kyiv. The checkpoint picker scored synthetic validation pairs, which kept improving while real clean text got worse. Now every run is --steps N.

2. Keep bias: a free precision dial

At decode time we add a constant to the KEEP logit:

score = logits.copy()
score[:, KEEP] += keep_bias   # higher = fewer, surer edits
labels = score.argmax(-1)

Adding phone slips (fat-finger neighbours, short words, dropped g) lifted the phone-typo F0.5 from 0.689 to ~0.90 but raised false positives. Raising keep bias 2 → 4 fixed that without retraining. v2.0's sweet spot is 3.

3. Seed variance, and soups that only work with a shared head init

Four seeds of one v1.9 recipe ranged from 1.2% to 4.3% suite false positives at the same keep bias. So we tried a model soup (uniform weight average). Result: 31% false positives (57% at keep bias 2). Useless.

The cause: torch.manual_seed(args.seed) ran before building the model, so each seed drew its own random tag head. The encoder shares pretraining, but four unrelated heads don't average. Model soups assume the fine-tunes start from one point.

torch.manual_seed(args.init_seed)   # same head for every run
model = TaggerModel(config)
model.encoder.load_state_dict(pretrained)
torch.manual_seed(args.seed)        # data order and dropout differ

With a shared head init:

Model (keep bias 4)	Suite FP	Real-message FP	Phone F0.5
seed 7 / 8 / 9 alone	1.41% / 1.23% / 1.47%	1.28% / 1.21% / 1.31%	0.902 / 0.897 / 0.904
soup of the three	1.16%	1.28%	0.907
EMA (0.9995) of each, souped	1.23%	1.42%	0.906

Every seed passes on its own now, and the soup is better than any of them. EMA added nothing. v2.0 is the soup of three runs on data with more casual messages (keep bias 3: suite F0.5 0.803, FP 1.9%, real-message FP 1.26%, phone F0.5 0.912).

Still weak

Multi-word rewrites (CoEdIT): 11% exact. One-edit-per-word tagging isn't built for them.
~2.4% of correct formal sentences still get an edit, mostly articles.
English only.

DEV Community