Новая SotA для генеративного моделирования - вероятностные модели с уменьшением шума диффузии

Автор: Себастьян Орбелл, аспирант Оксфордского университета, факультет материалов.

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

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

Вступление

Вероятностные модели рассеяния шумов - это класс генеративных моделей, вдохновленных статистической термодинамикой (J. Sohl-Dickstein et. al.) (для ясности я сейчас буду ссылаться на них как диффузионные модели). Общая концепция заключается в последовательном уничтожении информации в распределении данных посредством процесса прямого распространения. Затем модель обучается обращать вспять этот процесс распространения и таким образом восстанавливать информацию, как одноименный демон Максвелла. Этот метод позволяет модели эффективно обучаться и производить выборку из вероятностей в глубоких нейронных сетях, в дополнение к созданию условных и апостериорных распределений вероятностей из изученной модели. Выборка из этих моделей позволяет генерировать реалистичные данные с помощью процесса цепи Маркова, обусловленного белым шумом.

Модели диффузии - это параметризованные цепи Маркова, которые обращают процесс диффузии и таким образом генерируют информацию. Если эта параметризация построена так, что модель предсказывает гауссов шум, добавленный на этапе прямой диффузии, можно продемонстрировать эквивалентность между шумоподавляющей диффузией и отожженной динамикой Ланжевена. Эта параметризация также снижает вариационную нижнюю границу до цели, используемой для методов сопоставления оценок шумоподавления. Важно отметить, что использование этой специфической параметризации цепи Маркова (идентифицированной J. Ho et al.) Обеспечивает высочайшее качество выборки.

Общеизвестно, что при изучении генеративных моделей сложно оценить качество синтеза изображений. Однако популярной метрикой является Расстояние начала Фрехлета (FID), которое оценивает расстояние между векторами признаков реальных и поддельных изображений. Для извлечения этого набора функций используется предварительно обученная модель Inception V2.

Недавние публикации

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

(i) Прогнозирование дисперсии

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

(ii) Планирование шума

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

(iii) Цель обучения (параметризация)

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

(iv) Функция потерь

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

(v) Архитектура модели

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

(vi) Руководство классификатора

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

(vii) Уменьшение количества шагов выборки

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

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

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

Аппаратно-ориентированная оптимизация модели

IPU (Intelligence Processing Unit) отличается от своих конкурентов во многих отношениях. Наиболее важно то, что IPU был разработан с огромной встроенной памятью с высокой пропускной способностью, расположенной рядом с 1472 ядрами процессора на кремниевом кристалле. Поскольку каждое ядро ​​имеет 6 потоков, каждый IPU может одновременно запускать 8 832 процесса, что обеспечивает массовое мелкозернистое распараллеливание. Такое распараллеливание дает значительные преимущества при выборке из моделей, основанных на диффузии, где необходимо вычислять итеративную цепь Маркова. IPU также имеет встроенный генератор псевдослучайных чисел, который позволяет значительно ускорить обучение в моделях, основанных на диффузии, где временной шаг и добавленный шум должны дискретизироваться на каждой итерации обучения.

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

Как и во многих современных методах машинного обучения SotA, производительность диффузионных моделей зависит от размера модели и вычислительных ресурсов (см. Рисунок выше). Более крупные модели не умещаются на одном кристалле и требуют более совершенных решений для их эффективного распределения по нескольким микросхемам. Graphcore разработала стратегию конвейерной обработки для оптимального распределения модели по нескольким этапам вычислений, при этом каждый этап распределяется между разными IPU.

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

Использование памяти может стать недопустимым при обучении очень больших сетей, однако программная оптимизация с учетом аппаратного обеспечения может смягчить эти проблемы. Graphcore's Poplar SDK позволяет пересчитывать активации при обратном проходе, что значительно сокращает количество активаций, которые необходимо сохранять между прямым и обратным проходами. Этот компромисс между потребностями в вычислениях и памяти является фундаментальным для оптимизации современных рабочих процессов машинного обучения, ориентированных на специализированное оборудование. Программный стек Graphcore предоставляет разработчикам набор инструментов для точного контроля баланса между вычислительными ресурсами и ресурсами памяти.

Другое использование диффузионных моделей:

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

Многомерное вероятностное прогнозирование временных рядов

Расул и др. предложить TimeGrad, метод, который изучает градиент распределения данных путем оптимизации вариационной нижней границы, а затем генерирует выборки, пропуская белый шум через цепь Маркова, которая следует динамике Ланжевена. Используя архитектуру модели, определяемую LSTM, остаточными и расширенными сверточными слоями, авторы достигают результатов SotA для шести популярных наборов данных временных рядов.

Преобразование изображения в изображение

Сасаки и др. достигать самых современных результатов в задачах преобразования изображений в изображения в нескольких общедоступных наборах данных (SotA, превзойдя CycleGAN). Они одновременно обучают две сети шумоподавления, одна из которых зависит от исходного домена, а другая - от целевого. Этот подход обеспечивает стабильное (не состязательное) обучение, при котором совместное распределение вероятностей по обоим доменам изучается путем оптимизации цели сопоставления оценок шумоподавления, зависящей от другой области. Таким образом, целевое изображение может быть сгенерировано с помощью метода Монте-Карло с шумоподавляющими цепями Маркова, основанного на динамике Ланжевена и обусловленного исходным изображением.

Создание условного изображения класса

Хо и др. report SotA приводит к классическому условному синтезу изображений (превосходящему BigGAN-deep и VQ-VAE-2) путем разработки каскада диффузионных моделей. Каскадирование - это метод, с помощью которого распределение данных с высоким разрешением изучается конвейером отдельно обученных моделей с несколькими разрешениями. Они обнаружили, что качество синтеза изображений из каскадных конвейеров может быть эффективно улучшено путем обучения каждой модели сверхвысокого разрешения с использованием входных данных с низким разрешением с учетом некоторой формы увеличения данных. Они утверждают, что это уменьшает ошибку сложения в каскадных конвейерах (смещение экспозиции).

Сверхвысокое разрешение изображения

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

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

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

Будущее

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

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