Марковские методы Монте-Карло (MCMC)

Апостериорное распределение реальной модели, распределение молекулярных конфигураций или энергетическая модель (energy-based model) изображений имеют плотность, которую можно вычислить лишь с точностью до константы, $p(x) = \tilde p(x) / Z$, где $Z$ — вычислительно неразрешимый (intractable) интеграл, причём в пространстве высокой размерности. Прямым методам из заметки о методах Монте-Карло нужно предлагающее распределение, близкое к $p$, а такого нет. Марковские методы Монте-Карло (Markov chain Monte Carlo, MCMC) переворачивают эргодическую теорему: строим цепь Маркова, стационарное распределение которой равно $p$, запускаем её и используем её состояния как сэмплы. Для построения нужны только отношения $\tilde p(y)/\tilde p(x)$, в которых $Z$ сокращается. В этой последней заметке курса мы выводим алгоритм Метрополиса–Гастингса, показываем, как он даёт сбой, и заканчиваем динамикой Ланжевена, которая идёт по градиенту логарифма плотности и даёт score, который затем используют диффузионные модели.

Обозначения. $p(x) = \tilde p(x)/Z$ — целевое распределение; $q(y \mid x)$ — предлагающее распределение (proposal); $X_0, X_1, \dots$ — цепь. $\nabla \log p(x)$ — score (функция score, градиент логарифма плотности), который не зависит от $Z$.

Условие детального баланса#

Цепь с вероятностями переходов $P(x \to y)$ удовлетворяет условию детального баланса (detailed balance) относительно $p$, если

$$p(x)\, P(x \to y) = p(y)\, P(y \to x) \quad \text{для всех } x, y .$$

В равновесии поток вероятности из $x$ в $y$ равен обратному потоку. Суммирование по $x$ даёт $\sum_x p(x) P(x \to y) = p(y)$: распределение $p$ стационарно. Условие детального баланса локально, и его легко обеспечить; если цепь к тому же неприводима и апериодична, то заметка о цепях Маркова гарантирует сходимость к $p$ из любого начального состояния — для конечного числа состояний; в непрерывных пространствах состояний нужны аналогичные, более тонкие условия.

Алгоритм Метрополиса–Гастингса#

Из текущего состояния $x$ предложим $y \sim q(\cdot \mid x)$ и примем его с вероятностью

$$\alpha(x, y) = \min\Big(1, \frac{\tilde p(y)\, q(x \mid y)}{\tilde p(x)\, q(y \mid x)}\Big) ;$$

иначе останемся в $x$. При $y \ne x$ вероятность перехода равна $q(y \mid x)\,\alpha(x, y)$, а выражение $p(x) q(y \mid x) \alpha(x, y) = \min\big(p(x) q(y \mid x),\ p(y) q(x \mid y)\big)$ симметрично по $x$ и $y$: условие детального баланса выполнено. Константа $Z$ входит и в числитель, и в знаменатель и сокращается. При симметричном предлагающем распределении, например $y = x + \sigma \varepsilon$ с $\varepsilon \sim N(0, 1)$, отношение равно просто $\tilde p(y)/\tilde p(x)$: шаги вверх по плотности принимаются всегда, шаги вниз — иногда. Это случайное блуждание Метрополиса (random-walk Metropolis) — алгоритм Метрополиса и соавторов 1953 года.

Цепь корректна при любом размере шага $\sigma$, но от $\sigma$ зависит, насколько хорошо она работает:

Рисунок

Метрополис со слишком малым, хорошим и слишком большим шагом

За 20 000 шагов на стандартном нормальном целевом распределении при шаге $0.1$ принимается 97% предложений, но цепь равноценна всего 57 независимым сэмплам; при шаге $25$ принимается 5%, и цепь равноценна 708; при шаге $2.4$ принимается 44%, и цепь равноценна 4815. Для случайного блуждания Метрополиса в высокой размерности теория рекомендует подбирать шаг так, чтобы доля принятых (acceptance rate) была около 0.23; в одномерном случае — около 0.44.

Перемешивание и диагностика#

Последовательные состояния цепи коррелированы, поэтому $n$ состояний обычно несут меньше информации, чем $n$ независимых сэмплов. Эффективный размер выборки (effective sample size, ESS) относится к интересующей величине $g(X_t)$, а не к цепи в целом: если $\rho_k$ — автокорреляция $g(X_t)$ на лаге $k$ в стационарной цепи, то

$$\text{ESS} = \frac{n}{1 + 2\sum_{k \ge 1} \rho_k} ,$$

