Сэмплирование ненормированных плотностей: MALA, гамильтонов Монте-Карло и NUTS

Потоки дают точную плотность, но только для обратимых сетей с дёшево вычислимым определителем Якоби. Модель, основанная на энергии (energy-based model), снимает все ограничения на архитектуру: любая сеть со скалярным выходом $E_\theta(x)$ задаёт плотность

$$p_\theta(x) = \frac{e^{-E_\theta(x)}}{Z_\theta}, \qquad Z_\theta = \int e^{-E_\theta(x)}\,dx .$$

Плотность определена, когда $Z_\theta$ конечна; положительна она автоматически, поскольку $e^{-E_\theta} > 0$. Цена — $Z_\theta$, интеграл по всему пространству данных, который для изображений никто не умеет вычислить. Без него нельзя вычислить $p_\theta(x)$. Но само по себе её отсутствие не исключает точного сэмплирования: выборка с отклонением, о которой речь ниже, даёт точные независимые точки из ненормированной плотности. Прямым методам нужна дополнительная структура — обратная функция распределения или огибающая, плотно ограничивающая плотность сверху, — а на изображениях такой структуры нет или от неё нет толку. Вычислить можно отношение $p_\theta(y)/p_\theta(x) = e^{E_\theta(x) - E_\theta(y)}$, в котором $Z_\theta$ сокращается, и градиент $\nabla_x \log p_\theta(x) = -\nabla_x E_\theta(x)$, в который $Z_\theta$ вообще не входит. Поэтому для сэмплирования из модели, основанной на энергии, и для её обучения нужны сэмплеры, использующие только отношения или градиенты ненормированной плотности. Эта заметка строит такие сэмплеры — от случайного блуждания Метрополиса–Гастингса до сэмплеров, которые движутся по градиенту и применяются на практике: MALA, гамильтонова Монте-Карло (Hamiltonian Monte Carlo) и его адаптивной формы NUTS. При первом чтении можно пройти Метрополиса–Гастингса, MALA, интегратор leapfrog и одну итерацию HMC, затем настройку, диагностику прогрева и Pyro; доказательство правила принятия и внутреннее устройство NUTS помечены как дополнительный материал. Следующая заметка встраивает их в обучение.

Обозначения. Целевое распределение задано плотностью на $\mathbb R^d$, известной с точностью до константы, $\pi(x) = e^{-U(x)}/Z$; $U$ — потенциальная энергия (potential energy), и для модели, основанной на энергии, $U = E_\theta$. Градиент $\nabla \log \pi(x) = -\nabla U(x)$ — это score (функция score). $q(y \mid x)$ — плотность предлагающего распределения (proposal), $\alpha(x, y)$ — вероятность принятия, $\mathcal N(\mu, \Sigma)$ — нормальное распределение с ковариацией $\Sigma$ (вторым аргументом указывается дисперсия), а $z \sim \mathcal N(0, I)$ всегда обозначает новый независимый стандартный нормальный шум. В гамильтоновом Монте-Карло $p \in \mathbb R^d$ — импульс (momentum; никогда не плотность), $H(x, p)$ — полная энергия, $\varepsilon$ — размер шага, $L$ — число шагов. $d$ — размерность.

Отправная точка: выборка с отклонением и алгоритм Метрополиса–Гастингса#

В курсе теории вероятностей уже есть два сэмплера, которые работают с ненормированной плотностью; здесь мы лишь напомним, что они дают.

Выборка с отклонением (rejection sampling; Методы Монте-Карло) сэмплирует $x \sim q$ из предлагающего распределения и оставляет его с вероятностью $\pi(x)/(M q(x))$, где $M \ge \sup_x \pi(x)/q(x)$. Оставленные точки — точные независимые сэмплы из $\pi$: точка сэмплируется и оставляется с плотностью $q(x) \cdot \pi(x)/(M q(x)) = \pi(x)/M$, пропорциональной $\pi$. Для нормированной целевой плотности доля принятых (acceptance rate) равна $1/M$, так что в среднем один сэмпл обходится в $M$ предложений; если вместо $\pi$ взять $\tilde\pi = Z\pi$, условие принимает вид $\tilde\pi \le M q$, а доля принятых — $Z/M$. Наилучшая константа — $M = \sup_x \pi(x)/q(x)$, и хвосты предлагающего распределения должны быть не легче, чем у целевого, иначе супремум бесконечен. Эта же граница и губит метод в многомерном случае. Для целевого распределения $\mathcal N(0, 1)$ и предлагающего $\mathcal N(0, 2^2)$ отношение $\pi/q = 2 e^{-3x^2/8}$ максимально при $x = 0$, поэтому $M = 2$, и принимается половина предложений. Для $\mathcal N(0, I_d)$ и $\mathcal N(0, 4 I_d)$ отношение равно произведению по координатам, $M = 2^d$: в десятимерном случае принимается одно предложение из 1024, в стомерном — одно из $10^{30}$.

Алгоритм Метрополиса–Гастингса (Марковские методы Монте-Карло) отказывается от независимости. Из текущего состояния $x$ он предлагает $y \sim q(\cdot \mid x)$ и принимает его с вероятностью

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

иначе остаётся в $x$. Входит только отношение $\pi(y)/\pi(x)$, поэтому $Z$ сокращается. Цепь удовлетворяет условию детального баланса $\pi(x) P(x \to y) = \pi(y) P(y \to x)$, которое делает $\pi$ стационарным распределением; если цепь к тому же неприводима (может попасть в любую область положительной плотности) и апериодична, она сходится к $\pi$ из любого начального состояния — по эргодической теореме в конечных пространствах и по её аналогу в непрерывных. Состояния цепи коррелированы, и $n$ состояний равноценны $n/(1 + 2\sum_{k \ge 1} \rho_k)$ независимым сэмплам — это эффективный размер выборки (effective sample size, ESS), где $\rho_k$ — автокорреляция на лаге $k$. При симметричном предложении случайного блуждания (random walk) $y = x + \sigma z$ отношение сводится к $\pi(y)/\pi(x)$, и всё решает шаг $\sigma$: при слишком малом шаге цепь ползёт, при слишком большом предложения отклоняются.

Почему случайное блуждание работает медленно#

Случайное блуждание движется независимыми шагами, поэтому после $n$ принятых шагов размера $\sigma$ оно обычно уходит на расстояние $\sigma\sqrt n$, а не $\sigma n$. Шаг ограничен самым узким направлением целевого распределения: шаг, намного превышающий ширину распределения в каком-либо направлении, попадает в область низкой плотности и отклоняется. Тогда, чтобы пересечь самое широкое направление, нужно около $(\sigma_{\max}/\sigma_{\min})^2$ шагов. Для двумерного нормального распределения с единичными дисперсиями и корреляцией $0.95$ стандартные отклонения вдоль двух осей эллипса равны $\sqrt{1.95} = 1.40$ и $\sqrt{0.05} = 0.22$; их отношение — $6.2$, его квадрат — $39$.

С ростом размерности становится хуже. Для целевого распределения с $d$ независимыми координатами изменение $\log\pi$ — сумма $d$ слагаемых, и чтобы доля принятых не стремилась к нулю, дисперсия предложения $\sigma^2$ должна быть пропорциональна $1/d$. Тогда цепи нужно порядка $d$ шагов на один независимый сэмпл, а оптимальная доля принятых с ростом $d$ стремится к $0.234$, причём для стандартных нормальных координат $\sigma = 2.38/\sqrt d$ (Roberts, Gelman and Gilks, 1997). У обеих проблем одна причина: предложение ничего не знает о том, где сосредоточена вероятность. Остальная часть заметки сообщает ему градиент.

