Как NextStat делает пайплайн HistFactory дифференцируемым в PyTorch

1.Проблема: метрика появляется после инференса

Нейросеть выдаёт скалярный скор для каждого события. Затем скор биннится в гистограмму, которая становится достаточной статистикой для правдоподобия HistFactory. Статистика для открытия обычно определяется так:

q₀  =  2 · [ NLL(θ̂₀, μ=0) − NLL(θ̂, μ̂) ]₊
Z₀  =  √q₀

где μ это параметр интереса (POI, «сила сигнала»), θ это мешающие параметры; шляпки означают оценки максимального правдоподобия (MLE), а [·]₊ это односторонний клиппинг для открытия.

Это и есть цель, важная для физики. Но вычислительный граф исторически ломает дифференцируемость в двух местах:

• ›Гистограммирование — жёсткие границы бинов дают нулевые градиенты почти всюду
• ›Профилирование — θ̂(s) = argmin NLL(θ; s) это задача оптимизации внутри прямого прохода

2.Обозначения и целевая функция

Используем стандартные обозначения для бинированного правдоподобия HistFactory: основные бины индексируются i, наблюдаемые счёты n_i, ожидаемые выходы модели ν_i(μ,θ; s) представлены суммой по сэмплам с модификаторами. Гистограмма сигнала s ∈ ℝ₊ᴮ это вектор, по которому мы дифференцируем.

NLL(μ, θ; s)  =  Σᵢ [ νᵢ − nᵢ·log(νᵢ) + log(nᵢ!) ]  +  ограничения(θ)

q₀(s)  =  2 · [ NLL(μ=0, θ̂₀(s); s) − NLL(μ̂(s), θ̂(s); s) ]₊

Цель:  ∂Z₀/∂s  →  правило цепочки  →  ∂Z₀/∂w   (веса сети)

3.Архитектура: общий вид

Сквозной стек, который превращает выход сети в дифференцируемую цель Z₀:

Что происходит при вызове loss.backward()

4.Слой 1 — SoftHistogram: делаем биннинг дифференцируемым

HistFactory работает с бинированными выходами, а нейросеть выдаёт непрерывные скоры. Жёсткая гистограмма h_j = Σ 𝟙{e_j ≤ x_k < e_{j+1}} недифференцируема по x_k.

Gaussian KDE mode (по умолчанию)

Каждое событие «мягко» вносит вклад во все бины через гауссово ядро с центром в центре бина:

w_{kj}  ∝  exp( −½ · ((x_k − c_j) / σ)² )
h_j    =  Σ_k  w̃_{kj}

Нормировка по событию: Σ_j w̃_{kj} = 1

Режим sigmoid-edge (быстрее)

Pr(x ∈ [e_j, e_{j+1}))  ≈  σ((x−e_j)/τ) − σ((x−e_{j+1})/τ)

Дешевле, но обычно менее гладко, чем KDE. Оба режима дают дифференцируемый тензор выходов по бинам, который и является вектором сигналов s.

5.Слой 2 — слитое GPU-ядро: NLL и аналитические градиенты

Ключевой примитив: слитое ядро, которое вычисляет (с ограничениями) пуассоновский NLL и аналитические градиенты и по мешающим параметрам θ (для L-BFGS-B), и по выходам сигнала s (для backpropagation).

Ожидаемые выходы HistFactory

νᵢ(θ; s)  =  Σₐ (nom_{a,i} + Δ_{a,i}(θ)) · F_{a,i}(θ)

F = произведение мультипликативных модификаторов (NormFactor/Lumi/NormSys/ShapeSys/…)
Δ = аддитивный вклад (HistoSys)

Ключевая идентичность для градиента

Для одного бина i выпишем пуассоновский вклад в NLL и его производную:

ℓᵢ(νᵢ)  =  νᵢ − nᵢ·log(νᵢ) + log(nᵢ!)

