PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis

Тех. репорт в продолжение темы латентной трансформерной диффузии, показывавший топовое (на момент выхода) качество генерации. Свои успехи авторы объясняют тремя факторами. Рассмотрим их по порядку.

Декомпозиция обучения. Во время обучения, модель должна научиться находить взаимосвязи между пикселями, связь текст-картинка и из всех возможных генераций порождать только эстетичные. Идея в том чтобы разбить обучение на стадии, в ходе которых модель учится каждой из подзадач отдельно.

Модель сначала учат на class-conditional генерацию ImageNet256, после чего дообучают на датасетах с хорошими (длинными, подробными) описаниями и эстетичными картинками соответственно, не меняя исходного разрешения. В конце есть еще две стадии файтнюнинга, в рамках которых продолжается обучение на эстетическом датасете, но разрешение повышают сначала до 512х512, а потом до 1024х1024. Занятно, что эстетический датасет на треть состоит из генераций MidJourney.

Обновленный DiT. В статье по Diffusion Transformer уже предлагалась архитектура на основе ViT с добавлением кондишенинга на время и класс генерируемого изображения. Авторы предлагают кондишениться на текст с помощью cross-attention, а также еще раз модифицируют многострадальный AdaLN слой, поскольку обуславливание на класс больше не актуально. Помимо этого, предлагается репараметризация послойных MLP для кодирования t на один общий MLP с небольшой модуляцией, что уменьшает эффективное число параметров, не раняя качество.

Улучшенный датасет. В очередной раз подчеркивается важность подробных описаний к картинкам и польза синтетических данных из SAM и LLaVA. Идея концептуально похожа на CapsFusion, хоть детали и не раскрываются.

Итоговая модель на основе DiT-XL/2 с T5-XXL текстовым энкодером и VAE из LDM обучается на 64 v100 в течение 26 дней. Модель генерит очень похоже на MidJourney, выдавая качественные картинки за 20 шагов DMP. У модели остаются проблемы с руками и текстами. Код есть.

Scaling Up to Excellence: Practicing Model Scaling for Photo-Realistic Image Restoration In the Wild

Свежая работа по использованию generative prior из предобученной SD. Концептуально, работа очень напоминает подход полугодичной свежести DiffBIR и отличается, в основном:
1. Размером базовой модели: авторы этой работы используют SDXL)
2. Размером датасета: используется проприетарный 20М датасет общего назначения + 70к картинок из FFHQ-raw для улучшения качества восстановления лиц
3. Файнтюнингом VAE в Real-ESRGAN постановке для увеличения устойчивости к частотным искажениям на инференсе
4. Использованием текстовых описаний (в том числе, negative) и своей модификации EDM семлпинга
5. Использованием уменьшенного варианта ControlNet для обуславливания на исходную картинку, который назвали ZeroSFT

С практической точки зрения важно, что основной прогресс в качестве достигается преимущественно за счет масштабирования модели и данных, а так же дообучения VAE с аугментациями, встречающимися в использовавшихся бенчмарках RealSR/DRealSR/Real47 (первые три пункта).

Вероятно, важной деталью является добавление в претрен датасет 100к заведомо плохих генераций SDXL с негативными промтами, позволяющим использовать гайденс, эффект от использования которого показан.

Однако, сама польза использования текстов (п. 4) и предложенной модификации семплинга не исследуется. То же самое касается последнего пункта про модификацию ControlNet. Её приходится проводить скорее по техническим причинам: модель получается слишком объемной при прямом копировании энкодера SDXL.

У проекта есть страница с интерактивными ползунками, код тоже выложен.

Content-Oriented Learned Image Compression

Есть такая задача, картинки сжимать. Проводятся эксперименты с обучением различных энкодеров дабы побороть качество всяких JPEG и BPG. SOTA по качеству при низних битрейтах является VAE, который обучают с некоторым набором трюков дабы сохранить как можно больше высокочастотной информации.

Эта работа предлагает yet another набор трюков, ориентированных в первую очередь на семантической неравнозначности разных регионов изображений. Например, мы можем готовы пожертвовать сохранением мелких деталей одежды и интерьера ради качественного восстановления черт лиц и конечностей. С другой стороны, информация о текстуре и границах объектов тоже важна, но классический MSE лосс делает модели склонными к сглаживанию деталей.