и эта величина показывает, скольким независимым сэмплам равноценны $n$ состояний при оценке $E_p[g(X)]$. Ошибка среднего $g$ по цепи тогда примерно равна $\sigma / \sqrt{\text{ESS}}$, а не $\sigma/\sqrt n$, где $\sigma^2 = \operatorname{Var}_p g(X)$. Для этого нужно, чтобы цепь вышла на стационарный режим, ряд автокорреляций сходился и для цепи выполнялась центральная предельная теорема — она верна для «хороших» цепей, но не для всех. Положительные корреляции делают ESS меньше $n$, отрицательные могут сделать его больше: у цепи с автокорреляциями $\rho_k = (-0.5)^k$ получается $\text{ESS} = 3n$. Значения ESS в этой заметке получены простой оценкой, которая обрывает сумму там, где автокорреляции падают ниже $0.05$; они иллюстративны, а в рабочих инструментах используют более аккуратные оценки. На практике к этому добавляются три привычки: отбрасывать начальный участок цепи, прогрев (burn-in), до выхода на стационарный режим, запускать несколько цепей из разнесённых начальных точек и сравнивать дисперсию между цепями с дисперсией внутри них (статистика $\hat R$).

Никакая диагностика не обнаружит того, чего цепь ни разу не посещала. Для целевого распределения с двумя разделёнными модами цепь с малым шагом остаётся в той моде, в которой начала:

Рисунок

Цепь, застрявшая в одной моде

За 20 000 шагов цепь с малым шагом проводит в правой моде 0% времени, цепь с большим шагом — 50%. Сэмплы первой цепи выглядят совершенно здоровыми; проблему выявляет только вторая цепь, запущенная в другой моде, или знание целевого распределения. Мультимодальность (multimodality) — одна из главных практических проблем MCMC, и для преодоления барьеров между модами существуют методы темперирования (tempering), которые запускают цепи на уплощённых версиях целевого распределения.

Сэмплирование по Гиббсу#

Когда условное распределение каждой координаты при остальных фиксированных легко сэмплировать, сэмплирование по Гиббсу (Gibbs sampling) обновляет координаты по одной, каждую — из её условного распределения. Это алгоритм Метрополиса–Гастингса, в котором предлагающим распределением служит условное, и вероятность принятия всегда равна 1. Для двумерного нормального распределения из заметки о многомерном нормальном распределении с единичными дисперсиями и корреляцией $\rho$ условные распределения равны $X_1 \mid X_2 = x_2 \sim N(\rho x_2, 1 - \rho^2)$ и симметрично для второй координаты.

Рисунок

Сэмплирование по Гиббсу ходит зигзагом

Сильная корреляция замедляет сэмплирование по Гиббсу: каждый шаг идёт вдоль оси, а условная дисперсия $1 - \rho^2$ мала. Последовательные значения $X_1$ имеют корреляцию $\rho^2$, что равно $0.81$ при $\rho = 0.9$ и $0.98$ при $\rho = 0.99$; 20 000 проходов (sweeps) тогда равноценны примерно 2300 и 190 независимым сэмплам. Это исправляет репараметризация, устраняющая корреляцию, или совместное обновление коррелированных координат блоками. Сэмплирование по Гиббсу — рабочая лошадка для моделей с сопряжёнными условными распределениями, таких как тематические модели и многие иерархические байесовские модели.

Динамика Ланжевена#

Предложения случайного блуждания не учитывают, где сосредоточена вероятность. Динамика Ланжевена (Langevin dynamics) использует score, чтобы сдвигать каждый шаг вверх по плотности, и добавляет шум, чтобы продолжать исследовать пространство:

$$x_{t+1} = x_t + \frac{\varepsilon}{2} \nabla \log p(x_t) + \sqrt{\varepsilon}\, z_t, \qquad z_t \sim N(0, I) .$$

Это нескорректированный алгоритм Ланжевена (unadjusted Langevin algorithm, ULA). Для гладкого целевого распределения и достаточно малых шагов его стационарное распределение стремится к $p$ при $\varepsilon \to 0$; при конечном шаге оно смещено, а при слишком большом шаге или score, быстро растущем на хвостах, цепь может стать неустойчивой и вовсе не иметь стационарного распределения. Для стандартного нормального целевого распределения, score которого равен $-x$, обновление имеет вид $x_{t+1} = (1 - \varepsilon/2)\, x_t + \sqrt{\varepsilon}\, z_t$ и устойчиво только при $0 < \varepsilon < 4$; тогда стационарная дисперсия равна $1/(1 - \varepsilon/4)$: $1.026$ при $\varepsilon = 0.1$ и $1.143$ при $\varepsilon = 0.5$. При $\varepsilon \ge 4$ цепь расходится. Добавление шага принятия Метрополиса–Гастингса устраняет смещение от дискретизации (MALA, Metropolis-adjusted Langevin algorithm): $p$ становится точно стационарным при любом размере шага, и цепь сходится к нему при обычных условиях, хотя при большом шаге принимается мало предложений; моделирование гамильтоновой динамики с переменной импульса даёт гамильтонов Монте-Карло (Hamiltonian Monte Carlo, HMC); в адаптивном варианте NUTS это сэмплер по умолчанию в современных системах вероятностного программирования, таких как Stan. С градиентами по мини-батчам то же обновление превращается в стохастическую градиентную динамику Ланжевена (stochastic gradient Langevin dynamics, SGLD), которая делает из зашумлённого градиентного спуска приближённое сэмплирование из апостериорного распределения для нейросетей.

