NextStat v0.9.6: нулевой JIT-налог, ESS/grad и сходимость

0.MAMS и LAPS за одну страницу (для пользователей NUTS)

Если вы уже знаете NUTS, ключевая ментальная модель:

• ›NUTS — динамический HMC с адаптивным расширением дерева на каждом переходе (адаптивная длина пути на каждом шаге).
• ›MAMS (Metropolis Adjusted Microcanonical Sampler) использует микроканоническую/изокинетическую динамику с фиксированной длиной траектории в прекондиционированном пространстве.
• ›LAPS — массивно параллельный GPU/Metal-путь сэмплирования, который применяет MAMS-динамику на тысячах цепей с аппаратно-ориентированным исполнением.

Что нового относительно стандартных реализаций NUTS

• ›Ядра переходов фиксированной формы проще эффективно исполнять на SIMD/GPU-бэкендах, чем рекурсивное построение деревьев.
• ›Сэмплер построен для очень большого параллелизма цепей (4096 цепей — штатная рабочая точка в этом отчёте).
• ›В этом релизе MAMS/LAPS используют микроканонические диагностические гейты и явное раскрытие NCP-vs-centered для funnel-подобных геометрий.

Где этот дизайн сильнее всего

• ›Time-to-result в итеративных воркфлоу (нет повторной JIT-компиляции в нашем AOT-пути).
• ›Иерархические и мультимасштабные цели, где мы наблюдаем высокий ESS/grad в matched-target запусках.
• ›Локальные пути ускорения (Metal) и серверные GPU-пути (CUDA) с единой семантикой сэмплера.

Где NUTS или XLA-тяжёлые стеки могут выиграть

• ›Простые низкоразмерные цели после прогрева компиляции (возможна более высокая raw warm-path пропускная способность).
• ›Некоторые гладкие концентрированные постериоры на CPU (например, large-n логистическая регрессия в этом отчёте).

Поэтому этот отчёт не заявляет универсального победителя среди сэмплеров — он документирует, где каждая модель исполнения сильнее, с явными оговорками о честности.

1.Протокол и правила честности

Разделение бэкендов (без смешивания)

• ›Результаты LAPS Metal публикуются отдельно от LAPS CUDA.
• ›Результаты CPU (EPYC) публикуются отдельно от GPU.
• ›Ни одна таблица не смешивает значения Metal и CUDA.

Мульти-ран агрегация (против cherry-pick)

• ›Финальные таблицы сравнения CPU/GPU используют три независимых сида: 42, 123, 777.
• ›Мы публикуем медиану как основное число (устойчиво к выбросам), с mean ± std где полезно.
• ›Значения одиночных сидов хранятся только в сырых артефактах, не как заголовочные утверждения.

Time-to-result в итеративном моделировании

• ›Байесовская работа — это цикл редактирования (изменить приор, добавить ковариату, репараметризовать, перезапустить).
• ›В JIT/XLA-стеках изменения графа/формы часто инвалидируют скомпилированные исполняемые файлы и запускают рекомпиляцию, поэтому cold-start затраты повторяются при исследовании.
• ›NextStat использует AOT-компилированные Rust/CUDA-ядра, поэтому wall-clock латентность итерации остаётся близкой к warm-path даже при эволюции моделей.

Раскрытие параметризации funnel

Для std_normal, eight_schools и glm_logistic оба движка сэмплируют одну и ту же целевую плотность (идентичные log-density функции).

Для neal_funnel_10d параметризации различаются в V100 parity run:

• ›NS LAPS сэмплирует Non-Centered Parameterization (NCP): log p(v, z) = -v²/18 - 0.5 · Σ(z_i²).
• ›BlackJAX сэмплирует centered-параметризацию: log p(v, x) = -v²/18 - 0.5·exp(-v)·Σ(x_i²) - 0.5·(d-1)·v.

Это не одна и та же задача оптимизации. Centered funnel имеет позиционно-зависимую кривизну, фундаментально более сложную для сэмплеров с фиксированной метрикой. Строки neal_funnel в секции 3 и Приложении отражают различия как алгоритма, так и параметризации, и не должны интерпретироваться как сравнение «при прочих равных». Они сохранены для демонстрации поведения сходимости (NS сходится, BlackJAX нет), но исключены из заголовочных ESS/grad утверждений.

• ›На CPU теперь есть явная FunnelNcpModel для честных NCP-сравнений (секция 6).
• ›Centered FunnelModel остаётся отдельным контролем сложной геометрии.

Алгоритмические изменения в v0.9.6