В итоге, авторы предлагают немного инженерное решение: сначала выделять семантически разные регионы, после чего применять к ним разные функции потерь:
— Лица детектируют с помощью YOLO-face. Для этих регионов применяют MSE лосс;
— Регионы резких яркостных переходов получают считая карты градиентов, после чего применяют к ним Laplacian Loss;
— Текстуры содержат множество скоррелированных деталей. Лоссы учитывающие только попискельные расстояния не способны это учесть. Поэтому для сохранения текстур к MSE добавляют GAN и LPIPS лоссы. В качестве адверсариальной компоненты используют подход Realistic Discriminator, с архитектурой из PatchGAN. Метод делает относительную оценку реалистичности картинок вместо выдачи бинарного ответа.

Метод выдаёт визуально неплохие реконструкции и норм циферки PSNR, SSIM, FID. При битрейте порядка 0.2bbp удаётся сохранить достаточно мелкие детали, в том числе небольшие лица и конечности. Кода нет.

YaART: Yet Another ART Rendering Technology

Много уже было работ про большие диффузонки. Это еще одна, отличающаяся от предыдущих применением RLHF в production-grade системах и подробными экспериментами по масштабированию не только моделей, но и данных, причем в не очень распространенной для академии каскадной постановке.

Авторы учат Imagen-подобный каскад с разбивкой на три стадии 2.3В-700М-700М на собственном пофильтрованном 330М датасете. В статье подробно описан пайплайн сбора данных для претрена, основанный на большом количестве фильтраций классификаторами качества текстов и картинок + их релевантности. Для последней Super Resolution модельки используют отдельный датасет на 180М картинок с упором на высокое техническое качество.

После претрена идут две стадии файтнюна:
- Supervised на небольшом наборе пар исключительного качества. Его основной целью является улучшение следованию текстов, хотя красивость картинок тоже повышается;
- Дальше идет RL alignment, задизайненый так чтобы еще больше улучшить эстетичность и уменьшить дефектность генераций без ущерба текстовой релевантности.

Основные эксперименты посвящены изучению влияния размеров модели, данных и количества компьюта на качество претрена. Получается, что большие модели учатся быстрее и эффективнее, а вот размер датасета при равномерном семплировании мало влияет на качество. Интересно, что если семплировать данные не равномерно, а брать только топ лучших, то можно найти sweet spot, в котором модель работает даже лучше обученной на полном датасете.

Отдельно показано, что можно обучить меньшую модель до качества основной за счет большего компьюта. Результаты претренов довольно неплохо переносятся на файнтюны, а значит непосредственно связаны с итоговым качеством.

Важная деталь работы - использование асессорских замеров для надежности результатов вообще везде. Авторы замеряются на DrawBench, а также релизят свою корзинку замеров YaBasket-300.

Корзинка доступна на сайте работы, а еще есть упрощенная/укороченная версия папиры на Хабре.

Boosting Diffusion Models with Moving Average Sampling in Frequency Domain

Классические диффузионные семплеры, такие как DDPM и DDIM, генерируют объекты из шума в Марковской постановке: семпл на каждом следующем шаге зависит только от предыдущего. Не идеальность их работы приводит к ошибкам дискретизации, из-за которой модель осциллирует вокруг некоторой истинной траектории семплирования. Эффект был замечен в работе DPM где был предложен солвер более высокого порядка, учитывающий результат 1-3 предыдущих шагов. Есть и другие подходы, например UniPC и его вариации, использующие схожий подход.

Авторы этой работы мотивируются распространенным в прикладной оптимизации методом EMA, суть которого в сглаживании шумной траектории путём взвешенного усреднения нескольких последних шагов. На практике, напрямую брать взвешенную сумму результатов нескольких шагов семплирования нельзя, поскольку каждый из них приходит из своего распределения. Поэтому предлагается метод усреднения оценок на х_0 как модификация к любому существующему семплеру.

Суть метода заключается разложении оценок на x_0 на частотные компоненты используя wavelet transform и взвешивании этих компонент в зависимости от шага семплирования. Известно, что в ходе обратного диффузионного процесса сначала генерируются низкочастотные компоненты. По ходу семплирования доля высокочастотных компонент увеличивается, доходят до максимума при малых t. Авторы используют это знание для взвешивания вклада частотных карт в своей модификации для солверов, которую назвали Moving Average Sampling in Frequency (MASF).

