Важно, что в переводе статьи мы опустили некоторые математические подробности. Мы обобщили математические выражения текстом. Курсивом выделены комментарии, чтобы новичкам было проще читать, а выделенный уровень подготовки  оправдывал себя и под статьей не пришлось ставить графу “сложно”. 

Чтобы статья не вышла слишком большой, приводится ее основная часть с методом. Мы приводим краткий вариант перевода ресерча от Meta.

Внутри оригинала можно ознакомиться с результатами работы. И конкретными метриками в таблицах. 

Вся статья сохраняет письмо от первого лица. 

Резюме статьи:

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

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

Наш подход включает три ключевых компонента:  обратную дистилляция, которая смягчает расхождения обучения и вывода путем калибровки ученика по его собственной обратной траектории; смещенная потеря восстановлений, которая динамически адаптирует передачу знаний на основе текущего временного шага. И коррекцию шума – он повышает качество выборки за счет устранения сингулярностей в прогнозировании шума. 

Благодаря экспериментам мы покажем, что наш метод уже существующие подходы по количественным метрикам и человеческим оценкам.

Примечательно, что он достигает производительности, сопоставимой с моделью учителя, используя всего три шага шумоподавления. 

Введение

Генеративное моделирование претерпело значительный сдвиг с появлением моделей диффузионного шумоподавления (DMs). 

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

Однако последовательный характер процесса шумоподавления — серьезная проблема. 

Отбор образцов из DMs является трудоемким и затратным процессом, время выполнения которого в значительной степени зависит от двух факторов: (i) задержка оценки нейронной сети на каждом шаге и (ii) общее количество шагов шумоподавления. 

Значительные исследовательские усилия были направлены на ускорение процесса отбора образцов. 

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

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

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

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

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

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

Адверсариальные потери — это функция, измеряющая качество различения между реальными и сгенерированными данными. 

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

Общей тема среди упомянутых методов — попытка согласовать модель ученика с малым количеством шагов со сложными путями учителя. 

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

Вкратце, наш (Meta) вклад заключается в трех аспектах: 

– Во-первых, наш подход вводит обратную дистилляцию — процесс дистилляции, предназначенный для калибровки модели ученика на ее собственной восходящей обратной траектории, что снижает разрыв между распределениями обучения и вывода. Мы получаем нулевую утечку данных во время обучения на всех временных шагах. 

– Во-вторых, мы предлагаем смещенную функцию реконструкции, которая динамически адаптирует передачу знаний от модели учителя.

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

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

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

Комбинируя эти три новаторских компонента, мы применяем нашу систему дистилляции к базовой модели диффузии, Emu [4] – получаем Imagine Flash, достигающей высококачественной генерации в условиях крайне низкого числа шагов без ущерба для качества образцов или верности условий (см. рис. 2). 

Диффузионные модели, в отличие от предыдущих генеративных моделей (например, GANs), подходят к оценке плотности и выборке данных итеративным образом, постепенно обращая процесс наложения шума. 

Эта итеративная природа приводит к многократным запросам к основе нейронной сети, что приводит к высоким затратам на вывод. 

В результате значительное количество работ сосредоточено на разработке более быстрых и эффективных способов выборки из диффузионных моделей. 

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

Другие исследования

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

Другие исследования рассматривают реформулировки диффузионного процесса с целью минимизации кривизны как в прямых (шумовых), так и в обратных (шумоподавляющих) траекториях. 

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

Исправление кривизны в диффузионных моделях (диффузионках) связано с улучшением качества генерации данных путем учета геометрии пространства данных. 

Диффузионки, которые постепенно преобразуют простое распределение (например, нормальное) в сложное целевое распределение, кривизна отражает отклонения траекторий от линейного или плоского пути в пространстве высокой размерности. 

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

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

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

Несмотря на значительное сокращение числа шагов благодаря этим методам “выпрямления кривых”, существует предел того, насколько большим может быть шаг вывода, не снижая при этом качество изображения. 

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

Сокращение размера модели: Ряд  работ направлены на уменьшение стоимости шага. В этом русле несколько ресерчей сосредотачиваются на использовании меньших архитектур и воспроизведения “мобильных сетей”

