atselousov Sep 4 at 12:00

Как Яндекс создавал новую end2end-модель генерации коротких видео

Hard

25 min

5.1K

Яндекс corporate blogWorking with video*Image processing*Machine learning*Artificial Intelligence

Приложение Шедеврум начало использовать бета‑версию end2end‑модели YandexART (Vi). Она позволяет создавать пятисекундные видео по текстовому запросу и учитывает взаимосвязь между кадрами, делая видео более связными и плавными, чем при использовании предыдущей модели. Предыдущее решение было основано на использовании эвристик для добавления движения камеры, где видео создавалось кадр за кадром с применением модели генерации изображений, что приводило к значительным изменениям в содержании каждого кадра.

Чтобы было понятнее, можно сравнить примеры генераций двух моделей:

В этой статье поделимся нашим опытом разработки первых версий end2end‑модели YandexART (Vi):

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

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

Зачем нужна диффузионная модель и как она работает

Представим, что есть сложное распределение $\pi_1$ , которое невозможно смоделировать напрямую, однако хочется получить из него новые объекты. На помощь приходят диффузионные модели.

Диффузионная модель генерирует новые объекты путём преобразования одного распределения в другое. Формально результат применения диффузионной модели можно представить в виде функции $D: \mathbb{R}^d \rightarrow \mathbb{R}^d$ . Пусть $Z_0 \sim \pi_0$ , где $\pi_0$ — некоторое распределение, которое мы умеем моделировать. Тогда $Z_1 := D(Z_0) \sim \pi_1$ .

Ключевую роль в этом процессе играет выбор начального распределения $\pi_0$ :

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

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

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

Формулировка задачи text2video

Задача генерации видео (или text2video) заключается в следующем: по текстовому запросу пользователя нужно сгенерировать серию кадров. Эта задача является логическим продолжением задачи text2image, когда по текстовому запросу создаётся изображение. Результат работы text2video‑модели — кадры, связанные между собой: движения последовательны, без рывков, объекты и фон резко не меняются.

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

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

Как работает Deforum

Deforum относится к категории image‑based‑подходов генерации видео. Эти подходы используют обученную диффузионную text2image‑модель, которая создаёт один кадр за раз, используя входной текстовый запрос. Генерация видео происходит следующим образом:

Создаётся первый кадр будущего видео. Этот процесс аналогичен генерации изображения с помощью диффузионной модели.
Для полученного кадра предсказывается карта глубины. Для этого применяется отдельная модель, которая предсказывает глубину для входного монокулярного изображения. О методах решения этой задачи можно почитать в релевантных статьях.
С картой глубины каждый пиксель изображения может быть представлен в 3D‑пространстве камеры (у пикселей появляются три пространственные координаты), где затем происходит преобразование положения камеры. Это может быть смещение (добавление вектора смещения), поворот (умножение на матрицу поворота) или всё вместе. Есть доступное объяснение процесса формирования изображения.
После преобразования точки изображения из 3D‑пространства проецируются обратно на 2D‑плоскость изображения, но уже для нового положения камеры. Точки в 3D‑пространстве, полученные на предыдущем шаге, не являются идеальной моделью мира, поэтому при такой проекции на результирующем изображении могут возникнуть артефакты или какие‑то пробелы.
text2image‑модель принимает на вход зашумлённое изображение из прошлого шага — начинает расшумление не с самого начала траектории. Создавая следующий кадр последовательности, модель дорисовывает и исправляет проблемные места, появившиеся из‑за преобразования.
Процесс повторяется начиная с шага 2 до тех пор, пока не будет получено нужное количество кадров.

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

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

Отличия модели end2end-text2video

Задача end‑to‑end‑модели — создать последовательность кадров, которые вместе выглядят как полноценное видео. На вход модель получает текст, в котором могут быть указаны субъект, объект и дано описание действия. В отличие от подхода Deforum end‑to‑end‑модель явно учитывает связь между кадрами, поэтому можно ожидать, что результат будет более согласованным.

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

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

Генерация базовых кадров

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

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

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

То, как референсные кадры подаются в модель, повторяет подход, описанный в статье Align your latents… Вход модели содержит зашумлённые фреймы и переданные как условие референсные кадры. Если для какого‑то фрейма не передан референс, то вместо него передаются нули. Также передаётся маска, которая подсказывает модели, где есть условие.

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

