Naukowcy zaproponowali prosty ale skuteczny mechanizm correction training, który redukuje lukę między trenowaniem a wnioskowaniem w modelach dyfuzji. Problem exposure bias wynika z faktu, że model denoisera podczas treningu pracuje na rzeczywistych zaszumionych obrazach, ale przy generowaniu musi radzić sobie z przewidywaniami z poprzednich kroków denoizacji - objektami, które nigdy nie widział w zbiorze treningowym.
Rozwiązanie polega na dodaniu korekcyjnego modułu, który przyjmuje przewidywanie bazowego denoisera jako wejście i trenuje się na tym samym ground-truth. W praktyce oznacza to, że korektor musi nauczyć się poprawiać artefakty, które pojawią się w przewidywaniach z wcześniejszych etapów. Podczas generowania próbek, na każdym kroku denoizacji model iteracyjnie ulepśla przewidywanie czystej próbki. Co ciekawe, korektor można aplicować wielokrotnie do swoich własnych wyników podczas wnioskowania, co prowadzi do dalszych ulepszeń.
Metodę testowano na architekturach DiT, SiT i JiT, osiągając znaczne polepszenia w miarach jakości takich jak FID, zwłaszcza w generowaniu bezwarunkowym bez dodatkowych technik wspomagających jak Classifier-Free Guidance. Podejście zachowuje różnorodność próbek, co czyni je potencjalnie przydatnym rozwiązaniem do poprawy jakości generowania w różnych zastosowaniach bez konieczności dodatkowych kosztów obliczeniowych przy treningu.