MANGO: An Autodiff Neutrino Oscillation Engine for Differentiable Analysis Pipelines

기술 요약: MANGO – 미분 가능한 분석 파이프라인을 위한 자동 미분 뉴트리노 진동 엔진

문제 정의
뉴트리노 진동 확률을 계산하는 것은 잘 확립된 문제이며 성숙하고 고도로 최적화된 소프트웨어(예: Prob3++, OscProb, GLoBES, NuFast)가 존재하지만, 이들의 정확한 도함수를 계산하고 이를 엔드 투 엔드(end-to-end) 미분 가능한 분석 파이프라인에 통합하는 것은 해결되지 않은 과제로 남아 있다. 기존 도구들은 순방향 확률 계산에는 뛰어나지만, 자동화된 정확한 도함수가 부족하다. 이러한 구조적 한계는 현대적인 미분 가능 워크플로우(예: 해밀토니안 몬테카를로(HMC), 변분 추론, 신경망 대리 모델, 경사도 기반 실험 설계)에 진동 물리학을 통합하는 것을 방해한다. 또한, 해석적 도함수 공식은 특정 영역(예: 일정 밀도)에 국한되는 경우가 많으며, 복잡한 기하학적 구조(예: 층상 구조의 지구 모델)나 표준적인 매개변수(예: 검출기 깊이 또는 대기 생성 높이와 같이 폐쇄형 확률식에 나타나지 않는 매개변수)를 쉽게 수용할 수 없다.

방법론
본 논문은 jax로 네이티브하게 구현된 구성 가능한 진동 엔진인 mango(MANGO: A Neutrino Gradient Oscillator)를 제시한다. mango는 특정 사례에 대한 해석적 도함수를 유도하는 대신, 일반적인 전파 코어를 구현하고 **자동 미분(AD)**에 의존하여 모든 입력값에 대한 정확한 도함수를 얻는다.

• 전파 코어: 이 엔진은 진공, 일정 밀도, 층상 PREM, 임의의 프로파일, 그리고 단열 태양 전파를 지원한다. 표준 3-플래버 진동뿐만 아니라 비표준 상호작용(NSI), 3+N 스테릴 상태, 환경적 결맞음(decoherence), 비단일성 혼합(non-unitary mixing)을 포함한 BSM(Beyond-Standard-Model) 시나리오를 처리한다.
• 기하학 및 미분 가능성: 프레임워크는 전파 기하학(천정각, 생성 높이, 검출기 깊이)을 미분 가능한 입력값으로 취급한다. 경로 코드(trajectory chord)와 껍질 교차점(shell intersections)을 폐쇄형 미분 함수를 통해 구축함으로써, 엔진은 전통적인 해석 공식에는 존재하지 않는 기하학적 매개변수에 대한 도함수를 계산한다.
• 수치적 백엔드: 스펙트럼 퇴화(spectral degeneracies) 구간에서도 수치적 정밀도와 미분 가능성을 보장하기 위해, mango는 기본 생산 엔진으로 **분할 차분 확장(divided-difference expansion, Cayley–Hamilton)**을 채택한다. 이 방식은 명시적인 고유벡터 계산을 피함으로써, 표준 대각화 루틴에서 질량 준위 교차(mass-level crossing) 근처에서 발생하는 도함수 특이점($1/(\lambda_i - \lambda_j)$) 문제를 우회한다. 대안적인 백엔드로는 JAX로 전사된 NuFast 알고리즘, 조밀한 에르미트 대각화(dense Hermitian diagonalization), 그리고 검증을 위한 Padé 지수(Padé exponentials)가 있다.
• 미분 가능성 모델: 엔진은 순방향(forward-mode) 및 역방향(reverse-mode) AD를 모두 지원한다. 결정적으로, 역방향 모드의 비용은 입력 차원이 아닌 출력 차원에 따라 결정된다. 이를 통해 mango는 수백 개의 매개변수(예: 층상 지구 모델의 369개 매개변수)에 대한 민감도를 단 한 번의 순방향 패스 대비 2.5~3배의 일정한 오버헤드만으로 평가할 수 있다.

주요 기여

