Обучение VAE и нижняя вариационная граница / forpes.ru

20.03.2023 18:43

IvanRodkin 0 835 Источник

В машинном обучении для оптимизации каких-либо вероятностных моделей часто используется нижняя вариационная граница (evidence lower bound), о ней и пойдет речь.

В первую очередь задача. Допустим мы генерируем картинки следующим образом: сэмплируем из стандартного многомерного нормального распределения какой-то вектор потом его преобразуем некоторым алгоритмом $p_\theta(x|z)$ , параметриуемым $\theta$ и хотим чтобы на выходе было нормальное распределение с центром в точке , которое похоже на лицо человека, с дисперсией, например, равной $\frac{1}{dim(x)}$ , где длина вектора . процесс работает со скрытыми переменными из стандартного нормального распределения и с полученными переменными , распределение которых явно зависит от переменных . То есть наш алгоритм явно моделирует условное распределение: берет на вход реализацию и возвращает новое распределение, которое постулируется нормальным. Но фактически, так как у этого нормального фиксирована дисперсия, наш алгоритм просто возвращает среднее этого нормального распределения.

Совместное распределение такого процесса мы назовем $p_\theta(x, z).$

Как такое обучить, если у нас есть очень много лиц? Заведем процесс, обратный процессу , его назовем . Суть его та же, что и Это просто совместное распределение $q_\psi(x, z),$ где имеет распределение наших данных (реальных лиц), а это некоторые зависящие от переменные. Зависят они по нашему правилу $q_\psi(z|x)$ парметризуемому $\psi$ . $q_\psi(z|x)$ также берет на вход реализацию и возвращает нормальное распределение. Но уже не с фиксированной дисперсией, а меняющейся по каждому направлению. Т.е. алгоритм по переменной выдает не только среднее нормального в пространстве , но и вектор дисперсий.

Для простоты я буду опускать далее индексы $\theta$ и $\psi$ , подразумевая всегда, что зависит от первого, а от второго

Окей, а теперь формально запишем, что мы хотим от нашей модели. Мы хотим, чтобы в среднем мы нормальные картинки генерировали. А на языке статистики это значит, что мы хотим, чтобы правдоподобие наших реальных картинок было максимальным в нашей генерирующей модели . Что такое правдоподобие картинки в нашей модели? Это просто

$p(x) = \int p(x, z) dz$

То есть это значение маргинального распределения на наших картинках. Само маргинальное распределние -- просто распределение картинок, которые мы генерируем. И мы пытаемся максимизировать это распределение в точках реальных картинок.

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

$\log p(x) = \log \int p(x, z) dz$

Так как картинок у нас много, усредним по "всем реальным картинкам" эту величину, взяв по их распределению мат. ожидание и получим

$\mathbb{E}_{x \sim Real } {\log p(x)} \approx \frac{1}{n}\sum_i \log p(x_i)$

Но теперь вспомним, что для такого распределения реальных картинок у нас уже есть обозначение -- это просто , просто по определению (см. выше).

Так что мы по факту хотим оптимизировать

$\mathbb{E}_{x \sim q(x)} \log p(x) = \mathbb{E}_{x \sim q(x)} {\log \int p(x, z) dz}$

Теперь распишем аккуратно то, что находится справа под математическим ожиданием, зафиксировав конкретный

$\log \int p(x, z)dz = \log \int q(z|x)\frac{p(x, z)}{q(z|x)}dz = \log \mathbb{E}_{z \sim q(z|x)}\frac{p(x, z)}{q(z|x)}$

Для вогнутых функций , типо логарифма, выполняется неравенство Йенсена, которое выглядит так:

$f(\mathbb{E}g) \geqslant \mathbb{E}f(g)$

Получаем, применим его:

$\log p(x) = \log \mathbb{E}_{z \sim q(z|x)}\frac{p(x, z)}{q(z|x)} \geqslant \mathbb{E}_{z \sim q(z|x)}\log \frac{p(x, z)}{q(z|x)}$

Величина слева и называется нижней вариационной границей или ELBO (Evidence Lower BOund)

Почему она нам вообще интересна? На это есть две причины:

Максимизируя нижнюю границу (по нашим параметрам $\theta$ и $\psi$ ) мы можем быть уверены, что наша оцениваемая величина не меньше её значения, а значит тоже как бы при этом процессе максимизируется
На самом деле эта нижняя граница может стать точной оценкой нашей величины (т.е. не только снизу)

Если первое, наверное, очевидно, то второе требует пояснения. Для этого, раскроем совместное распределение по правилу

$\log p(x) \geqslant \mathbb{E}_{z \sim q(z|x)}\log \frac{p(x, z)}{q(z|x)} = \mathbb{E}_{z \sim q(z|x)}\log\frac{p(x)p(z|x)}{q(z|x)} = \\ = \mathbb{E}_{z \sim q(z|x)} \log p(x) + \mathbb{E}_{z \sim q(z|x)}\log \frac{p(z|x)}{q(z|x)}$

