Machine Learning for neutron source distributions

본 논문은 몬테카를로 입자 목록으로부터 중성자 소스 분포를 추정하기 위해 확률적 생성 모델을 활용한 새로운 접근법을 제안하고 평가하여, 모델 학습 후 효율적이고 메모리 독립적인 샘플링을 가능하게 합니다.

핵심

완벽하고 다층적인 케이크의 완벽한 레시피를 재현하려고 노력한다고 상상해 보세요. 중자 과학의 세계에서 이 "케이크"는 소스에서 방출되는 중성자(미세 입자)의 흐름이며, 각각은 고유한 속도, 방향, 에너지, 타이밍을 가지고 있습니다.

전통적으로 과학자들은 이 흐름을 재현하기 위해 두 가지 방법을 시도해 왔습니다:

1. "복사-붙여넣기" 방법: 그들은 거대한 단일 중성자 목록을 생성하기 위해 방대하고 느린 컴퓨터 시뮬레이션을 실행합니다. 그런 다음 이 목록(MCPL 파일이라고 함)을 저장하고 반복적으로 사용하려고 합니다. 문제는 목록에 있는 중성자보다 더 많은 중성자가 필요할 경우, 동일한 중성자를 반복해서 복사하고 붙여넣는다는 점입니다. 이는 시뮬레이션에서 "글리치"나 "핫스팟"을 생성하며, 마치 같은 빵 부스러기 패턴이 끝없이 반복되어 보이는 것과 같습니다.
2. "경험칙" 방법: 그들은 성분을 개별적으로 살펴봄으로써 레시피를 추측하려 합니다 (예: "얼마나 많은 것이 빠른가?", "얼마나 많은 것이 느린가?"). 문제는 성분들이 어떻게 섞이는지를 무시한다는 점입니다. 실제로는 빠른 중성자가 항상 특정 방향으로 이동할 수 있지만, 이 방법은 이를 서로 무관한 것처럼 취급하여 실제 데이터의 "맛"을 잃게 합니다.

새로운 접근법: "AI 셰프"
이 논문은 머신러닝을 사용하여 이 문제를 해결하는 새로운 방법을 제시합니다. 목록을 복사하거나 규칙을 추측하는 대신, 저자들은 중성자 레시피의 본질을 학습하도록 네 가지 다른 유형의 "AI 셰프"(생성 모델) 를 훈련시켰습니다.

다음은 논문이 이를 어떻게 설명하는지입니다:

1. 훈련 단계 (레시피 학습)

AI 셰프들은 원래의 느린 컴퓨터 시뮬레이션 샘플 ( "훈련 데이터") 을 공급받습니다. 그들은 단순히 목록을 외우는 것이 아니라, 모든 변수 간의 복잡한 관계를 학습합니다.

• 비유: 한 셰프에게 특정 유형의 구름 사진 천 장을 보여주는 상황을 상상해 보세요. 그들은 단순히 사진들을 외우는 것이 아니라, 그 구름이 그런 구름처럼 보이게 만드는 요소를 학습합니다. 가장자리가 어떻게 말리는지, 밀도는 어떠한지, 빛이 어떻게 비추는지 등을요. 일단 이를 학습하면, 그들은 존재해 본 적 없지만 정확히 올바른 모습을 한 새로운 구름을 그릴 수 있습니다.

2. 네 가지 AI 셰프

저자들은 어떤 모델이 레시피를 가장 잘 학습했는지 확인하기 위해 네 가지 다른 유형의 AI 모델을 테스트했습니다:

• 정규화 흐름 (Normalizing Flows, NF): 이는 반죽을 완벽하게 늘리고 구부릴 수 있는 셰프라고 생각하세요. 그들은 단순하고 균일한 반죽 덩어리 (무작위 노이즈) 로 시작하여 중성자 구름의 정확한 복잡한 형태로 늘립니다. 논문은 이것이 가장 뛰어난 셰프라고 발견했으며, 원래 데이터와 완벽하게 일치하는 가장 정확한 "새로운" 중성자를 생성했습니다.
• 변분 오토인코더 (Variational Autoencoders, VAE): 이 셰프는 레시피를 요약으로 압축한 다음 다시 재구성하려고 시도합니다. 이는 빠르고 복잡한 형태에 좋지만, 때때로 재구성된 케이크가 원래보다 약간 "흐릿"하거나 덜 선명하게 나옵니다.
• 생성적 적대 신경망 (Generative Adversarial Networks, GAN): 이는 두 명의 셰프 사이의 "줄다리기"입니다. 한 셰프는 가짜 케이크를 굽고 다른 셰프는 가짜를 찾아내려 합니다. 가짜 케이크가 진짜와 구별할 수 없을 때까지 계속 경쟁합니다. 이 논문은 이를 훈련하기가 다소 어렵고 "속임수"(동일한 몇 가지 패턴을 반복함) 에 취약하다고 발견했습니다.
• 확산 모델 (Diffusion Models, DM): 이 셰프는 노이즈가 많고 지저분한 케이크로 시작하여 완벽해질 때까지 단계별로 천천히 정리합니다. 이는 잘 작동하지만 매우 느리고 계산 비용이 많이 듭니다. 마치 한 알의 먼지씩 주워가며 방을 치우는 것과 같습니다.

3. 결과: 왜 중요한가

논문은 두 가지 실제 시나리오에서 이러한 AI 셰프들을 테스트했습니다:

• 시나리오 A (TDR 데이터셋): 복잡하고 고에너지인 중성자 소스입니다. AI 셰프들은 레시피를 매우 잘 학습하여, "복사-붙여넣기" 글리치 없이 원래 시뮬레이션과 통계적으로 동일한 수백만 개의 새로운 중성자를 생성할 수 있었습니다.
• 시나리오 B (벤치마크 데이터셋): 실제 실험실 측정값과 AI 가 생성한 중성자를 비교한 실제 실험입니다. AI(특히 정규화 흐름) 는 실제 데이터와 거의 완벽하게 일치했습니다.

주요 장점:
AI 셰프가 레시피를 학습하면, 원래 중성자의 거대하고 무거운 목록은 더 이상 필요하지 않습니다. AI 모델은 매우 작습니다 (수 킬로바이트 수준) 그리고 통계적으로 완벽한 무한한 새로운 중성자를 즉시 생성할 수 있습니다. 이는 막대한 양의 컴퓨터 시간과 메모리를 절약합니다.

논문이 말하지 않는 것

저자들은 신중하게 이러한 모델이 데이터 기반이라고 명시합니다. 그들은 제공된 데이터에서 엄격하게 학습합니다.

• 원래 시뮬레이션에 특정 유형의 중성자가 누락되어 있다면, AI 는 그것을 발명하지 않습니다 (데이터 밖을 추측하도록 모델이 특별히 조정되지 않는 한, 이는 논문에서 다른 방법의 특정 기능으로 언급되지만 여기서는 주요 목표가 아닙니다).
• 논문은 이러한 모델이 새로운 물리학을 예측하거나 나쁜 데이터를 수정할 수 있다고 주장하지 않습니다. 그들은 중성자 기기를 설계하는 데 사용하기 위해 기존 데이터 패턴을 효율적으로 재현하기 위한 도구일 뿐입니다.

요약하자면:
이 논문은 중성자 데이터의 무겁고 오류가 발생하기 쉬운 목록을 작고 똑똑한 AI 모델로 대체할 수 있음을 보여줍니다. 이러한 모델은 중성자 흐름의 "DNA"를 학습하여 필요할 때마다 신선하고 현실적인 중성자를 생성할 수 있으며, 이는 미래 중성자 실험의 설계를 더 빠르고, 저렴하며, 정확하게 만듭니다. 테스트된 네 가지 모델 중 **정규화 흐름 (Normalizing Flow)**이 명확한 승자였습니다.

기술 요약

기술 요약: 중성자 원천 분포를 위한 머신러닝

문제 제기
중성자 광학 및 계기 설계에서 특정 표면에서의 중성자 위상 공간 분포 (위치, 방향, 시간, 에너지, 가중치, 편광) 를 정확하게 특성화하는 것은 필수적입니다. 전통적으로 이러한 분포는 몬테카를로 (MC) 시뮬레이션 (예: MCNP, PHITS, OpenMC) 에서 유도되어 몬테카를로 입자 목록 (MCPL) 파일로 저장됩니다. MCPL 파일이 라이스 트레이싱 소프트웨어 (예: Vitess, McStas) 의 표준 입력으로 사용되지만, 이를 직접 재사용하는 데에는 상당한 한계가 존재합니다:

1. 통계적 인공물: 원래 MCPL 파일의 과소 표본 추출 및 통계적 노이즈가 하류 시뮬레이션으로 전파되어 인공 구조 (예: "핫스팟") 를 생성하고 수렴을 저해합니다.
2. 상관관계 손실: 개별 변수에 분석적 또는 준분석적 모델을 적합시키는 대안적 접근법은 종종 위상 공간 변수 간의 복잡하고 고차원적인 상관관계를 보존하지 못합니다. 이러한 상관관계는 특히 고휘도 원천에서 현실적인 수송에 필수적입니다.
3. 계산 비용: 서로 다른 계기 구성에 대한 새로운 원천 목록을 생성하기 위해 고통계량 MC 시뮬레이션을 다시 실행하는 것은 시간이 많이 소요되고 계산 비용이 높습니다.

기존의 머신러닝 솔루션인 커널 밀도 추정 (KDSource) 은 일부 문제를 해결하지만, 종종 중요도 가중치의 수동 조정이 필요하며 특히 입자 가중치가 변할 때 완전한 다변량 위상 공간 추정에 어려움을 겪습니다.

방법론
저자들은 MCPL 파일에서 직접 중성자 위상 공간 변수의 결합 확률 분포를 학습하기 위해 확률적 생성 모델 (PGM) 을 활용하는 프레임워크를 제안합니다. 일단 훈련이 완료되면, 이러한 모델은 원래 입자 목록이 필요 없이 원래 분포의 상관관계를 보존하는 새로운 통계적으로 독립적인 중성자 샘플을 생성할 수 있습니다.

• 데이터셋: 두 가지 서로 다른 데이터셋을 사용하여 모델을 평가했습니다:

  1. TDR 데이터셋: 고휘도 원천 (HBS) 파라수소 감속기에 대한 PHITS 시뮬레이션 (676,558 개 중성자). 이 데이터셋은 입자 가중치를 포함하며 넓은 에너지 스펙트럼 (냉각, 열, 고속) 을 다루므로 복잡하고 고차원 (7 개 변수) 의 모델링 과제를 제시합니다.
  2. 벤치마크 데이터셋: JULIC 중성자 플랫폼에 대한 PHITS 시뮬레이션 (118,219 개 중성자) 으로, 2023 년 실험 측정치와 비교 검증되었습니다. 이 데이터셋은 일정한 입자 가중치를 가지며 (차원성이 6 개 변수로 축소됨), 실험 스펙트럼과 직접 비교 검증할 수 있습니다.
  3. 전처리: 모든 변수는 MinMax 스케일링을 사용하여 정규화되었습니다. 시간 및 에너지 변수는 넓은 동적 범위를 처리하기 위해 로그 변환이 적용되었습니다.

• 생성 모델: 네 가지 서로 다른 PGM 아키텍처가 비교되었습니다:

  1. 정규화 흐름 (NF): 특히 커플링 흐름과 마스킹 자기회귀 흐름 (MAF) 입니다. 이러한 모델은 간단한 기준 분포 (가우시안) 와 복잡한 데이터 분포 사이의 가역적 변환을 학습하여 정확한 가능도 평가 및 샘플링을 가능하게 합니다.
  2. 변분 오토인코더 (VAE): 재구성 손실과 Kullback–Leibler (KL) 발산 정규화 항을 통해 최적화되는 인코더 - 디코더 아키텍처로 데이터의 잠재 표현을 학습합니다.
  3. 생성적 적대 신경망 (GAN): 실제 데이터와 구별할 수 없는 샘플을 생성하도록 판별자 네트워크와 대조하여 훈련된 생성자 네트워크입니다.
  4. 확산 모델 (DM): 무작위 노이즈에서 데이터를 재구성하기 위해 순방향 노이즈화 과정을 역전시키는 것을 학습하는 모델입니다.

• 구현: 모델은 JUWELS 부스터 클러스터에서 PyTorch 를 사용하여 훈련되었습니다. Vitess MC 소프트웨어를 위해 사용자 정의 C++ 모듈인 source AI 가 개발되어 훈련된 PyTorch 모델을 (TorchScript 를 통해) 로드하고 온더플라이 샘플링을 수행하도록 하여 MCPL 파일 입력을 효과적으로 대체했습니다.

