NUTS v10: прогрессивный сэмплинг,
ESS/leapfrog и воспроизводимые бенчмарки

Реализация NUTS в NextStat и протокол бенчмарков для воспроизводимых байесовских метрик. Диагностика ESS/leapfrog изолирует алгоритмическую эффективность от wall-time, а прогрессивный сэмплинг на верхнем уровне tree join удваивает алгоритмическую эффективность на GLM и иерархических апостериорных — давая до 3,21× ESS/sec на non-centered иерархическом бенчмарке.

Февраль 2026 · ~15 мин чтения

1.Введение

Заявления о байесовской производительности имеют смысл только когда апостериорное распределение и протокол инференса зафиксированы. На практике кросс-фреймворковые сравнения регулярно проваливаются из-за тонких различий в модели, настройках адаптации, расчёте диагностик или из-за того, что wall-time отражает оркестрацию процессов, а не качество сэмплера.

NextStat решает эти проблемы через:

• ›Явный контракт апостериорного распределения (модель + априорные + трансформации).
• ›Реализацию сэмплера с документированными настройками.
• ›Бенчмарк-систему, генерирующую валидные JSON-артефакты для каждого запуска, включая сбои.

Вклад

• ›Протокол бенчмарков и формат артефактов для ESS/sec сравнений с явными health gates и мульти-seed стабильностью.
• ›Декомпозиция эффективности через ESS/leapfrog для изоляции алгоритмического качества сэмплирования от wall-time эффектов.
• ›Stan-совместимая NUTS v10 реализация с прогрессивным сэмплингом на верхнем уровне tree-doubling join, существенно улучшающая ESS/leapfrog на GLM и иерархических апостериорных.

2.Предпосылки: NUTS с мультиномиальным сэмплированием траектории

Рассматриваем евклидов HMC с позицией q ∈ ℝᴰ и импульсом p ∈ ℝᴰ, с гамильтонианом:

2.1 Мультиномиальные веса

В мультиномиальном варианте NUTS каждому листу присваивается ненормированный вес w_i = exp(−ΔH_i). Предложение внутри дерева выбирается пропорционально этим весам (категориальная выборка по листьям), поддеревья объединяются через численно стабильный log-sum-exp.

2.2 Обобщённый U-turn критерий

Используется обобщённый no-U-turn критерий на основе сумм импульсов (Betancourt 2017). При каждом объединении поддеревьев три проверки (полная траектория и две граничных) останавливают рост, когда траектория начинает разворачиваться.

Три U-turn проверки при каждом объединении поддеревьев:

Проверка 1 (полное дерево):   is_turning(ρ_total, p_left, p_right)
Проверка 2 (начало→стык):    is_turning(ρ_init + p_junction, p_start, p_junction)
Проверка 3 (стык→конец):     is_turning(ρ_final + p_junction, p_junction, p_end)

Если ЛЮБАЯ проверка обнаруживает U-turn → рост останавливается.
Соответствует base_nuts.hpp в Stan.

3.Прогрессивный сэмплинг в цикле удвоения дерева

NUTS строит траекторию последовательным удвоением дерева в случайном направлении. На каждом шаге имеем существующее дерево (вес W_exist, образец z_sample) и новое поддерево (вес W_sub, предложение z_propose).

Внутри рекурсивного построителя дерева объединение поддеревьев использует стандартное мультиномиальное обновление:

На верхнем уровне (внешний цикл удвоения) используется прогрессивный сэмплинг:

Это обновление смещает выбор в сторону от удержания начальной точки по мере роста пула предложений. Схема описана Betancourt (Appendix A.3.2) и используется в CmdStan.

3.1 Почему это важно (режим равных весов)

На хорошо обусловленных апостериорных с адаптированным шагом веса поддеревьев часто сопоставимы. В пределе равных весов W_sub ≈ W_exist:

• ›Мультиномиальное объединение даёт вероятность замены ≈ 0.5 — текущий образец часто остаётся.
• ›Прогрессивное объединение даёт вероятность замены 1 — агрессивно продвигает предложение по мере исследования новой массы траектории.

Python: воспроизвести кейс 3,21× на одной странице

В v10 прогрессивный сэмплинг встроен в core NUTS переход (без дополнительного флага). Сниппет ниже показывает форму Python API, которую использует бенчмарк-харнесс.

import json
import nextstat
import nextstat.bayes as bayes

# Загрузка набора данных бенчмарка (x, y, group_idx, n_groups).
spec = json.loads(open("hier_logistic_ri.json").read())

model = nextstat.ComposedGlmModel.logistic_regression(
    spec["x"],
    spec["y"],
    include_intercept=True,
    group_idx=spec["group_idx"],
    n_groups=spec["n_groups"],
    random_intercept_non_centered=True,
)

idata = bayes.sample(
    model,
    n_chains=4,
    n_warmup=1000,
    n_samples=2000,
    seed=42,
    target_accept=0.8,
)

quality = idata.attrs["nextstat_diagnostics"]["quality"]
print("min ESS (bulk):", quality["min_ess_bulk"])
print("max R-hat:", quality["max_r_hat"])

4.Диагностики и декомпозиция ESS/leapfrog

4.1 Health gates

Производительность публикуется только для запусков, прошедших health-проверки: частота дивергенций, насыщение treedepth, ранг-нормализованный R̂, минимальный bulk/tail ESS, минимальный E-BFMI.

