칭화팀은 최초로 분자 생성과 속성 예측을 통합하는 성과를 달성했습니다. 이는 2단계 확산 생성 메커니즘을 제안했으며 ICLR 2025에 선정되었습니다.
인공지능 기술은 약물 개발 과정을 근본적으로 바꾸고 있습니다.그 중 분자 특성 예측과 분자 생성은 두 가지 핵심 과제로서 오랫동안 독립적인 기술 경로를 따라 발전해 왔습니다.분자 특성 예측의 목적은 분자 구조 정보가 주어졌을 때 분자의 다양한 화학적, 생물학적 특성을 예측하고 약물 스크리닝을 가속화하는 것입니다. 분자 생성은 분자 데이터의 분포를 추정하고, 잠재적으로 원자 상호작용과 구조 정보를 학습하고, 화학적으로 합리적인 새로운 분자를 처음부터 생성하여 약물 설계 가능성의 경계를 확장하는 것을 목표로 합니다. 최근 몇 년 동안 이 분야에 대한 많은 연구가 이루어졌지만, 대부분은 독립적으로 발전했습니다.이 두 가지 핵심 연결 고리 사이의 협력 채널은 효과적으로 열린 적이 없습니다.
이를 고려하여,청화대학교와 중국과학원 연구팀은 UniGEM 모델을 제안했는데, 이는 확산 모델을 기반으로 두 가지 과제의 협력적 향상을 처음으로 달성한 사례입니다.연구팀은 생성과 속성 예측이 높은 상관관계를 가지고 있으며 효과적인 분자 표현에 의존한다고 지적했습니다. 연구팀은 기존 관절 훈련의 비일관성을 극복하고 분자 생성 및 특성 예측 분야에서 새로운 길을 여는 2단계 생성 프로세스를 혁신적으로 제안했습니다. 이 업적은 "UniGEM: 분자의 생성 및 속성 예측에 대한 통합적 접근 방식"이라는 제목으로 ICLR 2025에 선정되었습니다.
서류 주소:
https://openreview.net/pdf?id=Lb91pXwZMR
QM9 양자화학 데이터 세트:
GEOM-Drugs 3D 분자 구조 데이터 세트:
오픈소스 프로젝트인 "awesome-ai4s"는 200개 이상의 AI4S 논문 해석을 모아 방대한 데이터 세트와 도구를 제공합니다.
https://github.com/hyperai/awesome-ai4s
생성 및 예측 작업 통합에 대한 동기
연구팀은 생성 및 예측 작업의 본질은 분자 표현을 학습하는 데 있다고 믿습니다.한편, 다양한 분자 사전 훈련 방법의 효과는 분자 특성 예측이 견고한 분자 표현을 기반으로 한다는 것을 보여줍니다. 반면, 분자 생성에는 생성 과정에서 좋은 표현을 만들어내기 위해 분자 구조에 대한 깊은 이해가 필요합니다.
최근 연구 결과는 이러한 견해를 뒷받침합니다. 예를 들어, 컴퓨터 비전 분야에서의 연구는 확산 모델 자체가 효과적인 이미지 표현을 학습할 수 있는 능력을 가지고 있음을 보여주었습니다. 분자 영역에서는 생성적 사전 학습이 분자 속성 예측 작업을 향상시킬 수 있다는 연구 결과가 있지만, 이러한 방법은 최적의 예측 성능을 달성하기 위해 추가적인 미세 조정이 필요한 경우가 많습니다. 게다가 예측 변수가 분류기 안내 방법을 통해 분자 생성을 안내할 수 있지만, 예측 변수의 교육이 생성 성능을 직접적으로 개선할 수 있는지 여부는 아직 불확실합니다.
따라서 기존 연구에서는 생성 작업과 예측 작업 간의 관계를 아직 완전히 밝히지 못했습니다.이는 핵심적인 질문을 제기합니다. 생성 및 예측 작업의 시너지 효과를 높이는 통합 모델을 구축할 수 있을까요?
기존 방법의 실패 이유 분석
이 두 가지 작업을 결합하는 간단한 방법은 기존의 다중 작업 학습 프레임워크를 사용하는 것입니다. 여기서 모델은 생성 손실과 예측 손실을 모두 최적화합니다. 하지만 연구팀이 실시한 실험 결과, 이 접근법은 발전 작업과 속성 예측 작업의 성능을 상당히 저하시키는 것으로 나타났습니다(발전 안정성은 6%만큼 떨어지고, 예측 오차는 1배 이상 증가했습니다). 생성 모델의 가중치를 고정하고 속성 예측 작업에 별도의 헤드를 추가하여 생성 성능을 유지한 후에도 연구자들은 처음부터 학습하는 것과 비교했을 때 속성 예측 성능에 개선이 없다는 것을 관찰했습니다.
연구자들은 기존 방법의 결과가 좋지 않은 이유는 생성 작업과 예측 작업 간의 본질적인 불일치 때문이라고 설명합니다.확산 생성 과정 동안 분자 구조는 무질서한 노이즈에서 미세 구조로 점진적으로 재구성되어야 합니다. 그러나 예측 작업에서는 분자 구조가 기본적으로 확립된 후에야 의미 있는 분자 특성을 정의할 수 있습니다. 따라서 단순한 다중 작업 최적화 접근 방식을 채택하면 초기 확산 단계에서 매우 무질서한 분자 구조가 속성 레이블과 잘못 연관되어 분자 생성과 속성 예측에 부정적인 영향을 미치게 됩니다.
이 점을 더욱 잘 설명하기 위해 연구진은 확산 훈련 동안 잡음 제거 네트워크 내의 중간 표현과 표적 분자 간의 상호 정보에 대한 이론적 분석을 수행했습니다.더욱이, 확산 모델은 중간 표현과 목표 분자 사이의 상호 정보의 하한을 암묵적으로 최대화한다는 것이 이론적으로 증명되었으며, 이는 확산 모델 표현 학습의 능력을 나타냅니다. 그러나 중간 표현과 표적 분자 간의 상호 정보는 단조롭게 감소하는 경향을 보이며 더 큰 시간 단계에서는 0에 가까워지는데, 이는 무질서한 단계의 중간 표현이 효과적인 예측을 지원할 수 없음을 의미합니다. 따라서 직관과 이론 모두 생성 및 예측 작업은 분자가 비교적 질서 있는 상태를 유지하는 더 작은 시간 단계에서만 정렬될 수 있음을 시사합니다.
2단계 확산 생성 메커니즘
위의 분석을 바탕으로,연구팀은 아래 그림에서 보듯이 분자 특성의 예측과 생성을 통합하는 것을 목표로 하는 새로운 2단계 생성 방법을 제안했습니다.
연구자들은 분자 생성 과정을 두 단계로 구분합니다.즉, "분자 핵 생성 단계"와 "분자 성장 단계"입니다.이러한 구분은 물리학의 결정 형성 과정에서 영감을 얻었습니다.
분자 핵 생성 단계에서 분자는 완전히 무질서한 상태에서 골격을 형성하고, 이 골격을 바탕으로 완전한 분자가 성장합니다. 이 두 단계는 "핵 생성 시간"에 의해 구분됩니다. 연구자들은 이 두 단계를 설명하기 위해 분자를 생성하는 새로운 방법을 도입했습니다. 그 중, "핵 생성 시간" 이전에 확산 모델은 점차적으로 분자 좌표를 생성합니다. 핵 생성 후, 모델은 특성과 원자 유형 예측 손실을 최적화하면서 분자 좌표를 계속 조정합니다.
일반적으로 원자 유형과 좌표의 공동 확산을 수행하는 기존 생성 모델과 달리, 이 혁신적인 방법은 좌표의 확산에만 초점을 맞추고 원자 유형을 별도의 예측 작업으로 처리합니다.연구자들은 형성된 분자의 좌표로부터 원자 유형을 종종 추론할 수 있다는 것을 관찰했습니다. 구체적으로, 핵 생성 전에 확산 과정은 좌표를 재구성하는 것을 목표로 합니다. 핵 생성 후 원자 유형과 특성의 예측 손실을 통합하여 통합된 학습 프레임워크를 구축합니다.
UniGEM 훈련 전략
연구진은 기존의 공동 확산 방법과의 비교를 용이하게 하기 위해 EGNN을 네트워크 구조 골격으로 사용하는 E(3) 등변 확산 모델(EDM)을 채택했습니다. 그중 성장 단계는 전체 훈련 과정의 약 1%만을 차지합니다. 표준 확산 학습 절차를 따르고 시간 단계를 균일하게 샘플링하는 경우 예측 작업에 대한 반복 횟수는 전체 학습 과정의 1%에 불과하므로 이 작업에서 모델의 성능이 크게 저하됩니다.따라서 예측 작업에 대한 적절한 훈련을 보장하기 위해 연구자들은 성장 단계 동안 시간 단계를 과도하게 샘플링했습니다.
그러나 연구자들은 과도한 샘플링으로 인해 시간 단계 범위에 걸쳐 훈련이 불균형해질 수 있으며, 이로 인해 생성 프로세스의 품질이 영향을 받을 수 있다는 것을 관찰했습니다. 이 문제를 해결하기 위해 다중 지점 네트워크 아키텍처가 제안되었습니다. 네트워크는 얕은 계층에서는 매개변수를 공유하지만, 더 깊은 계층에서는 두 개의 분기로 나뉘며, 각각은 독립적인 매개변수 세트를 갖습니다.이러한 분기는 훈련의 여러 단계에서 활성화됩니다. 한 분기는 핵 형성 단계에 초점을 맞추고 다른 분기는 성장 단계를 처리합니다.아래 그림과 같습니다. 이 설계는 예측 작업과 생성 작업이 서로 영향을 미치지 않고 효과적으로 학습될 수 있도록 보장합니다.
UniGEM의 추론 과정
UniGEM에서는분자 생성은 역확산 과정을 통해 원자 좌표를 재구성한 다음 생성된 좌표를 기반으로 원자 유형을 예측하여 수행됩니다.그림에서 보듯이. 속성 예측의 경우, 네트워크 입력 시간 단계는 0으로 고정되고 속성 예측 헤드가 사용됩니다. 이 접근 방식은 생성 작업과 예측 작업 모두에 추가적인 계산 오버헤드를 발생시키지 않으며, 총 추론 시간은 기준선과 동일하다는 점에 주목할 가치가 있습니다.
분자 생성 과제의 경우, 연구진은 UniGEM과 기존의 공동 생성 방법 간의 생성 오류 차이도 분석했습니다.첫째, UniGEM에서 원자 유형 예측 손실의 오차는 공동 생성에서의 원자 유형 잡음 제거 생성 손실보다 작은 것으로 관찰되었습니다. 둘째, 공동 생성 과정에서 원자 유형 예측 결과의 진동에 의해 좌표 생성이 영향을 받아 오류가 증가하게 됩니다. 마지막으로, 공동 생성 방법은 더 큰 초기 분포 오류와 이산화 오류를 초래합니다. 이러한 요소들은 UniGEM이 어떻게 탁월한 생성 결과를 달성하는지 설명합니다.
실험 결과: 분자 생성 및 속성 예측 작업 모두에서 기준 모델보다 우수한 성능 발휘
분자 생성: UniGEM, 벤치마크 모델보다 우수한 성능 발휘
연구진은 먼저 EDM 기반 UniGEM을 QM9 및 GEOM-Drugs 데이터 세트의 EDM 변형과 비교했습니다. 아래 그림에서 볼 수 있듯이 UniGEM은 거의 모든 평가 지표에서 기준 모델보다 우수한 성과를 보였습니다. 다른 EDM 변형과 비교했을 때 주목할 점은 다음과 같습니다.UniGEM은 사전 지식에 의존하지 않고 추가적인 자동 인코더 학습도 필요 없기 때문에 훨씬 간단하지만, EDM-Bridge와 GeoLDM보다 성능이 뛰어나 UniGEM의 장점을 잘 보여줍니다.
연구진은 UniGEM이 다양한 생성 알고리즘에 적응하는 데 얼마나 유연한지를 보여주기 위해 UniGEM을 베이지안 흐름 네트워크(BFN)에 적용했습니다. 이는 QM9 데이터 세트에서 좌표와 원자 유형을 함께 생성하는 GeoBFN을 능가하는 결과이며, SOTA 결과를 달성했습니다.
또한 연구진은 샘플링 과정에서 모델 자체의 속성 예측 모듈을 가이드로 사용하여 조건부 생성 모델을 다시 학습할 필요 없이 조건부 생성 작업에서 UniGEM의 성능을 테스트했습니다.
분자 속성 예측: UniGEM은 대부분의 사전 학습 방법을 능가합니다.
연구진은 테스트 세트의 평균 절대 오차(MAE)를 평가 지표로 사용하여 QM9 데이터 세트에 대한 UniGEM 속성 예측의 성능을 평가했습니다. 그림에서 보는 바와 같이,UniGEM은 처음부터 학습한 EGNN보다 훨씬 우수한 성능을 보이며, 이는 통합 모델링의 효과를 보여줍니다.놀랍게도 UniGEM은 대규모 사전 학습 데이터 세트를 추가로 활용했음에도 불구하고 여전히 이러한 최첨단 사전 학습 방법 대부분보다 우수한 성과를 보입니다. 이는 생성 및 예측을 위한 통합 모델의 장점을 강조하는데, 이를 통해 추가 데이터와 사전 학습 단계 없이도 생성 과정에서 분자 표현 학습의 힘을 효과적으로 활용할 수 있습니다.
결론
UniGEM 모델은 분자 생성과 속성 예측 작업을 통합하여 두 작업의 성능을 크게 향상시킵니다. UniGEM의 향상된 성능은 견고한 이론적 분석과 포괄적인 실험 연구에 의해 뒷받침됩니다. 우리는 혁신적인 2단계 생성 공정과 이에 상응하는 모델이 분자 생성 프레임워크 개발을 위한 새로운 패러다임을 제공하고, 더욱 진보된 분자 생성 프레임워크의 개발에 영감을 주어, 보다 구체적인 응용 분야에서 분자 생성에 도움이 될 것이라고 믿습니다.
이 연구는 ATOM 연구실이 주도하고 있습니다. 이 팀은 분자 사전 훈련, 분자 생성, 단백질 구조 예측, 가상 스크리닝 등의 분야에서 더 많은 연구 결과를 보유하고 있으니 주목해 주시기 바랍니다!
ATOM Lab 홈페이지에 오신 것을 환영합니다.
https://atomlab.yanyanlan.com/
저자 소개:
* 란 얀얀은 청화대학교 지능형 산업 연구소(AIR)의 교수입니다. 그녀의 연구 관심사로는 AI4Science, 머신 러닝, 자연어 처리 등이 있습니다.
* 펑 시쿤은 청화대학교 지능형 산업 연구소(AIR)의 박사과정 학생입니다. 그의 연구 관심사에는 표현 학습, 생성 모델, AI4Science가 포함됩니다.
* Yuyan Ni는 중국 과학 아카데미 산하 수학 및 시스템 과학 아카데미(AMSS)의 박사과정생입니다. 그녀의 연구 관심사로는 생성 모델, 표현 학습, AI4Science, 딥 러닝 이론 등이 있습니다.
본 논문의 주요 저자인 Shikun Feng 박사와 Yuyan Ni 박사는 현재 취업 기회를 찾고 있습니다. 관심 있는 친구는 연락하면 됩니다.
* Feng Shikun 이메일: fsk21@mails.tsinghua.edu.cn
* Ni Yuyan의 이메일 주소: niyuyan17@mails.ucas.ac.cn