[Статья][Кода нет, а и TPUv5 у вас]
Введение
Есть такие scaling laws, согласно которым, чем больше модель и чем больше количество данных, на которых она училась - тем лучше ее качество или некий эрзац качества. Для LLM уже вышло множество работ по данной теме и выработаны compute-optimal рецепты обучения, но для диффузионных моделей систематического исследования до сих по не проводилось.
И ребята с Google вооруженные TPU решили заполнить зияющую пустоту ⚫️.
Эксперименты
В качестве базовой модели берут UNet от Stable-Diffusion v1.5 (866M) и меняют его ширину (глубина постоянная). И таким образом получают семейство моделей от 39М до 5B параметров.
Модели обучают на проприетарном датасете из 600M пар изображений и описаний к ним, отфильтрованном по эстетике.
Базовые настройки сэмплера -
DDIM солвер с 50 шагами и cfg_scale=7.5.Исследуют следующие задачи:
1️⃣ text-2-image генерация
2️⃣ Superresolution через дообучение модели из 1) следуя постановке из StableSR.
3️⃣ Dreambooth поверх модели из 1)
1️⃣ Для оценки качества используют по традиции FID и CLIP. В области некоторого среднего бюджета обучения модели разных размеров выдают примерно одни и те же метрики. Но при наличии значительного объёма вычислительных ресурсов - размер начинает играть значение, и самая большая модель побеждает своих меньших собратьев. Для 5B модели пришлось прибегнуть к memory-efficient training, так как при стандартном обучении с
Adam+amp модель с состояниями оптимизатора не влезает в TPUv5.Совсем маленькие модели генерируют дефектные картинки 🤪, но начиная с нескольких сот миллионов параметров качество генераций моделей примерно на одном уровне по ощущениям.
2️⃣ На downstream задачах размер моделей играет уже более существенную роль. В 4x image super-resolution большие модели даже при заданном бюджете обучения заметно выигрывают по FID. Однако по LPIPS модели разного размера с фиксированным training compute сравниваются.
3️⃣ Большие модели ожидаемо выдают более качественные DreamBooth генерации.