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

Essence

Fig. 1. Multi-task learning training for MLIPs on multi-source, multi-fidelity data, with shared message-passing represe

본 논문은 16개 제일원리 데이터셋(5억 4400만 구조)을 활용한 다중 과제 학습 기반 PaiNN 그래프 파운데이션 모델을 Frontier·Aurora·Perlmutter 슈퍼컴에서 학습하여, 11억 구조를 50초에 스크리닝할 수 있는 지수급 가속을 달성하였다.

Motivation

Known: 기존 원자 모델링은 DFT 같은 제일원리 방법에 의존하지만 계산 비용이 높으며, 최근 GemNet·MACE·Equiformer 등 등변 그래프 모델과 UMA·DPA-2·MACE-MP-0 같은 범용 MLIP이 다중 화학 영역에서 전이 성능을 보였다.
Gap: 기존 연구는 (i) 충분하지 않은 다중 충실도 균형, (ii) 배포 제약 하에서의 체계적 정밀도 특성화 부재, (iii) 이기종 가속기 시스템 간 제한적 확장성 검증에 그쳤으며, 대규모 불균형 다중 원본·다중 충실도 데이터에서 견고한 교차 도메인 표현을 유지하면서 동시 학습하는 방법이 해결되지 않았다.
Why: 원자 소재 발견의 산업 규모 스크리닝은 제일원리 방법으로 불가능하며, 다중 원본 데이터의 불일치 에너지 기준·레이블 노이즈·최적화 불균형을 다루면서도 지수급 추론 성능을 달성하는 것이 중요하다.
Approach: HydraGNN·ADIOS2·DDStore 기반 워크플로우로 다중 과제 아키텍처(데이터셋당 독립 헤드)를 구현하고, DeepHyper로 FP64 하이퍼파라미터 최적화 캠페인을 수행한 후, PaiNN 리드 모델을 2,048 노드 학습으로 승격하여 BF16/FP32/FP64 정밀도 트레이드오프를 검증하고, 12개 화학적으로 다양한 다운스트림 과제에서 전이 성능을 측정하며, Frontier·Aurora·Perlmutter에서 강한·약한 스케일링을 입증한다.

Achievement

Fig. 5. Fused Gradient Inference Pipeline. The encoder executes once and its

제1원리 데이터셋 규모: 16개 데이터셋 5억 4400만 구조로 기존 연구 초과. 지수급 추론 성능: 1.1억 구조를 50초 내 평가(수년치 DFT 계산 압축). 다중 충실도 균형: 불균형 코퍼스에서 교차 도메인 견고성 유지. 정밀도 분석: BF16/FP32/FP64 트레이드오프 및 강제 미분 수치 오차 영향 정량화. 전이 학습: 12개 다운스트림 과제에서 데이터 부족 미세조정 성능 입증. 이기종 확장성: Frontier·Aurora·Perlmutter 간 강한·약한 스케일링 효율 입증.

How

Fig. 5. Fused Gradient Inference Pipeline. The encoder executes once and its

Multi-task learning: 데이터셋마다 독립적 예측 헤드를 유지하면서 공유 message-passing 백본 사용
확장 가능 데이터 파이프라인: ADIOS2/DDStore로 544M 구조 분산 처리
DeepHyper 기반 대규모 HPO: 6개 캠페인으로 최적 MPNN 백본 자동 탐색
정밀도 트레이드오프: 배포 시 FP64→BF16/FP32 감소로 성능-효율 균형
강제 미분 오차 분석: 에너지 그래디언트 학습 시 수치 오차 상호 작용 정량화
다운스트림 미세조정: 제한 데이터 체제에서 사전학습 정규화 활용

Originality

단순 데이터셋 혼합 대신 초기부터 이기종 다중 원본에 노출(catastrophic forgetting 회피): 기존 MACE-MH-1은 단일 원본 사전학습 후 다중 원본 미세조정
multi-task 헤드 기반 데이터셋·충실도 특화: UMA의 mixture-of-experts 화학 전문화와 상보적
강제 미분 기반 학습에서 기계 정밀도 영향의 체계적 특성화: 배포 제약 하 과학 ML 작업흐름에 처음 정량화
544M 구조의 불균형 다중 충실도 코퍼스에서 교차 도메인 견고성 검증: 기존은 단일 도메인 또는 소규모 다중 원본

Limitation & Further Study

세밀한 비교 불가: 데이터 혼합·정밀도·하드웨어·측정 경계가 문헌마다 상이하여 상호 비교 어려움
단순 혼합의 한계: 논문에서 모든 데이터셋을 단순 혼합할 경우 지배 도메인 과적합 관찰 → 다중 과제 헤드 필요성 강조하나, 최적 헤드 설계 원리는 미해결
12개 다운스트림 과제의 화학적 다양성 범위: 유기·무기·혼성 시스템 포함하나 각 도메인별 세부 성능 분석 부재
대규모 HPO의 계산 비용: 6개 DeepHyper 캠페인 실행 비용 미보고로 재현성·실용성 평가 제한
정밀도 트레이드오프의 일반화: BF16 추론 시 정확도 저하 수량화 미흡
후속 연구: (i) 다운스트림 과제별 세밀한 오류 분석, (ii) 헤드 설계 최적화 이론, (iii) 더 광범위 화학 공간 검증

Evaluation

Novelty: 4/5 Technical Soundness: 4/5 Significance: 5/5 Clarity: 4/5 Overall: 4/5

총평: 본 논문은 불균형 다중 충실도 원자 데이터에서 처음 지수급 다중 과제 파운데이션 모델을 구축하고, 이기종 슈퍼컴 간 이식성과 정밀도 특성화를 입증하여 과학 머신러닝의 실무적 배포 장벽을 크게 낮추었다. 544M 구조 학습과 11억 구조 50초 스크리닝은 소재 발견의 패러다임을 바꾸며, 정밀도 트레이드오프·강제 미분 오차 분석은 배포 제약 하 과학 ML의 신뢰성을 강화한다. 다만 다운스트림 과제별 세밀한 오류 분석과 헤드 설계 원리 정밀화가 필요하다.