NextStatNextStat
← Статьи

Компиляторный и гибридный подходы:небинированные фиты на GPU в ФВЭ

Работа MoreFit (Langenbruch, EPJC 86, 2026) показала, что символьная дифференциация и JIT‑компиляция графа вычислений могут дать порядка ~10× ускорение по сравнению с RooFit CUDA для «чистых» небинированных фитов. NextStat решает ту же задачу иначе: делает ставку на гибридный пайплайн, где есть нейросетевые нормализующие потоки (normalizing flows) через ONNX (опционально, режим neural), слитые CUDA‑ядра с аналитическими градиентами и систематики выхода (модификаторы нормы; rate modifiers) как часть контракта с первого дня. В статье мы сравниваем оба подхода по смыслу и по архитектуре и объясняем, почему в реальных физических анализах выигрывает гибрид.

GPUONNXCUDAНЕБИНИРОВАННЫЕ ФИТЫАНАЛИТИЧЕСКИЕ ГРАДИЕНТЫСИСТЕМАТИКИ

Февраль 2026 · ~20 мин


1.Расширенное небинированное правдоподобие: математическая формулировка

Небинированный (event-level) анализ строит правдоподобие непосредственно из набора наблюдённых событий{xi}i=1…N, без потери информации при биннинге. Расширенное правдоподобие для модели со P процессами:

ℒ(θ) = Pois(N | ν_tot(θ)) · ∏ᵢ₌₁ᴺ f(xᵢ | θ)

где:
  ν_tot(θ) = Σₚ νₚ(θ)                      — суммарный ожидаемый выход
  f(xᵢ | θ) = Σₚ νₚ(θ) · pₚ(xᵢ | θ) / ν_tot(θ)  — взвешенная смесь PDF

Отрицательный логарифм:
  NLL(θ) = ν_tot(θ) − Σᵢ log[ Σₚ νₚ(θ) · pₚ(xᵢ | θ) ] + C_constr(θ)

  C_constr(θ) = ½ Σⱼ ((θⱼ − μⱼ) / σⱼ)²    — гауссовы ограничения

Ключевая вычислительная задача — внутренняя сумма по процессам для каждого из N событий. ДляN = 50 000–500 000 (типичные unbinned анализы в LHCb/Belle II) это доминирует над всеми остальными затратами. GPU идеально подходит: каждое событие — независимая единица работы.

Для минимизации NLL(θ) нужен градиент ∇θNLL. Здесь подходы принципиально расходятся.


2.Подход MoreFit: компиляторные оптимизации графа вычислений

MoreFit (Langenbruch, EPJC 86, 2026) предлагает компиляторный подход к unbinned фитам. Основная идея: пользователь описывает PDF как символьный граф вычислений, фреймворк оптимизирует граф, а затем JIT-компилирует его в GPU-ядро (OpenCL) или CPU-код (LLVM/Clang с SIMD).

2.1 Computation Graph и оптимизации

PDF записывается как DAG (directed acyclic graph), где узлы — арифметические операции, листья — параметры θ и наблюдаемые x. Над этим графом выполняются три ключевые оптимизации:

  • Parameter-only term hoisting. Все подграфы, зависящие только от θ (без x), вычисляются однократно на CPU и передаются в GPU-ядро как константы. Для нормированных PDF это включает нормализационные интегралы — они зависят от (μ, σ), но не от x. Результат: GPU-ядро вычисляет только ненормированную PDF, деление на норму происходит через предвычисленную константу.
  • CSE (Common Subexpression Elimination). Если (x − μ) встречается в двух местах графа, оно вычисляется один раз. Для Crystal Ball PDF, где (x − μ)/σ используется и в гауссовом, и в степенном хвосте, это даёт 20–30% сокращение арифметических операций.
  • Символьная дифференциация. Частные производные ∂NLL/∂θⱼ вычисляются символьно прямо из графа. Это даёт аналитические градиенты, а не числовые (конечные разности) или AD (reverse-mode tape). Градиент интегрируется в то же GPU-ядро — один проход по данным вычисляет и NLL, и ∇NLL.