Интерполяция кадров

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

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

Повышение разрешения

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

Пространство, в котором происходит генерация

Генерация видео может происходить в пиксельном или в латентном пространстве.

В пиксельном пространстве генерация происходит напрямую, модели работают непосредственно с кадрами будущего видео. Процесс создания видео сложен, включает множество этапов повышения разрешения как по временной, так и по пространственным осям. Для каждого этапа требуется отдельная модель, которую нужно обучить. Для обучения моделей используются ground‑truth‑данные, но во время инференса каждая последующая модель получает на вход предсказания предыдущей модели. Распределение этих предсказаний может отличаться от ground‑truth‑распределения, что приведёт к накоплению различий в распределениях на этапе инференса по мере генерации. И чем больше этапов, тем больше будет ошибка.

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

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

Как обучать text2video-модель и какие данные использовать

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

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

Яндекс — крупная компания, но даже нам сложно сразу собрать большой набор данных для обучения. Сложность заключается не столько в объёме, сколько в качестве и разнообразии данных. Не каждое видео подходит для обучения — надо уметь отсеивать неподходящие.

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

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

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

Из‑за ограниченного объёма данных мы выбрали альтернативный способ получения text2video‑модели — дообучение text2image‑модели. У нас была проверенная, основанная на Imagen архитектура для генерации изображений в пиксельном пространстве. По этой причине, несмотря на упомянутые сложности, мы начали работу над моделью, которая генерирует видео именно в пиксельном пространстве.

Мы планировали добавить temporal‑блоки (motion‑блоки) к картиночному претрейну YandexART text2image‑модели. Похожие подходы описаны во многих работах и показали свою эффективность. Для обучения motion‑блоков требуется меньше данных и предполагается, что text2image‑модель, на основе которой строится text2video‑модель, уже знакома с частью концептов. Задача motion‑блоков — научиться выстраивать кадры в связное видео. Однако насколько такой подход оптимален — вопрос открытый, ведь неясно, насколько схоже понимание концептов, связанных со временем, у видео‑ и картиночной моделей.

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

Мы применяли следующие инженерные оптимизации:

обучение в формате f16 (mixed precision) с использованием torch.FSDP (Fully Sharded Data Parallel). FSDP позволяет распределить веса модели, градиенты и состояния оптимизатора между графическими процессорами (GPU), снижая таким образом нагрузку на память за счёт появления дополнительных синхронизаций;
замораживание пространственных параметров и обучение только добавленных motion‑блоков, что позволяет не хранить состояния оптимизатора для замороженной части параметров;
поэтапное обучение, позволяющее ещё больше сократить затраты вычислительных ресурсов: разрешение и количество фреймов повышается постепенно, по мере обучения модели. Для проверки гипотез необязательно проходить все этапы, а это ускоряет их валидацию.

Получение картиночного претрейна

Стоит отметить значительное сходство Deforum и нашей end‑to‑end‑модели. Несмотря на то что модель для работы с изображениями применяется по‑разному, в обоих случаях качественная text2image‑модель является основой, которая закладывается в базу подхода. Поскольку без надёжного фундамента невозможно создать что‑то масштабное, нашим первоначальным действием стала оптимизация процесса обучения модели для работы с изображениями.

Хотя работа над моделью для генерации изображений была запущена задолго до text2video, все эксперименты с картиночной моделью проводились без FSDP. Именно поэтому наше приключение началось с решения внедрить torch.FSDP при обучении text2image‑модели.

Сразу после добавления поддержки torch.FSDP мы обратили внимание на резкие увеличения значений функции потерь («спайков») на кривой функции потерь. Кроме того, процесс обучения начал расходиться.

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

На стабильность обучения влияли три параметра:

Размер батча. Меньше батч — стабильнее модель, однако это замедляет процесс обучения, поэтому мы не хотели уменьшать батч.
Размер модели. Также мы не хотели уменьшать размер модели, он напрямую влияет на обобщающие способности.
Размер learning rate (lr). Большой lr также обычно приводит к нестабильности.

Что общего у этих трёх параметров и как они связаны? Посмотрим на формулировку алгоритма оптимизации. Мы используем Adam‑оптимизатор, с помощью которого обновление весов $\pmb w_i$ на шаге можно записать следующим образом:

