Auto-Encoding Variational Bayes

Разбор статьи

Проблема и мотивация

Направленные вероятностные модели с непрерывными скрытыми переменными $\mathbf{z}$ полезны для генеративного моделирования, обучения представлений без учителя и маргинального вывода. Однако их обучение методами максимального правдоподобия или байесовскими методами сталкивается с серьезными вычислительными препятствиями, когда:

1. Трудновычислимость апостериорного распределения (posterior intractability): Для нелинейных зависимостей (например, в глубоких нейросетях) интеграл маргинального правдоподобия $p_{\mathbf{\theta }}(\mathbf{x})=\int p_{\mathbf{\theta }}(\mathbf{z})p_{\mathbf{\theta }}(\mathbf{x}|\mathbf{z})d\mathbf{z}$ и апостериорное распределение $p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$ не имеют аналитического решения в замкнутой форме, что делает неприменимыми стандартный алгоритм Expectation-Maximization (EM) и вариационный байесовский вывод среднего поля (mean-field variational Bayes).
2. Большие наборы данных: Традиционные подходы к сэмплированию, такие как Monte Carlo EM, требуют медленного итеративного сэмплирования MCMC для каждой точки данных по всему датасету, из-за чего оптимизация на мини-батчах для больших масштабов становится непомерно дорогой.

Стандартный наивный градиентный оценщик функции счета (score-function) ${\nabla }_{\mathbf{\varphi }}{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z})}[f(\mathbf{z})]$ обладает чрезвычайно высокой дисперсией, препятствующей надежной оптимизации с помощью стохастического градиентного подъема.

Ключевые идеи

• Трюк с репараметризацией (Reparameterization Trick): Вместо прямого сэмплирования $\mathbf{z}$ из параметризованного распределения $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ переменная переписывается как детерминированное преобразование независимой переменной шума $\mathbf{ϵ}\sim p(\mathbf{ϵ})$: $\mathbf{z}=g_{\mathbf{\varphi }}(\mathbf{ϵ},\mathbf{x})$. Это позволяет дифференцировать по $\mathbf{\varphi }$ с помощью стандартного обратного распространения ошибки с существенно более низкой дисперсией.
• Амортизированный вывод (модель распознавания): Вместо оптимизации вариационных параметров отдельно для каждого объекта энкодер-нейросеть $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ предсказывает вариационные параметры (такие как среднее и дисперсия латентного гауссовского распределения), обусловленные $\mathbf{x}$.
• Auto-Encoding Variational Bayes (AEVB): Совместно оптимизирует генеративные параметры $\mathbf{\theta }$ (вероятностный декодер) и вариационные параметры $\mathbf{\varphi }$ (вероятностный энкодер) с использованием стохастического градиентного подъема по репараметризованной вариационной нижней оценке.

Метод

Вариационная целевая функция максимизирует нижнюю оценку обоснованности (Evidence Lower Bound, ELBO) $\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$: $$\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)})\ge \mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=-D_{KL}(q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})\parallel p_{\mathbf{\theta }}(\mathbf{z}))+{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})}\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|\mathbf{z})\right]$$

Применяя трюк с репараметризацией ${\mathbf{z}}^{(i,l)}=g_{\mathbf{\varphi }}({\mathbf{ϵ}}^{(i,l)},{\mathbf{x}}^{(i)})$ с ${\mathbf{ϵ}}^{(i,l)}\sim p(\mathbf{ϵ})$, получаем общий оценщик SGVB (Оценщик A): $${\tilde{\mathcal{L}}}^A(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=\frac{1}{L}\sum _{l=1}^L\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)},{\mathbf{z}}^{(i,l)})-\log q_{\mathbf{\varphi }}({\mathbf{z}}^{(i,l)}|{\mathbf{x}}^{(i)})\right]$$

Когда априорное распределение $p_{\mathbf{\theta }}(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ и приближенное апостериорное $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})=\mathcal{N}(\mathbf{\mu },{\mathbf{\sigma }}^2\mathbf{I})$ являются гауссовскими, дивергенция Кульбака — Лейблера (KL) может быть вычислена аналитически, что дает Оценщик B: $$\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})\simeq \frac{1}{2}\sum _{j=1}^J\left(1+\log (({\sigma }_j^{(i)}{)}^2)-({\mu }_j^{(i)}{)}^2-({\sigma }_j^{(i)}{)}^2\right)+\frac{1}{L}\sum _{l=1}^L\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|{\mathbf{z}}^{(i,l)})$$ где ${\mathbf{z}}^{(i,l)}={\mathbf{\mu }}^{(i)}+{\mathbf{\sigma }}^{(i)}\odot {\mathbf{ϵ}}^{(l)}$ и ${\mathbf{ϵ}}^{(l)}\sim \mathcal{N}(0,\mathbf{I})$.

Для мини-батчей размера $M$ итоговая оценка нижней границы равна: $${\tilde{\mathcal{L}}}^M(\mathbf{\theta },\mathbf{\varphi };{\mathbf{X}}^M)=\frac{N}{M}\sum _{i=1}^M\tilde{\mathcal{L}}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$$ На практике оказывается достаточно установить $L=1$ сэмпл на объект при размере мини-батча, например, $M=100$.

Результаты

• Настройки моделей: Генеративные модели оценивались на MNIST (декодер MLP с распределением Бернулли, 500 скрытых нейронов) и Frey Face (декодер MLP с распределением Гаусса, 200 скрытых нейронов), оптимизация проводилась с помощью Adagrad со скоростями обучения из $\{0.01,0.02,0.1\}$ и затуханием весов (weight decay).
• Оптимизация нижней оценки: Для различных размерностей скрытого пространства ($N_{\mathbf{z}}\in \{2,5,10,20,400\}$ на Frey Face и MNIST) AEVB сходился значительно быстрее и достигал более высокой нижней границы, чем алгоритм wake-sleep.
• Устойчивость к переобучению: Увеличение размерности $\mathbf{z}$ не приводило к переобучению благодаря регуляризующему эффекту аналитического члена дивергенции KL по отношению к априорному распределению.
• Оценка маргинального правдоподобия: В экспериментах на MNIST с 3-мерным скрытым пространством с использованием оценки маргинального правдоподобия на основе HMC алгоритм AEVB сходился быстрее, чем wake-sleep, и многократно превосходил возможности Monte Carlo EM (который был слишком медленным для запуска на полном наборе MNIST).
• Время работы: Обучение занимало примерно 20–40 минут на миллион сэмплов на процессоре Intel Xeon (~40 GFLOPS).

Ограничения и открытые вопросы

• Непрерывные скрытые переменные: Трюк с репараметризацией напрямую опирается на дифференцируемость по параметрам, из-за чего его затруднительно напрямую применить к дискретным распределениям скрытых переменных.
• Зазор аппроксимации: Предположение о диагональном гауссовском распределении для $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ ограничивает выразительную способность модели, если истинное апостериорное распределение является мультимодальным или имеет сложную корреляционную структуру.
• Расширение до полного вариационного байесовского вывода: Экспериментальная оценка фокусировалась на оценке параметров $\mathbf{\theta }$ методами MAP/ML; вариационный вывод для глобальных параметров $\mathbf{\theta }$ был сформулирован в приложении, но оставлен для дальнейшей эмпирической проверки.