2.2 Результат: сходимость за 2–3 итерации

Числовые производные (central finite differences, h = 10⁻⁴) требуют 2P+1 вызовов функции на каждую итерацию оптимизатора. Для модели с P = 10 параметрами — 21 вызов NLL. Символьные градиенты дают точные производные за 1 вызов. С квазиньютоновским оптимизатором (L-BFGS) MoreFit демонстрирует сходимость за 2–3 итерации вместо ~85 вызовов функции у RooFit CUDA с числовыми производными.

Стоимость минимизации:
  MoreFit:  K_sym × C_kernel    где K_sym ≈ 2–3 (итерации L-BFGS с точным градиентом)
  RooFit:   K_num × (2P+1) × C_kernel   где K_num ≈ 15–30, P ≈ 10

  Ускорение от градиентов:  K_num × (2P+1) / K_sym ≈ 30 × 21 / 3 ≈ 210×
  Ускорение от ядра (CSE + hoisting): ≈ 2–5×
  Итого: ~10× над RooFit CUDA (авторские бенчмарки)

2.3 Ограничения архитектуры

Компиляторный подход оптимален для замкнутого набора аналитических PDF. Но:

  • Каждая новая PDF-семья требует написания формулы, доказательства корректности градиента, тестирования.
  • Систематические неопределённости (nuisance parameters, rate modifiers) отложены на «будущие версии».
  • Acceptance corrections и per-event weights/selection filters в v0.1 не описаны как часть контракта.
  • Бинированные фиты (HistFactory) не входят в scope.
  • В многомерном случае (>3D) аналитические PDF масштабируются плохо — нет замкнутых форм для произвольных корреляций.

3.Подход NextStat: гибридный параметрически-нейронный пайплайн

NextStat решает ту же задачу минимизации NLL(θ), но с принципиально другими архитектурными приоритетами: систематики, произвольные PDF, production-readiness. Для этого используется трёхуровневая система градиентов и два пути GPU-ускорения.

3.1 Трёхуровневая система градиентов

text
Tier 1: Reverse-mode AD (ns-ad crate, Tape)
  ├ Computation graph: Add, Sub, Mul, Div, Ln, Exp, Powf, Max
  ├ O(1) backward pass по числу параметров
  ├ Используется для: общего CPU пути (в т.ч. HistFactory/binned), где нужна гибкость графа
  └ CPU only (Rust, no GPU)

Tier 2: Аналитические CUDA/Metal-ядра
  ├ Hand-derived ∂logp/∂θ для каждой PDF (Gaussian, Exponential,
  │   Crystal Ball, Double CB, Chebyshev, Histogram)
  ├ Fused NLL+grad kernel: 1 проход по данным = NLL + ∇NLL
  ├ Grid-stride loop, shared memory, atomic grad accumulation
  ├ 1 CUDA block = 1 dataset (batch toy: 1 block = 1 toy)
  └ Lockstep L-BFGS-B: все toys на одной итерации → 1 kernel launch

Tier 3: Центральные конечные разности (h = 10⁻⁴)
  ├ Для flow PDF (ONNX = чёрный ящик, нет доступа к графу)
  ├ 2P+1 вызовов eval_logp на итерацию
  ├ NLL reduction на GPU (CudaFlowNllAccelerator)
  └ Trade-off: гибкость > скорость для flow-компонент

Ключевой инсайт: NextStat не выбирает один подход — он комбинирует все три в рамках одной модели. Параметрические компоненты получают аналитические градиенты (Tier 2), flow-компоненты — конечные разности (Tier 3), а HistFactory-обёртка использует reverse-mode AD (Tier 1).

3.2 Fused CUDA-ядро для параметрических PDF