$\pmb w_i = \pmb w_{i-1} - \frac{\alpha * \pmb m_i}{\sqrt{\pmb v_i} + \epsilon}, \tag{1}$

где

$\pmb m_i = \beta_1 * \pmb m_{i-1} + (1 - \beta_1) * \nabla_{\pmb w} L$ — экспоненциальное скользящее среднее градиента;
$\pmb v_{i} = \beta_2 * \pmb v_{i-1} + (1 - \beta_2) * (\nabla_{\pmb w} L)^2$ — экспоненциальное скользящее среднее квадрата градиента;
$\alpha$ — learning rate.

Коэффициенты $\beta_1$ и $\beta_2$ определяют, насколько учитываются градиенты с предыдущих шагов. Чем больше эти коэффициенты, тем больший вклад вносят градиенты истории. В реализации Adam в Pytorch по умолчанию установлены следующие значения: $\beta_1 = 0.9$ и $\beta_2=0.999$ . Это означает, что в выражении (1) в краткосрочной перспективе большие градиенты сильнее влияют на числитель, чем на знаменатель.

Как выражение (1) поможет понять связь? Все перечисленные факторы, влияющие на стабильность обучения, могут привести к значительному обновлению весов, что может быть причиной возникновения «спайков» на графике функции потерь:

Batch size. С batch size всё достаточно просто. Большой batch size уменьшает дисперсию, что может привести к небольшим значениям $\pmb v_{i}$ . Как результат, в какой‑то момент мы сможем сделать больший апдейт весов, чем рассчитывали.

Размер модели. В нашем понимании, нестабильности при увеличении модели могут быть косвенно связаны с Lottery Ticket Hypothesis. Если говорить простыми словами, то, согласно этой теории, только часть весов действительно полезна после случайной инициализации модели в начале обучения. Интуитивно кажется, что при увеличении размера модели вероятность встретить «ненужные» веса должна расти. Но что это значит с точки зрения выражения (1)? Для таких весов градиенты долгое время могут быть небольшими и апдейт будет вида $\text{small} / \text{small}$ . Если придёт какой‑то обучающий пример, из‑за которого появятся большие значения градиентов для неиспользуемых весов, то апдейт в выражении (1) превратится в $\text{big}/(\text{small} < 1) \rightarrow \text{large}$ . Это приведёт к большому изменению весов и может быть причиной появления нестабильностей при обучении. Подробнее об этом читайте в материале A theory on Adam instability in Large Scale Machine learning.

Learning rate. Большой lr непосредственно влияет на величину апдейта весов.

Во всех примерах выше проблема возникает из‑за устаревшего (outdated) второго несмещённого момента в выражении (1). Можно сказать, что в формуле (1) нет задержки добавления больших значений градиента, если большие значения появились относительно внезапно. Чтобы снизить устаревание знаменателя в выражении (1), можно уменьшить $\beta_2$ и тем самым снизить вероятность «спайков» лосса. Уменьшение $\beta_2$ уменьшает влияние истории и делает величину апдейта более адаптивной к резким изменениям градиентов. Подробнее про outdated second moment можно почитать в статье про Adafactor.

В экспериментах мы уменьшили $\beta_2$ до 0,98, обучение text2image‑модели стабилизировалось, и «спайки» ушли. По результатам сравнения с референсным обучением уменьшение $\beta_2$ не привело к замедлению обучения и не повлияло на финальное качество модели.

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

Что нам дал переход на torch.FSDP? Ниже приведены цифры сравнения обучения модели в разных режимах и на разном количестве нод (1 нода содержит 8 карт). Все замеры c torch.FSDP проводились с использованием стратегии full sharding (распределение между всеми GPU):

C hybrid sharding (распределение между GPU одной ноды) на 16 нодах наблюдается 3%-ное замедление по сравнению с DDP-бейзлайном из-за исключения междунодовой синхронизации для сбора весов перед forward-/backward-операциями. Это значит, что примерно за то же время с torch.FSDP мы можем проводить два параллельных эксперимента вместо одного на том же количестве карт.

Расхождение видеомодели