MALA: пусть предлагает градиент#

Диффузия Ланжевена (Langevin diffusion)

$$dX_t = \tfrac12 \nabla \log \pi(X_t)\, dt + dW_t ,$$

где $W_t$ — стандартное броуновское движение, сносится вверх по $\log\pi$ и диффундирует; при мягких условиях её стационарное распределение — $\pi$. Один шаг метода Эйлера на время $\varepsilon$ даёт предложение

$$\begin{gathered} y = x + \frac{\varepsilon}{2} \nabla \log \pi(x) + \sqrt{\varepsilon}\, z, \\ q(y \mid x) = \mathcal N\Big(y;\ x + \frac{\varepsilon}{2} \nabla \log \pi(x),\ \varepsilon I\Big) . \end{gathered}$$

Если использовать этот шаг сам по себе как обновление цепи, получится нескорректированная динамика Ланжевена (unadjusted Langevin dynamics; Динамика Ланжевена), и при любом конечном $\varepsilon$ её стационарное распределение немного неверно: на $\mathcal N(0, 1)$ его дисперсия равна $1/(1 - \varepsilon/4)$. Алгоритм Ланжевена с поправкой Метрополиса (Metropolis-adjusted Langevin algorithm, MALA; Roberts and Tweedie, 1996) использует тот же шаг как предложение алгоритма Метрополиса–Гастингса и исправляет ошибку проверкой принятия. Предложение несимметрично, потому что снос (drift) в точке $y$ отличается от сноса в точке $x$, поэтому в отношение Метрополиса–Гастингса входят оба направления:

$$\begin{gathered} \alpha(x, y) = \min\Big(1, \frac{\pi(y)\, q(x \mid y)}{\pi(x)\, q(y \mid x)}\Big), \\ q(x \mid y) = \mathcal N\Big(x;\ y + \frac{\varepsilon}{2} \nabla \log \pi(y),\ \varepsilon I\Big) . \end{gathered}$$

Каждый шаг требует одного нового градиента — в $y$; градиент в $x$ сохраняется с предыдущего шага. Всё вычисляется в логарифмах: $\log \pi(y) - \log \pi(x) = U(x) - U(y)$, а нормировочные константы двух гауссиан равны и сокращаются.

Разобранный пример: MALA на стандартном нормальном распределении#

Возьмём $\pi = \mathcal N(0, 1)$, так что $\nabla \log \pi(x) = -x$, и обозначим $a = 1 - \varepsilon/2$. Среднее предложения из $x$ равно $a x$, и

$$\begin{aligned} \log \frac{q(x \mid y)}{q(y \mid x)} &= \frac{(y - a x)^2 - (x - a y)^2}{2\varepsilon} \\ &= \frac{(1 - a^2)(y^2 - x^2)}{2\varepsilon} = \Big(1 - \frac{\varepsilon}{4}\Big)\frac{y^2 - x^2}{2}, \end{aligned}$$

где использованы равенства $(y - ax)^2 - (x - ay)^2 = (1 - a^2)(y^2 - x^2)$ и $1 - a^2 = \varepsilon - \varepsilon^2/4$. После прибавления $\log \pi(y) - \log \pi(x) = (x^2 - y^2)/2$ остаётся

$$\log \frac{\pi(y)\, q(x \mid y)}{\pi(x)\, q(y \mid x)} = \frac{\varepsilon}{8}\,(x^2 - y^2) .$$

У случайного блуждания логарифм отношения для той же пары равен $(x^2 - y^2)/2$, у MALA — той же величине, умноженной на $\varepsilon/4$. Предложение уже учитывает бо́льшую часть изменения плотности, и проверка лишь исправляет остаток. При $x = 2$ и $\varepsilon = 1$ предложение — $\mathcal N(1, 1)$. Для сэмпла $y = 2.8$, ушедшего от центра ещё дальше, отношение равно $\pi(y)/\pi(x) = e^{-1.92} = 0.15$, поэтому случайное блуждание приняло бы его с вероятностью $0.15$, а MALA принимает с вероятностью $e^{-0.48} = 0.62$: обратное предложение из $2.8$ имеет среднее $1.4$, всего в $0.6$ от $x = 2$, так что такой шаг легко отменить.

Как шаг зависит от размерности#

Ошибки MALA возникают только из-за дискретизации динамики, для которой $\pi$ инвариантно, и они меньше, чем у случайного блуждания. Для целевых распределений с независимыми координатами Roberts and Rosenthal (1998) показывают, что шаг может уменьшаться как $\varepsilon \propto d^{-1/3}$, а не как $d^{-1}$, так что цепи нужно порядка $d^{1/3}$ шагов на независимый сэмпл вместо $d$, и что оптимальная доля принятых стремится к $0.574$ против $0.234$ у случайного блуждания.

Игрушечная проверка на $\mathcal N(0, I_d)$, где каждый сэмплер настроен на свою асимптотически оптимальную долю принятых, считает эффективный размер выборки по $x_1$ на 1000 вычислений целевой плотности (по одному запуску на 40 000 вычислений для каждого сэмплера; гамильтонов Монте-Карло, о котором речь ниже, — с траекториями длины около 1 и долей принятых около $0.65$):

$d$	блуждание: $\sigma$	ESS на 1000	MALA: $\varepsilon$	ESS на 1000	HMC: $\varepsilon$	ESS на 1000
10	0.75	33	1.27	316	1.05	290
100	0.24	4.4	0.59	108	0.65	126
1000	0.075	0.4	0.27	43	0.37	66

Каждое десятикратное увеличение $d$ обходится случайному блужданию примерно в десятикратную потерю эффективности. Шаг MALA следует закону $\varepsilon \approx 2.7\, d^{-1/3}$, а его эффективность падает в три и в два с половиной раза; последний столбец — предмет следующих разделов.

Две оговорки. Во-первых, шаг ограничен самым жёстким направлением целевого распределения, потому что снос $\frac{\varepsilon}{2}\nabla\log\pi$ проскакивает через максимум плотности там, где кривизна велика: на гауссиане с наименьшей дисперсией $\sigma^2_{\min}$ один только снос устойчив лишь при $\varepsilon < 4\sigma^2_{\min}$. Поэтому на вытянутом целевом распределении MALA ограничен почти так же, как случайное блуждание (рисунок 2 ниже). Предобусловленное (preconditioned) предложение $\mathcal N(x + \frac{\varepsilon}{2} A \nabla \log \pi(x),\ \varepsilon A)$ с матрицей $A$, близкой к ковариации целевого распределения, снимает проблему, а матрица $A$, зависящая от положения, даёт вариант на римановом многообразии из работы Girolami and Calderhead (2011). Во-вторых, для целевых распределений с хвостами легче гауссовых нескорректированная цепь может расходиться, а MALA может не сходиться с геометрической скоростью (Roberts and Tweedie, 1996).

Гамильтонов Монте-Карло: импульс и энергия#