• 평가 지표:

  • 최대 평균 불일치 (MMD): 재생성 커널 힐베르트 공간 (RKHS) 에서 학습된 분포와 원래 MCPL 분포 간의 거리를 정량화하는 데 사용되었습니다.
  • 샘플링 속도: 20,000 개의 중성자를 생성하는 데 소요된 시간.
  • 물리적 검증: JULIC 플랫폼의 실험 데이터와 비교한 시뮬레이션 파장 스펙트럼.

주요 결과

• 분포 충실도: 모든 모델이 기본 분포를 성공적으로 학습했지만, 정규화 흐름 (NF) 이 우수한 성능을 보였습니다. NF 는 두 데이터셋 모두에서 가장 낮은 MMD 점수를 달성했습니다 (예: TDR 의 경우 $1.19 \times 10^{-4}$, 벤치마크의 경우 $1.20 \times 10^{-4}$). 이는 VAE, GAN, 확산 모델 및 KDSource 기준선보다 우월한 성능입니다.
• 상관관계 보존: 2 차원 히스토그램의 시각적 분석은 NF 가 시간 - 에너지, 각도 - 위치와 같은 변수 간의 복잡한 상관관계를 성공적으로 보존했음을 확인시켜 주었으며, 이는 단순한 적합 방법에서는 손실되었습니다.
• 샘플링 효율성: KDSource 가 가장 빠른 샘플링 시간을 제공했지만, NF 는 속도와 정확성 사이의 유리한 균형을 제공했습니다. VAE 와 GAN 도 빠르지만, 확산 모델은 반복적 탈노이즈 특성으로 인해 샘플링이 현저히 느렸습니다.
• 실험적 검증: 벤치마크 데이터셋에서 NF 가 생성한 원천은 실험 파장 스펙트럼을 높은 충실도로 재현하여 KDSource 와 동등한 성능을 보였습니다. 그러나 KDSource 와 달리 NF 모델은 시그모이드 활성화 계층으로 인해 훈련 데이터의 경계를 엄격히 준수하여, 명시적으로 설계되지 않는 한 가이드 외부의 중성자나 관찰되지 않은 파장과 같은 비물리적 외삽을 생성하지 않았습니다.
• 강건성: 동일한 모델 아키텍처가 두 가지 서로 다른 데이터셋에 적용될 때 최소한의 재조정으로 작동하여 다양한 차원성과 분포 복잡성 전반에 걸쳐 강건성을 입증했습니다.

의의 및 주장
본 논문은 현대적 생성 머신러닝 모델을 중성자 몬테카를로 파이프라인에 통합하고 실험 데이터로 검증한 최초의 비교 연구를 제시한다고 주장합니다. 주요 의의는 다음과 같습니다:

1. PGM 은 viable 한 대리 모델입니다: 확률적 생성 모델은 원시 MCPL 파일을 대체할 수 있으며, 메모리 효율적 (기가바이트 파일 대비 킬로바이트 크기 모델) 이고 통계적으로 견고한 원천 표현 방법을 제공합니다.
2. 정규화 흐름이 최적입니다: 테스트된 아키텍처 중 NF 는 이 특정 응용 분야에 대해 가장 좋은 균형을 제공하며, 경계 위상 공간과 같은 하드 물리적 제약을 아키텍처에 직접 통합할 수 있는 능력을 갖춘 고흡실도 분포 학습을 제공합니다.
3. 워크플로우 통합: Vitess 에서 개발된 source AI 모듈은 이러한 모델이 기존 중성자 시뮬레이션 워크플로우에 원활하게 통합되어 원래 훈련 데이터가 필요 없이 신속하고 독립적인 샘플링을 가능하게 함을 입증합니다.

저자들은 생성 모델이 초기 고흡실도 MC 시뮬레이션의 필요성을 대체하지는 않지만, 직접적인 MCPL 재사용이 비현실적인 반복 시뮬레이션이나 고통계량 샘플링이 필요한 시나리오에서 하류 계기 설계 및 최적화를 위한 효율적이고 유연하며 물리적으로 의미 있는 도구를 제공한다고 결론지었습니다.