Дифференцируемый слой HistFactory для PyTorch

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

Если вы хотите оптимизировать именно метрику открытия, а не суррогат вроде BCE/MSE, то целевой конвейер такой:оценки NN → дифференцируемая гистограмма → профилированное правдоподобие → q₀ → Z₀. На практике это доступно как SoftHistogram и SignificanceLoss в nextstat.torch.

Что дифференцируемо (а что нет)

• Слой дифференцируем по сигнальной гистограмме s: это одномерный тензор побиновых выходов (yields), который формирует ваша модель совместно с SoftHistogram.
• Для целевой функции с фиксированными параметрами nll_loss(...) мешающие параметры считаются константами: backward возвращает ∂NLL/∂s.
• Для профилированных целей (q₀, qμ, Z₀, Zμ) NextStat используеттеорему об огибающей, чтобы вычислять точный градиент по sв точке найденного оптимума, не дифференцируя через итерации оптимизатора.
• Пока не поддерживается: производные по величинам, которые получаются через решение уравнений, например полосы CLs (Brazil band) или μ_up(α). Обычно для этого требуется неявное дифференцирование через этап root-finding (см. «Фаза 2 (в будущем)» ниже).

Архитектура

PyTorch Training Loop (GPU)
    |
    +-- Neural Network --> signal_histogram (torch.Tensor, CUDA)
    |
    +-- NextStatNLL.forward(signal_histogram)
    |   +-- tensor.data_ptr() --> raw CUDA device pointer (u64)
    |   +-- Rust (PyO3): passes pointer to CUDA kernel
    |   +-- Kernel: reads signal bins from PyTorch memory (ZERO-COPY)
    |   +-- Kernel: writes dNLL/d(signal) into PyTorch grad tensor (ZERO-COPY)
    |   +-- Return: NLL scalar
    |
    +-- NextStatNLL.backward(grad_output)
        +-- grad_signal = cached from forward (already on CUDA)
        +-- return grad_output * grad_signal --> flows back to NN

Ключевая идея (CUDA, фаза 1 NLL): в цикле обучения нет копирований устройство↔хост для сигнальной гистограммы и ее градиента. CUDA-ядро читает сигнальные бины напрямую из памяти PyTorch на GPU и записывает градиенты туда же. Небольшие передачи остаются только для мешающих параметров и возврата скалярного значения NLL.

Для профилированных целей (q₀/qμ/Z₀/Zμ) сигнальная гистограмма также читается zero-copy, но финальный вектор ∂/∂signal сейчас возвращается в Python как небольшой Vec<f64> на хосте и затем материализуется как CUDA-тензор (порядка 100–1000 чисел).

Быстрый старт: end-to-end обучение по Z₀

import torch
import nextstat
from nextstat.torch import SoftHistogram, SignificanceLoss

model = nextstat.from_pyhf(workspace_json)

# Однократная инициализация GPU-сессии (предпочтительно CUDA, с запасным вариантом Metal)
loss_fn = SignificanceLoss(model, signal_sample_name="signal", device="auto")  # по умолчанию возвращает -Z0

soft_hist = SoftHistogram(
    bin_edges=torch.linspace(0.0, 1.0, 11),
    bandwidth="auto", mode="kde",
)

optimizer = torch.optim.Adam(nn.parameters(), lr=1e-3)
for batch in dataloader:
    optimizer.zero_grad()

    scores = nn(batch)                      # [N]
    signal_hist = soft_hist(scores).double() # [B]
    if torch.cuda.is_available():
        signal_hist = signal_hist.cuda()    # CUDA path expects float64

    loss = loss_fn(signal_hist)             # returns -Z0
    loss.backward()
    optimizer.step()

Быстрый старт: NLL при фиксированных параметрах

import torch
import nextstat
from nextstat.torch import create_session, nll_loss

model = nextstat.from_pyhf(workspace_json)
session = create_session(model, signal_sample_name="signal")

optimizer = torch.optim.Adam(nn.parameters(), lr=1e-3)
for batch in dataloader:
    optimizer.zero_grad()
    signal_hist = nn(batch).double().cuda()  # CUDA float64
    loss = nll_loss(signal_hist, session)
    loss.backward()                          # градиент проходит через NextStat обратно в сеть
    optimizer.step()

Справка API

DifferentiableSession (нативный API)

Доступен как nextstat._core.DifferentiableSession при сборке с CUDA:

Профилированная значимость (q₀ / qμ) на GPU

NextStat также предоставляет GPU-ускоренный профилированный слой (CUDA или Metal) для стандартных тестовых статистик профильного правдоподобия:

• Открытие: q₀ = 2·(NLL(μ=0, θ̂₀) − NLL(μ̂, θ̂))
• Верхние пределы: qμ = 2·(NLL(μ_test, θ̂_μ) − NLL(μ̂, θ̂)) с односторонним отсечением

