원자 특성 예측을 위한 사전학습 데이터 정렬의 중요성

초록

본 논문은 대규모 데이터와 연산량에 의존하던 기존 원자 특성 예측 패러다임을 뒤흔들며, 작업에 정렬된 소규모 데이터셋을 사전학습에 활용하면 24배 적은 예산으로도 동일하거나 더 높은 성능을 달성할 수 있음을 입증한다. 이를 위해 화학적 유사성 지수(CSI)를 제안해 상류 데이터와 하류 작업 간의 정렬 정도를 정량화하고, CSI가 최소인 데이터셋을 선택했을 때 성능 향상이 일관됨을 실험적으로 보여준다. 또한, 정렬되지 않은 데이터를 무작위로 추가하면 오히려 성능이 저하될 수 있음을 강조한다.

상세 분석

이 연구는 원자 수준의 물성 예측에서 사전학습(pretraining) 단계의 데이터 선택이 모델 성능에 미치는 영향을 체계적으로 조사한다. 기존의 JMP와 같은 대규모 혼합 데이터셋을 이용한 사전학습은 수백만 개의 샘플과 수천 GPU‑day를 요구했지만, 저자는 “데이터 양보다 질”이라는 가설을 검증하기 위해 세 가지 핵심 요소를 도입한다. 첫째, 화학적 유사성 지수(CSI)를 정의한다. CSI는 FID(Frechet Inception Distance)를 그래프 기반 분자 표현에 적용한 것으로, 사전학습 데이터셋과 목표 다운스트림 데이터셋의 임베딩 평균·공분산을 비교해 정렬 정도를 수치화한다. 낮은 CSI 값은 두 데이터 분포가 유사함을 의미한다. 둘째, CSI를 이용해 K개의 후보 사전학습 데이터셋 중 가장 정렬된 하나를 선택하고, 선택된 데이터셋만으로 제한된 연산 예산 C=E×N(에포크 수×샘플 수) 하에 사전학습을 수행한다. 셋째, 동일한 예산 조건에서 선택된 소규모 데이터셋(≈10M 샘플)과 기존 대규모 혼합 데이터셋(≈240M 샘플)을 비교한다. 실험 결과, CSI가 최소인 데이터셋을 사용했을 때 평균 절대오차(MAE)가 JMP 대비 동일하거나 더 낮았으며, 특히 에너지·힘 예측에서 24배 적은 연산량으로도 경쟁력을 유지했다. 흥미롭게도, 가장 정렬된 데이터셋을 포함한 혼합 데이터셋을 사용하면 성능이 오히려 감소했는데, 이는 비정렬 데이터가 모델의 표현 공간을 혼란시켜 최적화 경로를 방해하기 때문이다. 또한, 데이터 양을 무조건 늘리는 것이 항상 이득이 아니라는 점을 정량적으로 입증했다. 이와 같은 결과는 3D 분자·재료 그래프 학습에서 데이터 선택 전략이 연산 효율성과 성능을 동시에 개선할 수 있음을 시사한다. 마지막으로, 저자는 CSI가 데이터 선택뿐 아니라 사전학습-파인튜닝 파이프라인 전반에 걸쳐 정렬을 측정하는 일반화 가능한 메트릭으로 활용될 수 있음을 강조한다.