Efficient Cross-Functional Learning for Atomistic Modeling of Materials

Motivation

• Known: Atomistic foundation model을 대규모 저비용 functional(PBE) 데이터로 사전학습한 뒤 소규모 고비용 functional(r2SCAN) 데이터로 fine-tuning하는 것이 일반적인 워크플로우이며, transfer learning이 sample efficiency를 개선하고 multi-task 학습이 downstream robustness를 향상시킨다는 것이 알려져 있다.
• Gap: Cross-functional transfer에서 adaptation 전략의 선택이 target 데이터 가용성 및 compute 제약과 어떻게 상호작용하는지에 대한 이해가 부족하며, functional 간 에너지 스케일 shift와 약한 상관관계로 인해 adaptation이 강력한 사전학습 모델에서도 자명하지 않다는 문제가 있다.
• Why: 현실적인 재료 스크리닝 워크플로우에서는 고정밀 DFT 라벨이 희소하고 계산 비용이 크기 때문에, 제한된 데이터와 compute 자원 하에서 효율적으로 higher-fidelity functional로 적응하는 실용적 레시피를 제공하는 것은 대규모 재료 탐색의 실질적 가치를 좌우한다.
• Approach: 저데이터·저컴퓨트 환경에 적합한 경량 delta-learning과 다중 functional을 공유 표현으로 공동 학습하는 multi-head fine-tuning 두 가지 cross-functional adaptation 전략을 비교하고, layer-wise probing으로 어떤 표현이 functional 간 전이 가능성이 높은지를 분석한다.

Achievement

Figure 2. Top: MAE (eV/atom) as a function of training set size

1. Layer-wise probing을 통한 전이 가능 표현 식별: 사전학습된 Orb-v3의 초기 message-passing layer(MP1)가 깊은 layer보다 낮은 test MAE를 보이며 cross-functional 전이에 더 유용한 특징을 제공함을 Ridge 및 MLP probe 모두에서 일관되게 확인했다.
2. 경량 delta-learning의 효율성 입증: frozen 표현(평균 atom embedding)과 foundation model 자체 예측을 입력으로 사용하는 delta-learning이 저데이터·저compute 환경에서 최대 2자릿수 적은 compute로 경쟁력 있는 성능을 달성함을 보였다.
3. Multi-head fine-tuning의 양방향 전이 효과 확인: 공유 backbone과 functional별 head를 사용한 multi-head 학습이 higher-fidelity 감독 신호가 lower-fidelity 예측 성능도 향상시키는 양방향 전이를 가능케 함을 MatPES와 LeMatBulk 데이터셋에서 확인했다.

How

Figure 3. Top: MAE (eV/atom) as a function of compute on

• Orb-v3를 백본으로 사용하여 encoder output 및 각 message-passing block의 atom-level embedding을 mean-pooling해 graph-level 표현을 추출하고, 각 layer마다 Ridge 및 MLP probe를 학습해 formation energy per atom을 예측하는 layer-wise probing 수행
• Probing 결과로 선택된 layer(MP1)의 embedding 평균과 Orb-v3의 예측값(단일 스칼라)을 입력으로 받아 EDFT − EOrb 보정값을 예측하는 delta-learning head(Ridge regression 또는 MLP)를 설계
• 공유 Orb-v3 backbone 위에 functional별 prediction head를 부착하고, task-normalized target space에서 weighted multi-task loss L = Σ_k λ_k L_k로 여러 DFT functional(PBE, r2SCAN)을 공동 학습하는 multi-head fine-tuning 구현
• MatPES 및 LeMatBulk 데이터셋에서 training set size 및 compute 대비 MAE(eV/atom)를 측정하여 두 전략을 비교하고, delta-learning에 대한 ablation study(embedding 및 foundation model 예측 사용 여부) 수행

Originality

• 기존 delta-learning(quantum chemistry에서 유래)을 atomistic foundation model의 frozen intermediate representation과 모델 자체 예측을 결합하는 형태로 재구성하여 cross-functional adaptation에 적용
• Layer-wise probing을 단순 분석 도구가 아니라 delta-learning head의 layer 선택 기준으로 직접 활용하는 방법론적 연결을 제시
• 동일한 사전학습 backbone 위에서 delta-learning과 multi-head fine-tuning이라는 두 상반된 적응 전략을 데이터·compute 제약이라는 축을 기준으로 체계적으로 비교한 최초의 시도

Limitation & Further Study

• 실험이 Orb-v3라는 단일 foundation model 백본과 PBE/r2SCAN 두 functional 조합에 한정되어 있어, 다른 아키텍처나 더 많은 functional 조합으로의 일반화 가능성이 검증되지 않았다
• Delta-learning의 성능이 특정 저데이터 구간에서만 우위를 보일 가능성이 있으며, 데이터가 충분히 커질 때의 점근적 성능 한계에 대한 심층 분석이 제한적이다
• Multi-head 학습에서 task-normalization 방식이 특정 스케일링 방법에 의존하므로, 더 다양한 정규화 기법이나 functional 수가 늘어났을 때의 확장성에 대한 추가 연구가 필요하다
• 후속 연구로는 3개 이상의 functional을 포함한 다중 fidelity 환경, 그리고 delta-learning과 multi-head 학습을 결합한 하이브리드 전략에 대한 탐구가 유의미할 것이다

Evaluation

총평: 저데이터·저컴퓨트 제약 하에서의 cross-functional adaptation 문제를 실용적 관점에서 체계적으로 다루었으며, layer-wise probing을 통한 근거 제시와 두 전략의 명확한 trade-off 분석이 돋보이는 실용적이고 명료한 연구이다.