Université de Paris(파리대학교), 그래프코어 IPU로 우주학 분야에서의 활용 가속화

Université de Paris의 한 연구자가 그래프코어 IPU를 사용해 우주학 분야를 위한 신경망 훈련을 가속화했습니다. 

최근 발행한 논문에서 Bastien Arcelin은 결정론적 심층 신경망과 베이지안 신경망(Bayesian neural network, BNN)을 사용하여 우주학의 두 가지 딥 러닝 사용 사례인 훈련된 VAE 잠재 공간에서의 은하 이미지 생성과

은하 모양 추정에 대한 IPU 프로세서의 적합성 및 성능을 탐구했습니다.

빅데이터

앞으로 천문학 조사에서는 유례없는 양의 관찰 데이터가 도출될 것입니다. Vera C. Rubin 천문대의 우주 및 시간에 관한 레거시 조사(Legacy Survey of Space and Time, LSST)가 바로 그러한 예입니다. 이 조사를 통해 매일 밤 20테라바이트 상당의 데이터가 생성되며, 전체 조사 기간인 10년 동안에는 총 60페타바이트 가량의 데이터가 도출될 것으로 예상됩니다. 

우주학 연구자들이 이처럼 대규모의 복잡한 데이터 세트를 신경망을 활용해 관리하는 경우가 점차 늘어나고 있습니다.

그러나 실제 관찰에서는 알 수 없거나 통제할 수 없는 변수의 양이 너무나도 많아 신경망에 이러한 유형의 데이터를 학습시키기가 매우 어려울 수 있습니다.

따라서 여타의 여러 AI 애플리케이션과 마찬가지로 우주학에서도 정확한 매개변수를 알고 통제할 수 있는 시뮬레이션 데이터가 훈련에 가장 적합합니다.

이러한 시뮬레이션 데이터의 양과 품질은 LSST와 같은 광도 은하 조사에서 기록된 실제 데이터와 최대한 가까운 것이 좋습니다. 

이 시뮬레이션 데이터도 양이 많고 복잡하므로, 이에 연결된 신경망은 빠르고 정확해야 하며 일부 애플리케이션에서는 인식론적 불확실성의 특성을 정확하게 파악할 수 있어야 합니다.

이러한 연산 수요는 'AI용으로 설계된 하드웨어가 우주학 분야의 딥 러닝에서 우월한 성능을 제공할 수 있을까?'라는 질문으로 자연스레 이어집니다.

그래프코어의 IPU는 모든 AI 알고리즘의 기본 구조인 그래프를 효율적으로 처리할 수 있도록 설계된 프로세서이므로 이러한 이론을 테스트하기에 이상적인 예시입니다. 

이 연구에서는 그래프코어의 1세대 칩인 GC2 IPU 하나의 성능을 엔비디아의 V100 GPU와 비교했습니다.

소규모 배치로 신경망을 훈련하는 경우, IPU는 GPU에 비해 DNN 훈련 성능은 2배 이상, BNN 훈련 성능은 4배 이상 빠른 것으로 나타났습니다. GC2 IPU는 GPU 대비 불과 절반의 전력을 소비하여 이러한 성능을 달성했습니다.

XLA 백엔드 통합으로 IPU에서 TensorFlow 1과 2가 모두 지원되므로, 모든 실험에는 TensorFlow 2.1 프레임워크가 사용되었습니다.

은하 이미지 생성

은하와 같은 우주론적 출처의 시뮬레이션에 기반한 데이터는 

전통적으로 Sérsic 프로필과 같은 단순한 분석 프로필에 근거해 왔습니다. Sérsic 프로필은 저속 생성 기법으로서 모델 편향을 초래할 위험을 높입니다.

