Factorizable Normalizing Flows for parameter-dependent density morphing

기술 요약: 매개변수 의존적 밀도 모핑을 위한 인수 분해 가능한 정규화 흐름(Factorizable Normalizing Flows)

문제 정의
고에너지 물리학(HEP) 및 기타 과학 분야에서, 우도 기반 추론(likelihood-based inference)은 종종 연속적인 매개변수의 함수로서 확률 밀도 함수(PDF)가 어떻게 변형되는지를 모델링해야 합니다. 이러한 매개변수는 물리적 효과, 교정 상수 또는 계통 오차(nuisance parameters)를 나타낼 수 있습니다. 정규화 흐름(Normalizing Flows, NF)은 단일 고정 밀도를 모델링하는 데는 탁월하지만, $K$개의 매개변수에 대한 모든 가능한 구성을 위해 별도의 흐름을 학습하는 것은 $K$에 따라 지수적으로 증가하는 결합 구성 공간 때문에 불가능합니다.

전통적인 HEP 방식은 각 계통 매개변수에 대해 개별적으로 $\pm 1\sigma$인 시뮬레이션을 실행하고, 그 결과를 빈(bin)이 있는 저차원 데이터에 대해 저차 다항식을 사용하여 보간하는 "템플릿 변이(template variations)"에 의존합니다. 이 방법은 견고하지만 빈(binning)에 대한 의존성과 계통 효과가 인수 분해된다는 가정(상관관계를 무시함)으로 인해 제한적입니다. 이를 비빈(unbinned), 고차원 적합(fit)으로 확장하기 위해서는 다음과 같은 조건을 갖춘 방법이 필요합니다:

1. 충실도(Fidelitous): 밀도 변형을 정확하게 포착해야 함.
2. 프로파일링 가능성(Profiling-capable): 매개변수에 대한 매끄럽고 미분 가능한 의존성을 제공해야 함.
3. 확장성(Scalable): 전체 $K$차원 계통 공간을 샘플링하는 것을 피해야 함.
4. 해석 가능성(Interpretable): 각 매개변수가 분포를 어떻게 재형성하는지 명시적으로 보여주어야 함.

방법론: 인수 분해 가능한 정규화 흐름 (FNFs)
저자들은 복잡한 명목 밀도(nominal density) 모델링과 계통 변이를 분리하는 생성 프레임워크인 인수 분해 가능한 정규화 흐름(Factorizable Normalizing Flows, FNF)을 소개합니다. 핵심 아키텍처는 두 가지 구성 요소로 이루어집니다:

1. 고정된 명목 흐름(Fixed Nominal Flow): 참조 시뮬레이션($\nu = 0$)에 대해 한 번 학습되어 명목 밀도 $p_{\text{nom}}(y | x)$를 모델링하는 고충실도 가역 흐름(예: Neural Spline Flow)입니다. 이 구성 요소는 계통 모델링 단계 동안 동결됩니다.

2. 학습 가능한 인수 분해 가능한 변환 ($T_\nu$): 관측된 이벤트 $y$(계통 $\nu$의 영향을 받는)를 명목 참조 프레임 $y_{\text{nom}}$으로 매핑하는 변환을 학습합니다. 우도는 다음과 같은 변수 변환 공식을 통해 평가됩니다:
   $$p(y | x, \nu) = p_{\text{nom}}(T_\nu(y | x, \nu) | x) \cdot |\det \nabla_y T_\nu(y | x, \nu)|$$

   변환 $T_\nu$는 스케일($s$)과 시프트($t$) 매개변수가 계통 매개변수 $\nu$의 다항식인 자기회귀 아핀 맵(autoregressive affine map)으로 구축됩니다. 결정적으로, 이 다항식들은 개별 계통들에 대해 인수 분해됩니다:
   $$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$$

   계수($\alpha, \beta, \gamma, \delta, \phi, \psi$)는 입력 특징값에 조건화된 마스크드 다층 퍼셉트론(MLP)의 출력입니다.

   • 가산성(Additivity): 기본 구조는 각 계통으로부터의 가산적 기여를 가정하여, 특정 매개변수에 대한 $\pm 1\sigma$ 샘플만을 사용하여 각 계통을 독립적으로 학습할 수 있게 합니다.
   • 상호작용 항(Interaction Terms): 계통 간의 쌍별 상관관계를 포착하기 위해 선택적인 교차 항($\phi, \psi$)을 포함할 수 있으며, 이는 결합된 변이 샘플을 사용하여 학습해야 합니다.