В экспериментах, метод существенно улучшает FID и немного уменьшает артефактность генераций. Наибольший прирост заметен на малом числе шагов семплирования. Кода нет.

Scaling Laws for Data Filtering - Data Curation cannot be Compute Agnostic

Большинство статей, посвященных масштабированию DL систем, изучают вопрос взамосвязи размера модели, размера датасета и компьюта, потраченного на обучение. Данные в таких исследованиях, как правило, семплируют равномерно по качеству. На практике же данные обычно фильтруются перед использованием. Какую долю данных оставить - вопрос не праздный.

В статье Beyond Neural Scaling Laws было показано, что если выбирать нужное количество сложных примеров для обучения, то это ломает привычные степенные законы масштабирования, а при оптимальности выбора можно достигать экспоненциального масштабирования.

В этой работе авторы изучают quality-quantity trade-off (QQT) заключающийся в вопросе: взять побольше плохих данных или нафильтровать хороших, но поменьше. Вопрос изучают в контексте VLM, фильтруя DataComp-medium для обучения CLIP. Качество оценивают по zero-shot классификации ImageNet/CIFAR, и ретривале Flickr и MSCOCO.

Основным выводом статьи является наблюдение: оптимальное количество данных зависит от доступного компьюта: чем больше раз вы проходитесь по датасету во время обучения, тем полезнее оставлять в нём не самые качественные примеры. Авторы находят сетап, при котором даже LAION фильтрации для Common Crawl оказываются избыточными и мешают дальнейшему улучшению качества модели.

В статье много любопытных кривулек для ознакомления. Какого-то универсального закона (сколько точно данных мне брать если знаю выч. бюджет), к сожалению, нет. Кода тоже.

Input Perturbation Reduces Exposure Bias in Diffusion Models

В авторегрессионной генерации текста существует проблема exposure bias. Суть в том что модель учится генерировать следующий токен при условии набора реальных токенов (контекста). Во время инференса, для генерации всех токенов кроме первого используются ранее полученные, синтетические токены, распределение которых отличается ввиду не идеальности генеративной модели.

Подобный эффект происходит и во время семплирования из диффузионных моделей. В диффузионках также наблюдается накопление ошибки семплирования. Это приводит к известному парадоксу: семплирование DDPM работает лучше при некотором T’ < T. Авторы квантифицируют совместный эффект exposure bias и накопления ошибки подсчетом FID. Для этого они зашумляют исходные картинки до некоторого t, после чего расшумляют и показывают, что чем больше t, тем дальше от исходной картинки оказываются итоговые генерации.

Для борьбы с описанными проблема предлагаются методы:
1. Regularization with Input Perturbation. Эта идея преимущественно направлена на устранение exposure bias и состоит в том чтобы сделать модель более устойчивой к отличию шума на шаге t во время прямого и обратного процессов за счет… впрыскивания дополнительного шума! Для этого будем учить модель предсказывать исходный, базовый шум на каждом шаге по чуть более зашумленной картинке. То есть, если согласно q-sample на шаге t у нас должен быть шум ε, то предсказывать мы будем ε, но по картинке, на которую наложен ε + σ, где σ тоже распределен нормально.
2. Regularization based on Lipschitz Continuous Functions. Предлагается уменьшать накопление ошибки за счет добавок к лоссу (Gradient penalty, Weight decay).

Важно, что оба метода требуют, как минимум, дообучения, а, скорее всего, полного переобучения модели, что не всегда возможно и удобно.

Методы проверяют только для DDPM семплера, а замеры проводят только на игрушечных CIFAR10, ImageNet32 etc. Показывают, что работает лучше ванильного DDIM. Код есть.

Alleviating Exposure Bias in Diffusion Models through Sampling with Shifted Time Steps

Продолжаем рассматривать проблему exposure bias в диффузии. В предыдущей серии мы узнали, что проблема существует и может быть решена парой хаков: впрыском дополнительного шума и регуляризацией на Липшицевость.
Минус в том что требуется переобучение (возможно, дообучение, этот вопрос не рассматривался) модели.

Важным наблюдением является то что ошибки семлпирования на шагах t_i распределены нормально. Это значит, что вместо добавления шума во время обучения можно просто немного смещать значение t для следующего шага семплирования так чтобы минимизировать ошибку. То есть, вместо переобучения с аугментацией, можно просто модифицировать семплер!

