Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

이 논문은 Frontier 슈퍼컴퓨터를 활용하여 16 개의 다양한 원자 구조 데이터셋으로 학습된 엑사스케일 멀티태스크 그래프 기반 모델을 개발함으로써, 기존 양자 역학 계산에 수년이 걸릴 작업을 50 초 만에 처리할 수 있는 초고속 물질 발견 워크플로우를 제시합니다.

핵심

이 논문은 **"원자 세계의 거대한 도서관을 50 초 만에 훑어보는 초고속 검색 엔진"**을 개발한 이야기입니다.

기존의 과학적 방법으로는 새로운 재료를 찾기 위해 컴퓨터로 시뮬레이션을 돌리는 데 수십 년이 걸렸다면, 이 연구팀은 슈퍼컴퓨터를 이용해 50 초 만에 11 억 개나 되는 원자 구조를 검사할 수 있게 했습니다. 마치 도서관에서 책 한 권을 찾는 데 100 년이 걸리던 것을, 1 초 만에 찾아내는 것과 같습니다.

이 놀라운 성과를 어떻게 이뤘는지, 일상적인 비유로 설명해 드리겠습니다.

---

1. 문제: 너무 많은 책, 너무 다른 글씨체

과학자들은 새로운 배터리나 약품을 만들기 위해 '원자'가 어떻게 배열되는지 알아야 합니다. 하지만 기존 방법 (양자역학 계산) 은 너무 정확하지만 너무 느리고 비싸서, 책 한 권 (새로운 재료) 을 읽으려면 몇 년이 걸립니다.

또한, 기존에 쌓아둔 데이터들은 제각각이었습니다. 어떤 데이터는 '정밀한 손글씨' (고정밀 데이터), 어떤 것은 '대충 쓴 메모' (저정밀 데이터) 였고, 주제도 유기물, 무기물, 합금 등으로 나뉘어 있어 한 번에 모두 배우기 어려웠습니다.

2. 해결책: "모든 책을 한 번에 읽는 거대한 AI 선생님"

연구팀은 HydraGNN이라는 AI 모델을 만들어 16 개의 서로 다른 데이터베이스 (총 5 억 4 천만 개의 원자 구조) 를 한꺼번에 학습시켰습니다.

• 비유: imagine(상상해 보세요) 16 명의 서로 다른 전공을 가진 교수님 (화학, 물리, 공학 등) 이 한 교실에 모여 있습니다. 보통은 각자 따로 가르치지만, 이 연구팀은 **한 명의 '수석 교수님' (공통 학습 부분)**을 만들어 모든 교수의 지식을 공유하게 했습니다.
• 핵심 기술 (멀티태스킹): 수석 교수님은 모든 학생 (데이터) 에게 공통된 원리 (원자 간 힘) 를 가르치고, 각 전공별 **보조 교수님 (전용 헤드)**을 두어 각자의 특수한 문제만 해결하게 했습니다. 이렇게 하면 서로 다른 데이터가 섞여도 혼란스럽지 않고, 작은 데이터도 큰 데이터에 묻히지 않게 됩니다.

3. 슈퍼컴퓨터의 힘: "Frontier"라는 거대한 공장

이 학습은 미국의 Frontier라는 세계 최고 성능의 슈퍼컴퓨터 (엑사스케일) 에서 이루어졌습니다.

• 규모: 16,384 개의 그래픽 카드 (GPU) 를 동시에 가동했습니다. 이는 16,000 대의 최신 게임용 그래픽 카드를 한꺼번에 돌려 학습을 시킨 것과 같습니다.
• 데이터 파이프라인: 이렇게 많은 컴퓨터가 데이터를 처리할 때 병목 현상이 생기지 않도록, ADIOS2/DDStore라는 특수한 '데이터 트럭' 시스템을 도입했습니다. 데이터를 컴퓨터 바로 옆에 미리 쌓아두어, 컴퓨터가 "데이터 주세요!"라고 기다리는 시간을 없앤 것입니다.

4. 최고의 모델 찾기: "초고속 레이스"

어떤 AI 구조가 가장 잘할지 알 수 없었기에, 6 가지 다른 구조 (EGNN, PaiNN 등) 로 동시에 대규모 실험을 했습니다. 마치 F1 레이서 6 대를 동시에 내보내어 가장 빠르고 정확한 차를 고르는 과정이었습니다.
그 결과, PaiNN이라는 구조가 가장 뛰어난 성능을 보였고, 이를 '최종 우승자 (리드 모델)'로 선정했습니다.

5. 놀라운 결과: 50 초의 마법

이렇게 훈련된 AI 모델은 다음과 같은 기적을 이루었습니다.

• 속도: 11 억 개의 원자 구조를 50 초 만에 검사했습니다.
  • 기존 방법으로는 이 작업을 하려면 약 6 년 7 개월을 쉬지 않고 계산해야 했습니다.
  • 이는 수십 년 걸릴 일을 점심시간 동안 끝낸 것과 같습니다.
