ED-DiT: 물리학 기반 확산 사전 학습을 통한 전자기 밀도로부터 생성되는 이식 가능한 분자 표현

ED-DiT: Physics-Guided Diffusion Pretraining for Transferable Molecular Representations from Electron Density

사전 학습은 이식 가능한 표현을 학습하는 데 강력한 잠재력을 보여주었지만, 전자 밀도를 기반으로 하는 분자 학습 분야에서는 아직 덜 연구되었습니다. 전자 밀도는 분자의 전자 구조에 대한 연속적인 3차원 설명을 제공하며, 국소적인 공간 패턴과 전반적인 물리적 양을 모두 포착합니다. 이러한 점은 핵심 질문을 제기합니다: 전자 밀도 필드를 사용하여 다양한 전자 구조 관련 작업에 이식 가능한 공유 표현을 학습하는 자기 지도 사전 학습이 가능할까요? 우리는 ED-DiT라는 물리학 기반 확산 트랜스포머를 제안하며, 이는 전자 밀도 포인트 클라우드에 대한 자기 지도 사전 학습을 위해 설계되었습니다. ED-DiT는 다양한 확산 노이즈 수준에서 손상되고 부분적으로 마스크된 로그 밀도 필드를 재구성함으로써 재사용 가능한 표현을 학습합니다. 또한, 전체 전자 질량을 유지하기 위한 전자 수 일관성 제약 조건을 도입합니다. 사전 학습된 인코더는 속성 예측, 개방/닫힌 쉘 분류, 분자-전자 밀도 검색 및 분자를 조건으로 하는 전자 밀도 예측에 적용될 수 있습니다. 여섯 가지 EDBench 작업에 대한 실험 결과, ED-DiT는 동일한 아키텍처를 처음부터 학습하는 것보다 일관되게 우수한 성능을 보이며, 특히 제한적인 감독 환경에서 더욱 그렇습니다. 분자를 조건으로 하는 전자 밀도 예측의 경우, RMSE가 2.2474에서 1.3753으로 감소하고 기존 기준 모델을 능가합니다. 단 10%의 레이블만 사용하여 오비탈 에너지 예측 RMSE를 0.0293에서 0.0138로 개선했습니다. 이러한 결과는 물리학 기반 전자 밀도 사전 학습이 이식 가능한 분자 표현을 학습하는 데 효과적임을 보여줍니다.

Pretraining has shown strong potential for learning transferable representations, yet it remains underexplored for electron-density-based molecular learning. Electron density provides a continuous three-dimensional description of molecular electronic structure, capturing both local spatial patterns and global physical quantities. This raises a key question: can electron-density fields be used for self-supervised pretraining to learn a shared representation that transfers across diverse electronic-structure-related tasks? We propose ED-DiT, a physics-guided Diffusion Transformer for self-supervised pretraining on electron-density point clouds. ED-DiT learns reusable representations by reconstructing corrupted and partially masked log-density fields across diffusion noise levels. An electron-number consistency constraint is further introduced to preserve the total electronic mass. The pretrained encoder can be adapted to property prediction, open-/closed-shell classification, molecule-electron-density retrieval, and molecule-conditioned electron-density prediction. Experiments on six EDBench tasks show that ED-DiT consistently outperforms the same architecture trained from scratch, especially under limited supervision. For molecule-conditioned electron-density prediction, it reduces RMSE from 2.2474 to 1.3753 and surpasses the available baseline. With only 10% labels, it improves orbital energy prediction RMSE from 0.0293 to 0.0138. These results demonstrate the effectiveness of physics-guided electron-density pretraining for learning transferable molecular representations.