Latent Diffusion Pretraining for Crystal Property Prediction

Motivation

• Known: GNN 및 Transformer 기반 모델(CGCNN, Matformer, eComFormer, PDDFormer 등)이 결정 물성 예측에서 우수한 성능을 보이지만 labeled 데이터가 부족한 문제가 있으며, CrysXPP, CrysGNN, Crystal Twins 같은 pretraining-finetuning 방법과 CrysDiff, DPF 같은 diffusion 기반 pretraining 방법들이 이를 완화하기 위해 제안되어 왔다.
• Gap: 기존 diffusion 기반 pretraining 방법들은 원자 타입(discrete diffusion), 격자 파라미터(DDPM), 분수 좌표(wrapped normal distribution 기반 score-based modeling) 등 이질적인 고차원 feature space에서 직접 diffusion을 수행하여 복잡한 아키텍처와 많은 diffusion step이 필요하고, 이로 인해 비효율적인 학습과 non-smooth input space에 제약된 표현이 발생한다.
• Why: 결정 물성(형성 에너지, 밴드갭 등)은 원자 배열과 격자 구조에 근본적으로 좌우되므로, 원시 feature space를 직접 모델링하기보다 조성과 3D 기하를 충실히 포착하는 압축되고 표현력 있는 latent representation을 학습하는 것이 데이터가 부족한 재료과학 분야에서 downstream property prediction 성능과 데이터 효율성을 크게 개선할 수 있어 중요하다.
• Approach: VAE encoder로 결정 구조(원자 타입, 좌표, 격자)를 물리적 대칭성을 보존하는 매끄러운 latent space로 인코딩하고, 이 latent space 안에서 transformer 기반 Latent Diffusion Model(LDM)로 노이즈 추가/제거를 학습시켜 encoder가 더 풍부한 구조적·화학적 semantics를 포착하도록 유도한 뒤, 이렇게 pretrain된 encoder를 소량의 labeled 데이터로 finetuning하여 물성을 예측한다.

Achievement

Figure 2. Performance comparison (MAE) under limited training data. MAE on four JARVIS properties using 20%, 40%, 60%, a

1. 성능 우수성: JARVIS 및 MP 데이터셋에서 training-from-scratch 및 기존 pretrained baseline 대비 각각 4.26%, 4.90% 개선을 달성했으며, 강력한 diffusion 및 GNN 기반 baseline 대비 최대 19.34%까지 상대 오차를 감소시켰다.
2. 저데이터 환경에서의 강건성: 학습된 표현이 sparse-data 조건에서도 강건하며, 제한된 실험 데이터로 finetuning 시 DFT 오차를 보정할 수 있을 만큼 표현력이 뛰어남을 보였다.
3. Backbone-agnostic 설계: encoder architecture(Matformer, PDDFormer, CoGN, iComFormer, eComFormer 등)를 자유롭게 교체할 수 있어, 더 강력한 backbone(Matformer→PDDFormer)으로 교체 시 JARVIS에서 평균 10.46%, MP에서 12.39%의 추가 성능 향상을 확인했다.

How

Figure 1. Overview of our proposed pretrain–finetune framework, CrysLDNet. In the pretraining stage, a VAE encodes cryst

• Crystal material을 M = (A, X, L) (원자 타입 행렬, 좌표 행렬, 격자 행렬)로 표현하고 회전 불변성 및 주기적 병진 불변성 같은 물리적 대칭을 만족하는 인코더 설계
• VAE encoder Eϕ와 decoder Dψ를 통해 M을 노드 임베딩 Z ∈ RN×d의 latent space로 인코딩
• Latent space 상에서 diffusion process를 적용하는 transformer 기반 Latent Diffusion Model(LDM)을 학습시켜, 점진적으로 노이즈를 추가하고 다시 제거하는 학습을 통해 latent 분포를 학습
• Pretraining 이후, 소량의 property-labeled 데이터 Dp = {(Mi, yi)}로 encoder와 경량 property predictor(readout)를 함께 finetuning
• backbone encoder로 CoGN, Matformer, iComFormer, eComFormer, PDDFormer 등 다양한 아키텍처를 교체 가능하도록 설계하여 backbone-agnostic 특성을 실험적으로 검증
• JARVIS, MP 등 DFT benchmark 데이터셋에서 여러 물성에 대해 comprehensive 실험 수행 및 loss curve, 표현력 비교(Figure 3, 5) 등 ablation 분석 수행

Originality

• 기존 결정 pretraining 연구들이 이질적인 고차원 raw feature space(원자 타입, 좌표, 격자)에서 직접 diffusion을 수행한 것과 달리, VAE를 통해 매끄럽고 통합된 latent space를 먼저 구성한 후 그 안에서 diffusion을 적용하는 latent diffusion 개념을 결정 구조 pretraining에 처음 도입
• 서로 다른 물리적 특성(discrete, continuous, periodic)을 가진 이질적 변수들을 각각 다른 diffusion 방식(D3PM, DDPM, wrapped normal)으로 모델링해야 하는 복잡성을 latent space 통합을 통해 근본적으로 우회
• Encoder architecture에 구애받지 않는(backbone-agnostic) pretraining 프레임워크를 설계하여 미래의 더 강력한 encoder로 손쉽게 교체 가능하도록 한 실용적 확장성 제공

Limitation & Further Study

• VAE의 latent space 품질이 전체 파이프라인 성능에 결정적 영향을 미치는데, VAE 재구성 품질 저하나 posterior collapse 등의 문제에 대한 심층적 분석이 본문 발췌에서는 충분히 제시되지 않음
• latent diffusion의 step 수, 아키텍처 크기 등에 따른 계산 비용과 raw feature-space diffusion 대비 실제 학습/추론 효율성 비교가 명확히 제시될 필요가 있음
• 물리적 대칭성(회전/병진 불변성) 보존이 latent space에서 어떻게 엄밀하게 보장되는지에 대한 이론적 증명이나 정량적 검증이 추가로 필요
• 향후 연구로 더 다양한 물성 및 더 큰 규모의 unlabeled 데이터셋, 다른 최신 backbone architecture와의 결합 가능성을 확장할 수 있음

Evaluation

총평: 이질적 feature space에서의 diffusion 대신 통합된 latent space에서 diffusion을 수행하는 아이디어는 결정 구조 pretraining 분야에서 참신하고 실용적이며, backbone-agnostic 설계와 강력한 실험 결과로 향후 재료 발견 연구에 의미 있는 기여를 할 것으로 기대된다.