∂ℓᵢ/∂νᵢ  =  1 − nᵢ/νᵢ   ≡  wᵢ   ("вес")

∂NLL/∂sᵢ  =  wᵢ · F_{signal,i}  =  (1 − nᵢ/νᵢ) · F_{signal,i}

Именно это ядро и пишет в буфер градиентов по сигналу.

zero-copy для CUDA-тензоров PyTorch

• ›Python передаёт signal.data_ptr() (сырой адрес CUDA-памяти) в Rust как u64
• ›Rust передаёт указатель в ядро как g_external_signal
• ›Ядро читает бины сигнала напрямую из GPU-памяти PyTorch
• ›Ядро пишет ∂NLL/∂signal через atomicAdd во внешний буфер градиента

// Внутренний цикл differentiable_nll_grad.cu
double w = 1.0 - obs / expected_bin;
atomicAdd(&g_grad_signal_out[sig_local_bin], w * signal_factor);

Покрытие модификаторов

Слитое ядро реализует все 7 типов модификаторов HistFactory: NormFactor, Lumi, NormSys (code-4 polynomial/exponential), HistoSys (побиновый аддитивный, code-4p), ShapeSys, ShapeFactor, StatError, а также градиенты для дополнительных пуассоновских ограничений (параметры γ в стиле Barlow–Beeston) и гауссовых ограничений.

6.Слой 3 — Rust GPU-сессии: держим модель на устройстве

Ядро вызывается сотни раз на один прямой проход (для профильных целей). NextStat один раз сериализует модель HistFactory в плоские GPU-буферы:

• ›Номинальные выходы сэмплов
• ›Дескрипторы сэмплов (диапазоны бинов, оффсеты)
• ›Дескрипторы модификаторов, таблицы оффсетов и данные (коэффициенты NormSys, дельты HistoSys)
• ›Таблицы ограничений (constraints)

DifferentiableAccelerator (CUDA)

crates/ns-compute/src/differentiable.rs загружает PTX на этапе сборки, один раз загружает статические буферы модели, выделяет переиспользуемые device-буферы и запускает ядро с shared memory, размер которого равен params[n_params] + scratch[block_size].

Поддержка многоканального сигнала

s = [ s⁽⁰⁾ ‖ s⁽¹⁾ ‖ ⋯ ]

Каждая запись: (sample_idx, first_bin, n_bins)
Ядро отображает индекс основного бина → корректный оффсет в g_external_signal

7.Слой 4 — профильная значимость и теорема об огибающей

Профилирование вводит вложенную оптимизацию. Наивный путь, развернуть итерации оптимизатора и дифференцировать через них, дорог по памяти и времени и делает градиенты зависимыми от выбранного оптимизатора и критериев останова. NextStat использует теорему об огибающей.

Статистика открытия как разность двух профильных оптимумов

ProfiledDifferentiableSession находит свободный оптимум (μ̂, θ̂) и оптимум при ограничении μ=0 (θ̂₀), в обоих случаях используя ограниченный L-BFGS-B на CPU, при этом каждое вычисление NLL и ∂NLL/∂θ реализовано слитым запуском GPU-ядра.

Градиент по огибающей для q₀

Пусть V(s) = min_θ f(θ, s):

  dV/ds = ∂f/∂s |_{θ = θ̂(s)}

Набросок доказательства:  dV/ds = ∂f/∂s + (∂f/∂θ)·(dθ̂/ds)
                                        ↑
                                      = 0 в оптимуме

Применяем к обоим слагаемым в q₀:

  ∂q₀/∂s = 2 · ( ∂NLL/∂s |_{μ=0, θ=θ̂₀}  −  ∂NLL/∂s |_{μ=μ̂, θ=θ̂} )

На языке реализации это выглядит так: два профильных фита → вычислить ∂NLL/∂signal в каждом оптимуме → взять разность. Градиент точный при сходимости. Если оптимизатор не сошёлся, NextStat возвращает ошибку, а не выдаёт некорректный градиент.