И MALA, и случайное блуждание забывают своё направление после каждого шага. Гамильтонов Монте-Карло (HMC) добавляет скорость, которая сохраняется. Введём импульс $p \in \mathbb R^d$, независимый от $x$ и нормальный, $p \sim \mathcal N(0, M)$, с матрицей масс (mass matrix) $M$ (единичной, если не сказано иное). Совместная плотность равна

$$\pi(x, p) \propto e^{-H(x, p)}, \qquad H(x, p) = U(x) + K(p), \qquad K(p) = \tfrac12\, p^\top M^{-1} p ,$$

а её маргинальное распределение по $x$ — это $\pi$: сэмплируя $(x, p)$ и отбрасывая $p$, мы сэмплируем $\pi$. На языке физики $x$ — положение частицы, $U$ — её потенциальная энергия, $K$ — кинетическая энергия, $H$ — полная энергия (Neal, 2011, раздел 3.1). Частица движется по уравнениям Гамильтона

$$\frac{dx}{dt} = \frac{\partial H}{\partial p} = M^{-1} p, \qquad \frac{dp}{dt} = -\frac{\partial H}{\partial x} = -\nabla U(x) .$$

Импульс несёт частицу вперёд, а сила $-\nabla U$, то есть score, искривляет её путь в сторону высокой вероятности. У точного потока три свойства, и HMC опирается на все три (Neal, 2011, раздел 2.2).

• Он сохраняет $H$. По правилу дифференцирования сложной функции $\frac{dH}{dt} = \nabla U \cdot \frac{dx}{dt} + \nabla K \cdot \frac{dp}{dt} = \nabla U \cdot M^{-1}p - M^{-1}p \cdot \nabla U = 0$.
• Он сохраняет объём. Дивергенция векторного поля $(M^{-1}p,\ -\nabla U(x))$ равна $\sum_i \partial_{x_i}(M^{-1}p)_i - \sum_i \partial_{p_i}\partial_{x_i}U = 0$, потому что первое слагаемое не зависит от $x$, а второе не зависит от $p$; поток с нулевой дивергенцией сохраняет объём любой области (теорема Лиувилля).
• Он обратим. Если пройти по потоку время $t$, сменить знак импульса и снова пройти время $t$, мы вернёмся в начальную точку с импульсом противоположного знака.

Если бы поток можно было вычислить точно, предложение, полученное движением вдоль него в течение некоторого времени, имело бы то же $H$ и принималось бы всегда, как бы далеко оно ни уходило. Вся работа — в интеграторе.

Интегратор leapfrog#

Метод Эйлера обновляет $x$ и $p$ по производным в текущей точке. На $\mathcal N(0, 1)$, где $H = x^2/2 + p^2/2$, а точные траектории — окружности на плоскости $(x, p)$, один шаг Эйлера переводит $(x, p)$ в $(x + \varepsilon p,\ p - \varepsilon x)$, и

$$(x + \varepsilon p)^2 + (p - \varepsilon x)^2 = (1 + \varepsilon^2)(x^2 + p^2) :$$

каждый шаг умножает энергию на $1 + \varepsilon^2$, и траектория раскручивается по спирали наружу. Метод leapfrog (метод Штёрмера–Верле) делит обновление импульса на две половины вокруг полного шага по положению:

