Вариационный автокодировщик

Вариационный автокодировщик (англ. Variational Autoencoder, VAE) — автокодировщик^{[на 28.01.19 не создан]} (генеративная модель, которая учится отображать объекты в заданное скрытое пространство (и обратно)) основанный на вариационном выводе.

Предпосылки

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

Вариационный автокодировщик в свою очередь предлагает пользователю самому определить распределение скрытого вектора.

Описание

Порождающее моделирование (англ. Generative modelling) — область машинного обучения, имеющая дело с распределением [math]P(X)[/math], определенном на датасете [math]X[/math] из пространства (возможно многомерного) [math]X[/math]. Так, например, популярные задачи генерации картинок имеют дело с огромным количеством измерений (пикселей).

Также как и в обыкновенных кодировщиках у нас имеется скрытое вероятностное пространство [math]Z[/math] соответствующее случайной величине [math](z, P(z))[/math] (распределенной как-нибудь фиксированно, здесь [math]\sim N(0, 1)[/math]). И мы хотим иметь декодер [math]f(z, \theta) \colon Z \times \Theta \to X [/math]. При этом мы хотим найти такие [math]\theta[/math], чтобы после разыгрывания [math]z[/math] по [math]P(z)[/math] мы получили "что-то похожее" на элементы [math]X[/math].

Вообще, для любого [math]x \in X[/math] мы хотим считать [math]P(x) = \int P(x|z; \theta)P(z)dz[/math], здесь мы заменили [math]f(z, \theta)[/math] на [math]P(x|z; \theta)[/math], чтобы явно показать зависимость между [math]x[/math] и [math]z[/math] и после этого применить формулу полной вероятности. Обычно [math]P(x|z; \theta)[/math] около нуля почти для всех пар [math](x, z)[/math]. Основная идея в том, что мы хотим теперь генерировать [math]z[/math], который бы давали что-то около [math]x[/math] и только их суммировать в [math]P(x)[/math]. Для этого нам требуется ввести еще одно распределение [math]Q(z|X)[/math], которое будет получать [math]x[/math] и говорить распределение на [math]z[/math] которое наиболее вероятно будет генерировать нам такой [math]x[/math]. Теперь нам нужно как-то сделать похожими распределения [math]E_{z\sim Q}P(X|z)[/math] и [math]P(X)[/math].

Рассмотрим следующую дивергенцию Кульбака-Лейблера (Kullback–Leibler divergence, KL-div).

[math]D[Q(z)||P(z|X)] = E_{z∼Q} [log Q(z|X) − log P(z|X)][/math],

Распишем [math]P(z|X)[/math] как [math]P(X|z) * P(z) / P(X)[/math].

[math]D[Q(z)||P(z|X)] = E_{z∼Q} [log Q(z) − log P(X|z) - log P(z)] + log P(X)[/math],

Что эквивалентно:

[math]logP(x) - D[Q(z)||P(z|X)] = E_{z∼Q}[log P(X|z)] - D[Q(z)||P(z)][/math],

Рассмотрим эту штуку для [math]Q(z|X)[/math], тогда:

[math]logP(x) - D[Q(z|X)||P(z|X)] = E_{z∼Q}[log P(X|z)] - D[Q(z|X)||P(z)][/math],

Посмотрим, на это равенство. Правую часть мы можем оптимизировать градиентным спуском (пусть пока и не совсем понятно как). В левой же части первое слагаемое — то, что мы хотим максимизировать. В то же время [math]D[Q(z|X)||P(z|X)][/math] мы хотим минимизировать. Если у нас [math]Q(z|X)[/math] — достаточно сильная модель, то в какой-то момент она будет хорошо матчить [math]P(z|X)[/math], а значит их дивергенция Кульбака-Лейблера будет почти 0. Значит, при оптимизации можно исключить эту часть и стараться максимизировать только правую. В качестве бонуса мы еще получили более "податливую" [math]P(z|X)[/math], вместо нее можно смотреть на [math]Q(z|X)[/math].

Теперь разберемся как оптимизировать правую часть. Сначала нужно определиться с моделью для [math]Q(z|X)[/math]. Обычно ее берут равной [math]N(z|\mu(X, \theta), \sigma(X, \theta))[/math]. Где [math]\mu[/math] и [math]\sigma[/math] какие-то детерминированные функции на X с обучаемыми параметрами [math]\theta[/math], которые мы впредь будем опускать (обычно используются нейронные сети).

Нетрудно проверить, что для дивергенция Кульбака-Лейблера двух нормальных распределений имеет следующий вид:

