Мечтают ли диффузионки о 3D-алайнменте, или что мы планируем рассказать на грядущей ICLR

Привет! Меня зовут Нина, я работаю инженером исследователем в AIRI, где мы с моими коллегами активно исследуем возможности генеративного ИИ. Особое место в нашей рабочей повестке занимает применение диффузионных моделей к различным задачам.

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

Введение

Про рисование двумерных картинок с помощью диффузионных моделей слышал, пожалуй, каждый, кто достаточно глубоко погружался в машинное обучение. Но, помимо этого, диффузионки активно применяют для 3D‑генерации, например, в задачах text‑to-3D или image‑to-3D — то есть когда мы хотим получить 3D‑представление объекта на основе заданного текстового промпта или входной картинки.

Одной из популярных методик для решения этих задач является использование оптимизации лосса с помощью Score Distillation Sampling для переиспользования знаний, зашитых в хорошо обученную 2D‑диффузионку. Основная идея такого метода в том, чтобы рендеры 3D‑генератора (например, NeRF) были похожи на распределения, выученные диффузионной моделью. Обучение происходит с помощью оптимизации 3D‑модели, когда сэмплируется случайное расположение камеры вокруг объекта, рендерится картинка, зашумляется до какого‑то шага , а затем делается один шаг расшумления и считается обычный диффузионный лосс.

Однако при подсчете градиента для оптимизации градиент по диффузионной модели (UNet) не берется:

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

Пусть у нас есть какая‑то пара промптов, например, «running man» и «running robot». Мы хотим, чтобы объекты сгенерировались в одинаковых позах и, например, руки/ноги робота были в том же положении в пространстве, что и у человека. Однако если мы возьмём и попробуем сгенерировать независимо эти два объекта, то даже при одинаковых сидах результаты получатся разные.

Наша работа как раз решает проблему неконсистентности в подобной генерации. Давайте разберёмся за счёт чего это работает.

Описание метода

Пусть у нас есть 3D‑генератор, который на вход обычно принимает точку в пространстве и возвращает цвет и прозрачность:

Также пусть у нас есть текстовых промптов для генерации — то есть  желаемых объектов. Выберем латентный код на ‑мерном вероятностном симплексе:

причем вершины этого симплекса соответствуют предоставленным промптам. Далее этот засэмплированный код подаётся дополнительным параметром в 3D‑генератор.

Дополнительно, при подсчёте SDS‑лосса мы также берём взвешенную сумму промптов

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

Таким образом, наш метод, получивший лаконичное название A3D, основан на обучении переходов (трансформаций) между несколькими выровненными объектами. Накладывая дополнительные условия на архитектуру модели и функцию потерь, мы добиваемся плавности переходов и некоторой реалистичности 3D‑модели в промежуточных переходах.

За счёт обучения такой траектории получается генерировать достаточно выровненные объекты.

Зачем это нужно?

Данный подход можно использовать для различных задач.

Во‑первых, можно решать задачу редактирования объектов. Например, можно редактировать поданный на вход меш пользователя. Такого рода изменения могут быть довольно глобальными, включая даже добавление небольших объектов:

Также можно проводить редактирование небольших частей объекта за счёт парной генерации, но с небольшим изменением промпта (например, обезьяна и обезьяна в шапке):

Наконец, из‑за того, что объекты находятся в одинаковых положениях и части одного объекта соответствуют частям другого — можно проделывать гибридизацию. Например, мы можем зафиксировать так называемые «якорные» точки в пространстве, которые должны относиться тому или иному сгенерированному объекту, а затем проинтерполировать эти веса, чтобы получить гибрид сгенерированных объектов.

Вывод

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

Важно отметить, что мы тестировали A3D не только на NeRF‑based модели, а также на DMTeT‑based 3D‑генераторе. И в обоих случаях метод показал себя очень хорошо, что означает его устойчивость к различным 3D‑представлениям.

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