Снижение стоимости за шаг также решается путем минимизации стоимости условной генерации на каждой итерации или кэширования промежуточных активаций в основаниях сети. 

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

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

Здесь авторы подчеркивают, что уменьшение размера модели, сведение ее к совсем маленьким/мобильным версиям по итогу все равно приводило к увеличению скорости работы нейронки (вывода данных).

Сокращение количества шагов выборки: один из способов дальнейшего уменьшения задержки вывода — это дистилляция шагов. 

В этих работах авторы предлагают прогрессивный подход к дистилляции двух или более шагов в один. 

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

Чтобы компенсировать потерю качества, другой “фронт” исследователей предлагает дополнительные улучшения обучения во время дистилляции. 

А именно, ADD, Lightning и UFOGEN. Они добавляют адверсариальные потери для повышения качества образцов. 

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

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

Практически это переводится в методы, позволяющие небольшое увеличение шагов (от 2 до 4 шагов) со значительным улучшением качества. 

Мы применяем этот подход в нашем методе. Для достижения лучшего качества мы предлагаем дистиллировать по обратному пути ученика, а не по прямому пути. 

Иными словами, вместо того чтобы ученик имитировал учителя, мы используем учителя для улучшения ученика на основе его текущих знаний. 

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

Метология и Imagine Flash

Мы представляем Imagine Flash, новаторскую технику дистилляции, предназначенную для быстрого преобразования текста в изображение, основанную на модели EMU, но не ограниченную ею.

 В отличие от оригинальной модели Emu, которая требует как минимум 50 оценок нейронных функций (NFE) для создания крутых картинок, Imagine Flash достигает сопоставимых результатов всего за несколько таких оценок. 

EMU (Efficient Meta-University) — это модель машинного обучения, разработанная для более эффективного и универсального переноса знаний между задачами. 

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

Предлагаемый метод дистилляции включает три ключевых компонента:  обратная дистилляция, процесс дистилляции, который обеспечивает нулевую утечку данных во время обучения на всех временных точках t. 

Смещенная функция реконструкции (SRL), адаптивная функция потерь, предназначенная для максимизации передачи знаний от учителя. 

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

В дальнейшем мы предполагаем доступ к предварительно обученной модели диффузии, которая предсказывает оценки шума ​. 

Эта модель учителя может работать как в пространстве изображений, так и в латентном (скрытом) пространстве. 

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

Этот итеративный характер приводит к множественным запросам в магистральной нейронной сети, что приводит к высоким затратам на логический вывод. 

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

Общие положения о диффузионных моделях и обратная дистилляция 

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

Модели диффузии состоят из двух взаимосвязанных процессов: прямого и обратного. Процесс прямой диффузии постепенно искажает данные путем интерполяции между выборочными точками данных.

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

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

Затем модель обучается решать обратную задачу: начиная с сильно зашумленных данных, шаг за шагом уменьшать шум, восстанавливая исходные данные. 

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

SDE – то самое стохастическое дифференциальное уравнение, которое описывает эволюцию шума во времени. 

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

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

Математически это осуществляется путем минимизации разности между выходами наставника и ученика. 

Наставник генерирует прогнозы или вероятностные распределения, которые ученик пытается воспроизвести. 

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

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

Таким образом, дистилляция позволяет перенести сложные знания и представления от большой модели к меньшей, сохраняя при этом важные характеристики и производительность модели.

И наоборот, процесс обратной диффузии предназначен для устранения процесса зашумления и генерации образцов. 

Согласно теореме Андерсона, прямое СДУ, введенное ранее, удовлетворяет уравнению диффузии в обратном времени, которое можно переформулировать с помощью уравнений Фоккера-Планка

Теорема Андерсона говорит, что если у нас есть прямое СДУ, описывающее этот процесс, то существует соответствующее уравнение диффузии в обратном времени. 

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

Этот обратный процесс описывается другим СДУ, которое в определённых условиях может быть найдено и использовано для обучения модели, восстанавливающей данные.

Теперь, что касается уравнения Фоккера-Планка. Это уравнение описывает эволюцию плотности вероятности распределения состояния системы во времени. 

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