Для параметрических PDF NextStat использует тот же принцип, что и MoreFit — аналитические производные, fused в одно ядро. Ядро unbinned_nll_grad.cu (~3150 строк CUDA C) включает:

  • 6 семейств PDF с аналитическими градиентами: Truncated Gaussian, Bounded Exponential, Crystal Ball, Double Crystal Ball, Chebyshev (произвольный порядок), Histogram (bin lookup).
  • Rate modifiers: NormSys (log-normal), WeightSys (Code0 + Code4p polynomial interpolation) — систематики выхода, интегрированные прямо в gradient kernel.
  • Online logsumexp по процессам для численной стабильности mixture likelihood.
  • Batch toy architecture: 1 CUDA block = 1 toy dataset. Все toys выполняются параллельно за один kernel launch. Lockstep L-BFGS-B: все toys на одной итерации оптимизатора.
cuda
// Архитектура ядра (упрощённо):
// 1 block = 1 dataset, threads = grid-stride loop по событиям

extern "C" __global__ void unbinned_batch_nll_grad(
    const double* params_flat,    // [n_toys × n_params]
    const double* obs_flat,       // [total_events]
    const uint*   toy_offsets,    // [n_toys + 1]
    /* ... model descriptors ... */
    double* nll_out,              // [n_toys]
    double* grad_out,             // [n_toys × n_params]
    uint n_params, uint n_procs, uint n_toys
) {
    uint toy_idx = blockIdx.x;
    // Load params[toy_idx] → shared memory
    // Grid-stride loop over events[toy_offsets[toy_idx]..toy_offsets[toy_idx+1]]
    //   For each event:
    //     logsumexp over processes: log Σₚ νₚ·pₚ(x|θ)
    //     Accumulate ∂NLL/∂θⱼ via atomicAdd (analytical)
    // Add Gaussian constraints + yield term
}

3.3 GPU Flow Session: NLL-редукция для нейронных PDF

Для flow PDF (normalizing flows, экспортированных через ONNX) NextStat разделяет вычисление на два этапа:

text
Путь 1: CPU Flow + GPU Reduce (Phase 3 MVP)
  ┌──────────────────┐     ┌─────────────────────┐
  │ FlowPdf::log_prob│────▶│ H2D upload logp_flat │
  │    (CPU / ONNX)  │     │   [n_procs × N]      │
  └──────────────────┘     └──────────┬────────────┘
                                      ▼
                           ┌─────────────────────┐
                           │ CudaFlowNllAccelerator│
                           │   logsumexp + yields  │
                           │   + constraints       │
                           │   → NLL (scalar)      │
                           └─────────────────────┘

Путь 2: CUDA EP Flow + GPU Reduce (D4, zero-copy)
  ┌──────────────────┐     ┌─────────────────────┐
  │ ONNX Runtime     │────▶│ CudaSlice<f64>       │
  │  CUDA EP         │     │ (device-resident)    │
  └──────────────────┘     └──────────┬────────────┘
                                      ▼
                           ┌─────────────────────┐
                           │ nll_device()         │
                           │   (zero-copy path)   │
                           └─────────────────────┘

Gradient для flow: central finite differences, h = 10⁻⁴. Каждая итерация оптимизатора требует 2P+1 вызововeval_logp. Это медленнее аналитики, но позволяет использовать любую PDF, экспортированную как ONNX.

Деталь реализации (контракт): в коде NextStat ядра GPU компилируются в PTX на этапе сборки (nvcc --ptx, см. crates/ns-compute/build.rs). «Device-resident» путь для flow означает, что редукция умеет читать GPU-буфер logp напрямую (CudaFlowNllAccelerator::nll_device), но yields/params остаются небольшими host→device копиями. Это принципиально другой класс JIT по сравнению с DAG→kernel в MoreFit.


4.Стратегии градиентов: математическое сравнение

Обозначим NLL(θ) = L(θ). Нам нужен ∇θL для P параметров.

4.1 Символьная дифференциация (MoreFit)

Для каждого параметра θⱼ:

  ∂L/∂θⱼ = ∂ν_tot/∂θⱼ − Σᵢ [ Σₚ (∂νₚ/∂θⱼ · pₚ(xᵢ) + νₚ · ∂pₚ/∂θⱼ) ] / [ Σₚ νₚ · pₚ(xᵢ) ]
            + Σⱼ (θⱼ − μⱼ) / σⱼ²