После перехода на torch.FSDP мы получили картиночный претрейн. Следующим шагом было дообучение text2video‑модели, которая получается путём добавления motion‑блоков. Для первых экспериментов в text2image Unet, основанный на Imagen, после каждого spatial‑attention‑блока был добавлен time‑attention‑блок с absolute positional embedding, который задавал позицию кадра в видео:

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

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

Учитывая величины весов, применение weight decay в нашей конфигурации предсказуемо привело к снижению качества генераций. Например, функция потерь эксперимента с использованием weight decay 1e-6 достигает плато раньше (оранжевая кривая), чем без weight decay (синяя кривая), при этом полученные генерации содержали больше артефактов.

Если веса не растут, то можно предположить, что проблема связана с обучением со смешанной точностью (mixed precision).

Мы решили начать проверку этой гипотезы с attention‑слоёв, нестабильности в которых обычно возникают из‑за высоких значений attention‑весов: $\pmb A = \pmb Q \pmb K^T$ . Источником такого роста может служить entropy collapse. Entropy collapse в attention‑весах происходит из‑за того, что в стандартной формулировке attention‑слоя нет механизма пропуска слоя, поэтому модель должна всегда что‑то взять из входных токенов (входа слоя) и добавить в основной поток модели. Чтобы выучить пропуск слоя, модель может сконцентрировать всё своё внимание на каком‑то бесполезном токене (например, токен фона или токен пунктуации), что позволит не добавлять никакой новой информации. Если посчитать энтропию весов внимания для такой ситуации, то её значение будет близким к нулю.

Анализ показал, что величины весов в $\pmb A$ действительно были большими, но связано это не с entropy collapse, а с тем, что на вход attention‑слоя уже поступали большие значения, поэтому cosine attention или какие‑то другие модификации attention не помогали побороть нестабильности обучения, а лишь откладывали расхождение обучения.

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

Мы выяснили, что активации принимают большие значения в text2image‑претрейне и наши motion‑блоки тут ни при чём. Переполнение происходило в основном в Unet‑декодере для $t \rightarrow 1$ в самом начале расшумления, когда на вход сети приходит что‑то близкое к шуму.

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

В визуализации активаций «проблемных» и «хороших» картинок для блока в середине декодера Unet явно видна вырожденность активаций для $t \rightarrow 1$ . Визуализации активаций для «плохих» входов содержат блобы или вообще выглядят как квадраты:

А активации для входов, в которых не наблюдаются большие активации, содержат черты объекта:

Рост абсолютных величин активаций в зависимости от более заметен на следующей анимации:

При взгляде на визуализации сразу вспоминается статья о StyleGan2 и выводы, сделанные в ней: модель создаёт блобы в активациях, чтобы провести статистики сигнала через слои нормализации, и это зачем‑то нужно модели.

Почему такое вообще может происходить в рамках обучения диффузионной text2image‑модели?

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

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

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

Во время обучения с torch.FSDP c mixed precision активации постепенно вырастают до больших значений. В какой‑то момент значения становятся такими большими, что происходит переполнение и обучение расходится.

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

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

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

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

Как бороться с ростом активаций при обучении text2video-модели

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

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

Как сгладить эффект больших активаций? Мы пробовали делать torch.clamp активаций после каждого блока, обучение стабилизировалось, но сходилось к дефектным генерациям. Также мы пробовали масштабировать residual и skip connections, однако модель научилась компенсировать умножение на коэффициент в случае как с константными, так и с выученными коэффициентами, что не дало желаемого результата.

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

$\min_{\pmb{W}} \frac{1}{N} \sum_{i=1}^{N} ||\pmb{A}_i \pmb{w}_i||_2^2 \tag{2}$

где $\pmb{A}_i$ — тензор входных активаций на слое с весами $\pmb{w}_i$ , a $\pmb{W}$ — все веса модели.

Выражение (2) может быть добавлено к функции потерь с небольшим коэффициентом — аналогично тому, как это сделано с weight decay.

Мы уже говорили, что в наших экспериментах weight decay ухудшает результаты. Нет ли такого эффекта и при использовании выражения (2) и чем оно вообще отличается от weight decay? Отличие заключается в том, что предложенная регуляризация на норму активаций накладывается на результат применения весов, а не на сами веса. С weigth decay мы можем получить небольшие значения весов, но результат их применения всё равно может быть большим из‑за наложения значений с одинаковым знаком.

