Вариационные автоэнкодеры (VAE) выступают в качестве основных строительных блоков в современных современных генераторах преобразования текста в изображение, таких как DALL-E и Stable Diffusion. Раньше современные генераторы изображений, как правило, были GAN, однако недавно мы увидели сдвиг с введением диффузоров, где модели, содержащие VAE, взяли верх в этой области. Имея это в виду, мы хотели бы глубже понять основные концепции VAE и реализовать небольшую версию в PyTorch.
Статья представляет собой краткое изложение статьи и реализацию: Auto-Encoding Variational Bayes Дидерика П. Кингмы и Макса Веллинга, ссылка: https://arxiv.org/abs/1312.6114.
Ключевые понятия вариационного автоэнкодера
Цель VAE — реконструкция и генерация информации. Учитывая набор данных, взятый из некоторого неизвестного распределения, мы хотим, например, иметь возможность условно генерировать новые данные с тем же распределением. Этот принцип можно применять в нескольких областях, но в основном он используется для создания изображений. Игрушечный пример, который мы будем использовать, который также использовался в оригинальной статье, — это создание новых изображений MNIST.
Модель состоит из двух частей — кодера и декодера. Мы предполагаем, что наши данные имеют скрытое распределение, подробно описанное ниже. Кодер переводит входные данные в скрытое представление и выводит распределение этого представления. Затем декодер производит выборку из этого распределения и создает новую точку данных. Поскольку используется скрытое распределение входного пакета, копия самой точки входных данных имеет высокую вероятность того, что она будет сгенерирована декодером.
Скрытое распределение является ключевой концепцией, которая отличает VAE от автоэнкодера. Вместо простого сжатия и реконструкции входных данных VAE пытается смоделировать базовое распределение данных. Итак, что же представляет собой это скрытое распространение? В большинстве реализаций вариационного автоэнкодера делаются два сильных предположения/выбора моделирования. Мы выбрали стандартную гауссову матрицу, а ковариационная матрица была диагональной. Это создает довольно сильное предположение о том, что особенности распределения независимы друг от друга. Затем кодировщик будет выводить только вектор как для среднего значения, так и для стандартного отклонения скрытого распределения.
Математика VAE
Не стесняйтесь пропустить этот раздел, если вам нужно только более интуитивное понимание основных концепций. В противном случае, давайте углубимся в детали бумаги. На самом деле VAE — это только пример в исходной статье основных идей. Однако благодаря своей полезности он стал широко известен. Вместо этого реальная ценность статьи заключается в идеях, лежащих в основе вариационного байесовского алгоритма автоматического кодирования.
Проблемный режим
Проблема, которую пытается решить статья, заключается в том, что у нас есть большой набор данных одинаково распределенных независимых выборок стохастической переменной X. Эта переменная создается скрытым процессом, зависящим от скрытой переменной z, полученной из предыдущего распределения с параметрами θ. Это конкретная выборка Xсоздается из условного распределения (вероятности)
Наша цель, чтобы иметь возможность генерировать новые выборки из X, состоит в том, чтобы найти предельное правдоподобие p(x) Но мы вообще столкнулись с проблемами с неподатливостью. Например, во многих случаях мы не в состоянии вычислить интеграл
аналитически по всем значениям z или апостериорного p(z| x) неизвестно.
Чтобы сформулировать обходные пути для этой проблемы, мы делаем два простых предположения: априорное, p(z), и вероятность p(x| z )имеют PDF-файлы, которые дифференцируемы (почти везде) и зависят от параметров θ.
Теперь мы хотели бы сформулировать цель, которая находит оценку максимального правдоподобия параметров θ и аппроксимирует апостериорное распределение, чтобы иметь возможность эффективно представлять данные. Другими словами, мы хотим закодировать данные в представление распределения и иметь возможность генерировать выборки из этого распределения.
Вывод вариационной границы — цель обучения
Начиная с цели: генерировать изображения. Это означает, что при наличии скрытой переменной z мы хотим реконструировать и/или сгенерировать изображение x. Для этого мы хотим оценить условное распределение или вероятность, p(x| z)однако, поскольку скрытая переменная z нам неизвестно, мы не можем обучить модель, чтобы она соответствовала этому распределению напрямую. Вместо этого идея состоит в том, чтобы позволить сети декодера аппроксимировать вероятность, а затем использовать правило Байеса, чтобы найти предельное распределение, которому следуют данные, т.е. вычислить
Мы можем предположить гауссово априорное значение для z, но у нас все еще остается проблема, заключающаяся в том, что апостериорное значение неразрешимо. Предлагаемое решение состоит в том, чтобы аппроксимировать это распределение сетью кодировщика q с параметрамиϕ.
Теперь мы начнем с логарифма маргинального распределения, приведенного выше, и попробуем вывести цель, которую мы можем оптимизировать с помощью стохастического градиентного спуска. Сначала мы используем уловку и умножаем числитель и знаменатель на наш приблизительный апостериор.
Затем мы используем логарифмические правила, чтобы разделить члены так, как нам удобно.
Поскольку левая часть не зависит от z, мы можем использовать дополнительный прием и взять математическое ожидание вместо z и тогда будет затронута только правая сторона.
Теперь мы находимся в точке, где мы можем видеть, что первый член представляет собой математическое ожидание логарифма правдоподобия данных. Максимальное увеличение этого члена будет означать, что у нас есть высокая вероятность правильного восстановления данных x. Два других термина, которые мы можем определить из определения KL-дивергенции, измеряют, насколько близко наше аппроксимированное распределение соответствует априорному и истинному апостериорному.
Обратите внимание, что последний член трудноразрешим, поскольку апостериорное значение неизвестно, но мы можем использовать, что KL-дивергенция будет неотрицательной, и сформировать нижнюю границу предельного правдоподобия следующим образом:
Это наша конечная цель. Первый член будет показателем реконструкции, который измеряет, насколько хорошо декодер восстанавливает данные, а второй член будет конкурирующей целью, которая приближает приблизительную апостериорную вероятность к априорной. На практике мы часто выбираем априорную норму, а второй член будет иметь регуляризирующий эффект, упрощающий распределение, выдаваемое кодировщиком.
Трюк с репараметризацией — как сэмплировать z?
Чтобы все это заработало, необходимо учесть еще одну деталь. Мы обучаем декодер генерировать выборку из условного распределения при заданном значении z. Но как нам сгенерировать z в первую очередь? Нам нужно найти способ выборки дифференцируемого распределения, чтобы иметь возможность оптимизировать его с помощью стохастического градиентного спуска. Хитрость, представленная в статье, заключается в том, чтобы отделить стохастическую часть z и затем преобразовать ее с заданными входными данными и параметрами кодировщика с помощью функции преобразования г
На практике мы сначала генерируем
из распределения, не зависящего от параметров кодировщика. Затем преобразуем его с помощью функции к нужному распределению. В VAE мы используем простую репараметризацию
Архитектура модели VAE
Вариационный автоэнкодер — это только пример того, как можно использовать идеи, представленные в статье. Однако с небольшими изменениями он показал, что это очень полезный пример.
В VAE мы выбираем априор скрытой переменной как единицу Гаусса с диагональной ковариационной матрицей. Это говорит модели, что мы хотим, чтобы она изучила представление скрытой переменной с независимыми функциями, что на самом деле является довольно строгим предположением.
Мы реализуем кодировщик и декодер как простые MLP с несколькими уровнями. Кодер выводит среднее значение и стандартное отклонение приблизительного апостериорного значения. Поскольку мы предполагаем, что оно является гауссовским с диагональной ковариацией, мы используем описанный выше прием репараметризации для выборки скрытого распределения. Мы передаем это значение $z$ декодеру, который генерирует реконструированную точку данных.
Проигрыш состоит из двух конкурирующих целей. KL-дивергенция, которая подталкивает распределение скрытой переменной к единичному нормальному распределению, а потери при реконструкции подталкивают эту модель к точному восстановлению исходных входных данных. В этом случае мы можем аналитически вычислить KL-дивергенцию, и, проведя вычисления, получим следующую формулу
где J — размерность z, и если вы немного посмотрите на формулу, вы поймете, что она максимальна для стандартного нормального распределения. . Потеря реконструкции
реализован с потерей BCE в PyTorch, которая, по сути, делает выходные значения пикселей похожими на входные.
Код в PyTorch
Реализация вариационного автоэнкодера упрощена и содержит только основные части. Пример относится к набору данных MNIST, а для сети кодировщика и декодера мы используем простой MLP.
Импорт
Модель
Кодер и декодер представляют собой зеркальные сети, состоящие из двух слоев. В кодировщике мы передаем входные данные в скрытое измерение через линейный слой, а затем передаем скрытое состояние в два разных линейных слоя, выводя среднее значение и стандартное отклонение скрытого распределения соответственно.
Затем мы делаем выборку из скрытого распределения и вводим ее в декодер, который, в свою очередь, выдает вектор той же формы, что и вход.
Конфигурация обучения
Загрузка данных MNIST
Подобно примерам в статье, мы используем набор данных MNIST для демонстрации концепций модели.
Настройка обучения
Как описано выше, потеря состоит из двух разных терминов: потери при реконструкции, реализованной здесь с помощью BCE-потери, и KL-дивергенции. Мы делаем довольно строгие предположения, что априор $z$ является единичной нормалью, а апостериор приблизительно гауссовской с диагональной матрицей ковариаций, что означает, что мы можем упростить выражение для KL-дивергенции, как описано выше.
1875it [00:28, 65.04it/s, loss=5.5e+3]
1875it [00:25, 72.15it/s, loss=4.75e+3]
1875it [00:26, 71.89it/s, loss=4.82e+3]
1875it [00:26, 71.39it/s, loss=4.17e+3]
1875it [00:26, 70.06it/s, loss=4.17e+3]
1875it [00:26, 70.68it/s, loss=4.22e+3]
1875it [00:26, 70.44it/s, loss=4.14e+3]
1875it [00:27, 68.29it/s, loss=4.59e+3]
1875it [00:28, 66.01it/s, loss=4.26e+3]
1875it [00:26, 69.46it/s, loss=4.06e+3]Как генерировать новые изображения
Ниже мы видим, что вариационный автоэнкодер генерирует слегка различающиеся изображения при одинаковых входных данных благодаря выборке нового значения скрытой переменной в каждом поколении.
Имея некоторое представление о том, как работают VAE, и увидев пример их реализации, я надеюсь, что теперь вы лучше подготовлены для понимания и реализации более современных архитектур, включающих эти идеи! По крайней мере, так я себя чувствую после прочтения газеты.