Все производные ∂pₚ/∂θⱼ вычислены символьно из графа. Результат:
  — Точность: machine epsilon (~10⁻¹⁶)
  — Стоимость: 1 проход по данным (fused с NLL)
  — Сходимость L-BFGS: 2–3 итерации
  — Ограничение: только для PDF с известной аналитической формой

4.2 Аналитические CUDA-ядра (NextStat Tier 2)

Тот же принцип, но реализован вручную в CUDA C, а не через компилятор графа.

Пример: Truncated Gaussian на [a, b]

  log p(x | μ, σ) = −½((x−μ)/σ)² − log σ − log Z
  Z = Φ((b−μ)/σ) − Φ((a−μ)/σ)

  ∂log p/∂μ = (x−μ)/σ² − (φ(zₐ) − φ(z_b)) / (σ · Z)
  ∂log p/∂σ = ((x−μ)² − σ²) / σ³ − (zₐ·φ(zₐ) − z_b·φ(z_b)) / (σ · Z)

  где zₐ = (a−μ)/σ, z_b = (b−μ)/σ, φ = PDF стандартного нормального, Φ = CDF

Аналогично для Crystal Ball (5 параметров: μ, σ, α, n + нормировка),
Double Crystal Ball (7 параметров), Chebyshev (произвольный порядок).

Стоимость: идентична MoreFit для параметрических PDF
Преимущество: rate modifiers (NormSys, WeightSys) интегрированы в ядро

4.3 Reverse-mode AD (NextStat Tier 1, ns-ad Tape)

Запись вычислительного графа (forward pass), затем обратный проход:

  tape.var(θ₁), tape.var(θ₂), ...
  z = tape.mul(x, y)    // записывает z = x·y и ∂z/∂x = y, ∂z/∂y = x
  w = tape.add(z, x)    // записывает w = z+x и ∂w/∂z = 1, ∂w/∂x = 1
  tape.backward(w)       // обратный проход: O(|граф|)
  tape.adjoint(θⱼ)      // dw/dθⱼ

Стоимость backward: O(|граф|) ≈ O(N·P) — один обратный проход
Сравнение с конечными разностями: O(N·P) vs O(N·P²)
Ограничение: CPU only, нет GPU-версии (пока)
Используется для: HistFactory binned models, profile likelihood

4.4 Конечные разности (NextStat Tier 3)

Центральная схема:

  ∂L/∂θⱼ ≈ (L(θ + h·eⱼ) − L(θ − h·eⱼ)) / (2h)    h = 10⁻⁴

Стоимость: 2P+1 вызовов L(θ) на итерацию
Точность: O(h²) ≈ 10⁻⁸ (достаточно для L-BFGS, но не для Hessian)

Используется для: flow PDF (ONNX = чёрный ящик)
Компенсация: NLL reduction на GPU, flow eval на GPU (CUDA EP)
  → узкое место: flow inference, не reduction

5.Сравнительная таблица возможностей

ВозможностьMoreFit v0.1NextStat
Аналитические градиентыСимвольные из графаHand-derived CUDA + reverse-mode AD
GPU backendOpenCL (vendor-agnostic)CUDA + Metal (Apple Silicon)
JIT-компиляция ядраДа (граф → ядро)Нет (ядра компилируются в PTX на этапе сборки; возможен JIT драйвера при запуске на новых SM)
Parameter hoistingАвтоматический (из графа)Явное: закрытая нормировка у встроенных PDF + ручные константы; общего DAG-оптимизатора нет
CSEАвтоматическийРучной (в CUDA C коде)
Систематики (NormSys)Нет (отложено)Да, rate modifier в GPU-ядре
Систематики (WeightSys)Нет (отложено)Да (Code0 + Code4p interp)
Conditional flowsНетp(x|α): α как context vector
Selection / weightsОграниченоSelection (ns-root expr). Per-event weights в spec зарезервированы; weight systematics поддерживаются при построении шаблонов из TTree
Бинированные фитыНетПолная HistFactory + GPU
Произвольные PDFТолько встроенная библиотекаTrain flow → ONNX → drop-in
Морфинг шаблоновНетDCR surrogate (нейросетевой)
Batch toy fitsНе описанLockstep L-BFGS-B, 1 kernel
Device-resident toysНе описанToy sampling → batch fit без H2D obs_flat (H2D только toy_offsets)
ВалидацияБенчмарки скоростиGPU↔CPU parity tests (cuda/metal), CLI контрактные тесты, guardrails для toy-фитов
Multi-channelНе описанΣ NLL/grad по каналам; constraints привязаны к одному каналу

