TriForces: Augmenting Atomistic GNNs for Transferable Representations

Motivation

• Known: 대규모 DFT 데이터로 학습된 MLIP은 property prediction에서 뛰어난 정확도를 보이며, MACE, eSEN, Orb-v3와 같은 equivariant/non-equivariant geometric GNN이 이러한 예측 작업에 널리 쓰인다. 또한 denoising, masked reconstruction 등 SSL 기법이 atomistic 모델의 data efficiency와 fine-tuning 성능을 개선한다는 것도 알려져 있다.
• Gap: 그러나 MLIP은 새로운 chemistry나 DFT functional로 전이될 때 성능이 일관되지 않고, 대규모 사전학습 모델조차 crystal system 식별이나 majority element 예측 같은 간단한 진단 과제에서 fine-tuning에 실패하는 경우가 있다. 기존 SSL 연구는 예측 정확도 중심으로만 평가되어 representation이 retrieval이나 exploratory analysis 같은 재사용 목적에 얼마나 적합한지는 충분히 다뤄지지 않았다.
• Why: composition과 structure 정보가 뒤엉킨 단일 latent representation은 fine-tuning 시 catastrophic forgetting이나 불안정한 전이를 유발할 수 있으므로, 이를 분리하고 transferable하게 만드는 것은 제한된 task-specific 데이터로 MLIP을 실용적으로 적용하는 데 핵심적이다.
• Approach: 기존 geometric GNN을 백본으로 유지하면서 composition 전용 Transformer 스트림과 좌표 없는 invariant geometry 기반 structure 스트림을 추가하고, denoising·masking·LeJEPA·Barlow Twins 등 여러 self-supervised objective를 결합해 embedding space를 구조화한다.

Achievement

1. MatBench/QM9 성능 향상: DFT label 없이도 baseline 대비 MatBench 및 QM9에서 성능을 개선했다.
2. OMat24 저데이터 효율성: 20K 샘플의 limited-data 환경에서 energy MAE를 57% 감소시켰고, 다양한 sample size에서 force MAE도 개선했다.
3. 효율적 유사 구조 검색: 학습된 latent space를 통해 composition, structure, 또는 joint 기준으로 해석 가능한 similarity retrieval을 지원한다.
4. 다중 아키텍처 검증 및 공개: MACE, eSEN, Orb-v3 등 equivariant/non-equivariant 아키텍처에 걸쳐 검증하고 pretrained TriForces variants와 코드를 공개했다.

How

• 입력 구조 X={zi, xi}에 대해 원자 종류와 위치를 분리하여 node-level representation hi=(h_comp_i, h_struct_i, h_int_i)를 생성하는 세 개의 병렬 스트림 구성.
• Compositional stream: 원자 리스트를 unique element 집합 {(zt, ct)}로 압축하고 learnable element embedding 후 count-weighted attention을 적용한 Transformer로 좌표 없이 chemistry만 인코딩.
• Structural stream: element identity 없이 invariant geometry 정보만으로 구조를 인코딩.
• 기존 base geometric GNN(interaction stream)이 composition과 structure의 결합(coupling)을 포착.
• Self-supervised pretraining: denoising(geometric stability 강조), masking(compositional pattern 학습), LeJEPA·Barlow Twins 같은 non-reconstruction objective(latent separability 향상)를 결합.
• TriForces-Streams(랜덤 초기화 학습), TriForces(SSL 초기화), Base+SSL(백본만 SSL)로 구분해 stream factorization과 SSL pretraining의 기여를 분리 평가.

Originality

• 단일 latent vector에 auxiliary feature를 추가하는 기존 방식과 달리 composition·structure·interaction을 명시적으로 분리하는 3-stream 아키텍처를 제안.
• 여러 SSL objective(denoising, masking, LeJEPA, Barlow Twins)를 각 스트림의 역할에 맞게 상호보완적으로 결합해 representation을 구조화.
• 예측 성능뿐 아니라 retrieval·probing 등 representation 재사용성 평가에 초점을 맞춤.
• model-agnostic하게 MACE, eSEN, Orb-v3 등 다양한 base 아키텍처에 적용 가능함을 입증.

Limitation & Further Study

• Stream factorization과 SSL pretraining의 상대적 기여가 설정(대규모 지도학습 vs 저데이터 전이)에 따라 달라짐을 저자들도 인정하고 있어, 언제 어떤 요소가 지배적인지에 대한 이론적 설명이 부족하다.
• 본문 발췌만으로는 3-stream 구조 도입에 따른 추가 연산/메모리 비용, 다양한 DFT functional 간 전이, 더 큰 pretraining corpus로의 확장성에 대한 구체적 분석이 제한적이다.
• retrieval 성능 평가가 element-set 기반 recall@k 등 제한된 지표에 의존하는 것으로 보이며, 실제 재료 탐색 워크플로우에서의 정성적 검증이 더 필요하다.

Evaluation

총평: 구성과 구조 정보를 명시적으로 분리하는 아키텍처와 다중 SSL objective 결합을 통해 MLIP의 저데이터 전이와 표현 재사용성을 동시에 개선한 실용적이고 잘 검증된 연구로, atomistic representation learning 커뮤니티에 실질적인 기여를 한다.