В диффузионных моделях ухудшение обобщения начинается до запоминания данных

Авторы посмотрели на проблему с двух сторон. Сначала они обучали обычные U-Net — популярный тип нейросети для диффузионных моделей — предсказывать шум на чёрно-белых изображениях CelebA. Для каждого исходного изображения использовали несколько заранее зафиксированных вариантов шума, поэтому модель видела не просто набор картинок, а много пар «изображение плюс шум». Затем тот же вопрос проверили на упрощённой модели со случайными признаками: в ней первый слой случайно задают и больше не меняют, а обучают только второй. Такая конструкция позволяет получить формулы для кривых обучения и понять, откуда берётся ошибка.

Ключевая деталь — в диффузионном обучении у одного изображения бывает много шумовых вариантов. Поэтому точная подгонка всех обучающих примеров наступает намного позже, чем первые признаки ухудшения на новых данных. Авторы связывают это не с тем, что модель просто слишком «шумная», а с изменением того, что именно она оценивает: вместо общей закономерности начинает всё точнее воспроизводить особенности конечного набора изображений. Разбор ошибки на две части показал, что сначала растёт систематическая ошибка — модель всё хуже приближает закономерность, характерную для новых данных; после пика случайная составляющая уменьшается, но систематическая остаётся большой.

Это объясняет, почему увеличение числа вариантов шума не спасает от раннего ухудшения: оно отодвигает момент полной подгонки, но не начало роста ошибки. Практически авторы проверяли два способа не дать модели уйти в этот режим: штрафовали слишком большие веса в упрощённой модели и останавливали обучение U-Net раньше. В обоих случаях крупные модели использовали свою вместимость с пользой, тогда как без такого контроля увеличение размера не гарантировало лучшего результата.