딥 러닝 방법은 데이터를 훨씬 더 빨리 시뮬레이션할 수 있는 잠재력을 지닙니다. 이에 따라 우주학의 여러 애플리케이션을 위한 은하 모델링에 발생적 신경망을 활용하는 경우가 점점 증가하고 있습니다(예: Lanusse 등., 2020, Regier, McAuliffe 및 Prabhat, 2015 또는 Arcelin 등., 2021).

이미지는 격리된 은하 이미지로 훈련된 VAE(Variational AutoEncoder)의 잠재 공간 분포를 표본 추출하여 생성됩니다.

이 사례에서는 소규모 이미지 배치 생성 시 IPU가 GPU보다 성능이 더 우수하다는 사실을 확인할 수 있습니다.

이 워크로드를 그래프코어의 2세대 IPU인 GC200에서 실행하는 경우, 대폭 증가한 프로세서 칩 메모리 덕분에 규모가 큰 배치를 생성할 때에도 성능을 크게 개선할 수 있을 것으로 예상됩니다.

은하 모양 매개변수 추정

차세대 천문학 조사에서는 그 어느 때보다도 하늘을 더 자세히 관찰할 것이며, 이로 인해 개체 혼합(중복)의 확률도 높아질 것으로 예상됩니다. 물론 일부 가려진 은하의 모양을 측정하기란 훨씬 더 어렵습니다.

기존의 은하 모양 측정 방법으로는 이러한 중복된 개체를 정확하게 측정할 수 없으므로, 하늘을 더 자세히 관찰하기 위해서는 새로운 기법이 필요합니다. 

파리대학교의 연구자 Bastien Arcelin은 심층 신경망과 합성곱 레이어를 활용하여 고립된 은하와 혼합된 은하 모두의 모양 타원도 매개변수를 측정하는 새로운 기법을 개발하고 있습니다.

첫 번째 실험은 결정론적 신경망을 사용하여 수행했으며, 두 번째 테스트에서는 BNN을 사용했습니다. 두 신경망의 훈련 종료 후 결정론적 네트워크의 훈련 가능 매개변수는 특정한 값으로 고정되며, 네트워크에 같은 고립된 은하 이미지를 두 번 제시해도 변화하지 않습니다. 이는 가중치 자체에 단일 값이 아닌 확률분포가 할당되는 BNN과 대조됩니다. 즉, 확률분포에서 같은 이미지를 두 번 표본 추출하여 BNN 네트워크에 제시하면 2회의 결과가 약간 다르게 나타납니다. 이러한 분포의 표본을 여러 번 추출하면 결과의 인식론적 불확실성을 측정할 수 있습니다.

결정론적 신경망의 경우, 그래프코어 IPU는 특히 소규모 배치를 사용하는 경우 GPU에 비해 훈련 시간이 더 짧았습니다. IPU의 기술은 이미 컴퓨터 비전 애플리케이션에서 합성곱 신경망에 대한 효율적인 성능을 통해 널리 알려져 있으므로 이는 놀라운 결과가 아닙니다.

다음으로 이 기법을 베이지안 신경망에 테스트했습니다. 우주학 연구자들에게는 신경망 예측을 통해 얻을 수 있는 신뢰도 수준을 확인하는 것이 매우 중요합니다. 인식론적 불확실성은 대규모 측정 오류와 밀접한 연관 관계에 있으므로, 연구자들은 BNN을 활용해 예측의 인식론적 불확실성을 계산하고 이를 통해 이러한 예측의 신뢰도 수준을 파악할 수 있습니다.

BNN의 경우에도 IPU가 GPU보다 성능이 우수하며, 훈련 시간도 최소한 4배 더 빠릅니다.

IPU는 은하 모양 매개변수 추정을 위한 인공 신경망 훈련 시간을 상당 부분 단축할 수 있으며, 이 연구처럼 IPU 1개를 사용하는 경우 배치 규모가 작을 때 최적의 성능을 제공합니다. 배치 규모가 큰 경우 네트워크를 여러 개의 IPU로 나누어 성능과 효율성을 개선할 수 있습니다.

논문 읽기