기술 요약: Transport Novelty Distance (TNovD)

문제 정의

재료 발견을 위한 생성형 머신러닝의 급격한 발전은 엄격한 평가 지표의 발전을 앞질렀습니다. SUN 메트릭(Stability, Uniqueness, Novelty)과 같은 기존 방식은 생성된 구조를 개별 인스턴스 단위로 평가하며 분포적 특성을 포착하지 못합니다. 반대로, Fréchet Inception Distance (FID) 및 그 분자적 아날로그인 FCD와 같은 분포적 메트릭은 데이터를 평균과 분산으로 압축하는 가우시안 가정을 기반으로 하므로, 모델이 단순히 훈련 데이터를 복제하는 "암기(memorization)" 현상을 감지할 수 없습니다. 또한, 현재의 방법들은 품질(훈련 분포와의 통계적 유사성)과 신규성(특정 훈련 샘플과의 비유사성)을 동시에 평가하는 데 어려움을 겪고 있습니다. 본 논문은 고품질의 일반화와 위험한 암기를 구분하는 것이 매우 중요하다고 강조하며, 특히 확산 모델(diffusion models)과 같은 최신 모델들이 과적합(overfitting)에 더 취약해짐에 따라 이 구분이 더욱 결정적임을 밝힙니다.

방법론

저자들은 최적 운송(Optimal Transport, OT) 이론과 대조 학습(contrastive learning)에 기반한 분포적 메트릭인 **Transport Novelty Distance (TNovD)**를 제안합니다. 이 방법론은 세 가지 주요 구성 요소로 이루어집니다:

1. 등변 GNN을 통한 특징 추출 (Feature Extraction via Equivariant GNN):

   • 결정 구조(Crystal structures)는 $SE(3)$-등변 그래프 신경망(GNN)을 사용하여 특징 공간으로 임베딩됩니다.
   • GNN은 MP20 데이터셋을 사용하여 **InfoNCE 대조 손실(contrastive loss)**로 학습됩니다.
   • 증강 전략 (Augmentation Strategy): 동일한 결정 구조에 회전, 이동, 슈퍼셀 확장을 적용하여 양의 쌍(positive pairs)을 생성합니다. 음의 쌍(negative pairs)은 서로 다른 재료들로 구성됩니다. 이러한 학습은 특징 공간이 강체 변환(rigid transformations) 및 슈퍼셀 크기에 불변하면서도 서로 다른 화학적 조성을 구별할 수 있도록 보장합니다.

2. 최적 운송 결합 (Optimal Transport Coupling):

   • TNDV는 Wasserstein 거리 프레임워크를 사용하여 훈련 세트($P_D$)와 생성 세트($P_G$) 사이의 특징 분포 간 결합을 계산합니다.
   • 표준 Wasserstein 거리와 달리, TNovD는 비용 함수를 두 영역으로 나누기 위해 임계값 $\tau$를 도입합니다:
     • 암기 영역 (Memorization Regime, $x < \tau$): 생성된 샘플이 훈련 샘의 너무 가까운 경우(거리 $x < \tau$), 거듭제곱 함수 $(\tau - x)^M$를 사용하여 강력하게 페널티를 부여합니다. 이는 정확하거나 거의 동일한 복제본을 감지합니다.
     • 품질 영역 (Quality Regime, $x \ge \tau$): 샘플이 충분히 새로운 경우($x \ge \tau$), 선형 함수 $(x - \tau)$를 사용하여 해당 샘플이 훈련 분포에서 얼마나 벗어났는지 측정합니다.

3. 하이퍼파라미터 휴리스틱 (Hyperparameter Heuristics):

   • $\tau$ (임계값): 훈련 세트 내의 1-최근접 이웃(1-nearest neighbor) 및 2-최근접 이웃 거리를 분석하여, 어떤 점이 "암기"된 것으로 간주될지를 근사함으로써 휴리스틱하게 결정됩니다.
   • $M$ (지수): "좋은" 분포의 두 분할에 대해 $\tau$에 대한 TNovD의 미분이 0이 되는 평형 조건을 통해 도출되며, 이를 통해 암기에 대한 페널티와 낮은 품질에 대한 페널티 사이의 균형을 맞춥니다.

주요 기여

