Factorizable Normalizing Flows for parameter-dependent density morphing

Technische Samenvatting: Factoriseerbare Normalizing Flows voor Parameterafhankelijke Dichtheidsmorfing

Probleemstelling
In de hoge-energiefysica (HEP) en andere wetenschappelijke domeinen vereist likelihood-gebaseerde inferentie vaak het modelleren van hoe een waarschijnlijkheidsdichtheidsfunctie vervormt als functie van continue parameters. Deze parameters kunnen fysieke effecten, kalibratieconstanten of systematische onzekerheden (nuisance parameters) vertegenwoordigen. Hoewel Normalizing Flows (NFs) uitblinken in het modelleren van een enkele vaste dichtheid, is het leren van een aparte flow voor elke mogelijke configuratie van $K$ parameters onhandelbaar omdat de gezamenlijke configuratieruimte exponentieel groeit met $K$.

Traditionele HEP-benaderingen vertrouwen op "template variaties", waarbij simulaties worden uitgevoerd bij $\pm 1\sigma$ voor elke nuisance parameter afzonderlijk, en de resultaten worden geïnterpoleerd met laag-orde polynomen op gebinde, laag-dimensionale data. Deze methode is robuust maar beperkt door de afhankelijkheid van binning en de aanname dat nuisance-effecten factoriseren (waarbij correlaties worden genegeerd). Het uitbreiden hiervan naar unbinned, hoog-dimensionale fits vereist een methode die:

1. Getrouw is: Accuraat dichtheidsvervormingen vastlegt.
2. Profiling-geschikt is: Een gladde, differentieerbare afhankelijkheid van parameters biedt.
3. Schaalbaar is: Het vermijden van het samplen van de volledige $K$-dimensionale nuisance-ruimte.
4. Interpreteerbaar is: Expliciet laat zien hoe elke parameter de distributie vormgeeft.

Methodologie: Factoriseerbare Normalizing Flows (FNFs)
De auteurs introduceren Factorizable Normalizing Flows (FNFs), een generatief framework dat het modelleren van de complexe nominale dichtheid ontkoppelt van het modelleren van systematische variaties. De kernarchitectuur bestaat uit twee componenten:

1. Fixed Nominal Flow: Een hoog-getrouwe, inverse flow (bijv. Neural Spline Flow) wordt eenmalig getraind op een referentisimulatie ($\nu = 0$) om de nominale dichtheid $p_{\text{nom}}(y | x)$ te modelleren. Deze component wordt bevroren tijdens de fase van systematische modellering.

2. Learnable Factorizable Transformation ($T_\nu$): Een transformatie wordt geleerd die geobserveerde gebeurtenissen $y$ (beïnvloed door de systematiek $\nu$) terugmaakt naar het nominale referentiekader $y_{\text{nom}}$. De likelihood wordt geëvalueerd via de change-of-variables formule:
   $$p(y | x, \nu) = p_{\text{nom}}(T_\nu(y | x, \nu) | x) \cdot |\det \nabla_y T_\nu(y | x, \nu)|$$

   De transformatie $T_\nu$ wordt geconstrueerd als een autoregressieve affine kaart waarbij de schaal ($s$) en verschuiving ($t$) parameters polynoomfuncties zijn van de nuisance parameters $\nu$. Cruciaal is dat deze polynomen factoriseerbaar zijn over de individuele nuisances:
   $$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$$

   De coëfficiënten ($\alpha, \beta, \gamma, \delta, \phi, \psi$) zijn outputs van gemaskeerde Multi-Layer Perceptrons (MLP's) die geconditioneerd zijn op de inputfeatures.

   • Additiviteit: De primaire structuur gaat uit van additieve bijdragen van elke nuisance, waardoor elke parameter in isolatie getraind kan worden met enkel $\pm 1\sigma$ samples voor die specifieke parameter.
   • Interactietermen: Optionele cross-termen ($\phi, \psi$) kunnen worden toegevoegd om paarcorrelaties tussen nuisances te vangen, wat training op gecombineerde variatie-samples vereist.

Belangrijkste Bijdragen

• Schaalbare Architectuur: De FNF generaliseert de discrete template-variatie methode naar een continue, differentieerbare hoog-dimensionale setting. Het vermijdt de "vloek van dimensionaliteit" door te trainen op per-nuisance variaties in plaats van de gezamenlijke ruimte.
• Lineaire Schaling: De computationele kosten en de benodigde trainingsdata schalen lineair met het aantal nuisance parameters (wanneer cross-termen worden weggelaten), wat het haalbaar maakt voor analyses met honderden onzekerheden.
• Tractabele Likelihood: De autoregressieve structuur zorgt ervoor dat de Jacobian determinant driehoekig en goedkoop te berekenen is ($O(D)$), wat efficiënte evaluatie tijdens likelihood fits mogelijk maakt.
• Interpreteerbaarheid: Het model scheidt expliciet de lineaire, kwadratische en interactie-effecten van elke parameter, wat een duidelijke fysieke interpretatie geeft van hoe systematieken de dichtheid vervormen.

Experimentele Resultaten
De auteurs hebben de FNF gevalideerd op een gecontroleerde toy-dataset bestaande uit twee klassen en twee nuisance parameters ($\nu_{\text{shift}}$ en $\nu_{\text{squeeze}}$) die zowel kinematische variabelen als conditionele score-distributies vervormen.

• Reconstructie: De geleerde transformatie reproduceerde de opgelegde vervormingen succesvol. Voor een pure translatie-nuisance leerde het model een lineaire verschuiving met een eenheidsschaal; voor een volume-behoudende squeeze leerde het de passende niet-lineaire schaling en verschuiving.
• Closure: Het factoriseerbare model (zonder cross-termen) bereikte closure tot de optimale likelihood (binnen $\sim 10^{-3}$ nats/event) nabij de trainingsassen en het nominale punt.
• Correlaties: In regio's waar beide nuisances tegelijkertijd sterk varieerden, vertoonde het factoriseerbare model een degradatie in de kwaliteit van de fit door de verwaarloosde correlaties. Het inschakelen van de optionele cross-termen verminderde de Kullback-Leibler divergentie met meer dan een factor vier in deze gecorreleerde regio's, wat de capaciteit van de architectuur demonstreert om te interpoleren tussen een standaard factoriseerbare benadering en een volledig gecorreleerd model.

Betekenis en Claims
Het artikel claimt dat FNFs een algemeen instrument bieden voor inferentie-workflows die continue dichtheidsmorfing vereisen. Door het mogelijk te maken om systematische onzekerheden te modelleren in hoog-dimensionale, unbinned ruimtes zonder de combinatorische kosten van gezamenlijke sampling, verlagen FNFs een belangrijke barrière voor de adoptie van unbinned likelihood fits in de hoge-energiefysica. De auteurs stellen dat hoewel dit werk de parametrische dichtheidsmodellering valideert, de inzet van FNFs in een volledige, geprofileerde, unbinned likelihood fit om parameters van belang te extraheren, het onderwerp is van toegewijd toekomstig werk. De aanpak wordt gepresenteerd als een directe, continue generalisatie van het gevestigde template-variatie recept, waarbij de robuustheid behouden blijft terwijl de toepasbaarheid naar hoog-dimensionale observabelen wordt uitgebreid.