Для одного forward-прохода требуется два профилированных фита, поэтому это дороже, чем nll_loss, но зато оптимизируется напрямую физическая тестовая статистика.

from nextstat.torch import (
    create_profiled_session,
    profiled_q0_loss,    # тестовая статистика открытия q₀
    profiled_z0_loss,    # √q₀ с численной устойчивостью
    profiled_qmu_loss,   # тестовая статистика верхнего предела qμ
)

session = create_profiled_session(model, "signal")
signal = nn(batch).double().cuda().requires_grad_(True)

q0 = profiled_q0_loss(signal, session)
z0 = profiled_z0_loss(signal, session)
qmu = profiled_qmu_loss(signal, session, mu_test=5.0)

Формулы для градиентов

Теорема об огибающей (профилированные цели):

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

Фаза 1: градиент NLL по сигнальным бинам при фиксированных параметрах:

expected_i = (signal_i + delta_i) * factor_i + sum_{other samples}
dNLL/d(signal_i) = (1 - obs_i / expected_i) * factor_i

Здесь factor_i это произведение всех мультипликативных модификаторов (NormFactor, NormSys, ShapeSys и т. п.), примененных к сигнальному сэмплу в бине i. Односторонняя значимость: если μ̂ < 0 (или статистика отсечется к q₀=0), то возвращаемые q₀ и его градиент равны нулю.

Практические замечания (важно для корректности)

• Градиентный буфер должен начинаться с нулей (CUDA): CUDA-ядро накапливает вклад через atomicAdd, поэтому выходной тензор градиента нужно обнулить. Python-обертка делает это через torch.zeros_like(signal).
• Синхронизация CUDA-потоков: PyTorch и NextStat могут использовать разные CUDA streams. Обертки вызывают torch.cuda.synchronize()до и после нативного вызова.
• Многоканальная раскладка сигнала: если сигнальный сэмпл присутствует в нескольких каналах, внешний буфер сигнала задается как конкатенация[ch0_bins..., ch1_bins..., ...]. Общее число бинов можно получить через session.signal_n_bins().

# Многоканальный случай: конкатенируйте гистограммы по каналам в порядке модели
signal = torch.cat([signal_sr, signal_vr], dim=0).double().cuda()
loss = loss_fn(signal)

Metal (Apple Silicon)

create_profiled_session(..., device="auto") предпочитает CUDA, а при отсутствии может перейти на Metal (требуется сборка с --features metal).

• Вычисления на GPU идут в f32 (ограничения точности Apple GPU); входы и выходы конвертируются на границе API
• Сигнал загружается с CPU (нет прямого взаимодействия указателей с MPS-тензорами)
• Допуск L-BFGS-B ослаблен до 1e-3 (вместо 1e-5 на CUDA)

Валидация и доказательства

• CUDA zero-copy NLL + signal gradient (∂NLL/∂s): tests/python/test_torch_layer.py
• Профилированные градиенты q₀/qμ (теорема об огибающей): tests/python/test_differentiable_profiled_q0.py
• Finite-difference error: max FD error 2.07e⁻⁹ on benchmark fixtures

Архитектурные решения

Зачем отдельное CUDA-ядро? Существующее batch_nll_grad.cu оптимизировано под батч-фиттинг тоев (1 block = 1 toy). У дифференцируемого слоя другие требования: вычисление для одной модели, внешний указатель на сигнал, вывод градиента по сигналу. Разделение позволяет избежать ветвлений на горячем пути.

Зачем zero-copy? Традиционный путь выглядит так: PyTorch GPU → CPU → Rust → GPU → CPU → PyTorch GPU (четыре передачи по PCIe). В zero-copy варианте ядро читает сигнальные бины и пишет градиенты напрямую в памяти GPU. Единственная заметная H→D передача это небольшой вектор мешающих параметров (порядка 250 double, около 2 КБ).

Фаза 2 (в будущем): неявное дифференцирование

Для производных профилированных величин (например, интерполированный верхний предел μ_up(α) или полные полосы CLs) обычно нужно неявное дифференцирование через слой решателя или root-finding. Градиента по теореме об огибающей достаточно для q₀ и qμ, но не для всех downstream-величин.

dq/ds = dq/ds|_{θ fixed} - (d²NLL/ds dθ)ᵀ (d²NLL/dθ²)⁻¹ dq/dθ|_{s fixed}

Для этого нужна смешанная производная (кросс-гессиан) d²NLL/ds/dθ, которую можно оценивать конечными разностями от GPU-градиента по сигнальным бинам.

Конкурентный ландшафт

Требования

• PyTorch (опциональная зависимость, импортируется лениво)
• NVIDIA GPU с CUDA для zero-copy; Apple Silicon (Metal) как запасной
• CPU-тензоры поддерживаются, но требуют копирования устройство→хост