주요 기여

• 확장 가능한 아키텍처: FNF는 이산적인 템플릿 변이 방식을 연속적이고 미분 가능한 고차원 설정으로 일반화합니다. 이는 결합된 공간을 샘PLING하는 대신 매 계통별 변이에 기반하여 학습함으로써 "차원의 저주"를 피합니다.
• 선형 확장성: 계통 매개변수의 수와 관계없이 계산 비용과 학습 데이터 요구 사항이 선형적으로 확장됩니다(교차 항을 제외할 경우). 이는 수백 개의 불확실성이 있는 분석에서도 실행 가능하게 만듭니다.
• 다루기 쉬운 우도(Tractable Likelihood): 자기회귀 구조는 야코비안 행렬식(Jacobian determinant)이 삼각 행렬임을 보장하여 계산 비용이 저렴하며($O(D)$), 효율적인 우도 적합(fit)을 가능하게 합니다.
• 해석 가능성: 모델은 각 매개변수의 선형, 이차 및 상호작용 효과를 명시적으로 분리하여, 각 계통이 밀도를 어떻게 변형하는지에 대한 명확한 물리적 해석을 제공합니다.

실험 결과
저자들은 두 클래스와 두 개의 계통 매개변수($\nu_{\text{shift}}$ 및 $\nu_{\text{squeeze}}$)가 포함된 통제된 토이 데이터셋을 통해 FNF를 검증했습니다. 이 매개변수들은 운동학적 변수와 조건부 스코어 분포를 모두 변형합니다.

• 재구성(Reconstruction): 학습된 변환은 부과된 변형을 성공적으로 재현했습니다. 순수 이동(translation) 계통의 경우, 모델은 단위 스케일을 가진 선형 시프트를 학습했습니다. 부피 보존 압축(volume-preserving squeeze)의 경우, 모델은 적절한 비선형 스케일링과 시프트를 학습했습니다.
• 폐쇄성(Closure): 인수 분해된 모델(교차 항이 없는 경우)은 훈련 축 근처와 명목 지점 근처에서 최적의 우도에 근접한 폐쇄성(오차 약 $10^{-3}$ nats/event)을 달ato했습니다.
• 상관관계(Correlations): 두 계통이 동시에 강하게 변하는 영역에서는 무시된 상관관계로 인해 인수 분해 모델의 적합 품질이 저하되었습니다. 선택적인 교차 항을 활성화하면 상관된 영역에서 쿨백-라이블러 발산(KL divergence)이 4배 이상 감소하였으며, 이는 표준적인 인수 분해 근사법과 완전히 상관된 모델 사이를 보간하는 아키텍처의 능력을 입증합니다.

의의 및 주장
본 논문은 FNF가 연속적인 밀도 모핑을 요구하는 추론 워크플로우를 위한 일반적인 도구를 제공한다고 주장합니다. FNF는 결합 샘플링의 조합론적 비용 없이 고차원, 비빈 공간에서 계통 불확실성을 모델링할 수 있게 함으로써, 비빈 우도 적합 채택의 주요 장벽을 낮춥니다. 저자들은 본 연구가 매개변수 밀도 모델링을 검증했다면, FNF를 사용하여 관심 매개변수를 추출하기 위한 전체 프로파일 적합(profiled), 비빈 우도 적합에 배치하는 것은 향후 전용 연구의 대상이라고 밝힙니다. 이 접근 방식은 기존의 템플릿 변이 레시피를 직접적이고 연속적인 일반화로서 제시하며, 그 견고함을 유지하면서 고차원 관측량으로의 적용 가능성을 확장합니다.