8.PyTorch autograd: как «провести» градиенты в backprop

NLL при фиксированных параметрах (фаза 1)

Forward: выделить и занулить grad_signal (для atomicAdd нужен нулевой буфер) → синхронизировать CUDA-потоки → вызвать нативную сессию с сырыми device-указателями → синхронизировать → сохранить grad_signal. Backward это просто правило цепочки: ∂L/∂s = (∂L/∂NLL) · (∂NLL/∂s).

SignificanceLoss: внешний API для пользователя

from nextstat.torch import SoftHistogram, SignificanceLoss

loss_fn = SignificanceLoss(model, signal_sample_name="signal")
soft_hist = SoftHistogram(bin_edges, bandwidth="auto", mode="kde")

scores = classifier(batch)                 # [N]
signal_hist = soft_hist(scores).double()   # [B] float64 для CUDA

loss = loss_fn(signal_hist)               # возвращает −Z₀
loss.backward()                           # градиенты идут до весов NN
optimizer.step()

Зачем здесь явный torch.cuda.synchronize()?

PyTorch и NextStat запускают ядра в разных CUDA-потоках. Без барьера ядро NextStat может прочитать незаполненный тензор сигнала, или PyTorch может прочитать незаполненный градиент. Текущая реализация использует полный synchronize() ради корректности; будущая оптимизация это ожидание событий на уровне конкретных потоков.

9.Экспериментальная валидация

Для дифференцируемого слоя инференса градиент это и есть продукт. NextStat проверяет градиенты двумя способами:

• ›PyTorch CUDA слой NLL гейтится конечными разностями на GPU (pytest: tests/python/test_torch_layer.py, критерий: max_err < 1e-4).
• ›Профилированные q₀/qμ гейтятся конечными разностями на одном информативном бине (pytest: tests/python/test_differentiable_profiled_q0.py, критерий: rel < 0.1).

10.Metal-бэкенд (Apple Silicon)

Тот же алгоритм, другие ограничения:

• ›Слитое ядро работает в f32 (точность GPU у Apple)
• ›Входы/выходы приводятся на границе API (тензоры PyTorch остаются f64)
• ›Допуск L-BFGS-B ослаблен до 1e−3 (против 1e−5 на CUDA)
• ›Гистограмма сигнала загружается через CPU (нет raw-pointer interop с тензорами MPS)

Согласие NLL с CPU f64: ~1e−6 по относительному уровню для типичных воркспейсов.

11.Заметки о производительности

Дифференцируемый слой спроектирован так, чтобы на каждой итерации NLL и градиенты вычислялись в одном слитом запуске GPU-ядра, а работа на CPU ограничивалась обновлением состояния L-BFGS-B. Профильные цели требуют двух профильных фитов на один forward pass, но работают на небольших гистограммах (O(10–1000) бинов), и фит с тёплым стартом обычно сходится быстро.

12.Ограничения и будущая работа

• ›Производные величины инференса (μ_up, полосы CLs) требуют дифференцирования через нахождение корней; теорема об огибающей напрямую покрывает только q₀ и q_μ
• ›Интеграция со стримами CUDA может снизить накладные расходы на синхронизацию
• ›Более агрессивное слияние ядер: вычислять оба профильных градиента внутри Rust и писать их напрямую в буфер градиентов PyTorch

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

NextStat превращает исторически недифференцируемый HEP-пайплайн инференса в loss-функцию, совместимую с PyTorch, за счёт:

• ›замены жёсткого биннинга на дифференцируемую гистограмму (SoftHistogram)
• ›реализации слитого GPU-ядра для NLL HistFactory и аналитических градиентов
• ›удержания сериализованного состояния модели на устройстве через Rust-сессии
• ›применения теоремы об огибающей для точных профильных градиентов без AD через оптимизатор

Это позволяет обучать нейросети напрямую на физической значимости, оптимизируя то, что в итоге и важно для анализа, а не суррогатную прокси-метрику.