[math]D_{K}[N(\mu_1, \Sigma_0)||N(\mu_1, \Sigma_0)][/math], KLD есть [math]\frac{1}{2} (tr(\Sigma_1^{-1}\Sigma_0) + (\mu_1 - \mu_0)^T\Sigma_1^{-1}(\mu_1 - \mu_0) - k + log(\frac{det\Sigma_1}{det\Sigma_0})) [/math].

Это значит, что

[math]D[Q(z|X)||P(z)] = D[N(\mu(X), \Sigma(X))||N(0, I)] = \frac12 (tr(\Sigma(X)) + \mu(X)^T\mu(X) - k - log(det\Sigma(X)))[/math].

Теперь здесь можно считать градиенты, для BackPropagation. С первым слагаемым в правой части все немного сложнее. [math]E_{z∼Q}[log P(X|z)][/math] мы можем считать методом Монте-Карло(МК), но тогда такая штука (из-за того, что переменные спрятаны в распределении, из которого мы генерируем себе выборку, для МК) не является гладкой относительно них, а значит непонятно, как проталкивать через это градиент. Для того, чтобы все-таки можно было протолкнуть градиент, применяется так называемый трюк репараметризации, который базируется на простой формуле [math]N(\Sigma(X), \mu(X)) = \mu(X) + \Sigma^{\frac12}(X) * N(0, I) [/math].

[math]E_{z∼Q}[log P(X|z)] = E_{\epsilon \sim N(0, I)}[log P(X = f(\mu(X) + \Sigma^{\frac12}(X) * \epsilon), \theta)][/math].

В такой форме мы уже можем использовать BackPropagation для переменных из функций [math]\Sigma[/math] и [math]\mu[/math].

Следующая картинка лучше поможет осознать структуру VAE и, в частности, зачем нужен (и как работает) трюк репараметризации.

На левой части диаграмма без использования reparameterization trick. На правой части диаграмма с использованием reparameterization trick.

взято из https://arxiv.org/pdf/1606.05908.pdf

Пример реализации

Ниже приведена реализация частного случая VAE на языке Python с использованием библиотеки Pytorch. Эта реализация работает с датасетом MNIST. Размерность скрытого слоя — 2. Координаты в нем считаются независимыми (из-за этого, например, матрица [math]\Sigma[/math] диагональная, и формула для расчета KLD немного другая).

class VariationalAutoencoder(nn.Module):
   def __init__(self):
       super().__init__()
       self.mu = nn.Linear(32, 2)
       self.gamma = nn.Linear(32, 2)
       self.encoder = nn.Sequential(nn.Linear(784, 32), nn.ReLU(True))
       self.decoder = nn.Sequential(nn.Linear(2, 32), nn.ReLU(True), nn.Linear(32, 784), nn.Sigmoid())

   def forward(self, x):
       mu, gamma = self.encode(x)
       encoding = self.reparameterize(mu, gamma)
       x = self.decoder(encoding)
       return x, mu, gamma

   def reparameterize(self, mu, gamma):
       if self.training:
           sigma = torch.exp(0.5*gamma)
           std_z = Variable(torch.from_numpy(np.random.normal(0, 1, size=sigma.size())).float())
           encoding = std_z.mul(sigma).add(mu)
           return encoding
       else:
           return mu

   def encode(self, x):
       x = self.encoder(x)
       mu = self.mu(x)
       gamma = self.gamma(x)
       return mu, gamma
 
   def decode(self, x):
       return self.decoder(x)

   def latent(self, x):
       mu, gamma = self.encode(x)
       encoding = self.reparameterize(mu, gamma)
       return encoding

def loss_function(input, output, mu, gamma, batch_size=batch_size):
   BCE = F.binary_cross_entropy(output, input)
   KLD = -0.5*torch.sum(1 + gamma - mu.pow(2) - gamma.exp())
   KLD /= batch_size*784
   return BCE + KLD

Применение

Область применения вариационных автокодировщиков совпадает с областью применения обыкновенных автокодировщиков. А именно:

• Каскадное обучение глубоких сетей (хотя сейчас применяется все реже, в связи с появлением новых методов инициализации весов);
• Уменьшение шума в данных;
• Уменьшение размерности данных (иногда работает лучше, чем метод главных компонент^{[на 28.01.19 не создан]}).

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

См. также

• Автокодировщик
• Порождающие состязательные сети

Примечания

• Вариационные автокодировщики: теория и рабочий код
• Tutorial - What is a variational autoencoder?
• Intuitively Understanding Variational Autoencoders

Источники информации

• Tutorial on Variational Autoencoders
• Datalore презентация Дениса Степанова