6.Архитектурное обоснование: почему ONNX, а не символьный JIT

6.1 Проблема расширяемости

MoreFit: каждая новая PDF-семья → формула → символьные производные → тест корректности → интеграция в кодогенератор. Для Crystal Ball это ~200 строк, для Double Crystal Ball — ~400. Для произвольной многомерной корреляционной структуры — невозможно.

NextStat: обучи normalizing flow на Monte Carlo данных → torch.onnx.export() → файл. Новый GPU-код не нужен. Flow автоматически учитывает корреляции любой размерности.

6.2 Проблема систематик в unbinned

В реальных анализах PDF зависит от nuisance parameters α: p(x | θ, α). MoreFit не имеет механизма для морфинга PDF по α.

NextStat: conditional normalizing flow принимает α как context vector. Один flow обучается на семплах при разных значениях α. На этапе фита flow вычисляет log p(x | θ, α) для произвольных значений α — гладкий морфинг по всему пространству систематик без дискретных шаблонов.

6.3 Проблема высокой размерности

Аналитические PDF плохо масштабируются в >3D: нет замкнутых форм для сложных корреляционных структур. Normalizing flows работают в произвольной размерности D и автоматически моделируют корреляции через обучение. Для анализов с несколькими наблюдаемыми (minv, cos θ, φ, Δη, ...) это критическое преимущество.


7.NextStat R&D: что в разработке

7.1 Device-resident toy pipeline

Текущая реализация (Phase 1, 1D): sampling на GPU → toy events остаются в device memory → batch fit без H2D copy основного буфера событий. Контракт: оффсеты тоев считаются на CPU и загружаются как маленький буфер ((n_toys+1)×4 байта), а большой obs_flat остаётся device-resident. API: fit_unbinned_toys_batch_cuda_device(). Toy offsets (~4 байта × n_toys) — единственный host→device transfer.

Важно: текущий GPU toy sampler ограничен аналитическими 1D PDF (Gaussian/Exponential); это инженерный гейт против «тихих» деградаций корректности при генерации.

7.2 Multi-channel batch fits

Модели с несколькими каналами (категориями): NLL = Σch NLLch. Каждый канал имеет свой набор процессов, но общие параметры. Реализовано через CudaUnbinnedBatchAccelerator per channel, суммирование NLL/grad на хосте. Gaussian constraints привязаны к одному каналу во избежание двойного счёта.

7.3 Metal backend (Apple Silicon)

Полностью параллельная реализация на Metal Shading Language для M1/M2/M3. Тот же lockstep L-BFGS-B, те же аналитические градиенты, f32 (Metal не поддерживает native f64). Контракт точности фиксируется тестами паритета: Metal (f32, atomic accumulation) сравнивается с CPU с ослабленными допусками порядка 1e−2 по abs и ~1e−6 по rel (см. crates/ns-compute/tests/unbinned_gpu_parity.rs).

7.4 Reverse-mode AD на GPU (исследование)

Текущий ns-ad::Tape работает на CPU. Исследуется возможность GPU-версии для HistFactory gradient kernel — потенциально заменит конечные разности для профильного likelihood и даст O(1) gradient cost по числу систематик вместо O(P).

7.5 Гибридная аналитика + flow в одном ядре

Для моделей, где сигнал — flow PDF, а фон — параметрическая PDF (Chebyshev, Exponential), планируется гибридное ядро: аналитические градиенты для фоновых компонент + finite differences только для signal flow. Это сократит число flow evaluations с 2P+1 до 2Pflow+1, где Pflow ≪ P.


8.Анализ производительности: теоретическая сложность

Обозначения:
  N = число событий, P = число параметров, K = итерации оптимизатора
  C_pdf = стоимость одного log p(x|θ)
  C_flow = стоимость одного flow inference (ONNX)