Уравнение Фоккера-Планка связано с СДУ через соотношение между траекторией отдельного состояния (описанного СДУ) и общим поведением всего распределения вероятностей (описанного уравнением Фоккера-Планка).

Итак, когда говорится, что "прямое СДУ удовлетворяет уравнению диффузии в обратном времени", это означает, что существует обратный процесс, который может быть описан через СДУ, и который формально соответствует обратному уравнению Фоккера-Планка. 

То есть, зная, как изменяется распределение в прямом процессе (добавление шума), мы можем описать и моделировать, как это распределение должно изменяться в обратном процессе (удаление шума).

Авторы предлагают модель альтернативной, обратной дистилляции. 

Для получения качественных образцов, Imagine Flash достигает сопоставимых результатов с помощью всего нескольких NFE. 

Метод дистилляции включает три новых ключевых компонента: Обратная дистилляция, процесс дистилляции, который обеспечивает отсутствие утечки данных на всех временных точках t;  Смещенная функция потерь восстановления (SRL), адаптивная функция потерь, предназначенная для максимизации переноса знаний от модели-наставника. 

Смещенная функция потерь (Shifted Reconstruction Loss, SRL) – метод оценки, разработанный для оптимизации процесса дистилляции знаний между моделями. 

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

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

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

Но тут SLR работает таким образом, что мы получаем обратный процесс дистилляции. 

Широко признано, что традиционные схемы шума часто не достигают нулевого конечного соотношения сигнал/шум (SNR) на конечном шаге, что создает несоответствие между обучением и инференсом. 

В частности, схема шума обычно выбирается так, что конечное состояние не является чистым шумом, а содержит низкочастотную информацию, утекшую из исходных данных. 

Это несоответствие приводит к ухудшению производительности во время инференса, особенно если используются лишь несколько шагов. 

Чтобы решить эту проблему, некоторые исследователи предлагают пересчитать существующие схемы шума, чтобы обеспечить нулевое конечное SNR.

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

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

Чем ближе к началу процесса, тем больше информации о сигнале сохраняется и тем сложнее корректировать ошибки.

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

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

 Это позволяет избежать зависимости от истинного сигнала во время обучения, улучшая согласованность между обучением и инференсом.

Что происходит с SLR?

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

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

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

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

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

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

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

 Этот подход помогает ученику сосредоточиться на формировании структуры изображения в ранних этапах и на создании более тонких деталей на поздних этапах.

Коррекция шума

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

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

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

Это позволяет обойти проблему тривиальности предсказания шума на начальном этапе, но при этом может вносить предвзятости в первый шаг обновления.

Когда модель обучается предсказывать шум, её оценка сигнала основывается на самом шуме, что приводит к искажению результатов. В результате процесс денойзинга может накапливать ошибки, поскольку модель не должна просто предсказать шум, а должна учитывать текущие результаты.

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

Мы рассматриваем случай на начальном этапе как особый и заменяем шум модели на истинный шум, чтобы корректировать процесс обновления. Это небольшое изменение значительно улучшает оценку цветов, делая их более яркими и насыщенными. 

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

Заключение

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

Наш подход включает три ключевых компонента: Обратную Дистилляцию, которая снижает несоответствие между обучением и инференсом, Потерю Сдвинутой Реконструкции (SRL), динамически адаптирующую передачу знаний на каждом шаге времени, и Коррекцию Шума, улучшающую качество начального образца.

В результате обширных экспериментов Imagine Flash демонстрирует выдающиеся результаты, сопоставимые с производительностью предварительно обученной модели-учителя при использовании всего трех шагов денойзинга, и неизменно превосходит существующие методы. 

Приводим несколько табличек:

Cравнение методов по ключевым бенчмаркам:

Сравнение значимости отдельных частей технологии для всего Imagine Flashe (больше всего от обратной дистилляции и SLR)

Cравнение ImagineFlash с SOTA машинного обучения в этом направлении

Эта беспрецедентная эффективность выборки в сочетании с высоким качеством и разнообразием образцов делает нашу модель идеально подходящей для приложений генерации в real-time. Таким образом, Meta вышла вперед других корпораций, как минимум, по скорости.