Factorizable Normalizing Flows for parameter-dependent density morphing

Resumo Técnico: Fluxos Normalizantes Fatorizáveis para Morfismo de Densidade Dependente de Parâmetros

Declaração do Problema
Na física de altas energias (HEP) e outros domínios científicos, a inferência baseada em verossimilhança frequentemente requer a modelagem de como uma função de densidade de probabilidade (PDF) se deforma como uma função de parâmetros contínuos. Esses parâmetros podem representar efeitos físicos, constantes de calibração ou incertezas sistemáticas (parâmetros de incerteza/nuisance parameters). Embora os Fluxos Normalizantes (NFs) se destaquem na modelagem de uma densidade fixa única, aprender um fluxo separado para cada configuração possível de $K$ parâmetros é intratável porque o espaço de configuração conjunto cresce exponencialmente com $K$.

As abordagens tradicionais em HEP baseiam-se em "variações de template", onde simulações são executadas para $\pm 1\sigma$ para cada parâmetro de incerteza individualmente, e os resultados são interpolados usando polinômios de baixa ordem em dados binados e de baixa dimensão. Este método é robusto, mas limitado por sua dependência de binagem e por sua suposição de que os efeitos de incerteza fatorizam (negligenciando correlações). Estender isso para ajustes (fits) não binados e de alta dimensão exige um método que seja:

1. Fiel: Captura com precisão as deformações de densidade.
2. Capaz de Profiling: Fornece uma dependência suave e diferenciável dos parâmetros.
3. Escalável: Evita a amostragem do espaço de incerteza total de $K$ dimensões.
4. Interpretável: Mostra explicitamente como cada parâmetro remodela a distribuição.

Metodologia: Fluxos Normalizantes Fatorizáveis (FNFs)
Os autores introduzem os Fluxos Normalizantes Fatorizáveis (FNFs), um arcabouço generativo que desacopla a modelagem da densidade nominal complexa da modelagem das variações sistemáticas. A arquitetura central consiste em dois componentes:

1. Fluxo Nominal Fixo: Um fluxo de alta fidelidade e invertível (ex: Neural Spline Flow) é treinado uma única vez em uma simulação de referência ($\nu = 0$) para modelar a densidade nominal $p_{\text{nom}}(y | x)$. Este componente é congelado durante a fase de modelagem de sistemas.
2. Transformação Fatorizável Aprendível ($T_\nu$): Uma transformação é aprendida para mapear eventos observados $y$ (afetados por sistemáticas $\nu$) de volta ao referencial nominal $y_{\text{nom}}$. A verossimilhança é avaliada via a fórmula de mudança de variáveis:
   $$p(y | x, \nu) = p_{\text{nom}}(T_\nu(y | x, \nu) | x) \cdot |\det \nabla_y T_\nu(y | x, \nu)|$$

A transformação $T_\nu$ é construída como um mapa afim autorregressivo onde os parâmetros de escala ($s$) e deslocamento ($t$) são funções polinomiais dos parâmetros de incerteza $\nu$. Crucialmente, estes polinômios são fatorizados sobre as incertezas individuais:
$$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$$

Os coeficientes ($\alpha, \beta, \gamma, \delta, \phi, \psi$) são saídas de MLPs (Multi-Layer Perceptrons) mascaradas condicionadas às características de entrada.

• Aditividade: A estrutura primária assume contribuições aditivas de cada incerteza, permitindo que cada uma seja treinada isoladamente usando apenas amostras de $\pm 1\sigma$ para aquele parâmetro específico.
• Termos de Interação: Termos cruzados opcionais ($\phi, \psi$) podem ser incluídos para capturar correlações de pares entre incertezas, exigendo o treinamento em amostras de variação combinada.

Principais Contribuições

• Arquitetura Escalável: O FNF generaliza o método de variação de template discreto para um cenário contínuo, diferenciável e de alta dimensão. Ele evita a "maldição da dimensionalidade" ao treinar nas variações por incerteza, em vez do espaço conjunto.
• Escalabilidade Linear: O custo computacional e os requisitos de dados para treinamento escalam linearmente com o número de parâmetros de inciaza (quando os termos cruzados são omitidos), tornando-o viável para análises com centenas de incertezas.
• Verossimilhança Tratável: A estrutura autorregressiva garante que o determinante Jacobiano seja triangular e barato de computar ($O(D)$), permitindo a avaliação eficiente durante ajustes de verossimilhança.
• Interpretabilidade: O modelo separa explicitamente os efeitos lineares, quadráticos e de interação de cada parâmetro, fornecendo uma interpretação física clara de como as sistemáticas deformam a densidade.

Resultados Experimentais
Os autores validaram o FNF em um conjunto de dados de brinquedo (toy dataset) controlado envolvendo duas classes e dois parâmetros de incerteza ($\nu_{\text{shift}}$ e $\nu_{\text{squeeze}}$) que deformam tanto variáveis cinemáticas quanto distribuições de score condicionais.

• Reconstrução: A transformação aprendida reproduziu com sucesso as deformações impostas. Para uma incerteza de translação pura, o modelo aprendeu um deslocamento linear com escala unitária; para um esmagamento (squeeze) preservador de volume, o modelo aprendeu a escala não linear e o deslocamento apropriados.
• Fechamento (Closure): O modelo fatorizado (sem termos cruzados) alcançou fechamento em relação à verossimilhança ótima (dentro de $\sim 10^{-3}$ nats/evento) próximo aos eixos de treinamento e ao ponto nominal.
• Correlações: Em regiões onde ambas as incertezas variavam fortemente simultaneamente, o modelo fatorizado mostrou uma degradação na qualidade do ajuste devido às correlações negligenciadas. Habilitar os termos cruzados opcionais reduziu a divergência de Kullback-Leibler em mais de um fator de quatro nessas regiões correlacionadas, demonstrando a capacidade da arquitetura de interpolar entre uma aproximação fatorizada padrão e um modelo totalmente correlacionado.

Significância e Alegações
O artigo afirma que os FNFs fornecem uma ferramenta geral para fluxos de trabalho de inferência que requerem morfismo de densidade contínuo. Ao permitir a modelagem de incertezas sistemáticas em espaços de alta dimensão e não binados, sem o custo combinatório da amostragem conjunta, os FNFs reduzem uma barreira fundamental para a adoção de ajustes de verossimilhança não binados. Os autores afirmam que, embora este trabalho valide a modelagem paramétrica de densidade, a implementação de FNFs em um ajuste de verossimilhança completo e perfilado (profiled) para extrair parâmetros de interesse é objeto de um trabalho dedicado futuro. A abordagem é apresentada como uma generalização direta e contínua da receita estabelecida de variação de template, mantendo sua robustez enquanto estende sua aplicabilidade a observáveis de alta dimensão.