Авторы статьи предлагают для каждого t_i на пути расписания семплирования искать новую точку в окрестности t_w такую чтобы ошибка из-за exposure bias оказалась минимальной. Точку выбирают оптимизируя дисперсию шума так, чтобы она максимально совпадала со значением из расписания q-sample для данного шага. Утверждается, что процедура может быть добавлена не только в DDPM и DDIM, но и в солверы высокого порядка. Поскольку результатом работы алгоритма является буквально смещение параметра времени t на каждом шаге семплирования, метод назвали Time Shift (TS), а соответствующие модификации семплеров: TS-DDPM, TS-DDIM и т.д.

Метод проверяют на игрушечных датасетах с низким разрешением для одностадийной пиксельной диффузии (ADM). Показывают прирост качества в терминах FID в широком диапазоне числа шагов семплирования. Важным экспериментальным фактом является то что начиная с какого-то t < t_c метод можно не использовать потому что процесс семлпирования становится менее чувствительным к небольшим ошибкам. Это экономит доп. вычисления, что всегда приятно. Средний оверхед по компьюту обещают в районе 10%. Код есть.

Elucidating the Exposure Bias in Diffusion Models

Еще одна статья с последнего ICLR, изучающая проблему exposure bias. На этот раз авторы проводят более глубокий теоретический анализ проблемы и находят, что:
- Дисперсия шума во время обратного процесса всегда выше чем во время прямого и тем больше чем меньше значение t;
- Предлагается метрика для измерения exposure bias для данного t, идейно сильно похожая на FID. Её значения также сильно скоррелированы с FID.

Аналитический вывод о том что дисперсия шума во время обратного процесса всегда выше чем во время прямого согласуется с интуицией предыдущих работа, предлагающих добавлять шум во время обучения и сдвигать значения t во время инференса. Более того, аккуратно выбирая параметры семплирования, EDM также позволяет уменьшить проблему, но требует переобучения. Авторы предлагают свою модификацию семплирования, не требующую этого.

Авторы предлагают масштабировать предсказанный во время семплирования шум в меньшую сторону. Оказывается, что оптимальной коррекции можно достичь с помощью линейной функции вида λ(t) = kt + b, где k и b могут быть найдены экспериментально путём семплирования из уже обученной модели. На практике оказывается, что даже этого делать не обязательно и что константное правило уменьшения дисперсии вида λ(t) = b даёт близкие к оптимальным результаты, но при этом не требует дополнительной оптимизации. Метод назвали Epsilon Scaling (ES), а модификации методов ADM-ES, DDPM-ES и т.д.

Метод проверяют, в том числе, во вполне реалистичных сетапах вроде DiT + LDM и датасетах уровня CelebA-HQ. Показывают прирост качества в терминах FID и по картиночкам в широком диапазоне значений шагов семплирования, в том числе и на практически интересных <30 шагах. Есть код для ADM-ES.

Imagine Flash: Accelerating Emu Diffusion Models with Backward Distillation

Инферить диффузионки долго, для ускорения их дистиллят. Можно дистиллить гайденс и это ускорит примерно в два раза, а можно по шагам и тогда ускорение будет кратно уменьшению их количества.

В этой работе предлагают дистилляцию шагов для text-to-image модели Emu. Основной claim в уменьшении шагов до 3 с аж 50 (хотя в исходной работе про число шагов ни слова) с помощью трёх трюков:

1. Backward Distillation. В обычной, forward дистилляции по шагам модель учат предказывать x_0 по x_t минуя промежуточные шаги, а сам x_t получают прямым процессом. Вместо этого, авторы предлагают получать x_t семплированием из студента, начиная с x_T. Затем, студента учат имитировать предсказания x_0 учителя, но уже стартуя с нового, “синтетического” x_t. Таким образом, модель как бы учится нивелировать exposure bias, о котором мы говорили выше (раз, два, три).

2. Shifted Reconstruction Loss. Обычно для вычисления разницы предсказаний учителя и ученика используют один и тот же момент времени t. Авторы предлагают смещать t для учителя в меньшую сторону, то есть в качестве референсов брать менее шумные x_t. Таким образом, студент будет ориентироваться на более богатую контентом картинку. Утверждается, что это улучшает сигнал для ученика и положительно влияет на его качество.