• 정확도: 단순히 빠르기만 한 게 아니라, 과학적으로 매우 정확한 결과를 냈습니다. 특히 데이터가 부족한 새로운 분야에서도 기존 모델을 처음부터 다시 훈련시키는 것보다 훨씬 잘 작동했습니다.
• 유연성: 이 모델은 작은 컴퓨터에서도 작동할 수 있을 정도로 가볍고 (1200 만 개의 파라미터), 다양한 화학 반응에 맞춰 '미세 조정 (Fine-tuning)'이 가능합니다.

6. 요약: 왜 이것이 중요한가요?

이 연구는 **"원자 세계의 탐험"**을 완전히 바꿉니다.

• 과거: 새로운 재료를 찾으려면 "우연히 발견"하거나, 아주 천천히 하나씩 실험해야 했습니다.
• 현재와 미래: 이 AI 를 사용하면 전 세계의 모든 가능한 재료 조합을 순식간에 훑어보고, 가장 유망한 후보만 골라내어 실험실로 보낼 수 있습니다.

마치 천문학적인 양의 별 (원자 구조) 을 밤새도록 하나하나 세는 대신, AI 가 50 초 만에 가장 빛나는 별을 찾아주는 것과 같습니다. 이를 통해 배터리, 태양전지, 신약 개발 등 우리 삶의 질을 높이는 혁신적인 재료들을 훨씬 빠르게 찾아낼 수 있게 되었습니다.

기술 요약

1. 문제 정의 (Problem Statement)

재료 과학 및 분자 설계 분야에서 **밀도 범함수 이론 (DFT)**과 같은 1 차원적 (first-principles) 계산 방법은 높은 정확도를 제공하지만, 계산 비용이 매우 커서 산업적 규모에서의 광범위한 스크리닝에는 적용하기 어렵습니다. 이를 해결하기 위해 머신러닝 상호원자 포텐셜 (MLIP) 및 원자 단위 그래프 기반 모델 (GFM) 이 대안으로 제시되었으나, 다음과 같은 주요 도전 과제가 존재했습니다.

• 불균형 및 다중 정밀도 (Multi-Fidelity) 데이터의 통합: 유기물, 무기물, 하이브리드 시스템을 아우르는 다양한 데이터셋은 서로 다른 근사 이론 (DFT 함수형, 기저 집합 등) 으로 생성되어 에너지 기준이 불일치하고, 라벨 노이즈와 최적화 불균형이 발생합니다. 기존 단일 데이터셋 학습은 이러한 이질성을 처리하지 못했습니다.
• 확장성 및 모델 선택의 어려움: 단일 아키텍처가 모든 화학 공간에서 최적이 아니며, 대규모 하이퍼파라미터 최적화 (HPO) 와 효율적인 분산 학습을 통한 모델 선별이 필요했습니다.
• 이식성 및 데이터 부족: 다양한 슈퍼컴퓨터 (Frontier, Aurora, Perlmutter) 간 워크플로우 이식성 확보와, 레이블이 부족한 하위 작업 (downstream tasks) 에서의 정확한 적응 (fine-tuning) 이 필요했습니다.

2. 방법론 (Methodology)

이 연구는 HydraGNN을 기반으로 한 엑사스케일 (Exascale) 워크플로우를 제안하며, 다음과 같은 혁신적인 접근법을 사용합니다.

• 대규모 다중 작업 학습 (Multi-Task Learning, MTL):
  • 16 개의 공개 1 차원 데이터셋 (총 5.44 억 개 이상의 구조물, 85 개 이상 원소) 을 동시에 학습합니다.
  • 공유 메시지 전달 레이어 (Shared Message-Passing Layers): 모든 데이터셋에 공통적인 물리적 상호작용 특징을 학습합니다.
  • 작업별 헤더 (Per-Dataset Heads): 각 데이터셋의 고유한 라벨 특성 (에너지 기준, 정밀도 차이) 을 흡수하도록 설계하여, 데이터 간 기울기 간섭을 줄이고 불균형 데이터 학습을 안정화합니다.
• 시스템 및 데이터 파이프라인 최적화:
  • ADIOS2/DDStore: 비동기 데이터 이동 및 NVMe 스테이징을 활용하여 파일 시스템 병목 현상을 해결하고, 2,048 노드 이상의 대규모 학습을 지원합니다.
  • Sharded Distributed Training: HydraGNN 의 분산 데이터 병렬 처리 및 태스크 병렬 처리를 통해 메모리 부하를 분산하고 확장성을 극대화합니다.
• 대규모 하이퍼파라미터 최적화 (HPO):
  • Frontier 슈퍼컴퓨터 (16,384 GPU) 에서 6 가지 MPNN 백본 (EGNN, SchNet, DimeNet, MACE, PaiNN, PNAEq) 에 대한 대규모 HPO 를 수행하여 성능과 효율성을 균형 있게 갖춘 모델을 선정했습니다.