Предложенный метод помог стабилизировать обучение, но имеет один существенный недостаток. Сбор норм активаций на каждом шаге замедляет обучение примерно на 25%. Значения количества обрабатываемых во время обучения изображений в секунду можно увидеть в таблице ниже, где Baseline — эксперимент без какой‑либо регуляризации, а столбец All steps содержит цифры для эксперимента, где регуляризация применялась на каждом шаге.

Чтобы предотвратить замедление, можно использовать регуляризацию не на каждом шаге обучения, а раз в несколько шагов. Если добавлять нормы активаций к функции потерь раз в 8 шагов, то получится ограничить рост без значительного замедления — всего лишь 5% по сравнению с Baseline.

Архитектура motion-блока

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

Нашим первым решением было заменить absolut positional embedding на Rotary positional embedding (RoPE). Вдохновившись идеями NTK RoPE и YaRN, мы хотели научиться расширять временное контекстное окно модели без дополнительного обучения, что позволило бы генерировать последовательности различной длины. Однако мы обнаружили, что модель всё равно переобучается на размер контекстного окна, используемого во время обучения. Это неудовлетворительный результат, который мы пока исследуем.

Следующим улучшением motion‑блока было применение time attention не к отдельным токенам, а к блокам токенов.

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

Вышеупомянутые изменения архитектуры позволили повысить временную согласованность и снизить количество дефектов. Схематично финальная версия text2video‑Unet‑блока выглядит следующим образом:

Получили видеомодель. Что дальше?

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

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

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

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

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

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

Поясним, как происходит сравнение моделей и что означают представленные далее метрики. В ходе первых экспериментов мы поняли, что распространённые автоматические метрики, такие как FVD (Frechet Video Distance), не всегда отражают реальное улучшение качества генераций, поэтому нашим основным методом сравнения стала оценка, которую дают люди.

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

Beauty (красота генерации). Красочное выглядит лучше, чем тусклое или бледное. Более чёткая генерация предпочтительна. Также выбор зависит от стиля.
Softness (плавность). Объект должен двигаться без рывков, каждое следующее его положение должно плавно вытекать из предыдущего.
Movement (движение). Отражает количество и качество движений на сгенерированных видео. Плохо, когда движений слишком много или слишком мало. Оптимальный вариант — гармоничное и реалистичное движение.
Consistency (консистентность). Следит за естественностью изменения объектов: за превращениями, сменой цвета, исчезновением или появлением.
Defect (дефектность). Покрывает дефекты главного объекта и фона сцены, «кипение» и мерцание света.
Relevant (релевантность). Показывает уровень соответствия сгенерированного видео входному текстовому запросу. Побеждает видео, в котором есть все описанные детали и которое точно передаёт действие.
User preference (субъективное предпочтение асессоров). Отражает личное предпочтение асессора, не связанное ни с одним из других аспектов.

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

На примерах ниже заметна разница в качестве движения:

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

Генерация в латентном пространстве. Фон статичен, движется только верхняя часть хомяка.

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

Как боролись с малым количеством движения

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

Для этого мы решили использовать метод classifier‑free guidance (CFG) с негативным текстовым запросом, который задаёт статичное движение. Например, в качестве негативного промта можно взять static video, no movement.

В классической формулировке CFG направляет генерацию от моды безусловного распределения к областям с высокими значениями функции плотности желаемого условного распределения. Если задан негативный текстовый запрос, логика аналогична — генерация отталкивается от моды распределения с учётом заданного негативного запроса.

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

Для модели важно понимать концепты, используемые в негативных текстовых запросах. Если мы берём концепт, неизвестный модели, то для неё это out of distribution и она не может предсказать правильное направление для генерации.

Это станет более очевидным, если записать выражение для CFG через скор‑функцию. Для этого начнём с записи детерминированного обратного диффузионного процесса в виде обыкновенного дифференциального уравнения (ODE):

$d \pmb x_\sigma = - \sigma \nabla_{\pmb x_{\sigma}} \log p(\pmb x_\sigma; \sigma) d\sigma$

где