$$\begin{aligned} p_{1/2} &= p - \tfrac{\varepsilon}{2} \nabla U(x), \\ x' &= x + \varepsilon\, M^{-1} p_{1/2}, \\ p' &= p_{1/2} - \tfrac{\varepsilon}{2} \nabla U(x') . \end{aligned}$$

В цепочке из $L$ шагов последний полушаг одного шага и первый полушаг следующего используют один и тот же градиент, поэтому $L$ шагов стоят $L$ градиентов. Leapfrog не сохраняет $H$, но два других свойства сохраняет точно (Neal, 2011, раздел 2.3):

• Объём. Каждая из трёх строк — сдвиг (shear): она меняет один блок переменных на величину, зависящую только от другого блока. Её матрица Якоби блочно-треугольная с единичными блоками на диагонали, поэтому её определитель равен 1, и определитель композиции тоже равен 1. Это тот же довод о треугольной матрице Якоби, благодаря которому дёшевы слои связи; здесь определитель не просто дёшев, а равен единице.
• Обратимость. Шаг симметричен: если сменить знак $p'$ и применить тот же шаг, получится $(x, -p)$.

Ошибка энергии у leapfrog остаётся ограниченной, а не растёт. На $\mathcal N(0, 1)$ шаг — линейное отображение $x' = (1 - \varepsilon^2/2)x + \varepsilon p$, $p' = (1 - \varepsilon^2/2)p - \varepsilon(1 - \varepsilon^2/4)x$, и прямая проверка показывает, что оно точно сохраняет теневую энергию (shadow energy)

$$\tilde H(x, p) = \frac{p^2}{2} + \Big(1 - \frac{\varepsilon^2}{4}\Big)\frac{x^2}{2} ,$$

так что вдоль любой траектории leapfrog разность $H - \tilde H = \varepsilon^2 x^2/8$ лишь колеблется. При $\varepsilon \ge 2$ коэффициент при $x^2$ перестаёт быть положительным, и траектории перестают быть ограниченными; для гауссианы со стандартным отклонением $\sigma$ предел — $\varepsilon < 2\sigma$, и задаёт его самое узкое направление. В общем случае ошибка leapfrog по $H$ имеет порядок $\varepsilon^3$ после одного шага и порядок $\varepsilon^2$ на траектории фиксированной длины.

Разобранный пример. На $\mathcal N(0, 1)$ из $x = 1$, $p = 0.5$ с $\varepsilon = 0.5$: $p_{1/2} = 0.5 - 0.25 \cdot 1 = 0.25$, $x' = 1 + 0.5 \cdot 0.25 = 1.125$, $p' = 0.25 - 0.25 \cdot 1.125 = -0.03125$. Энергия меняется с $H = 0.5 + 0.125 = 0.625$ на $0.63330$, так что $\Delta H = 0.0083$, и шаг был бы принят с вероятностью $e^{-0.0083} = 0.992$. Шаг Эйлера из той же точки даёт $(1.25, 0)$ с $H = 0.78125$, то есть ошибку $\Delta H = 0.156 = \varepsilon^2 H$, в девятнадцать раз больше. Теневая энергия $\tilde H = 0.59375$ одинакова до и после шага leapfrog.

Одна итерация HMC#

Итерация HMC из текущего состояния $x$ состоит из трёх шагов, и цепь их повторяет (Neal, 2011, раздел 3.2):

Рисунок

Итерации HMC в фазовом пространстве: leapfrog против Эйлера

1. Обновить импульс. Сэмплировать $p \sim \mathcal N(0, M)$ независимо от всего остального. На плоскости $(x, p)$ это вертикальный сдвиг, и это единственный шаг, который меняет уровень энергии.
2. Смоделировать динамику. Сделать $L$ шагов leapfrog размера $\varepsilon$ из $(x, p)$ и получить $(x', p')$. Траектория идёт по линии уровня $H$ с точностью до малой ошибки leapfrog.
3. Принять или отклонить. Принять $x'$ с вероятностью $\min(1, e^{-\Delta H})$, где $\Delta H = H(x', p') - H(x, p)$; иначе оставить $x$. Импульс отбросить.
4. Повторить с шага 1 с новым импульсом.

На рисунке каждый новый импульс переводит цепь на другую окружность, точки leapfrog остаются близко к своей окружности, а точки Эйлера уходят с неё. Исследование разных уровней энергии обеспечивает шаг 1, длинные перемещения на фиксированном уровне — шаг 2.

Почему для принятия достаточно ошибки энергии#

Дополнительный материал: доказательство того, что шаг принятия HMC сохраняет целевое распределение. Остальная часть заметки на него не опирается.

Запишем $z = (x, p)$, пусть $\Lambda$ — один шаг leapfrog, $F(x, p) = (x, -p)$ — смена знака импульса, и определим отображение предложения $\Phi = F \circ \Lambda^L$: сделать $L$ шагов, затем сменить знак импульса. Два факта об $\Phi$ делают вероятность принятия простой.

• $\Phi$ — инволюция, $\Phi(\Phi(z)) = z$. Обратимость означает $F \circ \Lambda \circ F = \Lambda^{-1}$, поэтому $F \Lambda^L F = \Lambda^{-L}$ и $\Phi \circ \Phi = F\Lambda^L F \Lambda^L = \Lambda^{-L}\Lambda^L$ — тождественное отображение.
• $\Phi$ сохраняет объём, $\lvert\det \partial\Phi/\partial z\rvert = 1$, потому что его сохраняют каждый шаг leapfrog и смена знака.

Будем принимать детерминированное предложение $\Phi(z)$ с вероятностью $\alpha(z) = \min\big(1, \pi(\Phi(z))/\pi(z)\big)$. Для любых двух областей $A$ и $B$ фазового пространства поток вероятности из $A$ в $B$ равен

$$\begin{aligned} \int_A \pi(z)\, \alpha(z)\, \mathbf 1[\Phi(z) \in B]\, dz &= \int_B \pi(\Phi(w))\, \alpha(\Phi(w))\, \mathbf 1[\Phi(w) \in A]\, dw \\ &= \int_B \pi(w)\, \alpha(w)\, \mathbf 1[\Phi(w) \in A]\, dw , \end{aligned}$$

а это поток из $B$ в $A$. В первой строке сделана замена $z = \Phi(w)$: она допустима, потому что $\Phi$ — сама себе обратная, а её якобиан равен единице, потому что $\Phi$ сохраняет объём. Во второй строке использовано равенство $\pi(z)\,\alpha(z) = \min\big(\pi(z), \pi(\Phi(z))\big)$, симметричное относительно перестановки $z \leftrightarrow \Phi(z)$. Это условие детального баланса, поэтому шаг принятия оставляет $\pi(x, p)$ инвариантным. Само отношение равно

$$\frac{\pi(\Phi(z))}{\pi(z)} = e^{-H(x', -p') + H(x, p)} = e^{-\Delta H} ,$$

поскольку $K(-p') = K(p')$; по той же причине смена знака на практике ничего не меняет и её опускают — импульс всё равно сэмплируется заново. Обновление импульса — точный сэмпл из условного распределения $p$ при данном $x$, то есть шаг Гиббса, поэтому оно тоже оставляет $\pi(x, p)$ инвариантным, а значит, и вся итерация (Neal, 2011, раздел 3.2).

Именно сохранение объёма избавляет отношение от якобианов. Отображению предложения, которое растягивает пространство, понадобился бы в отношении множитель $\lvert\det \partial\Phi/\partial z\rvert$ — тот же член замены переменных, за который нормализующие потоки платят в своём правдоподобии. У метода Эйлера определитель на $\mathcal N(0,1)$ равен $1 + \varepsilon^2$ за шаг, и простое отношение $e^{-\Delta H}$ сэмплировало бы не то распределение.

Три сэмплера на вытянутой гауссиане#

На рисунке 2 случайное блуждание, MALA и HMC запущены на двумерном нормальном распределении с корреляцией $0.95$ из раздела о случайном блуждании, с одинаковым бюджетом в 3000 вычислений целевой плотности для каждого. Случайное блуждание использует $\sigma = 1.0$ и принимает 22% предложений, что близко к правилу $0.234$; MALA использует $\varepsilon = 0.17$, чуть ниже своего предела устойчивости $4\sigma^2_{\min} = 0.2$, и принимает 57%; HMC использует $\varepsilon = 0.3$ и $L = 10$, то есть 300 итераций, и принимает 86%.

Рисунок

Три сэмплера с одинаковым бюджетом на вытянутой гауссиане

Эффективные размеры выборки по $x_1$ — 92, 38 и 300. К этим числам три оговорки. Бюджет считает один градиент за одно вычисление, как одно значение плотности, хотя градиент при обратном автоматическом дифференцировании стоит в несколько раз больше; это одна модель учёта, а не равные затраты по времени. Оценка суммирует автокорреляции до первой, меньшей $0.05$, поэтому для HMC, у которого автокорреляция на лаге один отрицательна, она возвращает ровно число итераций, что здесь занижает результат HMC. Оценка Гейера по начальной последовательности (initial-sequence estimator), которая учитывает отрицательные лаги, даёт на этом прогоне 661 (Geyer, 1992). Наконец, каждое число получено с одним сидом: по 20 сидам та же оценка даёт от 46 до 107 для случайного блуждания, от 21 до 62 для MALA и от 299 до 783 (медиана 482) для HMC. Через десять вычислений автокорреляция $x_1$ всё ещё равна $0.59$ у случайного блуждания и $0.79$ у MALA, а у HMC, который за это время закончил одну итерацию, она равна $-0.41$. Каждая траектория HMC длится $\varepsilon L = 3$ единицы времени. Вдоль длинной оси эллипса, стандартное отклонение по которой равно $1.40$, десять шагов leapfrog продвигают колебание на $2.15$ радиана ($123°$), поэтому траектория обычно заканчивается по другую сторону от среднего; отрицательная автокорреляция и отражает этот переброс. На этом целевом распределении градиент ничего не даёт MALA, и он работает хуже хорошо настроенного случайного блуждания: его шаг ограничен узким направлением, а вдоль длинного направления он по-прежнему диффундирует. Его преимущество — масштабирование с размерностью из таблицы выше. HMC отличается качественно: импульс заставляет его двигаться вдоль длинного направления в течение всей траектории.

Настройка HMC: размер шага, длина траектории и матрица масс#

Размер шага. $\varepsilon$ должен оставаться ниже предела устойчивости — примерно удвоенного стандартного отклонения по самому узкому направлению ($0.45$ для целевого распределения выше), а на практике и несколько ниже, потому что доля принятых падает с ростом ошибки энергии. Для целевых распределений с $d$ независимыми координатами шаг должен уменьшаться как $\varepsilon \propto d^{-1/4}$, траектории фиксированной длины тогда нужно порядка $d^{1/4}$ шагов leapfrog, а оптимальная доля принятых стремится к $0.65$ (Neal, 2011, раздел 4.4; Beskos et al., 2013 дают $0.651$). Поскольку каждый градиент стоит порядка $d$ операций, полная стоимость одного независимого сэмпла растёт как $d^{5/4}$ против $d^2$ у случайного блуждания и $d^{4/3}$ у MALA (Neal, 2011, разделы 4.4 и 5.2).

Длина траектории. $\varepsilon L$ определяет, как далеко уходит предложение. Слишком короткая траектория — и HMC вырождается в случайное блуждание мелкими шагами (при $L = 1$ это MALA). Слишком длинная — и траектория поворачивает обратно к началу, а лишние градиенты тратятся впустую. Хуже того, фиксированная длина может войти в резонанс с периодами целевого распределения. На целевом распределении выше шаг leapfrog по узкому направлению после перемасштабирования — поворот на угол $\theta$, где $\cos\theta = 1 - \varepsilon^2/(2\sigma^2_{\min})$. При $\varepsilon = 0.25$ и $L = 10$ это $1.19$ радиана за шаг и $1.89$ оборота за траекторию, так что узкая координата возвращается почти туда, откуда начала (косинус полного угла равен $0.76$); эффективный размер выборки по $x_1 - x_2$ в четырёх запусках падает до значений от 46 до 78 из 300 итераций, тогда как $\varepsilon = 0.3$ даёт 300 во всех четырёх. Небольшая случайная вариация $\varepsilon$ или $L$ на каждой итерации разрушает такие резонансы (Neal, 2011, раздел 3.2), а NUTS вообще избавляет от выбора $L$.

Матрица масс. При $K(p) = \frac12 p^\top M^{-1} p$ шаг leapfrog по положению принимает вид $x' = x + \varepsilon M^{-1} p_{1/2}$. Выбрать $M^{-1}$ равной оценке $\Sigma$ ковариации $\pi$ — то же самое, что запустить HMC с единичной матрицей масс на линейно преобразованной переменной $\Sigma^{-1/2} x$ (Neal, 2011, раздел 4.1; Betancourt, 2017, раздел 4.2), для которой целевое распределение примерно изотропно: на гауссиане предел устойчивости перестаёт зависеть от её корреляций, а короткие траектории достигают почти независимых точек. В высокой размерности обычный компромисс — диагональная $M$, по одному масштабу на координату; её оценивают по цепи во время прогрева (warm-up).

NUTS: остановка на развороте#

Сэмплер без разворотов (No-U-Turn Sampler, NUTS; Hoffman and Gelman, 2014) сохраняет динамику HMC и убирает $L$: он удлиняет траекторию, пока та не начнёт возвращаться.

Критерий разворота#

При единичной матрице масс расстояние от начала траектории $x$ до её текущей точки $\tilde x$ меняется со скоростью

$$\frac{d}{dt}\, \frac{(\tilde x - x)\cdot(\tilde x - x)}{2} = (\tilde x - x) \cdot \frac{d\tilde x}{dt} = (\tilde x - x) \cdot \tilde p .$$

Пока эта величина положительна, продолжение моделирования уводит предложение дальше; как только она становится отрицательной, траектория разворачивается (при неединичной матрице масс $\tilde p$ заменяется скоростью $M^{-1}\tilde p$). Если просто остановиться при первом отрицательном значении, нарушится обратимость: траектория, запущенная обратно из точки остановки, обычно остановилась бы в другом месте, и условие детального баланса не выполнялось бы (Hoffman and Gelman, 2014, раздел 3.1).

Удвоение траектории#

NUTS строит траекторию как сбалансированное двоичное дерево. Он начинает с единственного состояния $(x, p)$ после обновления импульса. На удвоении $j = 0, 1, 2, \dots$ он равновероятно выбирает направление $v \in \{-1, +1\}$ и делает $2^j$ шагов leapfrog размера $v\varepsilon$ от соответствующего конца, вперёд или назад по времени, что удваивает число состояний. Он останавливается, когда для всей траектории или для любого из её сбалансированных поддеревьев крайнее левое и крайнее правое состояния $(x^-, p^-)$ и $(x^+, p^+)$ удовлетворяют условию

$$(x^+ - x^-) \cdot p^- < 0 \quad \text{или} \quad (x^+ - x^-) \cdot p^+ < 0 ,$$

то есть когда продолжение в любую сторону сблизило бы концы. Он останавливается и тогда, когда моделирование становится совсем неточным и ошибка энергии превышает $\Delta_{\max} = 1000$, — это расходимость (divergence; Hoffman and Gelman, 2014, раздел 3.1.1). Поскольку дерево строится случайными удвоениями, любое его состояние могло бы породить то же дерево с той же вероятностью $2^{-j}$ после $j$ удвоений, и эта симметрия вместе с сохранением объёма заменяет обратимость, которую потеряло бы простое правило остановки.

Какое состояние оставить#

Дополнительный материал: как NUTS выбирает следующее состояние из своего дерева. Остальная часть заметки на него не опирается.

У NUTS нет единого шага принятия или отклонения конечной точки. Он выбирает следующее состояние среди состояний траектории так, чтобы совместное распределение положения и импульса оставалось инвариантным.

• Сэмплирование срезами (slice sampling), исходный вариант: сэмплировать $u \sim U\big(0, e^{-H(x, p)}\big)$ и равновероятно выбрать одно из состояний траектории с $e^{-H} \ge u$, исключив состояния последнего удвоения, которые вызвали бы остановку раньше, чем было построено всё дерево (Hoffman and Gelman, 2014, раздел 3.1.1). Эффективный вариант хранит лишь $O(j)$ состояний и смещает выбор к более новой половине дерева, чтобы следующее состояние, как правило, было далеко от текущего (раздел 3.1.2).
• Мультиномиальное сэмплирование (multinomial sampling): выбрать одно из всех состояний траектории с вероятностью, пропорциональной $e^{-H}$, без переменной среза. Отказ от вспомогательной переменной повышает эффективность по рассуждению Рао–Блэкуэлла, и именно этот вариант Stan и Pyro используют по умолчанию (Betancourt, 2017, приложения A.2.2 и A.5).

Одна итерация NUTS после $j$ удвоений стоит $2^j - 1$ градиентов, и реализации ограничивают глубину дерева: в Pyro по умолчанию — 10 удвоениями.

Адаптация размера шага: двойственное усреднение#

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

NUTS по-прежнему нужен $\varepsilon$. Его настраивают во время прогрева так, чтобы статистика принятия $\alpha_t$ в среднем равнялась целевому значению $\delta$. Для NUTS $\alpha_t$ — среднее $\min(1, e^{-\Delta H})$ по состояниям последнего удвоения итерации $t$, то есть вероятность принятия, которая была бы у HMC для этих состояний. Двойственное усреднение (dual averaging), заимствованное из прямо-двойственного метода Нестерова, обновляет

$$\begin{aligned} \bar H_t &= \Big(1 - \frac{1}{t + t_0}\Big)\bar H_{t-1} + \frac{\delta - \alpha_t}{t + t_0}, \\ \log \varepsilon_t &= \mu - \frac{\sqrt t}{\gamma}\, \bar H_t, \\ \log \bar\varepsilon_t &= t^{-\kappa} \log \varepsilon_t + (1 - t^{-\kappa}) \log \bar\varepsilon_{t-1} . \end{aligned}$$

Когда цепь принимает реже, чем $\delta$, $\bar H_t$ растёт и шаг уменьшается; стягивание к $\mu = \log(10\,\varepsilon_1)$ поощряет шаги крупнее начального приближения $\varepsilon_1$, а $\bar\varepsilon_t$ — сходящееся скользящее среднее. Hoffman и Gelman используют $\gamma = 0.05$, $t_0 = 10$, $\kappa = 0.75$, а после прогрева фиксируют $\varepsilon$ равным $\bar\varepsilon$ (Hoffman and Gelman, 2014, раздел 3.2). Они нашли, что NUTS эффективен при любом $\delta$ от 0.45 до 0.65, и предложили $0.6$ против примерно $0.65$ для HMC (раздел 4.4). В Stan и Pyro по умолчанию $\delta = 0.8$: меньший шаг, меньше расходимостей на сложной геометрии и больше шагов leapfrog на траекторию.

Прогрев#

Первая часть запуска, прогрев, уходит на адаптацию и отбрасывается. В Pyro, вслед за Stan, прогрев начинается с буфера из 75 итераций, в котором адаптируется только размер шага, продолжается окнами по 25, 50, 100, … итераций, в конце каждого из которых матрица масс заново оценивается по состояниям цепи, и заканчивается буфером из 50 итераций, который подстраивает размер шага под окончательную матрицу масс. Если прогрев слишком короток для такого расписания, два буфера сокращаются до 15% и 10% его длины. После прогрева переходное ядро (kernel) замораживается. Ядро, которое продолжало бы адаптироваться, уже не задавало бы фиксированную цепь Маркова, и её стационарное распределение больше не было бы гарантировано.

Каждый запуск сопровождают две диагностики. Расходимости отмечают траектории, ошибка энергии которых взорвалась; они сигнализируют об областях высокой кривизны, которые цепь не может исследовать с данным размером шага, и устойчиво повторяющиеся расходимости означают смещённые оценки (Betancourt, 2017, раздел 6.2). Эффективный размер выборки и расщеплённая (split) статистика $\hat R$ по нескольким цепям измеряют перемешивание, как в курсе теории вероятностей.

Запуск сэмплеров в Pyro#

Библиотеки вероятностного программирования прячут всё это за двумя объектами: переходным ядром (HMC или NUTS), которое делает один переход, и MCMC, который запускает прогрев и сэмплирование. Плотность можно задать моделью — программой, которая вызывает pyro.sample:

import torch
import pyro
import pyro.distributions as dist
from pyro.infer import MCMC, NUTS

Sigma = torch.tensor([[1.0, 0.95], [0.95, 1.0]])

def model():
    pyro.sample("x", dist.MultivariateNormal(torch.zeros(2), covariance_matrix=Sigma))

mcmc = MCMC(NUTS(model), num_samples=1000, warmup_steps=500)
mcmc.run()
x = mcmc.get_samples()["x"]   # тензор формы (1000, 2)
mcmc.summary()                # среднее, sd, квантили, n_eff и r_hat по каждой координате

По модели Pyro строит потенциал $U$ — отрицательный логарифм совместной плотности, в котором каждая переменная с ограничениями отображается в переменную без ограничений, а к логарифму плотности добавляется логарифм якобиана этого отображения, так что leapfrog движется по всему $\mathbb R^d$. Настройки по умолчанию — те, что описаны выше: step_size — лишь начальное $\varepsilon_1$, которое двойственное усреднение подстраивает под target_accept_prob=0.8; во время прогрева адаптируется диагональная матрица масс (adapt_mass_matrix=True, full_mass=False); NUTS использует мультиномиальное сэмплирование и max_tree_depth=10; по умолчанию warmup_steps равно num_samples. Для обычного HMC вызов HMC(model, step_size=0.1, num_steps=10) фиксирует длину траектории $\varepsilon L = 1$, и $L$ пересчитывается как целая часть $1/\varepsilon$ при каждой адаптации размера шага.

Для модели, основанной на энергии, потенциал — сама энергия, и его можно передать напрямую:

energy = ...                       # обученная сеть E_theta: тензор формы (2,) -> скаляр

def potential(params):             # U(x) = E_theta(x); нормировочная константа не нужна
    return energy(params["x"])

mcmc = MCMC(NUTS(potential_fn=potential), num_samples=1000, warmup_steps=500,
            initial_params={"x": torch.zeros(2)})
mcmc.run()

С potential_fn нет модели, из которой можно было бы взять начальную точку, поэтому initial_params обязателен. Градиенты $U$ вычисляет автоматическое дифференцирование; каждый шаг leapfrog — один прямой и один обратный проход сети.

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

Сэмплер	Предложение	Что нужно	Шаг принятия	Настраивается на	Шагов на независимый сэмпл при $d$ координатах	Отказывает, когда
Выборка с отклонением	независимый сэмпл из $q$	$\pi$ с точностью до константы, граница $M$	оставить с вероятностью $\pi/(Mq)$	наилучшее $M = \sup \pi/q$	около $M$, часто экспоненциально по $d$	нет плотной огибающей
Случайное блуждание Метрополиса	$x + \sigma z$	отношения значений $\pi$	$\min(1, \pi(y)/\pi(x))$	доля принятых $\approx 0.234$	порядка $d$	вытянутые или многомерные целевые распределения
MALA	$x + \frac{\varepsilon}{2}\nabla\log\pi + \sqrt\varepsilon\, z$	отношения и score	отношение Метрополиса–Гастингса с $q$ в обе стороны	доля принятых $\approx 0.574$	порядка $d^{1/3}$	жёсткие направления, лёгкие хвосты
HMC	$L$ шагов leapfrog из нового импульса	отношения и score	$\min(1, e^{-\Delta H})$	доля принятых $\approx 0.65$; $L$ и $M$ вручную	порядка $d^{1/4}$ шагов leapfrog	неудачное $L$, расходимости
NUTS	траектория, удваиваемая до разворота	отношения и score	состояние траектории с весом $e^{-H}$	$\delta = 0.8$ двойственным усреднением; $M$ при прогреве	как у HMC, с $L$, выбираемым на каждой итерации	расходимости, несколько мод

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

• Переход между модами может быть крайне медленным. Ничто не запрещает перепрыгнуть область нулевой плотности: выборка с отклонением даёт независимые точки, а предложение случайного блуждания или MALA может оказаться по ту сторону провала, не заходя в него. Но такое предложение должно дотянуться до другой моды за один шаг, а с шагом, настроенным на локальную ширину, этого почти никогда не происходит. HMC движется по непрерывной траектории, энергия на которой почти постоянна, поэтому не может пройти через область почти нулевой плотности, где потенциал намного выше, чем позволяет его энергия. Поэтому на целевом распределении с разделёнными модами цепь обычно остаётся в той моде, которую нашла, как курс теории вероятностей показывает для случайного блуждания; HMC и NUTS — не исключение (Neal, 2011, раздел 4.2). Эту проблему решают методы темперирования и отжига (tempering, annealing) ценой дополнительных вычислений.
• Нужны градиенты. MALA, HMC и NUTS требуют дифференцируемого $\log\pi$ на непрерывном пространстве. Дискретные переменные приходится суммировать (маргинализировать) или обновлять другими шагами, а переменным с ограничениями нужна замена переменных.
• Один размер шага на всё пространство. Области высокой кривизны, например горлышку «воронки» иерархической модели, нужен шаг меньше, чем остальному пространству; глобальный $\varepsilon$ либо тратит градиенты впустую в других местах, либо порождает расходимости в этой области.
• Гарантии асимптотические. Стационарность, длину прогрева и перемешивание можно проверить только диагностиками, которые обнаруживают некоторые сбои, но никогда не все.
• Каждый градиент — проход по сети. Для модели, основанной на энергии, каждый шаг leapfrog или Ланжевена — прямой и обратный проход через $E_\theta$, а обучению нужны новые сэмплы после каждого обновления параметров. Цепи, доведённые до сходимости, там непозволительно дороги — с этого и начинается следующая заметка.

Итоги#

• Ненормированная плотность даёт отношения $\pi(y)/\pi(x)$ и score $\nabla\log\pi$ без $Z$; все сэмплеры этой заметки используют только их.
• Выборке с отклонением нужна граница $M \ge \sup \pi/q$, и она принимает долю $1/M$ предложений; алгоритму Метрополиса–Гастингса нужны только отношения, но случайное блуждание диффундирует: порядка $d$ шагов на независимый сэмпл при оптимальной доле принятых $0.234$.
• MALA предлагает один шаг Эйлера для диффузии Ланжевена, $\mathcal N(x + \frac{\varepsilon}{2}\nabla\log\pi(x), \varepsilon I)$, и исправляет его несимметричным отношением Метрополиса–Гастингса; его шаг масштабируется как $d^{-1/3}$, а оптимальная доля принятых равна $0.574$. На $\mathcal N(0,1)$ логарифм его отношения равен $\frac{\varepsilon}{8}(x^2 - y^2)$.
• HMC добавляет импульс $p \sim \mathcal N(0, M)$ и следует уравнениям Гамильтона для $H = U + K$ с интегратором leapfrog. Leapfrog сохраняет объём и обратим, поэтому вероятность принятия равна $\min(1, e^{-\Delta H})$ без якобиана; его ошибка энергии остаётся ограниченной, тогда как у Эйлера энергия на гауссиане умножается на $1 + \varepsilon^2$ за шаг.
• Шаг HMC должен оставаться ниже примерно удвоенного наименьшего стандартного отклонения, масштабируется как $d^{-1/4}$ и оптимален при доле принятых $\approx 0.65$; длина траектории должна избегать и поведения случайного блуждания, и резонанса; матрица масс, близкая к обратной ковариации, делает целевое распределение изотропным.
• NUTS удваивает траекторию в случайных направлениях до разворота, $(x^+ - x^-)\cdot p^- < 0$ или $(x^+ - x^-)\cdot p^+ < 0$, выбирает следующее состояние из всей траектории (сэмплированием срезами или мультиномиальным) и во время прогрева настраивает $\varepsilon$ двойственным усреднением, а также матрицу масс.
• В Pyro всё это запускает MCMC(NUTS(model), num_samples, warmup_steps); для модели, основанной на энергии, NUTS(potential_fn=...) принимает энергию в качестве потенциала.

Ноутбук#

Ноутбук о MCMC из репозитория HSE-LAMBDA Deep Generative Models реализует пять сэмплеров на одном целевом распределении — одномерном $\mathcal N(0, 1)$ (ячейки нумеруются с 1, считая все ячейки подряд, включая текстовые). Сохранённых выходов у него нет, поэтому он документирует процедуры, а не результаты.

• Ячейка 1: введение. Pyro как библиотека вероятностного программирования на PyTorch и список сэмплеров: выборка с отклонением, Метрополис–Гастингс, MALA, HMC и NUTS.
• Ячейки 2–3: подготовка. Импорт torch, pyro, pyro.distributions, pyro.infer.mcmc, matplotlib и seaborn, начальное значение генератора (seed) 42 и вспомогательная функция time_sampling, которая замеряет время работы сэмплера.
• Ячейка 4: выборка с отклонением, теория. Граница $\pi \le Mq$, доказательство того, что принятые точки имеют плотность $\pi$, доля принятых $1/M$ при геометрически распределённом числе попыток со средним $M$, оптимальное $M = \sup \pi/q$, а также достоинства и ограничения, включая экспоненциальный рост $M$ с размерностью.
• Ячейки 5–6: выборка с отклонением, код. Целевое распределение $\mathcal N(0, 1)$, предлагающее $\mathcal N(0, 2^2)$ (dist.Normal(0, 2), масштаб 2), $M = 1.5$, 1000 сэмплов, по одному предложению за раз, и гистограмма с ядерной оценкой плотности.
• Ячейка 7: Метрополис–Гастингс, теория. Отношение принятия, эргодичность и детальный баланс, отбрасывание начального участка цепи (burn-in), влияние ширины предложения и предложение случайного блуждания $\mathcal N(x, \sigma^2)$ с $\sigma = 0.5$.
• Ячейки 8–9: Метрополис–Гастингс, код. 1000 шагов из $x = 0$ с симметричным отношением $\pi(y)/\pi(x)$; гистограмма.
• Ячейка 10: MALA, теория. Диффузия Ланжевена, предложение и несимметричное отношение, почему MALA превосходит случайное блуждание, и советы по выбору шага. Шаг $s$ в ячейке — стандартное отклонение: её предложение $\mathcal N(x + \frac{s^2}{2}\nabla\log\pi(x), s^2)$ совпадает с предложением этой заметки при $\varepsilon = s^2$.
• Ячейки 11–12: MALA, код. Градиенты $\log\pi$ через autograd в текущей и предложенной точках, плотности прямого и обратного предложений, проверка принятия, step_size = 0.1, 1000 сэмплов из $x = 0$; гистограмма.
• Ячейка 13: HMC, теория. Положение, импульс, $H = U + K$ с $K = p^2/2$, уравнения Гамильтона, шаг leapfrog, вероятность принятия $\min(1, e^{-\Delta H})$ и влияние $\varepsilon$ и $L$.
• Ячейки 14–15: HMC в Pyro. Модель с одной случайной переменной (site) x ~ Normal(0, 1), mcmc.HMC(hmc_model, step_size=0.1, num_steps=10), MCMC(..., num_samples=500, warmup_steps=100), сэмплы из get_samples()["x"]; гистограмма.
• Ячейка 16: NUTS, теория. Почему $L$ трудно задать, разворот, удвоение траектории, адаптация размера шага двойственным усреднением во время прогрева.
• Ячейки 17–18: NUTS в Pyro. Та же модель с mcmc.NUTS(model), MCMC(nuts_kernel, num_samples=500, warmup_steps=100); гистограмма.
• Ячейка 19: итоги и сравнительная таблица пяти методов, которую заменяет таблица из раздела «Что изменилось» выше.

Ошибки и неточности, с исправлениями:

• Ячейка 5: $M = 1.5$ не ограничивает $\pi/q$. Отношение равно $2e^{-3x^2/8}$, его максимум 2 достигается при $x = 0$, и оно превышает 1.5 при $\lvert x\rvert < 0.876$. Там $\pi/(Mq)$ больше 1, каждое предложение принимается, и принятые точки имеют плотность, пропорциональную $\min(\pi, 1.5\,q)$, а не $\pi$: это нормальное распределение со срезанной вершиной, с плотностью $0.337$ в нуле вместо $0.399$, дисперсией $1.107$ вместо 1 и долей принятых $0.593$ вместо $1/M = 0.667$, которую предсказала бы теория. Правильная константа — $M = 2$, с долей принятых $1/2$.
• Ячейка 4: предложение для тяжёлых хвостов. Нормальное предлагающее распределение не может покрыть целевое с тяжёлыми хвостами: тогда $\pi/q$ не ограничено в хвостах. Хвосты предлагающего распределения должны быть не легче, чем у целевого (например, распределение Стьюдента). Замечание ячейки о том, что $\pi$ можно не нормировать, верно, но тогда в границе $\pi$ заменяется на $\tilde\pi$, а доля принятых равна $Z/M$, а не $1/M$.
• Ячейки 7–8: шаг случайного блуждания. Рекомендованная доля принятых 25–50% смешивает два результата: оптимум равен $0.234$ в высокой размерности и около $0.44$ в одномерном случае. При $\sigma = 0.5$ на $\mathcal N(0, 1)$ стационарная доля принятых равна $(2/\pi)\arctan(2/\sigma) = 0.84$, так что цепь делает мелкие, сильно коррелированные шаги; $\sigma = 2.4$ даёт $0.44$. Утверждение, что алгоритм Метрополиса–Гастингса работает для мультимодальных целевых распределений, вводит в заблуждение: для них он корректен, но может так и не перейти из одной моды в другую. Кроме того, ячейка 8 вычисляет отношение как частное значений exp(log_prob); разности логарифмов плотностей избавляют от переполнения и от обращения в 0 при потере точности (underflow) на других целевых распределениях.
• Ячейки 10–11: шаг MALA слишком мал. step_size = 0.1 означает $\varepsilon = 0.01$. На $\mathcal N(0,1)$ предложение — это $0.995\,x$ плюс шум со стандартным отклонением $0.1$, практически каждое предложение принимается, а интегральное время автокорреляции равно $(1 + 0.995)/(1 - 0.995) \approx 400$ шагам, так что 1000 шагов из $x = 0$ стоят двух-трёх независимых сэмплов, и гистограмма не может походить на $\mathcal N(0, 1)$. Для этого одномерного целевого распределения step_size от 1 до 1.4 принимает 79–92% предложений и даёт примерно от 370 до 780 эффективных сэмплов на 1000 шагов. Эмпирическое правило 50% принятых соответствует оптимуму $0.574$ для высокой размерности. «RMSprop MALA» — не устоявшийся метод: предобусловливание в духе RMSprop встречается в предобусловленной стохастической градиентной динамике Ланжевена (preconditioned SGLD; Li et al., 2016), в которой нет шага принятия, а предобусловленный MALA с поправкой Метрополиса — это MALA на римановом многообразии из работы Girolami and Calderhead (2011).
• Ячейка 13: сохранение энергии. Ячейка объясняет высокую долю принятых тем, что гамильтонова динамика сохраняет объём и энергию. Точная динамика принималась бы всегда; leapfrog точно сохраняет объём и обратимость, но энергию — лишь приближённо, и вероятность принятия меньше 1 измеряет его ошибку энергии.
• Ячейка 14: настройки HMC перекрываются адаптацией. HMC в Pyro во время прогрева адаптирует размер шага двойственным усреднением под target_accept_prob=0.8 и диагональную матрицу масс; обе адаптации включены по умолчанию. Поэтому step_size=0.1, num_steps=10 фиксирует лишь длину траектории $\varepsilon L = 1$ и начальный шаг; после прогрева $L$ — целая часть $1/\varepsilon$ для адаптированного $\varepsilon$. Чтобы запустить HMC ровно с этими значениями, нужно передать adapt_step_size=False и adapt_mass_matrix=False.
• Ячейка 16: как NUTS выбирает следующее состояние. Ячейка описывает принятие или отклонение конечной точки траектории по Метрополису. У NUTS такого шага нет: он выбирает следующее состояние среди состояний траектории — сэмплированием срезами в исходном алгоритме и мультиномиальным сэмплированием в Pyro по умолчанию (use_multinomial_sampling=True). Каждое удвоение продлевает траекторию в одном случайно выбранном направлении, а не в обоих. «Доля принятых 60–90%» — это целевое значение адаптации $\delta$, по умолчанию 0.8, а не свойство метода.
• Ячейка 19: сравнительная таблица. Её диапазоны долей принятых — эмпирические правила; асимптотически оптимальные значения равны $0.234$, $0.574$ и $0.65$, выборка с отклонением принимает ровно $1/M$, а доля принятых NUTS равна тому $\delta$, которое задано. Ни один сэмплер из таблицы не справляется с хорошо разделёнными модами, а NUTS — выбор по умолчанию только для непрерывных дифференцируемых целевых распределений.
• Все ячейки: одно целевое распределение. На одномерном $\mathcal N(0, 1)$ пять сэмплеров выглядят одинаково. Коррелированная двумерная гауссиана с рисунка 2 их различает.

Дальше#

Сэмплеры этой заметки предполагают фиксированное целевое распределение и достаточно времени, чтобы сойтись. Модели, основанной на энергии, после каждого обновления $\theta$ нужны сэмплы из её собственной плотности, чтобы оценить градиент $\log Z_\theta$, и доводить цепи до сходимости она себе позволить не может. Следующая заметка выводит этот градиент с его негативной фазой (negative phase), которой нужны сэмплы из модели, и обучает модели, основанные на энергии, с помощью коротких цепей Ланжевена, персистентных цепей и буфера воспроизведения (replay buffer). В ней разобраны и два родственника этих сэмплеров, которые там используются: стохастическая градиентная динамика Ланжевена и выборка по значимости с отжигом (annealed importance sampling), которая оценивает нормировочную константу, чтобы сообщить правдоподобие.

Литература#

• R. M. Neal. MCMC Using Hamiltonian Dynamics. In S. Brooks, A. Gelman, G. Jones, X.-L. Meng (eds.), Handbook of Markov Chain Monte Carlo, Chapter 5, Chapman & Hall/CRC, 2011. Разделы 2.2–2.3 (свойства динамики, leapfrog), 3.1–3.2 (алгоритм HMC и его инвариантность), 4.1–4.2 и 4.4 (линейные преобразования, настройка, масштабирование с размерностью), 5.2 (метод Ланжевена как один шаг leapfrog).
• M. D. Hoffman, A. Gelman. The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(47), 1593–1623, 2014. arXiv:1111.4246. Разделы 3.1 (критерий разворота, удвоение, сэмплирование срезами), 3.2 (двойственное усреднение), 4.4 (выбор $\delta$).
• M. Betancourt. A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv, 2017. Разделы 4.2 (евклидово-гауссовы кинетические энергии, матрица масс), 4.3 (время интегрирования), 6.2 (расходимости), приложения A.2.2 и A.5 (мультиномиальное сэмплирование, текущая версия NUTS в Stan).
• G. O. Roberts, R. L. Tweedie. Exponential Convergence of Langevin Distributions and Their Discrete Approximations. Bernoulli, 2(4), 341–363, 1996.
• G. O. Roberts, J. S. Rosenthal. Optimal Scaling of Discrete Approximations to Langevin Diffusions. Journal of the Royal Statistical Society: Series B, 60(1), 255–268, 1998.
• 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), 110–120, 1997.
• A. Beskos, N. Pillai, G. Roberts, J.-M. Sanz-Serna, A. Stuart. Optimal Tuning of the Hybrid Monte Carlo Algorithm. Bernoulli, 19(5A), 1501–1534, 2013. arXiv:1001.4460.
• M. Girolami, B. Calderhead. Riemann Manifold Langevin and Hamiltonian Monte Carlo Methods. Journal of the Royal Statistical Society: Series B, 73(2), 123–214, 2011.
• C. Li, C. Chen, D. Carlson, L. Carin. Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural Networks. AAAI, 2016. arXiv:1512.07666.
• C. J. Geyer. Practical Markov Chain Monte Carlo. Statistical Science, 7(4), 473–483, 1992.
• Документация Pyro: MCMC, HMC and NUTS.
• HSE-LAMBDA, репозиторий курса Deep Generative Models: MCMC.ipynb.