• 정밀도 민감도 분석 및 추론 최적화:
  • BF16, FP32, FP64 정밀도 간의 정확도 - 성능 트레이드오프를 정량화했습니다.
  • 추론 파이프라인 최적화: 인코더 재사용, 브랜치 스킵핑 (Branch Skipping), Fused Gradient(모든 헤드의 에너지를 한 번의 역전파로 통합), torch.compile 적용 등을 통해 추론 속도를 극대화했습니다.
• 컴포지션 조건부 브랜치 가중치 학습:
  • 16 개의 디코딩 브랜치 출력을 화학 조성 벡터에 기반한 MLP 를 통해 동적으로 가중치 부여하여, 훈련되지 않은 새로운 화학 시스템에도 적용 가능하도록 했습니다.

3. 주요 기여 (Key Contributions)

1. 첫 번째 엑사스케일 다중 작업 원자 GFM 워크플로우: 이질적이고 불균형한 16 개 데이터셋 (5.44 억 개 구조물) 을 16,384 개의 GPU 에서 성공적으로 학습한 최초의 사례입니다.
2. 초대규모 스크리닝 달성: 학습된 모델은 50 초 만에 11 억 개의 원자 구조물을 평가할 수 있으며, 이는 기존 1 차원적 계산으로 수 년이 걸릴 작업을 가능하게 합니다.
3. 이식 가능한 워크플로우: Frontier (AMD), Aurora (Intel), Perlmutter (NVIDIA) 등 이기종 슈퍼컴퓨터 환경에서 재현 가능한 소프트웨어 환경과 데이터 파이프라인을 구축하여 확장성을 입증했습니다.
4. 데이터 효율적인 파인튜닝: 사전 학습된 GFM 은 데이터가 부족한 환경에서도 훈련부터 시작하는 모델보다 검증 오차를 최대 10 배까지 낮추는 성능을 보였습니다.

4. 성능 결과 (Performance Results)

• 학습 및 HPO: Frontier 에서 6 개 MPNN 백본에 대한 HPO 를 수행하여 PaiNN 기반의 리드 (Lead) 모델을 선정했습니다. 이 모델은 2,048 노드에서 100 에포크 학습을 완료하며 수렴했습니다.
• 확장성 (Scaling):
  • Strong Scaling: Perlmutter(2,048 GPU), Aurora(6,144 GPU), Frontier(1,024 GPU) 에서 거의 선형적인 확장성을 보였습니다.
  • Weak Scaling: 시스템 규모가 커짐에 따라 효율성이 점진적으로 감소했으나, 세 플랫폼 모두에서 안정적인 성능을 유지했습니다.
• 하위 작업 (Downstream Tasks) 성능:
  • QM9, MD17, OQMD 등 12 가지 화학적으로 다양한 작업에서 **파인튜닝 (Unfrozen)**이 훈련부터 시작하는 모델보다 훨씬 낮은 오차 (최대 10 배 개선) 를 보였습니다.
  • 특히 에너지 표면 (PES) 과 관련된 작업에서 파인튜닝의 효과가 두드러졌습니다.
• 추론 속도:
  • Frontier (9,300 노드) 에서 초당 2,180 만 개 (전체), GPU 당 초당 293 개의 구조물을 처리했습니다.
  • 최적화 기법 (Fused Gradient, torch.compile 등) 을 적용하여 베이스라인 대비 최대 33 배의 속도 향상을 달성했습니다.
  • FP64 정밀도 유지 시 과학적 정확도를 보장하면서도, FP32/BF16 사용 시 속도 대비 정확도 손실을 정량화했습니다.

5. 의의 및 중요성 (Significance)

이 연구는 엑사스케일 과학적 머신러닝의 새로운 지평을 열었습니다.

• 재료 발견의 패러다임 전환: 기존에 접근 불가능했던 거대한 화학 설계 공간 (Chemical Design Space) 을 수 초 내에 스크리닝할 수 있게 하여, 희귀하고 고부가가치인 후보 물질을 찾는 과정을 가능하게 했습니다.
• 실용적인 과학 도구: 단순히 큰 모델을 만드는 것을 넘어, 데이터 효율성, 이식성, 그리고 생산 환경 (Production) 에 적합한 높은 추론 속도를 갖춘 실용적인 과학 도구로 자리 잡았습니다.
• 미래 방향 제시: 불균형하고 다중 정밀도의 데이터를 통합하는 다중 작업 학습 전략과, 엑사스케일 HPO 를 통한 모델 선정 프로세스는 향후 대규모 과학적 AI 시스템 개발의 표준이 될 것입니다.

결론적으로, 이 논문은 엑사스케일 컴퓨팅 자원을 활용하여 복잡한 물리 법칙을 학습하고, 이를 통해 전통적인 계산 방법으로는 불가능했던 규모의 재료 탐색을 실현한 획기적인 성과입니다.