TriForces: Augmenting Atomistic GNNs for Transferable Representations

TriForces 는 자기지도 학습을 분리된 조성 및 구조 표현과 결합하여 기계 학습 원자간 퍼텐셜을 위한 원자 단위 그래프 신경망의 전이성과 데이터 효율성을 크게 향상시키는 모델 무관한 3 스트림 프레임워크입니다.

핵심

로봇 셰프가 요리를 하는 법을 가르치려 한다고 상상해 보세요.

문제: "일률적" 셰프
현재 과학자들은 원자의 거동 (예: 물질의 에너지 양이나 원자를 움직이는 데 필요한 힘) 을 예측하기 위해 강력한 AI 모델 (MLIPs) 을 사용합니다. 이러한 모델들은 슈퍼컴퓨터 (DFT) 에서 생성된 방대한 양의 데이터를 기반으로 훈련됩니다.

그러나 이러한 모델에는 결함이 있습니다. 특정 요리의 정확한 맛을 외웠지만, 그 맛이 왜 그런지 그 이유를 잊어버린 셰프와 같다는 점입니다. 만약 약간 다른 요리 (새로운 유형의 물질) 를 만들어 달라고 요청하면 그들은 어려움을 겪습니다. 그들은 재료 (조성) 와 냄비 모양 (구조) 을 혼동합니다. 재료를 바꾸면 모양에 대해 혼란을 겪고, 그 반대의 경우도 마찬가지입니다. 이로 인해 특히 학습할 데이터가 많지 않을 때 새로운 작업을 빠르게 습득하는 데 서툴게 됩니다.

해결책: TriForces (세 개의 흐름을 가진 주방)
저자들은 이러한 AI 셰프를 구축하는 새로운 방법인 TriForces를 소개합니다. 모든 것을 한 번에 기억하려는 거대한 뇌 하나 대신, 뇌를 세 가지 전문화된 "흐름" 또는 부서로 나눕니다:

1. 재료 흐름 (조성): 이 부서는 냄비 안에 무엇이 들어 있는지 (예: "수소 2 개와 산소 1 개가 있습니다") 만 봅니다. 모양은 완전히 무시합니다. 화학을 배웁니다.
2. 모양 흐름 (구조): 이 부서는 원자들이 공간에 어떻게 배열되어 있는지 (예: "그들은 삼각형 모양입니다") 만 봅니다. 원자가 실제로 무엇인지는 무시합니다. 기하학을 배웁니다.
3. 상호작용 흐름: 이는 재료 및 모양 부서에서 받은 메모를 받아 최종 결과 (에너지 또는 힘) 를 예측하기 위해 이를 결합하는 메인 셰프입니다.

비밀 소스: 자기지도학습
모델이 특정 속성을 예측하도록 요청받기 전에, 저자들은 "자기지도학습"이라는 게임을 통해 모델을 훈련시킵니다. 이는 AI 가 다음을 수행해야 하는 연습 세션과 같습니다:

• 잡음 제거: 약간 손상되거나 잡음이 섞인 분자 이미지를 보고 이를 수정합니다.
• 마스킹: 재료를 가리고 이웃을 바탕으로 그것이 무엇인지 추측합니다.
• 매칭: 동일한 분자의 약간 다른 두 버전을 보고 그것이 같은 것임을 파악합니다.

이러한 훈련은 AI 가 지식을 깔끔하게 조직하도록 강제합니다. "재료"는 한 폴더에, "모양"은 다른 폴더에 속한다는 것을 배우고, 이를 뒤섞지 않도록 합니다.

중요성 (결과)
이 논문은 이러한 새로운 "세 흐름" 주방이 기존의 "단일 뇌" 주방보다 훨씬 잘 작동함을 보여줍니다:

• 빠른 학습: 소량의 새로운 데이터 (수백만 개 대신 2 만 개 정도의 예시 등) 를 제공받았을 때, TriForces 는 훨씬 더 빠르게 학습하고 실수를 줄입니다. 이는 천 번 요리해 보아야 새로운 레시피를 익히는 셰프가 아니라, 한 번 맛만 보고도 새로운 레시피를 익히는 셰프와 같습니다.
• 향상된 기억: AI 는 배운 것을 잊지 않습니다. 혼란을 겪지 않고 한 유형의 물질에서 다른 유형의 물질로 지식을 이전할 수 있습니다.
• 검색 가능한 지식: AI 가 "재료"와 "모양"을 분리해 두기 때문에, 재료가 다르지만 모양은 같은 물질을 찾거나, 모양은 다르지만 재료가 같은 물질을 찾아달라고 요청할 수 있습니다. 기존 모델들은 지식이 너무 뒤섞여 있어 이를 수행할 수 없었습니다.