4.2 ESS/sec

4.3 ESS/leapfrog — ключевая диагностика

Wall time смешивает алгоритмическую эффективность сэмплирования с вычислительными затратами на leapfrog (автодифф, оркестрация процессов, runtime языка, BLAS). Для изоляции алгоритмического поведения:

5.Детали реализации

Реализация NUTS в NextStat:

crates/ns-inference/src/nuts.rs     — NUTS сэмплер, build_tree, nuts_transition
crates/ns-inference/src/adapt.rs    — Оконная адаптация, dual averaging, Welford
crates/ns-inference/src/hmc.rs      — Leapfrog интегратор, метрика (Diag/DenseCholesky)
crates/ns-inference/src/chain.rs    — Мульти-цепная оркестрация (Rayon параллелизм)
crates/ns-inference/src/diagnostics.rs — R̂, ESS, E-BFMI

5.1 Код (top-level join)

Ключевое изменение локализовано: объединение внутри дерева использует стандартное мультиномиальное обновление, а внешний join — прогрессивный выбор:

// crates/ns-inference/src/nuts.rs
// exp(logW_sub - logW_exist) clamped to [0,1] (progressive sampling).
let p = prob_select_outer_progressive(tree.log_sum_weight, subtree.log_sum_weight);
if rng.random::<f64>() < p {
    tree.q = subtree.q;
}

Паритет со Stan:

6.Протокол бенчмарков

6.1 Модели

• ›GLM логистическая регрессия (6 параметров) — хорошо обусловленная базовая модель.
• ›Иерархическая логистическая со случайным интерсептом (22 параметра, non-centered) — воронкообразная геометрия.
• ›Eight Schools (10 параметров, non-centered) — классический иерархический бенчмарк.

6.2 Настройки

Все модели:
  4 цепи, 1000 warmup, 2000 выборок
  Диагональная метрика
  target_accept = 0.8 (Eight Schools: 0.95)
  Seed датасета: 12345 (фиксированный)
  Seed цепей: {42, 0, 123}

6.3 Артефакты

Каждый запуск генерирует JSON-артефакты по схеме nextstat.bayesian_benchmark_result.v1, индекс suite и агрегированный мульти-seed отчёт. Все v10 результаты чекинены в benchmarks/nextstat-public-benchmarks/suites/bayesian/results_v10/.

7.Результаты (февраль 2026)

Оборудование: AMD EPYC 7502P (32C/64T), 128 ГБ RAM. Базовая линия CmdStan: 2.38.0 (релиз 2026-01-13). Все запуски проходят health gates: 0 дивергенций, макс R̂ < 1.01.

7.1 ESS/sec (среднее по 3 seed)

7.2 Health gates (наихудшие по seed)

7.3 ESS/leapfrog (алгоритмическая эффективность)

Диагностика изолирует эффективность сэмплера на единицу гамильтоновой интеграции — именно эта метрика выявила несоответствие proposal selection в v9.

7.4 Декомпозиция ESS/sec (подразумеваемая)

Используя тождество: (ESS/sec ratio) ≈ (ESS/LF ratio) × (LF/sec ratio).

8.Обсуждение

8.1 Что изменилось в v10

В v9 диагностики уровня траектории (размер шага, acceptance, глубина дерева) совпадали точно, но ESS/leapfrog был существенно ниже на GLM апостериорной. Это указало на проблему proposal selection, а не интеграции или адаптации. Переход на прогрессивный сэмплинг существенно улучшил ESS/leapfrog и ESS/sec на всех протестированных геометриях.

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

Даже с фиксированным протоколом, кросс-фреймворковые сравнения чувствительны к:

• ›Различиям реализации градиентов (аналитический vs AD).
• ›Оверхеду оркестрации процессов (один процесс vs мульти-процессные цепи).
• ›Конфигурации CPU-векторизации и BLAS.
• ›Параметризации (centered vs non-centered).

Мы рассматриваем это как часть экспериментального контекста и публикуем полные артефакты для воспроизведения.

9.Заключение

ESS/leapfrog — эффективная диагностика для разделения алгоритмического качества сэмплирования от wall-time затрат реализации. При фиксированном мульти-seed протоколе единственное изменение proposal selection (прогрессивный сэмплинг верхнего уровня) существенно улучшило эффективность сэмплера и даёт сильные ESS/sec результаты на иерархических апостериорных.

Для полной воспроизводимости pinned протокол, JSON-артефакты и фигуры находятся в benchmarks/nextstat-public-benchmarks/suites/bayesian/results_v10/.

Литература

• ›Hoffman, M. D., Gelman, A. The No-U-Turn Sampler. JMLR 15(1), 1593–1623 (2014).
• ›Betancourt, M. A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv:1701.02434 (2017).
• ›Carpenter, B., et al. Stan: A Probabilistic Programming Language. J. Stat. Soft. 76(1) (2017).
• ›Stan Development Team. Stan Reference Manual. Version 2.38 (2026).
• ›Stan Development Team. Release of CmdStan 2.38. January 13, 2026.
• ›Vehtari, A., et al. Rank-normalization, folding, and localization: An improved R̂. Bayesian Analysis (2021).
• ›Langenbruch, C. MoreFit: A More Optimised, Rapid and Efficient Fit. arXiv:2505.12414 (2025).