Как оживить Кандинский матрицами вращений для генерации видео — модель Splitter (часть 2-я)
В первой части, я познакомил вас c подходом для генерации видео, который основан на матрицах вращений. К нему меня привела интуиция, а затем я начал работать на формализацией идеи, после первичного погружения в теорию групп. Далее я уже был готов двигаться в решение задачи на основе машинного обучения.
Гипотеза
За основу методики был взят подход, используемый в видеокодеках.
опорный кадр
алгоритм, описывающий изменения, от опорного кадра далее на определенную глубину.
Широко применяемый подход I-Frame, который использует опорный кадр и последующие кадры изменений, подталкивает к мысли, что нужно научиться передавать в модель, работающей на границе двух латентных пространства, информацию об изменениях в латентном просранстве изображений!
На слайде далее представлены схематично 2 латентных пространства разной модальности.
латентные пространства разной модальности
Кроме основных векторов, есть вектора изменений. Так как прошлые опыты строились на изменениях, то возникла гипотеза.
Гипотеза для задачи text2video
Можно построить обучение модели, предсказывающей вектор для генерации i — го кадра через передачу в модель информации только об изменениях от 0-го кадра посредством функции потерь.
Функция потерь
Проведенные ранее эксперименты, показали, что имеет смысл построить обучение нейронной сети на изменениях. Сами изменения могут быть разными, и с точки зрения векторного пространства это могут быть угловые изменения вектора, так и сама длина вектора. Изменения в пикселях в кадрах, через преобразование энкодером будут отражаться в эмбеддингах изображений. Поэтому для обучения условной модели этим малым изменениям в эмбеддингах кадров, как основной вариант, взята комбинация двух функций ошибки. Основываясь на предварительных экспериментах и анализе изменений векторов в латентном пространстве, было решено использовать комбинированный подход, включающий две различные метрики потерь: CosineEmbeddingLoss и MSEloss.
На слайде ниже представлена функция потерь, связывающая изменения в пространствах через угловые и метрические данные.
Комбинированная функция потерь
С виду, простая формула, но с особенностью — выходом функции потерь для подсчета градиентов должно быть не число, а вектор размерности латентного пространства. По сути, при обучении, мы обращаемся не к среднему градиентов в точке, а к полю градиентов. И мы будем учить на изменениях в разных пространствах, то есть изменения изменений поля.
Тут есть также схожесть с подходом «Компенсация видео» в видеокодеках. Там картинка разбивается на блоки и ищутся изменения блоков, при переходах между кадрами, которые кодируется как вектора изменений.
Модель, которую нужно применить к латентному пространству текстов, будет учится на в чем-то подобных векторах изменений, но уже в латентном пространстве изображений.
Splitter — что за зверь?
Для реализации обучающего алгоритма я решил использовать модель Kandinsky 2.2, которую уже использовал для первых тестов (см. часть 1). Kandinsky 2.2 построена на подходе схожим с unclip, который применяется в некоторых версиях Stable Diffusion и DALL-E2 . Если точнее, то в Kandinsky 2.2 используется подход diffusion_mapping, для преобразования высокоразмерных текстовых эмбеддингов в латентные эмбеддинги с сохранением геометрических свойств и связности. На выходе этого процесса эмбеддинги такого же размера, что и эмбеддинг от энкодера изображений.
Генерация в диффузионной модели Кандинского 2.2 построена по принципу Image-2-Image, в ней вторая часть — Decoder, содержит диффузионную модель Unet Image-2-Image и модель MOVQ для перевода изображения в более высокое разрешение. Первая часть модели содержит модель Prior, которая обучена сближать unclip эмбеддинги текстов и эмбеддинги изображений. Они имеют высокую косинусную близость, а Unet в декодере учится хорошо восстанавливать из шума картинки из эмбеддингов изображений. Эта функциональность модели Kandinsky 2.2 в задачах, где необходимо соблюдение непрерывности и динамики изменений сцен, может обеспечивать более глубокое понимание преобразования текста в слегка изменяющихся кадрах и крайне удобна для экспериментирования в моем подходе.
Близость unclip эмбеддингов текстов и эмбеддингов изображений позволят в том числе и комбинировать картинки и теты для создания новых изображений с помощью Kandinsky 2.2. Структура модели, возможность раздельного обучения и использования модулей, качество генерации картинок — оказались вполне удобным тестовым алгоритмом для моих экспериментов по добавлению своего модуля, который я назвал Splitter (сплиттер) для получения из модели Kandinsky 2.2 составной модели Kandinsky 2.2 + Splitter, которая будет способна уже создавать видеоряд.
Kandinsky 2.2 + Splitter
Для ее адаптации под задачу Text2Video между Prior и Декодером применяется модель Splitter, которая должна изменять вектора, используемые для генерации изображений Декодером. Тут важно заметить, что модель Splitter учится только в латентном пространстве, а Декодер для обучения не нужен.
Splitter принимает на вход:
порядковый номер предсказываемого вектора,
полные текстовые эмбедденинги от модели CLIP-ViT‐G используемой в Кандинском
и стартовый эмбедденинг из модели Prior.
Базовый сценарий обучения Splitter
Если предполагаемый сценарий — рабочий, то первые результаты должны быть и на простой модели.
Splitter имеет простую конфигурацию из входных эмбеддинг слоев и далее каскад понижающих линейных слоев и слоев регуляризации и нелинейности. На выходе предсказываемый измененный эмбеддинг, который может быть использован в последствии для генерации изображения декодер модели Кандинский 2.2.
Сценарий обучения самописный и о нем будет далее. Все учится на карте T4, что существенно для того объема экспериментов, который нужно было провести.
Датасет
Для создания требуемых в обучении данных, использован простой и удобный датасет TGIF.
плюсы — это удобство, разнообразие, большое количество роликов и длина роликов в пределах 100 кадров.
минусы — в основном низкое разрешение, попадаются частично статичные ролики и короткое текстовое описание.
Для фильтрации проблем, набор данных для будущего обучения модели Splitter готовился отдельным скриптом. Получаемые скриптом данные представляет собой пайторч датасет уже векторизованных данных. Ввиду ограниченности ресурсов, для первичных тестовых обучений, автоматически было отобрано 200 отфильтрованных и векторизованных скриптом роликов.
Первичные обучения и тесты
Обучение ведется путем случайного выдергивания батча кадров из фильма, если фильм длиннее размера батча, или берутся все кадры, если фильм короче размера батча.
Схема инференса модели от номера генерируемого кадра
примеры первых генераций за счет простой модели Splitter
Первые результаты простого подхода, которые вдохновили меня для дальнейших исследований. И особенно это заметно в танцующей паре, где при наличие неустойчивости фона и одежды, все же имеется сложная связь в их совместных движениях.
Поиски улучшений
Вроде данных много, в каждом ролике от 15 до 100 кадров, но текст один на все кадры.
Если проиграть ролик назад — то часто подойдет тоже описание. Будем учить модель в обе стороны, от начального кадра и конечного, дав модели на вход лейбл направления.
добавление в датасет и модель лейбла направления
Другая гипотеза - НОЖНИЦЫ.
Новый стартовый кадр
От ролика можно сначала или с конца отрезать кадры и стартовать уже с нового кадра. И описание часто подойдет тоже.
Но для корректной работы и обогащения данных, мы должны изменить стартовые текстовые вектора, через матрицы вращений, полученные от векторов изображений первоначального стартового кадра и нового стартового кадра.
Применение матриц вращений для аугментации
Так в обучения появляются матрицы вращений. Для работы в процессе обучения потребовалось реализовать скрипт их применения на тензорах, чтобы все подсчеты велись на видеокарте. В первой части статьи, где не было обучения, матрицы вращений были реализованы на numpy.
Подсчет матриц вращений
Полученной матрицей вращений от старого и нового опорного кадра мы действуем на изначальный unclip вектор для получения измененного unclip вектора. Так же моступаем и для изменения полных текстовых векторов на основе матрицы вращения получаемой из старого и нового unclip вектора.
Так как обучение происходит на случайных кадрах из фильма, то нужно на лету учитывать номер «нового нулевого» кадра и выбирать только те кадры, что есть в случайной выборке далее в фильме, так как те, что до «нового нулевого» кадра будут использоваться для обучения при развороте направления. Был дополнен класс по учету порядка кадров и их перемешивания и контроль места кадра в батче после перемешивания. Учет и контроль места и расстояния от опорного кадра и псевдоопорного кадра до кадра в случайном батче очень важный момент, так как при ошибке происходит усреднение и «замораживание» вместо выучивания особенностей в изменениях.
пример после первичного дополнения обучения
Марковская цепочка
Выходом модели Splitter является вектор той же природы, что приходит на вход. Поэтому возникла идея дополнительного регрессионного шага модели при обучении, то есть сделать предсказание от предсказания, подав предсказанный эмбеддинг вновь в модель Splitter.
Цель:
авторегрессионный шаг
Тут также применяются матрицы вращений для передачи изменений в текстовое пространство от изменений первоначального стартового кадра на предсказанный стартовый кадр.
возможности модели к авторегрессии
На слайде выше представлены примеры генераций от векторов модели полученных авторегрессионно, в зависимости от того, как модель учили.
Видно, что возможности у модели разные, от того был ли в ее обучении шаг с вращением матриц и регрессионный шаг.
Трейнер обучения
Схематичное представление кастомного трейнера обучения, который сочетает в себе:
работу со случайной малой выборкой роликов на каждой эпохе, для более устойчивого и обобщенного обучения модели.
случайный батч кадров из ролика, для устойчивости к разнородностям в кадрах.
шаг обычного обучения в обе стороны.
шаг обучения с ножницами в обе стороны.
авторегрессионный шаг в обе стороны.
применение матриц вращений реализовано для работы на видеокарте.
периодическое обновление модели лучшими весами, на основе разных слагаемых в функции потерь, от застревания в локальном минимуме.
каждый сохраняемый чекпоинт модели имеет историю обучения модели для возможности использования статистики при новом обучении с измененными параметрами.
расширенный трейнер обучения модели Splitter
Построенный трейнер позволяет комфортно обучать и дообучать модель на больших объемах данных, используя карту Т4. Причем длина роликов может быть разной. Для обучения отбиралось уже 500 примеров.
Сравнение обученных весов Splitter
Сравнительные генерации с одним seed
Сравнительные генерации с одним шумом на основе векторов полученных от моделей с разными типами обучений и их сочетаниями.
Вполне заметно, что дополнительные шаги в обучении, как просто с матрицами вращений, так с регрессионным шагом вносят дополнительную информацию в модель и она делает генерацию интереснее.
Интересные результаты
Интересные примеры генерации от векторов по номеру шагу от стартового вектора. Они демонстрируют вполне связанные кадры сгенерированного видеоролика.
генерация по шагу
На слайде представлены примеры генераций от векторов модели полученных авторегрессионно, в зависимости от того, как модель учили.
Видно, что возможности у модели разные, от того был ли в ее обучении шаг с вращением матриц и регрессионный шаг.
генерация авторегрессией
по номеру шага
авторегрессионно
Исследования данного этапа также представлены в моем репозитории.
Это был этап предзащиты и по его результатам я уже углубился в усложнение самой модели Splitter чтобы понять как улучшить ее генеративные возможности. Продолжение следует.