$\pmb x_{\sigma} = \pmb y + \pmb n_{\sigma}$ — зашумлённый семпл;
$\pmb n_{\sigma} \sim N(0, \sigma^2 \pmb I)$ , a $\sigma$ — стандартное отклонение;
$\pmb y \sim p_{\text{data}}$ — семпл из распределения, которое мы хотим смоделировать;
$\nabla_{\pmb x_{\sigma}} \log p(\pmb x_\sigma; \sigma)$ — скор-функция, которая задаёт направление во время движения вдоль траектории и которую мы учимся аппроксимировать нашей моделью.

Для примера: с EDM-формулировкой диффузионного процесса скор-функцию можно записать следующим образом:

$\nabla_{\pmb x_\sigma} \log p(\pmb x_\sigma; \sigma) \approx \frac{D_\theta(\pmb x_{\sigma}; \sigma) - \pmb x_{\sigma}}{\sigma^2} \tag{3}$

где $D_\theta(\pmb x_\sigma; \sigma)$ выход модели, который аппроксимирует $\pmb y$ .

Также, используя EDM-формулировку, guidance-процесс диффузионной модели может быть рассмотрен как экстраполяция с коэффициентом между двумя разными предсказаниями — $D_0(\pmb x_\sigma; \sigma, \pmb c)$ и $D_1(\pmb x_\sigma; \sigma, \pmb c)$ — и может быть записана следующим образом:

$D_w(\pmb x_\sigma; \sigma, \pmb c) = w D_1(\pmb x_\sigma; \sigma, \pmb c) + (1 - w) D_0(\pmb x_\sigma; \sigma, \pmb c) \tag{4}$

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

$D_0(\pmb x_\sigma; \sigma, \pmb c) = D_\theta(\pmb x_\sigma; \sigma, \emptyset)$

Используя выражение (3), мы можем записать скор-функцию через $D_w(\pmb x_\sigma; \sigma, \pmb c)$ :

$D_w(\pmb x_\sigma; \sigma, \pmb c) \approx \pmb x_\sigma + \sigma^2 \nabla_{\pmb x_\sigma} \log p_w(\pmb x_\sigma| \pmb c; \sigma) \tag{5}$

где

$p_w(\pmb x_\sigma| \pmb c; \sigma) \propto p_0(\pmb x_\sigma| \pmb c; \sigma) \left[ \frac{p_1(\pmb x_\sigma| \pmb c; \sigma)}{p_0(\pmb x_\sigma | \pmb c; \sigma)} \right]^w$

Используя последнее выражение, мы можем записать искомую скор-функцию для CFG появившуюся в (5):

$\nabla_{\pmb x_\sigma} \log p_w(\pmb x_\sigma | \pmb c; \sigma) = \nabla_{\pmb x_\sigma} \log p_1(\pmb x_\sigma | \pmb c; \sigma) + (w - 1) \nabla_{\pmb x_\sigma} \log \left( \frac{p_1(\pmb x_\sigma | \pmb c; \sigma)}{p_0(\pmb x_\sigma | \pmb c; \sigma)} \right) \tag{6}$

Интуиция выражения (6) заключается в том, что под логарифмом в знаменателе находится предсказание с негативным текстовым запросом. Если модель его не знает, то происходит деление на ноль, и поэтому CFG не может направить генерацию в нужную сторону.

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

Мы взяли случайные кадры из видео и повторили их, чтобы получить статичную последовательность. Затем провели два эксперимента: в первом заменили оригинальные промты на static video, no movement, во втором негативный промт добавили как префикс. Оба эксперимента показали увеличение количества движения согласно нашим метрикам, но второй вариант оказался более эффективным. Скорее всего, это можно объяснить тем, что модель выучивала совместное распределение, а не просто концепт статики в изоляции.

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

Результаты и дальнейшие шаги

Согласно нашим метрикам, мы превзошли наш предыдущий метод генерации видео — Deforum. Генерации end-to-end-модели стали более плавными (softness) и последовательными (consistency), что также больше нравится пользователям (user preference).

Судя по нашим замерам, у end-to-end-модели не изменилась красота (beauty) генерации. Это, вероятно, связано с тем, что мы используем кадр, сгенерированный той же text2image-моделью, которая применяется для Deforum. Кроме того, не изменилось и соответствие генерации текстовому запросу (relevant). Это указывает на то, что у нашей end-to-end-модели явно есть проблемы с пониманием концептов движения.

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

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

Tags:

Hubs: