Grammar correction with a diffusion LoRA on a frozen model

A small adapter turns a frozen left-to-right model into a one-pass, both-directions corrector.

This started as a small test, prompted by the recent diffusion language models — a diffusion version of Gemma among them. A diffusion model predicts every position of a sequence from both directions at once, instead of one token at a time, left to right. The question was where that bidirectionality actually helps. Grammar correction is a clean case: correcting a word usually depends on the words that come after it — fixing the article in "a apple" means looking ahead to "apple", and agreement and tense hinge on downstream cues. An autoregressive model predicts each token from the ones before it, so it can't use the right-hand context a correction needs; a diffusion model, re-predicting every position from the whole sentence at once, can.

The experiment: instead of training a diffusion model from scratch, add a small LoRA adapter to a frozen autoregressive model and run it bidirectionally. The adapter (rank 64 on a frozen ~335M base) is trained with a diffusion objective specialized to grammar — a substitution-heavy noise kernel, plus "copy supervision" that teaches it to leave correct words alone. The result is a one-pass bidirectional corrector built on a model that was never trained to be one. Try it below.

Caveat. This is a small, deliberately under-trained prototype, built to test whether the approach works at all — not to be accurate. It misses errors and sometimes introduces new ones. Treat the outputs as illustrative of the method, not as a usable corrector.

Sentence to correct

Or try one

He do not know which of the …I seen a elephant and a owl …She have a apple on she desk…These books is more interest…We was waiting for they to a…

Corrected

Input

Method

The (possibly-erroneous) sentence is the input. Bidirectional attention predicts every position from the entire sentence at once; the argmax at each position is the corrected token. It is one forward pass regardless of how many errors are present, and each token is decided with its right-hand context in view — which a left-to-right model cannot do.

Only the LoRA adapter and the norm gains are trained; the base model’s weights stay frozen. Training uses a grammar-specific diffusion kernel — mostly realistic substitutions, some masking — and supervises the clean positions as well as the corrupted ones, so the model learns to keep what is already right instead of rewriting the sentence. Behaviour is approximate: it can paraphrase or miss a named entity.