• ›MAMS использует eps_jitter=0.1 по умолчанию (±10% равномерный шум размера шага на переход), ломая периодичность фиксированного L и улучшая tail ESS на периодических целях вроде std_normal.
• ›Длина траектории по умолчанию: L = √d в прекондиционированном пространстве (Robnik et al. 2023).

Конфигурация BlackJAX (V100 parity run)

Для предупреждения вопросов о мисконфигурации конкурента, полная конфигурация BlackJAX:

• ›Сэмплер: blackjax.adjusted_mclmc с интегратором isokinetic_mclachlan.
• ›Прогрев: встроенная blackjax.adjusted_mclmc_find_L_and_step_size (500 итераций, single-chain прогрев, target_accept=0.9, diagonal_preconditioning=True).
• ›Длина траектории: настроена встроенным прогревом BlackJAX (L, step_size), затем n_steps = round(L / step_size).
• ›Mass matrix: сэмплирование использует настроенную inverse_mass_matrix из прогрева BlackJAX.
• ›Multi-chain: 4096 цепей, jax.vmap(run_chain), block_until_ready() + device_get() для честного host-side тайминга.
• ›Cold/warm: cold = первый вызов vmap (включает XLA-компиляцию); warm = второй вызов с кэшированным JIT.
• ›Init: цепи инициализированы вокруг прогретого single-chain состояния (warmed_state.position + N(0, 0.5)).
• ›Seed: 42 (cold), 1042 (warm).
• ›Seeds: 42, 123, 777 (для каждого сида, warm run использует seed + 1000 key path).
• ›Исходник: benchmarks/gpu_triple_bench.py, функции _blackjax_builtin_warmup() и bench_blackjax().

Конфигурация V100 parity run (NS LAPS, 3 seeds)

• ›n_chains=4096, n_warmup=500, n_samples=1000, report_chains=256, seeds=42/123/777.
• ›Секции 3/4 публикуют медиану по 3 сидам.
• ›R̂ вычислена по 256 report chains (512 полу-цепей), давая существенно более точные диагностики по сравнению с предыдущим 64-chain репортингом.

2.Каноничные результаты LAPS Metal (финальные)

Железо: Apple M5, 10 GPU-ядер, 24 ГБ unified memory.

Примечание: строка 256-chain std_normal_10d (R̂ 1.175) демонстрирует минимально жизнеспособное количество цепей; строка 4096-chain — каноничная конфигурация бенчмарка.

На практике это показывает, что локальный Apple Silicon способен запускать массивно параллельные inference-нагрузки дата-центрового масштаба с надёжной диагностикой сходимости — без настройки CUDA и без задержки JIT-компиляции.

Политика quality gates для этой матрицы:

• ›MAMS/LAPS: QualityGates::microcanonical() (EBFMI — только предупреждение).
• ›NUTS: строгий дефолтный gate сохранён (EBFMI fail < 0.20).

3.CUDA V100 parity run (LAPS vs BlackJAX, 3-seed медиана)

Железо: Tesla V100-PCIE-16GB.

Как читать эту таблицу

• ›Нулевой JIT-налог: NS LAPS cold остаётся близким к warm (AOT-компилированный Rust/CUDA). Cold-start BlackJAX существенно выше в этой конфигурации (11.8–90.6 с).
• ›Warm-start пропускная способность (canonical прогон): NS LAPS выше на всех совпадающих целях в этой конфигурации.
• ›neal_funnel — не сравнение «при прочих равных» (см. секцию 1: NS сэмплирует NCP, BlackJAX — centered). В этих 3 сидах R̂ centered-funnel BlackJAX в диапазоне 1.260–1.275 — это ожидаемо из-за сложности параметризации, а не дефект сэмплера.

4.ESS/grad на V100 (фаза сэмплирования, только совпадающие цели, 3-seed медиана)

neal_funnel исключён из таблицы, так как два движка сэмплируют разные параметризации (см. секцию 1).

Основной вклад в изменение по сравнению с ранними драфтами — нормализация знаменателя: оба движка теперь вычисляют ESS/grad на одном бюджете report_chains.

Практическая интерпретация для этого canonical прогона:

• ›NS LAPS достигает более высокого ESS/grad на всех совпадающих целях в этом отчёте.
• ›glm_logistic остаётся самой дорогой целью для обоих движков по абсолютному wall time.

5.Верификация качества LAPS на V100 (report_chains=256)

Отдельный запуск с более строгой диагностикой (report_chains=256 → 512 полу-цепей → SE(R̂) ≈ 0.015).

Это подтверждает, что сходимость LAPS надёжна при достаточном количестве диагностических цепей. Значения R̂ из parity-run (секция 3, report_chains=256) напрямую сопоставимы с quality run.

6.CPU EPYC (MAMS vs NUTS) и исправление паритета funnel