Первое слагаемое от не зависит, поэтому оно просто $\log p (x)$ (как раз то, что мы оцениваем!). Второе слагаемое, если поменять числитель и знаменатель под логарифмом (и следовательно вынести из него минус), будет просто минус дивергенцией Кульбака — Лейблера между и , обозначаемой как.

$\log p(x) \geqslant \mathbb{E}_{z \sim q(z|x)}\log \frac{p(x, z)}{q(z|x)} = \log p(x) -KL[q(z|x)||p(z|x)]$

Что это вообще значит?

Максимизируя ELBO, мы не только максимизируем $\log p(x)$ но и минимизируем,
Если мы доминимизируем до нуля, то мы будем в аккурат максимизировать $\log p(x)$ (потому что ELBO будет попросту совпадать с ней)

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

Отлично! Теперь мы просто забиваем на лог-правдоподобие $\log p(x)$ и оптимизируем ELBO, которое потом все равно сравняется с первым

Следующая проблема, которая перед нами встает: как оценить это самое ELBO

Что мы вообще знаем в нашей модели? Мы умеем с помощью наших алгоритмов (нейросеток) и оценивать и, так что хотелось бы с помощью них как-то посчитать ELBO. Для этого вместо того, чтобы раскрывать совместное распределение в формуле ELBO как раскроем его так:. Что как бы хорошо для нас, потому что подобно тому, как это, по определению, распределение наших данных, мы тоже знаем и по определению это стандартное многомерное нормальное! (см. определение выше)

$ELBO = \mathbb{E}_{z \sim q(z|x)}{\log \frac{p(z)p(x|z)}{q(z|x)}} = \mathbb{E}_{z \sim q(z|x)}\log p(x|z) - \mathbb{E}_{z \sim q(z|x)}\log \frac{q(z|x)}{p(z)} = \\ =\mathbb{E}_{z \sim q(z|x)}\log p(x|z) - KL[q(z|x)||p(z)]$

Осталось понять, почему то, что мы получили -- это как бы хорошо. Для этого осознаем один факт про математическое ожидание:

$\mathbb{E}_{q(x)}\mathbb{E}_{q(z|x)} = \mathbb{E}_{q(x, z)}$

Доказать этот факт довольно просто:

$\mathbb{E}_{q(x)}\mathbb{E}_{q(z|x)} f(x, z) = \int q(x)\int q(z|x)f(x, z)dzdx = \int \int q(x)q(z|x)f(x, z)dzdx =\\ = \int \int q(x, z)f(x, z)dzdx = \mathbb{E}_{q(x, z)} f(x, z)$

А теперь вспомним, что на самом-то деле мы хотим не саму ELBO оптимизировать, а его мат. ожидание по , которое по сути является распределением наших изначальных картинок-лиц.

И тут нам как раз пригождается это свойство матожей:

$\mathbb{E}_{x\sim q(x)}ELBO = \mathbb{E}_{q(x)}(\mathbb{E}_{z \sim q(z|x)}\log p(x|z) - KL[q(z|x)||p(z)]) = \\ = \mathbb{E}_{x, z \sim q(x, z)}\log p(x|z) - \mathbb{E}_{q(x)} KL[q(z|x)||p(z)]$

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

Вот у нас есть картинка . Как посчитать ELBO?

Первое слагаемое (уменьшаемое) можно самым тупым образом оценить: чтобы оценить мат. ожидание, можно взять среднее. А среднее для одного элемента это сам этот элемент :). Так что мы просто сэмплируем $z \sim q(z|x)$ , прогнав через наш алгоритм и получив таким образом это условное распределение (оно у нас нормальное, а мы предсказали по сути среднее и вектор дисперсий). После чего прогоняем полученный через генератор , получив распределение $p(\cdot|z).$ В которое мы будем подставлять . Оно у нас тоже нормальное, причем с фиксированной дисперсией (т.е. предсказали мы нашим алгоритмом только среднее $\mu(z)$

Так как плотность многомерного нормального (оно же в данном случае правдоподобие) описывается формулой

$p(x) = \frac{1}{\sqrt{(2\pi)^n |\Sigma|}}e^{-\frac{1}{2}(x-\mu)^T\Sigma^{-1}(x-\mu)}$

Получается, что лог-правдоподобие:

$\log p(x|z) = C -\frac{1}{2dim(x)}(x-\mu(z))^T(x-\mu(z)) = C - \frac{1}{2}MSE(x, \mu(z))$

Второе слагаемое (вычитаемое) -- KL-дивергенция двух нормальных с известными средними и дисперсиями. Здесь мы уже можем явно записать, чему она равна:

$KL[N(a, \Sigma_1)||N(b, \Sigma_2)] = \frac{1}{2}(tr(\Sigma_2^{-1}\Sigma_1-I) +(a - b)^T\Sigma_2^{-1}(a-b) +\log\frac{|\Sigma_2|}{|\Sigma_1|})$

Таким образом мы можем явно оценить два слагаемых: первое через MSE реального изображения и реконструированного, второе по формуле KL-дивергенции. А значит оптимизировать наши параметры

Обучение VAE и нижняя вариационная граница +9

Комментарии (0)