Teaching Diffusion Models to Correct Their Own Work: Correction Training (arXiv 2610.14136)

Forum topic · 小凯 · 2026-10-22

Summary

A new paper from UC Berkeley and Reve, Correction Training (arXiv 2610.14136), addresses exposure bias in diffusion models: during training the denoiser sees real images plus noise, but at inference it consumes its own imperfect outputs, letting small errors snowball across sampling steps. The authors attach a corrector branch to a pretrained denoiser; the corrector learns to repair the denoiser's draft toward the ground-truth clean image, jointly fine-tuned for 50,000 steps with minimal overhead. At inference the corrector is applied repeatedly (up to six times) within each denoising step. On ImageNet at ~250 function evaluations, unguided FID drops from 9.58 to 2.07 for DiT-XL/2, 8.55 to 1.92 for SiT, and 5.53 to 1.31 for SiT-REPA, beating CFG guidance while achieving higher recall (0.603 vs 0.581), reducing oversaturation and AI-look artifacts. Controls confirm gains come from the draft-then-refine structure, not extra capacity, and benefits can be distilled back into the base model with no architecture change.

Teaching Diffusion Models to Correct Their Own Work

Diffusion models carry a built-in mismatch: during training, the denoiser only ever sees real photos plus noise. But at inference time, every step consumes the model's own previous output — imperfect images it never saw during training. One small error becomes the input for the next step, and flaws snowball down the sampling chain. This is exposure bias, and it is one source of the telltale "AI look."

A paper posted on October 9 by a UC Berkeley and Reve team (including Alexei Efros and pix2pix author Taesung Park) offers a strikingly simple fix: since the model was never taught how to fix its own mistakes, make error correction an explicit lesson. The paper is titled Correction Training (arXiv 2610.14136).

The method: a dual-branch structure

A corrector branch is attached alongside the pretrained denoiser. The denoiser predicts a clean image from noisy input; the corrector takes that "draft" as input and learns to repair it toward the ground-truth clean image. Both tasks share the same ground-truth target, and joint fine-tuning runs for 50,000 steps — a small additional cost compared to pretraining.

The real surprise comes at inference: the corrector was trained only on single-step corrections, yet it can be applied repeatedly — within the same denoising step, correct once, then again, up to six times. No one trained it to "revise its own revisions," but it keeps improving. Each step's error is absorbed on the spot instead of being cemented into structural artifacts later.

Results

On ImageNet, at roughly 250 function evaluations:

DiT-XL/2: unguided FID drops from 9.58 to 2.07
SiT: from 8.55 to 1.92
SiT-REPA (with representation alignment): from 5.53 to 1.31

CFG guidance yields FIDs of 2.51, 2.11, and lower for these models — the corrector beats all of them. Just as valuable is the diversity column: CFG trades distribution spread for quality, at the cost of oversaturation, neon colors, and an unmistakably AI look. The corrected model's recall is higher than CFG's (0.603 vs 0.581) — quality matches or exceeds CFG while images look less "plastic." Correction also stacks with CFG: correct first, then guide, further improving diversity and reducing oversaturation. Text-to-image settings (BLIP3o-based models) show gains on GenEval and DPG-Bench too.

Not just more compute

The paper's most rigorous move rules out a "brute force" explanation: merging the denoiser and corrector layers into a single, deeper denoiser with the same total layer count and inference budget produces a model that ties with single-pass correction, while repeated correction clearly wins. The gains come from the draft-then-refine structure itself, not from extra parameters.

Deployment is also covered: the correction gains can be distilled back into the base model, so a student model learns to "draw with corrective intuition" in a single forward pass — same architecture, same parameter count, same inference speed, no corrector attached.

Calibrating the claims

Two caveats on how this has been framed elsewhere: "first" is an interpreter's addition — the paper never claims to be first, only that it "introduces a correction step." "Fundamentally mitigates exposure bias" is also too strong — the paper hedges, saying exposure bias "may be one reason diffusion models need CFG." Still, the core conclusion stands: it fills a long-missing piece for unguided generation, making "drawing well without guidance" a practical path that can be distilled into weights. The new JiT architecture (Li & He, 2026) is also among the verified cases.

The most appealing aspect is not the FID tables but the transplant of an old writing principle into generative models: finish the draft, then reread it — the model learns not just to draw, but to revise.