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

Всем привет!

Всё началось с того, что мы в Sber AI решили немного поизучать/почитать подробнее про хайповую нейронную сеть DALL·E и понять её потенциал возможностей, а также в чём заключается боттлнек – что же мешает генерить картинки хорошего качества и как можно попытаться улучшить работу модели?

Все автоэнкодеры важны, все автоэнкодеры нужны

О том, что такое автоэнкодеры (AE, VAE, CVAE и др.), уже много писали на Хабре и не только (можно почитать здесь, здесь и здесь), так что углубляться не будем. В двух словах: это нейронная сеть, состоящая из двух частей, – энкодера и декодера. Энкодер сжимает входной сигнал (например картинку) в какое-то скрытое состояние меньшей размерности. Декодер переводит скрытое состояние в исходный сигнал.

Таким образом, есть возможность «обучить» свой архиватор. Сами автоэнкодеры уже существуют давно, ими никого не удивишь. Однако фантазия исследователей позволяет и по сей день ставить занимательные эксперименты. Так, например, идея работы DALL·E пришла после создания так называемого VQ-VAE и последующего его улучшения в VQ-VAE 2, а потом ещё и в VQ-GAN).

Причём тут вообще DALL-E?

Краткая вводная для тех, кто ещё не в курсе. DALL·E – это нейронная сеть, которая создаёт изображения, основываясь на текстовом описании на естественном языке. Например, вот что она может создать по тексту «a snail made of harp» (улитка, сделанная из арфы).

Принцип работы DALL·E устроен довольно просто, его можно сравнить с хорошо известной генерацией текста «Sequence to Sequence» с помощью Transformer. Для генерации вместо привычных текстовых токенов используются токены пространства codebook, которые отображают скрытое дискретное состояние изображений после процедуры энкодинга VAE. В качестве трансформера используется одна из версий GPT-3 на 12 миллиардов параметров, а подают в него эмбеддинги текстового запроса на естественном языке и эмбеддинги codebook текущего состояния генерации изображения, конкатенированные между собой. Затем сгенерённая последовательность codebook подаётся в декодинг VAE, который восстанавливает изображение.

Таким образом, очевидным bottleneck’ом для данной архитектуры является автоэнкодер – если он не способен восстановить исходное изображение процедурой encode-decode с приемлемым качеством, то и DALL·E никогда не сможет сгенерить изображение хорошего качества.

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

А чем оригинальный энкодер не угодил?

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

• “16384”: VQGAN ImageNet (f=16), 16384

• “VAE”: DALL-E dVAE (f=8), 8192, GumbelQuantization

• “gumbelf8”: VQGAN OpenImages (f=8), 8192, GumbelQuantization

• “SBER-gumbelf8”: VQGAN SberData (f=8), 8192, GumbelQuantization

Последний из списка мы дообучали на собственных собранных данных. Стоит отметить, что все выбранные модели не видели датасет COCO в процессе обучения, и оценка на этом датасете достаточно честная и независимая. Естественно, для улучшения качества наши данные были дополнены примерами из этих доменов. Мы лишь дофайнтюнили предобученные модели из оригинального репозитория на одну эпоху с warmup learning rate до 1.5e-06.

Вы всё ещё обучаете с нуля? Тогда мы идём к вам!

Для оценки качества восстановленных изображений мы использовали метрики Inception Score (IS) и Fréchet inception distance (FID), к сожалению, первая статья по FID куда-то пропала, поэтому предлагаем посмотреть такой вариант.

Что такое IS и как его посчитать, можно почитать тут или тут. Если вкратце, то для расчёта используется всем известная модель нейронной сети Inception v3. С помощью неё оценивается вероятность принадлежности изображений каждому из 1000 классов. Затем все распределения вероятностей по одинаковым классам суммируются, считается KL-дивергенция между этими распределениями каждого индивидуального класса и средним общим распределением. Оценка принимает значения от 1 до N, где N – количество классов. Чем больше значения вероятностей принадлежности к каким-то классам и чем больше разнообразие этих классов, тем больше IS.

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

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

IS (больше – лучше):

FID (меньше – лучше):

Также мы посмотрели, как Inception Score оценивает оригинальные изображения (без процедуры encode-decode). Интересно, что для данной метрики оригиналы выглядят менее естественными, чем то, что сгенерили некоторые автоэнкодеры (!). Вероятнее всего, мы уже упёрлись в возможности данной метрики, оценка с её помощью уже становится менее объективной. Предлагаем «дедовскими» методами (глазами) посмотреть на генерации. И вот что получилось с проблемными доменами:  визуально качество заметно улучшилось для SBER-gumbelf8:

«Текст»

«Несколько людей»

«Лица»

Не моё, а общее

Мы хотим поделиться с вами тем, что у нас получилось. А именно – дообученными чекпоинтами и примерами кода, который поможет вам использовать эти модели. Всё это есть на Google-диске и в колабе, так что приятного использования!  

Колаб с инференсом моделек.

Колаб с расчётом метрик.

Общая папка с моделями и ноутбуками.

Ссылка на гитхаб

На самом деле очень интересно, какие примеры у вас получатся. Спасибо, что дочитали до конца, и вперёд, к новым генерациям!

Команда: @denis_karachev, @shonenkov, @dendimitrov