Для score не нужны ни $Z$, ни значения самой плотности — только градиент её логарифма, и в этом мост к генеративным моделям. По этому мосту переходит score, а не сэмплер целиком, и три процедуры нужно различать:

• Равновесная динамика Ланжевена и MALA — эта заметка: одно фиксированное целевое распределение $p$, цепь запускают достаточно долго, чтобы она забыла начальную точку, и получают сэмплы из её стационарного распределения $p$ (приближённо, а для MALA точно).
• Отжиг динамики Ланжевена (annealed Langevin dynamics, Song and Ermon, 2019): нейросеть учит score $\nabla \log p_\sigma$ данных, размытых гауссовским шумом нескольких величин $\sigma$, и на каждом уровне, от самого большого к самому малому, делают несколько шагов Ланжевена. У каждого уровня своё целевое распределение, ни на одном из них цепь не доводят до равновесия, и результат — не MCMC-сэмпл из распределения данных (заметка о score matching).
• Обратная диффузия (DDPM, score-модели на основе стохастических дифференциальных уравнений, Song et al., 2021): генерация моделирует обращение во времени процесса зашумления — от чистого шума в конечный момент обратно к данным — за фиксированное число шагов. Её снос содержит score $\nabla \log p_t$ зашумлённых данных в момент $t$, поэтому правило перехода меняется от шага к шагу; это неоднородная по времени цепь Маркова без стационарного целевого распределения (сэмплирование в DDPM). На каждом шаге можно добавить шаги Ланжевена как корректор, но сам обратный процесс — не MCMC.

Что изменилось#

	Прямые методы Монте-Карло	Марковские методы Монте-Карло
Что нужно	сэмплы из $p$ или близкое предлагающее распределение	$\tilde p$ с точностью до константы (или её score)
Сэмплы	независимые	коррелированные; нужно считать эффективный размер выборки
Высокая размерность	предлагающие распределения экспоненциально ухудшаются	локальные шаги по-прежнему работают ценой времени перемешивания
Типичные алгоритмы	метод обратной функции, выборка с отклонением, выборка по значимости	Метрополис–Гастингс, Гиббс, Ланжевен, HMC
Главная проблема	веса с бесконечной дисперсией	медленное перемешивание, цепи, застрявшие в одной моде

Ограничения#

• Сходимость асимптотическая, и её трудно проверить. Цепь может выглядеть стационарной, исследовав лишь часть целевого распределения; диагностика обнаруживает некоторые сбои, но никогда не все.
• Настройка важна. Размеры шага, форма предлагающего распределения, параметризация и число цепей меняют эффективность на порядки.
• Это дорого для больших данных и больших моделей. Каждый шаг вычисляет целевую плотность, часто по всему набору данных; для нейросетей вариационный вывод (variational inference), ансамбли или MCMC со стохастическими градиентами — более дешёвые приближения.

Что нужно запомнить#

• Условие детального баланса $p(x) P(x \to y) = p(y) P(y \to x)$ делает распределение $p$ стационарным.
• Алгоритм Метрополиса–Гастингса: предложить $y \sim q(\cdot \mid x)$, принять с вероятностью $\min\big(1, \tilde p(y) q(x \mid y) / (\tilde p(x) q(y \mid x))\big)$; нужны только отношения значений $\tilde p$.
• При слишком малом шаге цепь ползёт, при слишком большом предложения отклоняются; в высокой размерности стоит ориентироваться на долю принятых около 0.23.
• Для величины $g(X_t)$ стационарной цепи коррелированные сэмплы равноценны выборке эффективного размера $n/(1 + 2\sum \rho_k)$, обычно меньшего $n$; нужно запускать несколько цепей и помнить о модах, которые цепь ни разу не посетила.
• Сэмплирование по Гиббсу обновляет координаты из их условных распределений; сильная корреляция его замедляет.
• Динамика Ланжевена движется по score с добавлением шума; без коррекции она смещена из-за размера шага и может быть неустойчивой ($0 < \varepsilon < 4$ для стандартного нормального распределения), MALA устраняет смещение, а HMC использует тот же score. Отжиг динамики Ланжевена и обратная диффузия используют выученные score зашумлённых данных, но их целевые распределения меняются по уровням шума или во времени, так что это не равновесный MCMC.

Код#

