Fast(er)PM and Moving Mesh: JAX-native Geometric Multigrid Methods

우주를 수십억 광년에 걸쳐 펼쳐진 거대하고 보이지 않는 중력의 그물망이라고 상상해 보세요. 이 그물이 어떻게 형성되는지 이해하기 위해 과학자들은 거대한 컴퓨터 시뮬레이션을 실행합니다. 이 시뮬레이션은 마치 "점 잇기" 게임과 같지만, 점 대신 암흑 물질을 나타내는 수십억 개의 입자가 있고, 그 "점"들은 중력이라는 힘에 의해 연결됩니다.

이 시뮬레이션에서 가장 큰 골칫거리는 **포아송 방정식(Poisson equation)**을 푸는 것입니다. 이것은 모든 입자가 다른 모든 입자를 어떻게 끌어당기는지를 알려주는 규칙서라고 생각하면 됩니다. 이는 물질이 어디에 있는지(밀도)를 중력이 얼마나 강하게 작용하는지(중력장)의 지도로 변환하는 수학입니다.

수십 년 동안 이 방정식을 푸는 표준적인 방법은 FFT(고속 푸리에 변환)라고 불리는 도구를 사용하는 것이었습니다. FFT를 초고속 글로벌 번역기라고 상상해 보세요. FFT는 전체 지도를 가져와서 비밀 코드로 번역한 뒤, 빠른 수학적 기교를 부리고 다시 원래대로 번역합니다. 이는 변화가 없는 평평한 격자 위에서는 믿을 수 없이 빠릅니다. 하지만 함정이 있습니다. FFT가 마법을 부리려면 모든 컴퓨터가 동시에 자신의 데이터를 모든 다른 컴퓨터에게 외쳐서 전달해야 합니다. 이는 마치 모든 학생이 동시에 서로에게 쪽지를 전달해야 하는 교실과 같습니다. 소규모 규모에서는 괜찮습니다. 하지만 시뮬레이션이 커지고 컴퓨터 수가 늘어남에 따라, 이 "전체 대 전체(all-to-all)" 방식의 외침은 속도를 늦추고 엄청난 양의 메모리를 잡아먹습니다.

논문의 주요 발견: 더 똑똑한 국소적 접근법

벤자민 호로비츠(Benjamin Horowitz)의 논문은 **기하학적 멀티그리드(Geometric Multigrid)**라는 방법을 사용하여 게임을 플레이하는 다른 방식을 제안합니다. 글로벌 번역기 대신, 신문사의 로컬 편집자 팀을 상상해 보세요.

1. "웜 스타트(Warm-Start)" 기법: 시간 단계별로 진행되는 시뮬레이션에서 우주는 매 초마다 급격하게 변하지 않습니다. 어제의 중력 지도는 오늘과 거의 비슷합니다. 논문은 만약 어제의 지도를 오늘의 계산을 위한 "웜 스타트"(헤드 스타트)로 사용한다면, 처음부터 무거운 작업을 다시 할 필요가 없음을 보여줍니다.
2. 멀티그리드 편집자: 멀티그리드 방식은 서로 다른 돋보기를 가진 편집자 팀처럼 작동합니다. 한 편집자는 전체 페이지를 보며 큰 뭉툭한 얼룩(장거리 중력)을 수정합니다. 다른 편집자는 확대하여 미세하고 거친 가장자리(단거리 중력)를 수정합니다. 이들은 수정 사항을 서로 주고받습니다.
3. 결과: 고정된 격자에서 이 국소적 접근 방식은 기존의 글로벌 FFT 번역기보다 더 빠르고 메모리를 적게 사용합니다. NVIDIA A100 GPU를 사용한 테스트에서, 새로운 방식은 큰 격자(예: 1024³ 메쉬)에 대해 기존 FFT 방식보다 최대 2.4배 더 빨랐습니다. 또한, 방대한 임시 데이터를 저장할 필요가 없었기 때문에 시뮬레이션을 실행하는 데 필요한 컴퓨터(노드)의 수를 절반으로 줄일 수 있었습니다.

거대한 도약: 움직이는 메쉬(Moving Mesh)

진정한 마법은 논문이 **움직이는 메쉬(Moving Mesh)**를 도입할 때 일어납니다.

표준적인 시뮬레이션 격자를 딱딱한 정사각형 낚시 그물이라고 상상해 보세요. 만약 당신이 광활한 바다(빈 공간) 속에 있는 작은 물고기 떼(밀집된 물질 클러스터)를 잡으려고 한다면, 그물의 대부분은 빈 물을 잡는 데 낭비됩니다. 물고기를 잡기 위해 거대한 그물이 필요하거나, 아니면 세부 사항을 놓치게 됩니다.

논문은 **준-라그랑지안 움직이는 메쉬(quasi-Lagrangian moving mesh)**를 제안합니다. 이것은 늘어나고 줄어들 수 있는 그물입니다.

• 물고기가 밀집된 곳(은하와 클러스터)에서는 그물이 수축하여, 셀(cell)들을 서로 밀착시켜 매우 날카롭고 높은 해상도의 뷰를 제공합니다.
• 물이 비어 있는 곳(보이드)에서는 그물이 확장되어, 빈 공간에 노력을 낭비하지 않습니다.