요약
TriForces 는 원자를 이해하는 복잡한 작업을 세 가지 더 간단한 작업으로 분해하는 프레임워크입니다. 즉, 재료를 알고, 모양을 알고, 그리고 그들이 어떻게 함께 작동하는지 아는 것입니다. AI 가 이러한 작업을 분리하여 유지하도록 훈련시키고 "추측 게임"(자기지도학습) 으로 연습시킴으로써, 모델은 새로운 물질을 발견하기 위한 훨씬 더 유연하고 효율적이며 정확한 도구가 됩니다.

저자들은 다른 과학자들이 물질 과학을 위한 더 나은 AI 를 구축하기 위해 이 "세 흐름 주방"을 사용할 수 있도록 코드와 사전 훈련된 모델을 공개했습니다.

기술 요약

기술 요약: TriForces

문제 정의

기하학적 그래프 신경망 (GNN) 기반의 머신러닝 원자간 포텐셜 (MLIPs) 은 대규모 밀도 범함수 이론 (DFT) 데이터셋으로 학습되었을 때 원자 수준의 속성을 예측하는 데 높은 정확도를 달성했습니다. 그러나 두 가지 주요 한계로 인해 실제 활용이 저해됩니다:

1. 일관되지 않은 전이성 (Transferability): MLIPs 는 종종 새로운 화학적 환경이나 작업으로 효과적으로 전이되지 못합니다. 특히 작고 비싼 작업별 데이터셋으로 파인튜닝할 때 그렇습니다. 사전 학습된 모델은 종종 결정계 식별이나 주성분 원소 확인과 같은 간단한 진단 작업에서 어려움을 겪으며, 도메인 간에 "파국적 망각 (catastrophic forgetting)"이나 불안정한 전이 성능을 보입니다.
2. 낮은 표현 재사용성: 현재 MLIPs 는 에너지와 힘과 같은 특정 예측 목표를 위해 표현을 최적화할 뿐, 일반적인 재사용을 위해 최적화되지 않습니다. 결과적으로 이러한 표현은 종종 조성 (화학) 과 구조 (기하학) 가 얽히게 되어 탐색적 분석, 최근접 이웃 검색, 또는 분해된 유사도 검색에 적합하지 않습니다. 표준 지도 학습 목적 함수는 회귀에 충분한 표현을 장려하지만, 접근 가능한 조성과 구조 정보를 보존하도록 조직화되지는 않습니다.

방법론: TriForces 프레임워크

저자들은 기존 기하학적 GNN 에 세 가지 스트림 아키텍처와 다목적 자기지도 학습 (SSL) 을 보강하는 모델 무관한 프레임워크인 TriForces를 제안합니다. 핵심 혁신은 원자 수준 표현을 세 가지 명확한 구성 요소로 명시적으로 분해하는 것입니다:

1. 세 가지 스트림 아키텍처

단일 잠재 벡터 대신 TriForces 는 노드 수준 표현 $h_i$를 세 개의 연결된 스트림으로 분해합니다:

• 조성 스트림 ($h^{comp}$): 좌표 없이 화학 정보를 인코딩합니다. 이 스트림은 카운트 가중 어텐션을 사용하는 Transformer 를 통해 고유한 원소들의 집합과 그 화학량론적 개수를 처리합니다. 이 스트림은 기하학과 무관하게 시스템 크기와 에너지 규모를 인코딩하기 위해 절대 원소 개수를 보존합니다.
• 구조 스트림 ($h^{struct}$): 원소 정체성 (유형 무관) 없이 기하학을 인코딩합니다. 이 스트림은 SOAP(원자 위치의 매끄러운 중첩) 에서 영감을 받은 회전 불변 국소 기술자를 구성하며, 방사 기저 함수, 구면 조화 함수, 그리고 다중 스케일 컷오프를 활용합니다. 이 스트림은 불변 메시지 전달에 의해 강화된 재사용 가능한 기하학적 모티프와 위상적 패턴을 포착합니다.
• 상호작용 스트림 ($h^{int}$): 조성과 기하학 사이의 결합을 포착하고 원래 아키텍처의 표현력을 보존하는 표준 기반 기하학적 GNN(예: MACE, eSEN, Orb-v3) 입니다.

2. 자기지도 사전 학습

잠재 공간을 조직화하고 전이성을 향상시키기 위해 TriForces 는 확률적 증강 (위치 노이즈, 원소 마스킹, 그래프 변형, 회전) 을 사용하는 다목적 자기지도 학습 (SSL) 사전 학습 전략을 채택합니다. 이 프레임워크는 세 가지 상호 보완적인 목적 함수를 결합합니다:

• 비재구성 (Non-Reconstruction, LeJEPA): 동일한 구조의 두 가지 증강된 뷰에서 노드 수준과 그래프 수준 모두의 임베딩을 정렬합니다. 이는 증강에 대한 불변성을 강제하고 스톱-그레이디언트나 모멘텀 인코더 없이 글로벌 잠재 공간을 조직화합니다.
• 디노이징 (Denoising): 노이즈가 있는 입력으로부터 깨끗한 원자 위치를 복원하도록 모델을 학습시킵니다. 이는 기하학적 표현을 안정화하고 암묵적으로 회전 증강을 제공합니다.
• 마스킹 (Masking): 주변 기하학과 조성에 기반하여 마스킹된 원소 유형을 예측합니다. 이는 모델이 조성 패턴과 맥락을 학습하도록 장려합니다.

