Factorizable Normalizing Flows for parameter-dependent density morphing

Resumen Técnico: Flujos Normalizantes Factorizables para el Morfismo de Densidad Dependiente de Parámetros

Planteamiento del Problema
En la física de altas energías (HEP) y otros dominios científicos, la inferencia basada en verosimilitud requiere a menudo modelar cómo una función de densidad de probabilidad se deforma como función de parámetros continuos. Estos parámetros pueden representar efectos físicos, constantes de calibración o incertidumbres sistemáticas (parámetros de molestia o nuisance parameters). Aunque los Flujos Normalizantes (NF) son excelentes para modelar una densidad fija única, aprender un flujo separado para cada configuración posible de $K$ parámetros es intratable porque el espacio de configuración conjunto crece exponencialmente con $K$.

Los enfoques tradicionales en HEP dependen de las "variaciones de plantillas" (template variations), donde se ejecutan simulaciones en $\pm 1\sigma$ para cada parámetro de molestia individualmente, y los resultados se interpolan mediante polinomios de bajo orden en datos discretizados (binned) de baja dimensionalidad. Este método es robusto pero limitado por su dependencia de la discretización y su suposición de que los efectos de molestia se factorizan (descuidando las correlaciones). Extender esto a ajustes no discretizados (unbinned) y de alta dimensionalidad requiere un método que sea:

1. Fiel: Capture con precisión las deformaciones de la densidad.
2. Capaz de perfilado (profiling-capable): Proporcione una dependencia suave y diferenciable de los parámetros.
3. Escalable: Evite el muestreo de todo el espacio de molestia de $K$ dimensiones.
4. Interpretable: Muestre explícitamente cómo cada parámetro reformatea la distribución.

Metodología: Flujos Normalizantes Factorizables (FNF)
Los autores introducen los Flujos Normalizantes Factorizables (FNF), un marco generativo que desacopla el modelado de la densidad nominal compleja del modelado de las variaciones sistemáticas. La arquitectura central consta de dos componentes:

1. Flujo Nominal Fijo: Un flujo de alta fidelidad e invertible (por ejemplo, Neural Spline Flow) que se entrena una sola vez sobre una simulación de referencia ($\nu = 0$) para modelar la densidad nominal $p_{\text{nom}}(y | x)$. Este componente permanece congelado durante la fase de modelado de sistemáticas.

2. Transformación Factorizable Aprendible ($T_\nu$): Se aprende una transformación que mapea los eventos observados $y$ (afectados por las sistemáticas $\nu$) de vuelta al marco de referencia nominal $y_{\text{nom}}$. La verosimilitud se evalúa mediante la fórmula de cambio de variables:
   $$p(y | x, \nu) = p_{\text{nom}}(T_\nu(y | x, \nu) | x) \cdot |\det \nabla_y T_\nu(y | x, \nu)|$$

   La transformación $T_\nu$ se construye como un mapa afín autorregresivo donde los parámetros de escala ($s$) y desplazamiento ($t$) son funciones polinómicas de los parámetros de molestia $\nu$. Crucialmente, estos polinomios están factorizados sobre las molestias individuales:
   $$s_j(y_{<j}, x, \nu) = \sum_{k=1}^K (\nu_k \alpha^k_j + \nu_k^2 \beta^k_j) + \sum_{k<\ell} \nu_k \nu_\ell \phi^k_\ell$$
   $$t_j(y_{<j}, x, \nu) = \sum_{k=1}^K (\nu_k \gamma^k_j + \nu_k^2 \delta^k_j) + \sum_{k<\ell} \nu_k \nu_\ell \psi^k_\ell$$
   Los coeficientes ($\alpha, \beta, \gamma, \delta, \phi, \psi$) son salidas de Perceptrones Multicapa (MLP) enmascarados condicionados a las características de entrada.

   • Aditividad: La estructura primaria asume contribuciones aditivas de cada molestia, lo que permite entrenar cada una de forma aislada utilizando solo muestras de $\pm 1\sigma$ para ese parámetro específico.
   • Términos de Interacción: Se pueden incluir términos cruzados opcionales ($\phi, \psi$) para capturar correlaciones de pares entre las molestias, lo que requiere el entrenamiento sobre muestras de variación combinadas.

Contribuciones Clave

• Arquitectura Escalable: El FNF generaliza el método de variación de plantillas discreto a un entorno continuo, diferenciable y de alta dimensionalidad. Evita la "maldición de la dimensionalidad" al entrenarse en variaciones por molestia en lugar de en el espacio conjunto.
• Escalamiento Lineal: El coste computacional y los requisitos de datos escalan linealmente con el número de parámetros de molestia (cuando se omiten los términos cruzados), lo que lo hace viable para análisis con cientos de incertidumbres.
• Verosimilitud Tratable: La estructura autorregresiva asegura que el determinante del Jacobiano sea triangular y de bajo coste de computación ($O(D)$), permitiendo una evaluación eficiente durante los ajustes de verosimilitud.
• Interpretabilidad: El modelo separa explícitamente los efectos lineales, cuadráticos e de interacción de cada parámetro, proporcionando una clara interpretación física de cómo las sistemáticas deforman la densidad.

Resultados Experimentales
Los autores validaron el FNF en un conjunto de datos de juguete controlado que involucra dos clases y dos parámetros de molestia ($\nu_{\text{shift}}$ y $\nu_{\text{squeeze}}$) que deforman tanto las variables cinemáticas como las distribuciones de puntuación (score distributions) condicionales.

• Reconstrucción: La transformación aprendida reprodujo con éxito las deformaciones impuestas. Para una molestia de traslación pura, el modelo aprendió un desplazamiento lineal con escala unitaria; para un estrujamiento (squeeze) de volumen preservado, aprendió el escalamiento y el desplazamiento no lineal apropiados.
• Cierre (Closure): El modelo factorizado (sin términos cruzados) logró un cierre respecto a la verosimilitud óptima (dentro de $\sim 10^{-3}$ nats/evento) cerca de los ejes de entrenamiento y del punto nominal.
• Correlaciones: En regiones donde ambas molestias variaban fuertemente de forma simultánea, el modelo factorizado mostró una degradación en la calidad del ajuste debido a las correlaciones omitidas. Habilitar los términos cruzados opcionales redujo la divergencia de Kullback-Leibler en más de un factor de cuatro en estas regiones correlacionadas, demostrando la capacidad de la arquitectura para interpolar entre una aproximación factorizada estándar y un modelo totalmente correlacionado.

Significado y Reivindicaciones
El artículo afirma que los FNF proporcionan una herramienta general para flujos de trabajo de inferencia que requieren el morfismo de densidad continua. Al permitir el modelado de incertididades sistemáticas en espacios de alta dimensionalidad y no discretizados, los FNF reducen una barrera clave para la adopción de ajustes de verosimilitud no discretizados en la física de altas energías. Los autores declaran que, si bien este trabajo valida el modelado paramétrico de la densidad, el despliegue de los FNF en un ajuste de verosimilitud completo y perfilado para extraer parámetros de interés es objeto de un trabajo dedicado futuro. El enfoque se presenta como una generalización directa y continua de la receta establecida de variación de plantillas, manteniendo su robustez mientras extiende su aplicabilidad a observables de alta dimensionalidad.