Железо: AMD EPYC 7502P, 32 ядра / 64 потока, 128 ГБ RAM (Hetzner dedicated).

Сводка EPYC multi-seed (42/123/777, 3-run агрегат)

Конфигурация: n_chains=4, n_warmup=1000, n_samples=1000, eps_jitter=0.1.

Наблюдаемый паттерн в этой real-run матрице:

Почему large-n логистическая регрессия благоприятствует NUTS

• ›Стоимость градиента масштабируется как O(n·p) на leapfrog-шаг; при n=5000, p=20 каждый дополнительный шаг дорог.
• ›NUTS может завершать траектории раньше через U-turn, тогда как MAMS использует фиксированную длину траектории в прекондиционированном пространстве.
• ›По мере роста n геометрия постериора приближается к хорошо обусловленному гауссиану; это сильный режим для NUTS с адаптивной длиной пути.

Практическая рекомендация

• ›Предпочитайте MAMS для иерархических / мультимасштабных геометрий.
• ›Предпочитайте NUTS для large-n GLM-подобных постериоров на CPU.
• ›Разумное продуктовое направление — явная эвристика method="auto" (напр.: GLM с большим n → NUTS; иерархические/funnel-подобные цели → MAMS), с сохранением ручного оверрайда.

Контроль параметризации funnel (EPYC, 3 сида)

Конфигурация: n_chains=4, n_warmup=1000, n_samples=1000.

MAMS:

NUTS:

Интерпретация:

• ›NCP — 6/6 ok по всем сидам для MAMS и NUTS. ESS_tail в диапазоне 1 604–2 516 (NUTS) и 1 897–1 924 (MAMS).
• ›Centered — 3/3 fail для NUTS и 1/3 fail для MAMS.
• ›Предыдущее несоответствие CPU funnel было методологическим (centered vs NCP), а не проблемой «CPU слабый».
• ›FunnelNcpModel — рекомендуемая параметризация бенчмарка для CPU/GPU паритета.
• ›Centered FunnelModel сохранён как известный патологический контроль; это демонстрация ограничений, а не регрессия продукта.
• ›В этих артефактах EPYC funnel-control EBFMI не экспортирован (n/a в таблицах), поэтому pass/fail основан на R̂/ESS quality gates.

7.Воспроизводимость и метаданные окружения

• ›JSON бенчмарка V100 содержит top-level снимок environment (python, jax, cuda, gpu, версии пакетов).
• ›Сьют EPYC хранит метаданные железа/конфига/сида и метрики по кейсам; полный снимок окружения на уровне пакетов сейчас только в V100 parity JSON.

Артефакты (все в docs/blog/artifacts/v096-zero-jit-tax/):

• ›Матрица V100 3-seed (canonical): v100-multi-seed-matrix-canonical.json
• ›Графические данные V100 (canonical): v100-parity-chart-data-canonical.csv, v100-essgrad-ratio-canonical.csv
• ›Сырые V100 3-seed (canonical): v100_v096_builtinwarmup_3seed_20260218T224654Z/seed_42/gpu_triple_bench.json, seed_123/..., seed_777/...
• ›V100 funnel addendum: v100_ns_funnel_3seed_20260218T231337Z/*, v100_bj_funnel_builtin3seed_20260218T231204Z/*
• ›V100 quality run: v100-quality-report256-5models.json
• ›V100 + EPYC refresh note: 2026-02-17-v096-refresh-v100-epyc.md
• ›EPYC multi-seed матрица: epyc-multi-seed-matrix.json
• ›Выход EPYC сьюта: epyc-mams-suite.json
• ›EPYC funnel-control: epyc-funnel-control-3seed.json

A.Приложение: V100 neal_funnel (разные параметризации)

Сохранено для прозрачности. Эти строки сравнивают NS LAPS (NCP) с BlackJAX (centered) — это не сравнение «при прочих равных».

Несходимость BlackJAX на centered funnel ожидаема (см. секцию 6: даже NUTS проваливается 3/3 на centered funnel с 4 цепями и стандартным бюджетом). Это сравнение прежде всего демонстрирует, что дефолтная NCP-диспатч NS производит сходящиеся результаты там, где centered-параметризация не сходится.

Ссылки

• ›Robnik, Cohn-Gordon, Seljak. Metropolis Adjusted Microcanonical Hamiltonian Monte Carlo (MAMS). arXiv:2503.01707
• ›BlackJAX. Composable Bayesian inference in JAX. arXiv:2402.10797

Связанное чтение

• ›NUTS v10: Progressive Sampling
• ›Байесовские бенчмарки, которые имеют смысл
• ›JAX Compile vs Execution