3. Noise Correction. Авторы говорят о том что в точке T, то есть на первом шаге семплирования, модели трудно сделать предсказание из-за того что в латенте пока что нету сигнала, что верно. Утверждается, что рассмотрение этого случая как отдельного и отказ от использования предсказания шума от студента стабилизирует процесс семплирования. Тем не менее, оценка на x_0 все же используется. Более того, судя по ablation, этот пункт не сказывается на FID/CLIP score.

Модель сравнивают с ADD-LDMXL и Lightning-LDMXL, а так же с ADD версией Emu. Аблейтят, что Backward Distillation даёт наибольший профит, а в сапмате показывают, что модель не коллапсирует по разнообразию. Кода нет.

Scaling (Down) CLIP: A Comprehensive Analysis of Data, Architecture, and Training Strategies

Большинство работ по масштабированию, в том числе мультимодальных моделей, направлены на исследование влияния увеличения модели, данных и компьюта. DataComp и основанные на нём работы, а также Beyond Neural Scaling Laws и YaART показали, что качество - важный фактор, влияющий на оптимальное количество данных необходимых для данной модели, а в Scaling Laws for Data Filtering показали, что качество данных связано еще и с тем сколько выч. ресурсов оптимально тратить на обучение.

Здесь авторы изучают quality-quantity trade-off, а также влияние конкретных методов фильтрафии и претрена, обучая CLIP на пофильтрованных подмножествах WebLI 3.4В. В статье есть несколько занятных выводов:

1. Можно получить буст качества претреном на 40% лучших (по текстовой релевантности) данных относительно обучения на всех данных в терминах zero-shot качества на ImageNet.

2. Наибольший ViT не всегда оптимален. Например, если учить одну эпоху на датасетах < 100М, то ViT-L/16 оказывается худшим из всей линейки ViT’ов. Оптимальный размер зависит от размера данных даже при фиксированном числе итераций обучения. Также на небольших выборках CNN (ResNet-50) могут быть даже лучше. Объясняется это сильным inductive bias, который, видимо, решает на малых выборках.

3. Продвинутые техники предобучения вроде SLIP, FLIP и добавления аугментаций не всегда побеждают стандартный CLIP. С увеличением масштаба эксперимента методы выравниваются по качеству, при этом SLIP становится не выгодным из-за в два раза большей выч. сложности. На максимальном тестировавшемся размере (3.4В) базовый CLIP оказывается лучшим вариантом.

Ни кодом ни обученными модельками авторы не делятся.

CatLIP: CLIP-level Visual Recognition Accuracy with 2.7× Faster Pre-training
on Web-scale Image-Text Data

Бывает так, что у статьи настолько кликбейтное название, что сначала испытываешь отторжение, а потом видишь, что авторы из Apple и все-таки читаешь.

Интересно, что изначальной мотивацией перехода на self-supervised обучение был уход от предобучения на классификационных задачах, обычно требующих явной разметки. Здесь авторы предлагают вернуться к классификации в угоду вычислительной эффективности, при этом все еще не требуя явной разметки. Повышение эффективности достигается за счет исключения большого числа мало информативных негативных сравнений в contrastive loss.

Метод основан на использовании иерархии концептов WordNet в качестве классов. Для этого из текстовых описаний к картинкам выделяют существительные. Каждое уникальное существительное - класс, при этом длинный хвост редко встречающихся существительных выбрасывается. На практике хорошо работает количество классов от 100 до 500. Каждый текст может содержать несколько существительных, поэтому классификация многоклассовая.

Основным результатом работы является вывод, что такой претрен в общем случае не хуже, а на малых выборках даже лучше контрастива. Важно, что использующийся для этого метод верификации (zero-shot классификация) в этом случае не является совсем уж zero-shot. Тем не менее, полученные таким способом модели хорошо масштабируются по размеру, компьюту и данным, а также не хуже контрастив претренов обобщаются и на другие задачи, например сегментацию и детекцию.

Для интересующихся есть детальная инфа про протоколы обучения и тестирования, а также есть код и веса обученных моделей.

SeeSR: Towards Semantics-Aware Real-World Image Super-Resolution

Диффузионки умеют генерировать картинки из шума: семплируем произвольный шум, получаем произвольную картинку. На практике получать произвольную картинку не интересно, генерацией хочется управлять. Для этого сначала придумали обуславливание на классы, а потом перешли на тексты как на самое удобное и понятное представление.