MoreFit (параметрическая PDF, символьный gradient):
  T_total = K_sym × N × P × C_pdf
  K_sym ≈ 2–3 (точный gradient → быстрая сходимость L-BFGS)

NextStat Tier 2 (параметрическая PDF, аналитический CUDA gradient):
  T_total = K_ana × N × P × C_pdf
  K_ana ≈ 2–5 (аналитический gradient, но без CSE/hoisting автоматизации)

NextStat Tier 3 (flow PDF, finite differences):
  T_total = K_fd × (2P+1) × N × C_flow
  K_fd ≈ 10–30 (менее точный gradient → больше итераций)

NextStat гибрид (P_param параметрических + P_flow flow компонент):
  T_total = K_hyb × [N × P_param × C_pdf + (2P_flow+1) × N × C_flow]
  K_hyb ≈ 5–15

Вывод: для чистых параметрических unbinned фитов без систематик MoreFit и NextStat Tier 2 дают сопоставимую производительность. MoreFit выигрывает за счёт автоматической оптимизации графа (CSE, hoisting). Но как только модель включает rate modifiers (NormSys/WeightSys), нейронные PDF (flow/ONNX) или требуется воспроизводимый toy-пайплайн с batch/lockstep фитом — компиляторно-символьный подход перестаёт быть универсальным, а NextStat остаётся применимым за счёт гибридной архитектуры.


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

MoreFit доказывает фундаментальный тезис: аналитические производные на GPU дают огромный выигрыш для простых unbinned фитов. Символьная дифференциация + JIT-компиляция — элегантный и мощный подход. Для чистых параметрических моделей без систематик MoreFit достигает ~10× ускорения над RooFit CUDA.

Но реальные физические анализы в LHCb, Belle II, ATLAS, CMS требуют:

  • Систематик. В unbinned — как минимум rate modifiers выхода (NormSys/WeightSys). В binned/HistFactory — shape/templating систематики (ShapeSys/HistoSys и т.п.).
  • Детерминированного определения выборки: selection и выражений наблюдаемых. Per-event weights для observed data распространены на практике, но в unbinned_spec_v0сейчас зарезервированы и не являются частью контракта модели.
  • Произвольных PDF для сигнала — аналитические формы не всегда существуют.
  • Бинированных фитов — для комбинаций каналов и legacy-анализов.
  • Валидации — GPU↔CPU паритет NLL/∇NLL, контрактные CLI-тесты, а также toy-based проверки bias/pull/coverage.

NextStat через гибридную архитектуру (аналитические CUDA-ядра для параметрических PDF + ONNX flow для произвольных PDF + reverse-mode AD для binned models) покрывает все эти требования. Цена — частичная потеря скорости для flow-компонент (конечные разности вместо аналитики). Но для production-анализов гибкость и полнота важнее пиковой скорости на бенчмарках.

Trade-off:
  MoreFit   = max скорость × min функциональность
  NextStat  = (0.3–1.0)× скорость MoreFit × полная функциональность

  Для production-анализа: NextStat покрывает класс задач, который выходит за рамки
  «закрытой библиотеки аналитических PDF без систематик» — ценой потери скорости
  на flow-компонентах (finite differences).

Ссылки

  • C. Langenbruch, «MoreFit: A More Optimised, Rapid and Efficient Fit», arXiv:2505.12414 (2025)
  • NextStat GPU kernels: crates/ns-compute/kernels/unbinned_nll_grad.cu (3150 строк, 6 PDF + rate modifiers + аналитические градиенты)
  • NextStat GPU Flow Session: crates/ns-inference/src/gpu_flow_session.rs (CPU/CUDA EP flow + GPU NLL reduction)
  • NextStat Reverse-mode AD: crates/ns-ad/src/tape.rs (computation graph + backward pass)
  • NextStat Batch Toy Fits: crates/ns-inference/src/unbinned_gpu_batch.rs (lockstep L-BFGS-B)
  • RooFit CUDA backend: L. Moneta et al., ROOT v6.32+ GPU acceleration