[논문 리뷰] Bridging the gap between machine learning and particle accelerator physics with high-speed, differentiable simulations

Jan Kaiser, Chenran Xu외 2명|arXiv (Cornell University)|2024. 01. 01.

Particle accelerators and beam dynamics참고 문헌 24인용 수 3

한 줄 요약

Cheetah는 PyTorch 기반의 고속, 미분 가능 선형 빔 동역학 시뮬레이터로, 여러 주요 차수의 속도 향상을 이루면서도 자동 미분을 통한 효율적인 기울기 계산을 가능하게 한다. 한 벤치마크에서 강화학습 학습 시간을 12일 이상에서 1시간 이내로 단축시켜 입자 가속기 물리학 분야에서의 확장 가능한 기계학습 응용을 가능하게 한다.

ABSTRACT

Machine learning has emerged as a powerful solution to the modern challenges in accelerator physics. However, the limited availability of beam time, the computational cost of simulations, and the high-dimensionality of optimisation problems pose significant challenges in generating the required data for training state-of-the-art machine learning models. In this work, we introduce Cheetah, a PyTorch-based high-speed differentiable linear-beam dynamics code. Cheetah enables the fast collection of large data sets by reducing computation times by multiple orders of magnitude and facilitates efficient gradient-based optimisation for accelerator tuning and system identification. This positions Cheetah as a user-friendly, readily extensible tool that integrates seamlessly with widely adopted machine learning tools. We showcase the utility of Cheetah through five examples, including reinforcement learning training, gradient-based beamline tuning, gradient-based system identification, physics-informed Bayesian optimisation priors, and modular neural network surrogate modelling of space charge effects. The use of such a high-speed differentiable simulation code will simplify the development of machine learning-based methods for particle accelerators and fast-track their integration into everyday operations of accelerator facilities.

연구 동기 및 목표

• 제한된 빔 시간과 비싼 시뮬레이션으로 인해 발생하는 기계학습의 데이터 부족 문제를 해결한다.
• 기존 가속기 시뮬레이션 코드로는 대규모 데이터 생성이 계산적으로 불가능한 문제를 해결한다.
• 고차원 매개변수 공간에서 기울기 기반 최적화를 가능하게 하여 빔 튜닝 및 시스템 식별을 지원한다.
• 물리 기반 시뮬레이션과 신경망 대체 모델을 원활하게 통합하여 재사용 가능하고, 미분 가능한 가속기 모델링을 가능하게 한다.
• 기계학습 도구와 가속기 물리 워크플로우를 연결하는 사용자 우아하고 확장 가능한 프레임워크를 제공한다.

제안 방법

• PyTorch를 활용해 고속 선형 빔 동역학 시뮬레이션을 위한 미분 가능한 텐서 기반 빔 동역학 엔진을 구현한다.
• PyTorch의 자동 미분 기능을 활용해 모든 시스템 매개변수에 대한 빔 관측값의 기울기를 효율적으로 계산한다.
• 벡터화된 연산과 GPU 가속을 최적화하여 시뮬레이션 시간을 여러 주요 차수만큼 감소시킨다.
• 스페이스 차지와 같은 복잡한 효과의 대체 모델링을 위해 신경망과의 모듈러 통합을 지원한다.
• 신경망 출력에서부터 미분 가능한 물리 시뮬레이터를 거쳐 기울기 흐름을 종단에서 종단까지 보장한다.
• 강화학습 환경(Gymnasium 등)과 같은 표준 기계학습 도구와의 호환성을 설계한다.

실험 결과

연구 질문

• RQ1 미분 가능하고 고속인 빔 동역학 시뮬레이터는 가속기 튜닝을 위한 강화학습 학습 시간을 수일에서 수시간으로 줄일 수 있는가?
• RQ2 미분 가능한 시뮬레이션은 고차원 가속기 매개변수 공간에서 기울기 기반 최적화를 어느 정도 가능하게 하는가?
• RQ3 미분 가능한 물리 모델은 가속기 시스템 식별을 위한 베이지안 최적화에서 사전 정보로 얼마나 효과적으로 기능하는가?
• RQ4 모듈러한 신경망 대체 모델은 미분 가능한 물리 시뮬레이터와 함께 훈련되고 통합되어 스페이스 차지와 같은 복잡한 효과를 모델링할 수 있는가?
• RQ5 기계학습 프레임워크와의 통합을 통해 미분 가능한 시뮬레이터가 대체 모델의 재사용성과 확장성에 어떻게 기여하는가?

주요 결과

• Cheetah는 Ocelot를 사용할 경우 12일 이상이 소요되던 횡방향 빔 튜닝 작업의 시뮬레이션 시간을 1시간 이내로 단축시켜 대규모 강화학습 학습의 실용성을 확보했다.
• 이 프레임워크는 모든 빔라인 매개변수를 통해 효율적인 기울기 계산을 지원하여, 전통적인 수치 방법의 한계를 초월한 고차원 최적화를 가능하게 했다.
• Cheetah는 다양한 사전 정보를 제공함으로써 물리 기반의 베이지안 최적화를 가능하게 하여 시스템 식별 작업에서 수렴 속도를 향상시켰다.
• 스페이스 차지 효과를 위한 모듈러한 신경망 대체 모델은 Cheetah와 함께 훈련되고 원활하게 통합되어 빠르고, 미분 가능하며 재사용 가능한 모델링을 가능하게 했다.
• Gymnasium 및 PyTorch와 같은 기계학습 생태계와의 통합을 통해 Cheetah는 기울기 가능 물리 모델과 신경망을 결합한 종단 간 학습 파이프라인을 가능하게 했다.
• 벤치마킹 결과, Cheetah는 기존의 Ocelot와 같은 레거시 시뮬레이션 코드보다 여러 주요 차수의 속도 향상을 달성하여 대규모 데이터 생성을 실용적으로 가능하게 했다.