Отношение к текстам в диффузонном Super Resolution (SR) не однозначное. С одной стороны, Imagen, YaART, SR3/SR3+ и многие другие не видят улучшения от использования текстов для SR в высоком разрешении. С другой стороны, в SUPIR утверждается, что тексты полезны.

Суть этой статьи - свой механизм подготовки текстов для SR, который, вроде как, бустит качество. Авторы предлагают:
— Обучить модель тегирования (по сути, классификации) ассоциировать low-rez картинки с тегами так же хорошо как их high-rez версии;
— С помощью неё предсказывать теги (ключевые слова через запятую) по low-rez картинке;
— Подавать строку, состоящую из тегов, в SR модель в качестве условия.

Большинство последних SR моделей полагаются на генеративный prior Stable Diffusion (SD). Для этого модели вроде DiffBIR и SUPIR берут предобученные варианты SD и адаптируют их: low-rez картинку подают через обучаемые отдельно варианты ControlNet, основная модель и VAE остаются замороженными.

Здесь авторы делают то же самое: полученные тегированием тексты подаются через cross-attention, а low-rez картинка через ControlNet в предобученную SD. Обучаются на мелких датасетах (DIV2k, Flickr etc.), с RealSR аугментациями, эвалятся на их тестах с теми же аугментациями + сборной корзинке из RealSR/DRealSR и собственных картинок.

Метод чуть напоминает натягивание совы (text-cond diffusion generation) на глобус (задача SR). Прирост качества, скорее всего, достигается за счет того что модель ломается меньше чем в работах, пытающихся оторвать ненужные тексты. Код есть.

CADS: Unleashing the Diversity of Diffusion Models through Condition-Annealed Sampling

Classifier-free guidance - важный элемент инференса диффузионных моделей, существенно повышающий качество генераций. Суть метода в том чтобы из всех возможных выученных мод направить генерацию в сторону той, которая лучше всего соответствует кондишену (обычно, тексту). Чем выше guidance scale (а для максимизации качества он, как правило, высокий), тем больше “сужается” пространство траекторий семплирования и, следовательно, разнообразие.

Авторы предлагают увеличивать разнообразие возможных генераций с помощью Condition-Annealed Diffusion Sampler (CADS) - метода, не требующего переобучения и модификаций параметров семплирования. Интуитивно, авторам хочется меньше ограничивать множество генераций промтом на ранних шагах семплирования. Ближе к финальным шагам роль промта можно снова увеличить для получения высокого качества и хорошего соответствия запросу.

Суть метода заключается во впрыскивании нормального шума непосредственно в текстовый эмбеддинг, подающийся в модельку в качества кондишена. Уровень шума задаётся некоторым расписанием, где базовый вариант - кусочно-линейная функция:
- На промежутке от Т до t1 - полностью заменяем эмбед на шум;
- От t1 до t2 линейно понижаем уровень шума;
- От t2 до 0 используем чистенький эмбед.

Чтобы эмбеды не ломались, используется дополнительная процедура нормализации эмбедингов перед их зашумлением.

Внимательный читатель спросит почему бы нам просто не варьировать силу guidance scale в ходе семплирования вплоть до полного его отключения промежутке от Т до t1, ведь кажется, что это эквивалентно. Такой подход является бейзлайном этой работы и авторы значительно его превосходят. В том ли причина, что авторы плохо завели бейзлайн - вопрос открытый.

В статье аблейтится метод в приложении к нескольким популярным семплерам (включая DMP++ и UniPC), рассматриваются вариации расписания зашумления. Все эксперименты на ImageNet 256 и 512, качество в терминах FID и Precision даже растёт. Кода нет.

Automatic Data Curation for Self-Supervised Learning: A Clustering-Based Approach

Современные Self-Supervised Learning (SSL) методы качественно и робастно решают широкий круг задач. SSL модели обучают на больших наборах данных, часто являющихся пофильтрованными подмножествами кусков интернета. Большой прогресс в области достигнут за счет совершенствования процедур фильтрации данных (DataComp), позволяющих отбирать более качественные семплы.

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

Авторы предлагают автоматический метод фильтрации данных с акцентом на сохранении редких концептов из длинного хвоста их распределения. В работе авторы извлекают DINOv2 эмбеды их картинок и SBERT эмбеды из текстов, а далее играются с кластеризацией.