최종 사전 학습 손실은 이러한 세 가지 목적 함수의 가중 합입니다.

주요 기여

1. 아키텍처 분해: 조성, 구조, 상호작용을 명시적으로 분리하는 세 가지 스트림 디자인으로, 두 요소가 얽히는 것이 아니라 설계에 의해 보존되도록 보장합니다.
2. 하이브리드 사전 학습 전략: 재구성 기반 목적 함수 (디노이징, 마스킹) 와 잠재 예측 학습 (LeJEPA) 을 결합한 자기지도 접근법으로, 하위 작업 전이를 개선하기 위해 임베딩 공간을 구조화합니다.
3. 해석 가능한 검색: 조성, 구조, 또는 결합 임베딩 공간에서 표적 유사도 검색을 수행할 수 있어, 화학만 또는 구조만과 같은 특정 기준에 기반한 재료 비교를 가능하게 합니다.
4. 실증적 검증: 여러 아키텍처 (Orb-v3, eSEN, MACE) 와 벤치마크 (OMat24, MatBench, QM9) 에 걸친 광범위한 실험을 통해 향상된 데이터 효율성, 전이 성능, 그리고 표현 품질을 입증했습니다.

결과

• 전이 성능 (OMat24): 제한된 데이터 환경에서 TriForces 는 베이스라인을 크게 능가합니다. 20K 샘플에서 에너지 평균 절대 오차 (MAE) 를 베이스 모델 대비 57% 감소시킵니다. 모든 샘플 크기에서 힘의 MAE 를 개선하고 응력 오차를 감소시킵니다.
• 데이터 효율성: TriForces 는 모든 데이터셋 크기 (20K~2M 샘플) 에서 더 낮은 오차를 달성하며, 가장 큰 개선은 저데이터 환경에서 관찰됩니다.
• 벤치마크 성능:
  • MatBench: TriForces 변형체는 8 개 작업 중 6 개에서 전체적으로 최상의 결과를 달성하여 자기지도 학습 및 DFT 라벨링 사전 학습 베이스라인 모두를 능가합니다. 예를 들어, 포논 (Phonons) MAE 는 57.8 에서 19.5 cm$^{-1}$로 개선되었습니다.
  • MatBench Discovery: TriForces eSEN-sm 은 60% 적은 파라미터를 사용하며 최대 5$\times$ 빠르게 학습하면서도 훨씬 더 큰 eSEN-30M-OAM 모델과 비교 가능한 에너지 MAE 를 달성합니다.
  • QM9: 다양한 화학 입력 (벌크 + 분자) 에 대한 사전 학습은 벌크 전용 또는 SSL 없는 베이스라인에 비해 일관되게 MAE 를 낮춥니다.
• 표현 품질: 고정된 임베딩에 대한 선형 프로빙은 TriForces 가 표준 MLIPs 가 잃어버리는 기본 정보 (결정계, 주성분 원소, 배위수) 를 보존함을 보여줍니다. TriForces 는 결정계 및 주성분 원소 분류에서 96–100% 의 정확도를 달성하는 반면, 베이스라인은 55–73% 로 어려움을 겪습니다.
• 검색: 이 프레임워크는 조성 스트림이 원소 집합 회상에 뛰어나고 구조 스트림이 공간군 회상에 뛰어난 효과적인 k-NN 검색을 가능하게 하며, 이는 단일 스트림 모델에는 없는 기능입니다.

의의 및 주장

이 논문은 TriForces 를 단순히 자기지도 학습 방법이 아닌, SSL 에 의해 표현이 더욱 강화된 아키텍처 프레임워크로 위치시킵니다.

• 영역 의존성: 저자들은 스트림 분해가 대규모 지도 학습 환경에서 지배적인 이득을 제공한다고 주장하는 반면, SSL 은 저데이터 전이, 표현 조직화, 그리고 검색 작업에서 가장 가치 있다고 주장합니다.
• 분리: 조성과 구조를 분리함으로써 TriForces 는 현재 기반 모델의 "취약한 전이"와 "재사용이 어려운" 문제를 해결합니다. 이는 예측뿐만 아니라 분석 (예: 프로빙, 검색) 을 위해 조직화된 표현을 학습할 수 있게 합니다.
• 실용성: 이 프레임워크는 모델 무관하고 플러그인 방식으로, 기존 또는 새로운 원자 수준 아키텍처에 즉시 적용 가능합니다. 저자들은 하위 작업 원자 수준 모델링에서의 재사용을 촉진하기 위해 사전 학습된 체크포인트와 코드를 공개합니다.

이 연구는 미래의 원자 수준 기반 모델이 단일 스트림 예측 아키텍처를 넘어 화학 시스템의 고유한 물리적 요인을 명시적으로 보존하는 분해된 표현으로 나아가야 함을 시사합니다.