[논문 리뷰] pmwd: A Differentiable Cosmological Particle-Mesh $N$-body Library

Yin Li, Libin Lu외 8명|arXiv (Cornell University)|2022. 11. 18.

Galaxies: Formation, Evolution, Phenomena인용 수 14

한 줄 요약

pmwd는 적분 역전파 방법과 JAX의 자동 미분 기능을 활용하여 메모리 효율적이고 GPU 가속화된 대규모 구조 형성의 미분 가능 입자-메쉬(N-body) 시뮬레이션 라이브러리이다. 시간을 거꾸로 적분하여 시뮬레이션 이력을 복원함으로써 표준의 시간 역행 역전파 방법에 비해 메모리 비용을 크게 줄이며 완전한 미분 가능성을 달성한다.

ABSTRACT

The formation of the large-scale structure, the evolution and distribution of galaxies, quasars, and dark matter on cosmological scales, requires numerical simulations. Differentiable simulations provide gradients of the cosmological parameters, that can accelerate the extraction of physical information from statistical analyses of observational data. The deep learning revolution has brought not only myriad powerful neural networks, but also breakthroughs including automatic differentiation (AD) tools and computational accelerators like GPUs, facilitating forward modeling of the Universe with differentiable simulations. Because AD needs to save the whole forward evolution history to backpropagate gradients, current differentiable cosmological simulations are limited by memory. Using the adjoint method, with reverse time integration to reconstruct the evolution history, we develop a differentiable cosmological particle-mesh (PM) simulation library pmwd (particle-mesh with derivatives) with a low memory cost. Based on the powerful AD library JAX, pmwd is fully differentiable, and is highly performant on GPUs.

연구 동기 및 목표

• 천체 물리적 파rameter 제약 조건을 위한 기반 기반 추론을 지원하는 미분 가능 천체역학 시뮬레이션 라이브러리 개발
• 표준의 시간 역행 역전파 방법에서 발생하는 메모리 병목 현상을 해결하기 위해 적분 역전파 방법을 활용한 미분 가능 시뮬레이션 개선
• 장영역, GPU 가속화된 시뮬레이션을 지원하며, 필드 수준의 전방 모델링에 대한 완전한 미분 가능성을 확보
• 딥 러닝 및 반분석 모델과의 통합을 통해 천체 물리적 파rameter와 초기 조건을 함께 추론할 수 있도록 지원
• 차세대 천체역학 데이터 분석을 위한 확장 가능하고 성능이 우수하며 오픈소스인 도구 제공

제안 방법

• 시간 역행 적분을 통해 역전파를 계산함으로써 전방 진화 이력을 전부 저장하지 않도록 함
• JAX를 사용하여 자동 미분 및 GPU 가속화를 구현함으로써 고성능의 미분 가능 시뮬레이션 지원
• 전방 적분과 동일한 이산 경로를 따르도록 보장하기 위해 '이산화-그리고-최적화' 접근법 적용
• 역전이 가능한 시간 적분을 갖는 입자-메쉬(PM) 중력 해법을 구현하여 일관된 적분 역전파 계산 지원
• 디자인에 C2 대칭성을 적용하여 시간 역행성 반영 및 메모리 효율적인 기울기 계산 지원
• 에뮬레이터 및 시공간 최적화기와 같은 다양한 미분 가능 구성요소와의 통합 지원

실험 결과

연구 질문

• RQ1 적분 역전파 방법이 천체역학적 N-body 시뮬레이션에 효과적으로 적용되어 메모리 비용을 줄이면서도 미분 가능성을 유지할 수 있는가?
• RQ2 전방 상태 전체를 저장하는 기존의 미분 가능 시뮬레이션 코드에 비해 pmwd의 메모리 효율성은 어떻게 비교되는가?
• RQ3 pmwd는 미분 가능 천체역학 시뮬레이션에서 더 높은 시간 해상도와 질량 해상도로 어떻게 확장 가능한가?
• RQ4 pmwd는 초기 조건과 물리적 파rameter를 포함한 천체 관측치의 종단 간 미분 가능 모델링을 가능하게 하는가?
• RQ5 pmwd는 통합 추론 파이프라인 내에서 딥 러닝 및 기타 다양한 미분 가능 구성요소와 얼마나 잘 통합되는가?

주요 결과

• pmwd는 전방 상태 전체를 저장하지 않고도 적분 역전파 방법을 통해 메모리 사용을 크게 줄이며 완전한 미분 가능성을 달성한다.
• 라이브러리는 JAX와 완전히 호환되며 GPU에서 효율적으로 실행되어 고성능 시뮬레이션을 가능하게 한다.
• 적분 역전파 방법 덕분에 pmwd는 이전의 AD 기반 시뮬레이션보다 더 많은 시간 단계와 더 높은 해상도로 확장 가능하다.
• pmwd는 필드 수준의 전방 모델링을 지원하여 천체 물리적 파arameter와 초기 조건을 함께 추론할 수 있다.
• 라이브러리는 미분 가능 딥 러닝 및 반분석 모델과의 상호 운용을 고려하여 설계되어 복잡한 추론 파이프라인을 가능하게 한다.
• pmwd는 오픈소스이며 공개적으로 아카이브되어 있어 재현성과 커뮤니티 기반 확장성을 지원한다.