k-means кластеризация хороша, но только когда размер потенциальных кластеров концептов сопоставим. На практике, у распределения концептов длинные хвосты, поэтому предлагается иерархический k-means.

Суть метода в том чтобы:
- Делать k-means на эмбедах, получать кластеры;
- Для каждого кластера считать центроиды;
- Снова делать k-means, но уже на центроидах;
- Асимптотически процесс сходится к равномерному покрытию произвольного распределения.

Получив иерархию, авторы семплируют нужное количество примеров из семпла интернет картинок/текстов и обучают свои версии DINOv2 на этих данных. Довольно убедительный эвал и ablation показывает высокую робастность, код обещают тут.

LiteVAE: Lightweight and Efficient Variational Autoencoders for Latent Diffusion Models

Автоэнкодер (AE) является важнейшей частью латентных диффузионок (LDM), а качество декодера задаёт верхнюю оценку качества генераций. При этом, большинство современных работ по LDM либо уделяют вопросам обучения AE не более пары абзацев либо вовсе опускают относящиеся к AE детали.

В этой работе авторы модифицируют архитектуру генератора и дискриминатора, внедряют wavelet фичи и меняют процедуру обучения для получения либо более эффективного VAE (то же качество за меньший компьют), либо более качественного VAE (лучшее качество за тот же копьют).

Основные составляющие метода:

1. Wavelet processing (основной элемент на приложенной схеме): к входной картинке применяется multi-scale wavelet transform, возвращающий wavelet карты разного разрешения (1/2, 1/4, 1/8 от исходного и тд). На практике, авторы берут 3 скейла, которые далее кодируются отдельными энкодерами. Латенты просто конкатенируются и декодируются.

2. Self-modulated convolution: авторам не нравится распределение значений активаций после group norm блоков декодера, поэтому они меняют их на свои свертки с модуляцией. Эти блоки выучивают скейлинг активаций, что приводит к меньшему разбросу значений.

3. Training improvements: учатся на картинках низкого разрешения для эффективности, меняют архитектуру дискриминатора на U-Net, добавляют частотный лосс.

Внимания заслуживает сапмат, в котором обсуждается важность выучивания латентного представления для реконструкции и последующего обучения LDM. Также обсуждаются более тонкие архитектурные моменты, выбор конкретного adversarial loss, детали имплементации. Кода нет.

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

В стандартной DDPM постановке, диффузионки требуют много шагов семплирования (>100) для получения качественных генераций. Это жуть как не удобно, поэтому делают: дистилляцию по шагам (ADD, Consistency etc.) и используют продвинутые солверы (часто высокого порядка) для решения ODE/SDE (DPM++, ER-SDE-Solvers), т.е. диффузионный дифур разбивается на интервалы (дискретизируется), далее численно решаемые солверами.

Замеченная авторами проблема заключается в том что сами интервалы выбираются исходя из общих эвристик и никак не адаптируются под модель, датасет и другие факторы. Получается, что в результате не оптимальной дискретизации возникает ошибка, которая может накапливаться и существенно влиять на итоговое качество генерации.

Для того чтобы оценить ошибку дискретизации достаточно узнать и решить генеративное SDE без дискретизации. Только вот на практике это невозможно из-за нелинейности нашей модели, приближающей скор функцию. На помощь приходит следующий трюк: давайте разбивать наше SDE на большое число малых линейных отрезков и решать их. Тогда в пределе это приблизит истинное решение. На практике достаточно выбрать достаточно больше число отрезков и сказать, что такое решение - достаточно близкая аппроксимация.

Теперь, когда у нас есть эвристическое расписание шагов t и некоторая аппроксимация идеального решения обратного SDE, давайте адаптировать наше расписание так чтобы минимизировать ошибку дискретизации. Авторы показывают, что делать это лучше иерархически. Допустим вы хотите лучшее расписание на 32 шага. Тогда нужно:
- Найти оптимальное расписание на 10 шагов минимизируя ошибку дискретизации;
- Добавить экспоненциально количество шагов (сначала до 20, потом до 40) и найти оптимальное решение для них;
- Аппроксимировать расписание какой-нибудь нелинейной функцией, после чего найти точки для нужного числа шагов (32) интерполяцией.

Метод показывает видимый прирост качества на картиночных и видео версия Stable Diffusion для ряда ODE/SDE солверов. Ссылки на потыкать (colab, diffusers) есть на сайте проекта.