[논문 리뷰] Branches of a Tree: Taking Derivatives of Programs with Discrete and Branching Randomness in High Energy Physics

M. Kagan, L. Heinrich|arXiv (Cornell University)|2023. 08. 31.

Distributed and Parallel Computing Systems인용 수 5

한 줄 요약

이 논문은 고에너지물리학(HEP)에서 이산적이고 분기적인 난수를 포함하는 프로그램을 미분 가능하게 하는 기법—특히 기저선을 사용한 스코어 함수 추정기와 확률적 자동미분(Stochastic AD)—을 소개한다. 본 연구는 이산적이고 분기적인 난수를 포함하는 첫 번째 완전히 미분 가능한 분기 프로그램을 제안하며, 작은 배치 크기에서도 확률적 AD와 기저선을 사용한 스코어 함수 추정기가 효과적인 기울기 기반 검출기 설계 최적화를 가능하게 한다.

ABSTRACT

We propose to apply several gradient estimation techniques to enable the differentiation of programs with discrete randomness in High Energy Physics. Such programs are common in High Energy Physics due to the presence of branching processes and clustering-based analysis. Thus differentiating such programs can open the way for gradient based optimization in the context of detector design optimization, simulator tuning, or data analysis and reconstruction optimization. We discuss several possible gradient estimation strategies, including the recent Stochastic AD method, and compare them in simplified detector design experiments. In doing so we develop, to the best of our knowledge, the first fully differentiable branching program.

연구 동기 및 목표

• 이산적이고 분기적인 난수를 포함하는 프로그램—예를 들어 입자 쇼워와 검출기 시뮬레이션—에 대해 고에너지물리학(HEP)에서 기울기 기반 최적화를 가능하게 하기.
• HEP 소프트웨어에서 매개변수에 의존하는 이산적 난수성에 대해 표준 자동미분(AD)이 처리할 수 없는 한계를 해결하기.
• 입자 쇼워 시뮬레이션을 위한 완전히 미분 가능한 분기 프로그램을 개발하고 검증하며, 이는 HEP의 미분 가능 프로그래밍 분야에서의 새로운 기여이다.
• 간단한 검출기 설계 최적화 작업에서 스코어 함수, 수치 기울기, 확률적 AD 등의 기울기 추정 기법을 비교하기.
• 검출기 설계, 시뮬레이터 캘리브레이션, 재구성 최적화 등의 애플리케이션을 위해 종단 간 미분 가능한 HEP 시뮬레이터의 실현 가능성을 입증하기.

제안 방법

• 이산적 난수성을 포함한 확률적 프로그램에서 기울기를 추정하기 위해 기저선(컨트롤 변수)을 사용한 스코어 함수 기울기 추정기를 적용한다.
• 최근 개발된 이산 난수 변수를 포함한 프로그램을 미분할 수 있는 기법인 확률적 자동미분(Stochastic AD; Arya 등, 2022)을 구현하고 HEP 워크로드에 적용한다.
• 검출기 반지름을 조정 가능한 매개변수로 설정하고 손실 함수를 반사적 쇼워 깊이로 정의한 단순화된 검출기 시뮬레이터를 설계한다.
• 작은 배치 크기(배치 크기 = 2)로 기울기 추정을 수행하기 위해 Adam 최적화기를 사용하여 검출기 반지름 매개변수에 대해 500단계의 최적화를 수행한다.
• 각 기법에 대해 10회 반복 최적화를 수행하고, 최적화 단계 동안 평균 손실와 표준편차를 평가하여 성능을 평가한다.
• 확률적 AD 추정기에서 원본 프로그램과 대체 프로그램에서 난수 변수를 재사용하여 분산을 감소시키며, 재사용이 적용되었을 경우 표준편차가 1.5배 이상 감소하는 것을 관찰한다.

(a) Single particle energy loss

실험 결과

연구 질문

• RQ1 기저선을 사용한 스코어 함수 기울기 추정기는 이산적 난수성을 포함한 HEP 검출기 설계 최적화에 효과적으로 적용될 수 있는가?
• RQ2 확률적 AD는 HEP에서의 분기 과정—예를 들어 파르톤 쇼워—에 성공적으로 적용되어 미분 가능 프로그래밍을 가능하게 할 수 있는가?
• RQ3 낮은 샘플 수의 HEP 최적화 환경에서 스코어 함수, 수치 기울기, 확률적 AD 등의 기울기 추정기 간 분산과 수렴 속도는 어떻게 비교되는가?
• RQ4 확률적 AD에서 난수 변수 재사용이 HEP 응용 분야의 기울기 추정기 분산에 어떤 영향을 미치는가?
• RQ5 기울기 추정기의 높은 분산을 감안할 때, 기울기 단계당 두 번의 시뮬레이션만으로도 HEP에서 기울기 기반 최적화를 수행하는 것이 가능한가?

주요 결과

• 확률적 AD와 기저선을 사용한 스코어 함수 추정기는 테스트된 모든 기울기 추정기 중에서 가장 낮은 표준편차를 보이며, 단 두 번의 시뮬레이션만으로도 안정적인 최적화를 가능하게 한다.
• 기저선이 없는 스코어 함수 추정기와 수치 기울기 추정기는 너무 높은 분산을 보이며, 주어진 낮은 샘플 환경에서는 최적화가 불가능하다.
• 확률적 AD와 기저선을 사용한 스코어 함수 추정기는 손실 최소값에 도달하는 데 유사한 수렴 진전을 보이며, 향후 HEP의 미분 가능 프로그래밍 응용 분야에서 모두 유망한 후보로 간주된다.
• 확률적 AD의 대체 프로그램에서 난수 변수를 재사용하면 기울기 추정기의 분산이 1.5배 이상 감소하며, 이는 결합 전략의 중요성을 강조한다.
• 본 연구는 HEP에서 이산적 분기 과정을 통해 기울기를 추정할 수 있는 첫 번째 알려진 미분 가능한 분기 프로그램을 제시하며, 입자 쇼워의 종단 간 미분 가능한 시뮬레이션을 가능하게 한다.
• 확률적 AD와 기저선을 사용한 스코어 함수 추정기는 최소한의 데이터로도 기울기 기반 검출기 설계 최적화에 효과적임을 입증하며, 더 넓은 HEP 응용 분야에서 큰 잠재력을 지닌다.

(b) Particle shower with splitting