• 통합된 분포적 메트릭: TNovD는 가우시안 가정에 의존하지 않고 품질과 신규성을 공동으로 평가하도록 설계된 분포적 메트릭입니다. 기존 연구[8]가 GNN 기반의 분포적 메트릭을 제안했지만, TNovD는 신규성 탐지를 위해 GNN을 명시적으로 학습시키고 가우시안 가정을 피함으로써 암기 탐지 능력을 차별화했습니다.
• 암기 탐지: 비선형 비용 함수를 활용함으로써, TNovD는 훈련 데이터를 재생산하는 것을 명시적으로 처벌하며, 이는 FID나 FCD에는 없는 기능입니다.
• 도메인 불가지론적 프레임워크: 결정 구조에 적용되었으나, 저자들은 특징 추출기를 교체함으로써 다른 도메인(예: 이미지, 분자)에도 적응 가능하다는 점을 언급합니다.
• 강건한 특징 학습: 특정 결정학적 증강(슈퍼셀, 회전)을 포함한 대조 학습을 사용하여, 물리적 대칭성을 존중하면서도 진정한 신규성과 동일 물질의 구조적 변이를 구별하는 특징 공간을 생성합니다.

결과

저자들은 다음과 같은 실험을 통해 TNovD를 검증했습니다:

1. 토이 실험 (MP20 검증 세트):

   • 암기: 검증 데이터가 점진적으로 훈련 데이터로 대체될 때, TNovD는 거의 선형적인 증가를 보이며 암기를 성공적으로 감지했습니다.
   • 노이즈 주입: 원자 좌표에 가우시안 노이즈를 추가하거나 격자를 전단(shearing)했을 때 TNovD 점수가 상승하여, 저품질 구조를 감지함을 보여주었습니다.
   • 불변성: 슈퍼셀이 생성되어도 메트릭이 안정적으로 유지되어, GNN이 크기 변화를 신규성으로 오인하지 않고 처리할 수 있음을 입증했습니다.
   • 치환: 무작위 및 그룹 내 원자 치환이 저품질 재료로 올바르게 식별되었으며, 무작위 치환이 더 강한 페널티를 받았습니다.

2. WBM 치환 데이터셋:

   • 반복적인 원자 치환을 통해 생성된 구조인 WBM 데이터셋에 적용했을 때, TN0VD는 구조가 MP20 훈련 세트에서 멀어짐에 따라 암기 성분은 0에 가깝게 떨어지고 품질 성분은 약간 증가하여, 멀어지는 정도를 반영함을 보여주었습니다.

3. 생성 모델 벤치마킹:

   • 10,000개의 생성 샘플에 대해 6개의 유명한 재료 생성 모델(MatterGen, CDVAE, DiffCSP, ADiT 포함)을 평가했습니다.
   • ADiT는 높은 품질에 힘입어 전반적으로 가장 좋은 점수를 기록했으나, 저자들은 이 모델이 메트릭이 완전히 포착하지 못하는 방식으로 데이터를 암기하고 있을 가능성이 있다고 가설을 세웠습니다.
   • MatterGen은 검증 세트 베이스라인과 유사한 성능을 보였습니다.
   • CDVAE는 품질 페널티(낮은 품질)가 높고 암기가 없는 모습을 보여, 훈련 분포를 맞추는 데 실패했음을 나타냈습니다.
   • Chemeleon은 가장 높은 암기 성분을 보였습니다.

의의 및 주장

본 논문은 TNovD가 기존의 인스턴스 수준 메트릭에 대한 "보완적 관점"을 제공한다고 주장합니다. 주요 의의는 현대적인 생성 모델의 평가에서 중요한 공백인 암기 탐지와 분포적 품질 평가를 동시에 수행할 수 있다는 점에 있습니다.

저자들은 다음과 같은 점을 겸허히 언급합니다:

• TNDV는 훈련 목적 함수(training objective)가 아닙니다. 생성 모델이 이론적으로 이 메트릭을 "속일(game)" 수 있으므로, 이는 엄격히 평가 도구로서 기능합니다.
• 메트릭은 재료의 안정성(예: 에너지 델타)이나 합성 가능성을 직접 포착하지는 않지만, 향후 사전 학습된 에너지 모델을 통해 확장될 수 있습니다.
• 결과는 선택된 특징 표현(MP20 기반 GNN)과 증강 방식에 크게 의존합니다.
• 많은 모델에서 관찰된 낮은 암기 점수는, 현재의 재료 생성 모델의 주요 과제가 암기를 피하는 것보다 올바른 데이터 분포를 맞추는 것(품질)임을 시사합니다.

이 연구는 홀리스틱한 평가를 위한 모듈형 청사진을 구축하며, 생성 모델이 발전함에 따라 진정한 발견과 훈련 세트의 복제 사이를 구분하는 능력이 무엇보다 중요해질 것임을 강조합니다.