Диффузионную модель сделали звездой

Учёные из Центра искусственного интеллекта и факультета компьютерных наук НИУ ВШЭ, а также Института искусственного интеллекта AIRI и Sber AI разработали новую структуру диффузионной нейросети, для которой возможно задать восемь видов распределения шума. Вместо классической структуры модели в виде цепи Маркова и применения нормального распределения учёные предложили звездообразную модель, где можно выбирать тип распределения. Это поможет решать задачи в разных геометрических пространствах. Результаты работы были представлены на конференции NeurIPS 2023.

Исследование поддержано грантом для исследовательских центров в области искусственного интеллекта, предоставленным Аналитическим центром при Правительстве РФ.

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

Диффузионные нейросети основаны на вероятностной модели шумоподавления и диффузии (Denoising Diffusion Probabilistic Models, DDPM). Работает это так: на каждом этапе к данным добавляются случайные изменения. Например, с каждым шагом могут изменяться цвета или яркость. Эти изменения постепенно уменьшают шум и делают данные более похожими на нужный результат — до тех пор, пока из хаоса не получится конечное изображение.

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

Процесс диффузии в действии

Karsten Kreis, Ruiqi Gao, Arash Vahdat / CVPR 2022

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

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

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

Андрей Охотин
Один из авторов статьи, стажер-исследователь Центра глубинного обучения и байесовских методов Института искусственного интеллекта и цифровых наук факультета компьютерных наук НИУ ВШЭ

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

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

Дмитрий Ветров
Научный руководитель Института искусственного интеллекта и цифровых наук ФКН НИУ ВШЭ, научный консультант Института искусственного интеллекта AIRI

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

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

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

Okhotin et al. / NeurlPS 2023

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

Сравнение оригинального изображения и результата звездообразной диффузионной модели

Okhotin et al. / NeurlPS 2023

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

Айбек Аланов
Один из авторов статьи, младший научный сотрудник Центра глубинного обучения и байесовских методов Института искусственного интеллекта и цифровых наук ФКН НИУ ВШЭ, научный сотрудник Института искусственного интеллекта AIRI