1. 구성 가능한 미분 가능 엔진: mango는 계산된 모든 물리량이 모든 입력값(전파 기하학, 검출기 깊이, 개별 껍질 밀도 포함)에 대해 미분 가능한 최초의 진동 엔진을 제공한다.
2. 범위 및 도달 범위: 엔진은 처리하는 모든 구성(층상 PREM, 연속적인 태양 프로파일, NSI, 스테릴 상태)에 대한 도함수를 지원한다. 이는 해석 공식에 존재하지 않는 양, 즉 코드 기하학 매개변수($\cos \theta_z$, $h_{atm}$, $d_{det}$) 및 껍질당 밀도에 대한 미분을 가능하게 한다.
3. 엔드 투 엔드 미분 가능성: AD 그래프는 확률 계산에서 종료되지 않는다. 이는 검출기 응답, 이벤트 가중치, 히스토그램 빈닝(binning), 그리고 가능도(likelihood)까지 연속적으로 흐른다. 이를 통해 단 한 번의 jax.grad 호출로 어떤 상위 물리 분석 지표(figure of merit)에 대해서도 그 경사도(gradient)를 반환할 수 있다.
4. 수치적 견고성: 분할 차분을 사용함으로써, 엔진은 정확한 질량 준위 교차 및 MSW 공명 구간에서도 $C^\infty$ 급으로 매끄러운 AD 도함수를 보장하며, 고유벡터 미분과 관련된 수치적 불안정성을 방지한다.

결과

• 정확도: 3-플래버 및 층상 지구 확률이 외부 벤치마크(OscProb, NuFast-Earth)와 $10^{-9}$에서 $10^{-5}$ 이내로 일치한다. 모든 순방향 모델, BSM 한계, 그리고 미분 경로는 정확한 해석적 해 및 유한 차분법(finite differences)을 통해 검증되었으며, 상대적 편차는 통상적으로 $10^{-10}$에서 $10^{-5}$ 범위 내에 있다.
• 성능: NVIDIA A100 GPU에서, constant-density NuFast 포트는 에너지 포인트당 61ns에 실행되며, 일반적인 분할 차분 커널은 116ns에 실행된다. 약 50개의 세그먼트를 포함하는 층상 PREM 경로는 포인트당 약 158$\mu$s의 비용이 든다.
• 확장성: 6개의 표준 진동 매개변수를 계산하는 데 드는 비용과 비교했을 때, 층상 지구의 369개 밀도, 전자 분율, 껍질 반지름 매개변수에 대한 민감도를 평가하는 비용은 거의 차이가 없다(순방향 패스 대비 약 3.0배의 비용 발생). 반면, 유한 차분법을 사용할 경우 약 250배 더 많은 평가가 필요하다.
• 시연: 본 논문은 정형화된 지구 토모그래피 분석을 보여준다. 단 한 번의 패스로, mango는 6개 구역의 방사형 밀도 모델에 대한 한계 불확실성(marginalized uncertainty)을 계산한다. 역 피셔 행렬(inverse Fisher matrix)을 통해 미분함으로써, 검출기의 각도 분해능에 대한 이 불확실성의 민감도를 평가한다. 이는 전통적인 해석적 확률 공식으로는 도달할 수 없는 실험 설계 지표(예: 각도 분해능 개선이 코어 밀도 정밀도 향상에 얼마나 기여하는지)를 제공한다.

의의
본 논문은 mango가 단순한 순방향 계산기를 넘어 정확한 엔드 투 엔드 미분 가능 엔진으로 뉴트리노 진동 워크플로우를 전환하는 기초적인 변화를 제공한다고 주장한다. 주요 의의는 물리 모델, 기기 설계 매개변수, 그리고 분석 파이프라인을 공동으로 미분하고 최적화할 수 있는 미분 가능한 추론 및 실험 설계를 가능하게 한다는 점에 있다. 전체 파이프라인을 통해 정확한 민감도 신호를 유지함으로써, mango는 다음과 같은 작업을 가능하게 한다:

• 실험 설계: 경사도 기반 최적화를 통해 빈닝 체계를 최적화하고, 하드웨어 트레이드오프(예: 각도 분해능 대 토모그래피 도달 범위)를 평가한다.
• 고차원 피팅: 경사도 기반 최소화 도구 및 HMC를 사용하여 수백 개의 계통 오차(nuisance parameters)에 대한 한계 불확실성을 빠르게 계산한다.
• 지구물리학적 역산: 뉴트리노 데이터로부터 지구 내부의 연속적인 밀도 프로파일을 직접 재구성하기 위한 경사도 기반 역산 기법을 용이하게 한다.

이 연구는 근본적인 양자 역학을 실험적 발견 잠재력과 직접 연결하는 구성 가능한 미분 가능 기초를 구축하며, 현대 고에너지 물리학 분석의 핵심적인 병목 현상을 해결한다.