이는 동작에 맞춰 자동으로 줌인(zoom-in)하고 배경에 대해 줌아웃(zoom-out)하는 카메라와 같습니다. 논문은 이 방식이 동일한 크기의 정적 메쉬보다 밀집된 영역에서 훨씬 더 많은 세부 정보를 복구해 낸다는 것을 보여줍니다. 예를 들어, 256³ 입자 시뮬레이션에서 움직이는 메쉬는 동일한 크기의 정적 메쉬가 놓친 미세한 구조들을 포착했습니다.

"미분 가능한" 초능력

가장 흥rix한 부분은 전체 시스템이 JAX라는 프로그래밍 프레임워크를 사용하여 구축되었다는 점입니다. JAX는 컴퓨터가 단순히 답을 계산하는 것을 넘어, 그로부터 학습할 수 있게 해줍니다.

보통 현재의 모습을 바탕으로 과거의 우주가 어떠했는지 알아내려면 추측하고 확인하는 과정을 반복해야 합니다. 하지만 이 새로운 솔버는 "미분 가능(differentiable)"하기 때문에, 시뮬레이션을 역방향으로 돌릴 수 있습니다. 당신은 컴퓨터에게 "과거를 이렇게 보이게 만들어줘"라고 말할 수 있고, 그러면 컴퓨터는 그 결과를 얻기 위해 초기 조건을 어떻게 조정해야 하는지 정확히 계산해 냅니다.

논문은 코스믹 웹(cosmic web)을 재구성함으로써 이를 입증합니다. 움직이는 메쉬는 중력이 가장 강한 곳에 해상도를 집중시키는 법을 성공적으로 "학습"했으며, 우주의 복잡한 웹 구조를 높은 충실도로 재현해 냈습니다.

이 방법이 배제하거나 제한하는 것들

이 방법이 무엇을 하지 않는지, 혹은 무엇에 반대하는지 아는 것도 중요합니다.

• 모든 것에 대한 마법의 해결책은 아닙니다: 논문은 매우 작은 격자의 경우 기존의 FFT 방식이 여전히 더 빠르다는 점을 명시적으로 밝히고 있습니다. 새로운 방식은 격자가 커질 때(예: 512³ 이상) 그리고 FFT의 "외침"이 너무 비용이 많이 들 때만 승리합니다.
• 전통적인 의미의 완전한 적응형 메쉬(Adaptive Mesh)는 아닙니다: 격자를 작고 불규칙한 조각(보로노이 셀 등)으로 나누는 일부 방식과 달리, 이 방식은 격자를 규칙적이고 늘어나는 직사각형 형태로 유지합니다. 그물을 찢는 것이 아니라 단순히 늘리는 것입니다. 따라서 그물이 겹쳐질 정도의 극단적인 왜곡은 처리할 수 없으며, 이 때문에 저자들은 셀이 너무 찌그러지는 것을 막기 위해 "리미터(limiters)"를 추가해야 했습니다.
• 시뮬레이션이지 물리적 발견이 아닙니다: 결과는 컴퓨터 시뮬레이션(CAMELS 스위트 및 AREPO 코드 참조)에 기반합니다. 논문은 이 방식이 미래 시뮬레이션을 위한 실용적인 가교가 될 수 있음을 시사하지만, 우주의 신비 자체를 풀었다고 주장하는 것은 아닙니다.

얼마나 확신하는가?

저자들은 성능 수치에 대해 매우 확신하고 있습니다. 그들은 실제 하드웨어(Perlmutter 슈퍼컴퓨터 및 A100 GPU)에서 단계를 실행하는 데 걸리는 시간을 측정했으며, 격자 크기와 컴퓨터 수에 따라 1.3배에서 2.4배 사이의 일관된 속도 향상을 발견했습니다.

또한 그들은 정확도에 대해서도 확신합니다. 움직이는 메쉬의 결과는 고해상도 참조 시뮬레이션과 매우 밀접하게 일치했으며, 교차 상관 계수(cross-correlation coefficient)는 약 0.97(1.0이 완벽한 일치를 의미)을 기록했습니다.

하지만 미래의 응용 분야에 대해서는 더 신중합니다. 그들은 이 방법이 필드 수준의 추론(field-level inference, 우주의 역사를 파악하는 것)이나 현재의 표준 컴퓨터로는 너무 큰 규모의 시뮬레이션을 실행하는 데 있어 게임 체인저가 될 수 있다고 제안합니다. 다만, 실제 과학적 연구를 위해서는 "리미터"를 튜닝하고 노이즈 모델을 추가하는 등의 추가 작업이 필요하다고 인정합니다.

요약하자면

이 논문은 "글로벌한 외침"(FFT)에서 "로컬한 속삭임"(멀티그리드)으로 전환하고, 격자가 살아있는 그물처럼 늘어나게 함으로써, 우리가 더 크고, 빠르고, 상세한 우주 시뮬레이션을 실행할 수 있음을 시사합니다. 이는 도시 크기의 슈퍼컴퓨터 없이도 가장 중요한 곳에서 고해상도의 중력을 얻는 방법입니다. 그리고 가장 멋진 점은, 이제 컴퓨터가 시뮬레이션으로부터 학습할 수 있어 우주의 역사를 역설계하는 데 도움을 줄 수 있다는 것입니다.