Случайное блуждание Метрополиса с тремя размерами шага, целевое распределение с двумя модами, сэмплирование по Гиббсу и нескорректированная динамика Ланжевена:

import numpy as np, math

rng = np.random.default_rng(0)

def metropolis(logp, x0, step, n):
    """Random-walk Metropolis: propose x + step * N(0, 1), accept with probability min(1, p(y) / p(x))."""
    x, lp, chain, accepted = x0, logp(x0), np.empty(n), 0
    for t in range(n):
        y = x + step * rng.normal()
        ly = logp(y)
        if math.log(rng.random()) < ly - lp:
            x, lp, accepted = y, ly, accepted + 1
        chain[t] = x
    return chain, accepted / n

def ess(chain):
    """A rough effective sample size for illustration: autocorrelations summed until they drop below 0.05 (production tools use more careful estimators)."""
    c = chain - chain.mean()
    f = np.fft.rfft(c, 2 * len(c))
    ac = np.fft.irfft(f * np.conj(f))[: len(c)]
    ac /= ac[0]
    cut = np.argmax(ac < 0.05)
    return len(c) / (1 + 2 * ac[1:cut].sum())

# a standard normal target with three step sizes, 20,000 steps each: acceptance rate and effective sample size
for step in (0.1, 2.4, 25):
    chain, rate = metropolis(lambda x: -x * x / 2, 0.0, step, 20_000)
    print(step, round(rate, 3), round(ess(chain)))
# 0.1 0.971 57 | 2.4 0.443 4815 | 25 0.053 708

# two modes at -3 and 3: the share of time in the right mode with a small and a large step
mix = lambda x: math.log(0.5 * math.exp(-(x + 3) ** 2 / 0.72) + 0.5 * math.exp(-(x - 3) ** 2 / 0.72))
for step in (0.5, 4.0):
    chain, rate = metropolis(mix, -3.0, step, 20_000)
    print(step, round(float((chain > 0).mean()), 3))
# 0.5 0.0 | 4.0 0.502

# Gibbs for a bivariate normal: successive draws of x1 have correlation rho^2
for rho in (0.5, 0.9, 0.99):
    x1 = x2 = 0.0
    s, xs = math.sqrt(1 - rho ** 2), np.empty(20_000)
    for t in range(20_000):
        x1 = rho * x2 + s * rng.normal()
        x2 = rho * x1 + s * rng.normal()
        xs[t] = x1
    print(rho, round(float(np.corrcoef(xs[:-1], xs[1:])[0, 1]), 3), round(ess(xs)))
# 0.5 0.264 12075 | 0.9 0.807 2292 | 0.99 0.978 188

# unadjusted Langevin on N(0, 1), whose score is -x: the step size biases the stationary variance to 1 / (1 - eps / 4)
for eps in (0.1, 0.5):
    x, xs = 0.0, np.empty(200_000)
    for t in range(200_000):
        x = x + eps / 2 * (-x) + math.sqrt(eps) * rng.normal()
        xs[t] = x
    print(eps, round(float(xs[1000:].var()), 3), round(1 / (1 - eps / 4), 3))
# 0.1 1.026 1.026 | 0.5 1.13 1.143

Дальше#

Это последняя заметка курса. Курс начался с монет, игральных костей и вопроса о том, что такое вероятность, а заканчивается сэмплерами, на которых держится байесовская статистика, и score, который используют диффузионные модели. Курс генеративных моделей опирается почти на каждую заметку этого курса: метод максимального правдоподобия и KL-дивергенцию, многомерное нормальное распределение, формулу замены переменных, оценки математических ожиданий методом Монте-Карло и динамику Ланжевена вместе с её score.

Литература#

• N. Metropolis, A. W. Rosenbluth, M. N. Rosenbluth, A. H. Teller, E. Teller. Equation of State Calculations by Fast Computing Machines. The Journal of Chemical Physics, 21(6), 1953.
• W. K. Hastings. Monte Carlo Sampling Methods Using Markov Chains and Their Applications. Biometrika, 57(1), 1970.
• S. Geman, D. Geman. Stochastic Relaxation, Gibbs Distributions, and the Bayesian Restoration of Images. IEEE Transactions on Pattern Analysis and Machine Intelligence, 6(6), 1984.
• G. O. Roberts, A. Gelman, W. R. Gilks. Weak Convergence and Optimal Scaling of Random Walk Metropolis Algorithms. The Annals of Applied Probability, 7(1), 1997.
• R. M. Neal. MCMC Using Hamiltonian Dynamics. In Handbook of Markov Chain Monte Carlo, CRC Press, 2011.
• M. Welling, Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. ICML, 2011.
• Y. Song, S. Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. NeurIPS, 2019. Отжиг динамики Ланжевена.
• Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, B. Poole. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR, 2021. Стохастические дифференциальные уравнения с обращённым временем и сэмплирование «предиктор–корректор».