ICML 2025에 선정된 Meta/Cambridge/MIT는 주기적 및 비주기적 원자 시스템의 통합 생성을 최초로 실현하는 전원자 확산 변환기 프레임워크를 제안했습니다.
오늘날 과학 연구 및 산업 응용 분야의 최전선에서 원자 시스템의 3차원 구조에 대한 생성적 모델링은 파괴적인 잠재력을 보여주고 있으며, 새로운 분자 및 재료의 역설계 환경을 완전히 바꿀 것으로 기대됩니다. 정밀한 구조 예측부터 유연한 조건부 생성에 이르기까지, 최첨단 확산 모델과 유동 매칭 모델은 생체 분자 분석, 신소재 연구 개발, 구조 기반 약물 설계와 같은 핵심 과제에서 활용되어 연구자들이 기술적 병목 현상을 극복하는 핵심 도구로 자리 잡았습니다.
그러나 이 붐비는 분야 뒤에는기술적 도약을 항상 제한해 온 주요 문제는 기존 모델이 시스템 전반에 걸쳐 보편성이 부족하다는 것입니다.모든 원자계는 3차원 구조와 상호작용을 결정하는 데 동일한 물리적 원리를 따르지만, 소분자, 생체분자, 결정 및 이들의 복합 시스템에 대한 모델링은 오랫동안 "분할 정복"의 상태에 있었습니다. 대부분의 확산 모델은 특정 시스템의 고유한 특성에 크게 의존하며, 범주형 데이터(예: 원자 유형)와 연속형 데이터(예: 3차원 좌표)가 얽혀 있는 복잡한 제품 다양체에 대한 다중 모드 생성을 요구하기 때문에 서로 다른 시스템 간의 모델 호환성이 어렵습니다.
특정 시나리오를 예로 들어 보겠습니다. 소분자의 신생(de novo) 생성은 원자 유형(분류)과 3차원 좌표(연속)라는 두 가지 독립적인 확산 과정으로 분리되어야 합니다. 잡음 제거 모델은 두 과정의 공진화 법칙을 학습해야 하지만, 중간 상태의 왜곡으로 인해 샘플링 효율이 저하되는 경우가 많습니다. 생물학적 분자 모델링은 회전 다양체를 추가로 도입해야 하며 원자 그룹을 강체로 취급합니다. 결정 및 재료의 확산 과정은 주기적 특성과 호환되어야 하며 원자 유형, 분수 좌표, 격자 매개변수와 같은 다차원 매개변수로 구성된 조인트 다양체에서 실행되어야 합니다. 이러한 차이로 인해 시스템 간 통합 모델링은 해당 분야에서 오랫동안 해결되지 않은 과제입니다.
이러한 맥락에서,메타 기초 인공지능 연구소(FAIR), 케임브리지 대학, 매사추세츠 공과대학의 공동 연구팀은 획기적인 솔루션인 전원자 확산 변압기(ADiT)를 제안했습니다.
Transformer를 기반으로 하는 통합된 잠재 확산 프레임워크로서,ADiT의 핵심 장점은 주기 시스템과 비주기 시스템 간의 모델링 장벽을 허물 수 있다는 것입니다. 모든 원자의 통합 잠재 표현과 변압기 잠재 확산이라는 두 가지 주요 혁신을 통해 단일 모델로 분자와 결정을 생성할 수 있습니다.이 설계는 유도 편향을 거의 발생시키지 않아, 자동 인코더 및 확산 모델을 기존의 등변 확산 모델보다 학습 및 추론 측면에서 훨씬 효율적으로 만듭니다. 동일한 하드웨어 조건에서 10,000개의 샘플을 생성하는 데 걸리는 시간이 2.5시간에서 20분 미만으로 단축됩니다. 더 중요한 것은, 모델 매개변수를 5억으로 확장했을 때 성능이 예측 가능한 선형적 향상을 보인다는 것입니다. 이 기능은 생성 화학의 보편적인 기본 모델을 구축하는 데 중요한 토대를 마련하며, 원자 시스템 모델링의 보편성과 대규모 적용에 있어 이정표를 세웁니다.
관련 연구 결과는 "모든 원자 확산 변환기: 분자 및 재료의 통합 생성 모델링"이라는 제목으로 ICML 2025에 선정되었습니다.
연구 하이라이트:
* ADiT는 주기적 물질과 비주기적 분자 시스템에 대한 생성 모델을 통합한 최초의 제품입니다.
* ADiT는 모든 원자의 통합 잠재 표현에 의존하고 잠재 확산을 위해 변압기를 사용하는데, 이는 생성 프로세스를 효과적으로 단순화하고 유도 편향이 거의 없습니다.
* ADiT는 뛰어난 확장성과 효율성을 가지고 있으며, 학습 및 추론 속도가 등변 확산 모델보다 훨씬 뛰어납니다.
서류 주소:
더 많은 AI 프런티어 논문:
https://go.hyper.ai/owxf6
데이터 세트: 주기적 데이터부터 비주기적 데이터까지 다양한 분야의 실험 데이터를 포함합니다.
이 연구에서 연구팀은 먼저 실험을 수행하기 위해 몇 가지 대표적인 데이터 세트를 선택했습니다.
* MP20 데이터 세트,최대 단위 셀에 원자 20개를 포함하여 89가지의 다양한 원소를 포함하는 Materials Project의 45,231개 준안정 결정 구조가 포함되어 있어 주기적 물질 시스템을 잘 표현할 수 있습니다.
* QM9 데이터 세트,이는 최대 9개의 무거운 원자(C, N, O, F)와 수소 원자를 포함하는 130,000개의 안정적인 작은 유기 분자로 구성되며, 비주기적 분자 시스템의 전형적인 대표입니다.
* GEOM-DRUGS 데이터 세트,최대 180개의 원자를 포함하는 430,000개의 대형 유기 분자
* QMOF 데이터 세트,14,000개의 금속 유기 골격 구조를 포함합니다.
안에,MP20과 QM9는 서로 다른 유형의 원자 시스템에 해당합니다.이는 주기적 및 비주기적 시스템에 대한 모델의 공동 학습을 위한 기반을 제공하며, 연구팀은 다른 모델과의 공정성을 보장하기 위해 기존 연구 방법에 따라 데이터를 나누었습니다. GEOM-DRUGS와 QMOF는 모델 테스트의 범위를 더욱 확장하고 모델의 일반화 능력을 보다 포괄적으로 테스트할 수 있습니다.
ADiT: 듀얼 코어 아이디어 기반 통합 원자 시스템 생성 모델 구축
잠재 확산 모델로서 ADiT의 핵심 설계는 주기적 및 비주기적 원자 시스템의 통합 생성 모델링을 달성하기 위한 두 가지 핵심 아이디어를 중심으로 이루어집니다.
첫 번째 핵심 아이디어는 모든 원자의 통합된 잠재 표현입니다.연구팀은 주기적 원자계와 비주기적 원자계를 모두 3차원 공간의 원자 집합으로 보고, 각 원자의 범주형 속성(예: 원자 유형)과 연속형 속성(예: 3차원 좌표)을 포함하는 통합 표현을 개발했습니다. 전체 원자 재구성을 위해 변분 자동 인코더(VAE)를 학습시킴으로써,인코더는 분자와 결정을 공유 잠재 공간에 내장할 수 있습니다.이는 다양한 유형의 원자 시스템을 통합적으로 처리하기 위한 기본적인 프레임워크를 제공합니다.
두 번째 핵심 아이디어는 잠재 확산을 위해 변압기를 사용하는 것입니다.연구팀은 VAE 인코더가 구축한 잠재 공간에 확산 변환기(DiT)를 도입하여 생성 모델링을 수행했습니다. 추론 과정에서분류자 없는 부트스트래핑 기술의 도움으로 새로운 잠재 변수를 샘플링할 수 있습니다.이러한 잠재 변수는 VAE 디코더를 통해 유효한 분자나 결정으로 재구성될 수 있으며, 이를 통해 잠재 공간에서 실제 원자계로의 변환이 완료됩니다.
ADiT의 실험 방법은 이 두 가지 핵심 아이디어를 바탕으로 두 단계로 나뉘며 체계적으로 진행됩니다.
첫 번째 단계에서 연구진은 재구성을 위한 자동 인코더를 구축했습니다.VAE를 통해 분자와 물질의 전체 원자적 표현이 공동으로 재구성되어 공유 잠재 공간을 학습하고 구성합니다. 이는 서로 다른 원자 시스템의 통합 모델링을 위한 전제 조건이며 후속 생성 프로세스의 기반을 마련합니다.
두 번째 단계에서 연구자들은 잠재 확산 생성 모델을 구축했습니다.DiT는 잠재 공간에서 새로운 샘플을 생성하는 데 사용되며, 이 샘플들은 분류기의 안내 없이 유효한 분자 또는 결정으로 디코딩됩니다. 이 잠재 확산 설계의 중요한 장점은 분류 및 연속 속성 처리의 복잡성이 자동 인코더로 이전되어 잠재 공간에서의 생성 과정을 더 간단하고 확장 가능하게 만들어 다양한 원자 시스템을 처리할 때 모델의 효율성과 적응성을 효과적으로 향상시킨다는 것입니다.
ADiT, 결정 및 분자 생성 분야에서 선도적 성과 달성
연구팀은 ADiT의 성능상의 이점을 최대한 강조하기 위해 여러 유형의 기준 모델을 선택하여 비교를 진행했습니다.결정 형성 분야에서비교 대상에는 다중 모드 제품 매니폴드를 기반으로 하는 CDVAE, DiffCSP, FlowMM 및 기타 등가 확산 및 흐름 매칭 모델과 비등가 확산 모델 UniMat 및 2단계 프레임워크 FlowLLM이 포함됩니다.분자생성 분야에서는ADiT는 등가 확산 모델(Equivariant Diffusion Model), GeoLDM, Symphony 등 다른 모델들과 비교됩니다. 이러한 분야의 고급 기준 모델들과의 체계적인 비교를 통해 ADiT의 성능 이점이 명확하게 입증됩니다.
구체적인 실험 결과로부터,ADiT는 결정 및 분자 생성 작업 모두에서 SOTA 수준을 달성합니다.결정 생성 측면에서 ADiT로 생성된 결정은 효과성, 안정성, 고유성, 그리고 신규성과 같은 주요 지표에서 우수한 성과를 보였습니다. 분자 생성 과제에서 ADiT는 10,000개의 샘플링된 분자 중 효과성과 고유성 측면에서 상위권을 차지했습니다.
ADiT의 관절 훈련 메커니즘은 또한 상당한 성능 향상을 가져옵니다. 실험 데이터는QM9와 MP20 데이터 세트로 학습한 ADiT는 하나의 데이터 세트로만 학습한 버전보다 물질 및 분자 생성 작업 모두에서 더 우수한 성능을 보였습니다.
모델 크기의 확장은 ADiT의 성능 향상을 예측할 수 있습니다. 아래 그림에서 볼 수 있듯이, DiT 디노이저 매개변수의 수가 3,200만 개(ADiT-S, 파란색)에서 1억 3,000만 개(ADiT-B, 주황색), 그리고 4억 5,000만 개(ADiT-L, 녹색)로 증가함에 따라, 약 13만 개의 샘플로 구성된 중간 크기의 데이터셋에서도 확산 학습 손실은 지속적으로 감소하고 효과비는 꾸준히 증가하여 유의미한 규모 효과를 나타냅니다. 모델 크기와 성능 간의 이러한 강력한 상관관계는 모델 매개변수와 데이터 볼륨의 확장을 통해 ADiT가 추가적인 혁신을 달성할 수 있을 것으로 기대됨을 시사합니다.
효율성 측면에서 ADiT는 등가 확산 모델에 비해 상당한 속도 이점을 보여줍니다.아래 그림에서 볼 수 있듯이, NVIDIA V100 GPU에서 10,000개의 샘플을 생성할 때, 표준 트랜스포머 기반 ADiT는 계산 집약적인 등변 네트워크를 사용하는 FlowMM 및 GeoLDM보다 통합 단계 측면에서 훨씬 더 뛰어난 확장성을 보입니다. ADiT-B는 등변 기준선보다 매개변수 크기가 100배 더 크지만, 추론 속도는 여전히 더 빠르며, 이는 트랜스포머 아키텍처가 확장성 측면에서 얼마나 우수한지를 보여줍니다.
또한, 대규모 시스템에서 ADiT의 확장성이 입증되었습니다. 최대 180개의 원자를 포함하는 430,000개의 분자를 포함하는 GEOM-DRUGS 분자 데이터 세트에서,ADiT은 효율성과 PoseBusters 지표 측면에서 최첨단 등가 확산 및 흐름 매칭 모델과 비슷한 성능을 보입니다.ADiT는 표준 Transformer 아키텍처를 기반으로 하며, 분자 유도 바이어스를 거의 도입하지 않고, 원자 결합에 대한 명시적 예측을 요구하지 않지만, 등가 모델과 비슷한 성능을 달성할 수 있다는 점이 주목할 만합니다. 이는 ADiT의 설계가 다재다능하고 폭넓게 적용될 수 있음을 다시 한번 입증합니다.
산업계와 연구계는 원자계의 3차원 구조 생성에 있어 획기적인 혁신을 공동으로 촉진합니다.
실제로, 원자계의 3차원 구조를 생성적으로 모델링하는 최첨단 연구 분야에서 학계와 기업은 끊임없이 노력하여 많은 주목할 만한 성과를 달성했습니다.
학계에서캘리포니아 대학교 버클리, 마이크로소프트 연구소, 제넨텍의 연구팀은 PLAID라는 다중 모드 단백질 생성 방법을 출시했습니다.이 방법은 사전 훈련된 가중치의 구조적 정보를 현명하게 활용하여 DiT로 노이즈 제거 작업을 수행하며, 다양한 길이의 단백질에 대한 구조적 품질과 다양성 분석에서 다른 벤치마크 방법보다 뛰어난 성능을 보여줍니다.
기업계 역시 이 분야를 적극적으로 탐색하고 혁신을 통해 개발을 추진하고 있습니다.중국의 생성적 AI 단백질 설계 혁신 기업 BioGeo가 세계 최초의 전방위 단백질 기반 모델인 GeoFlow V2를 출시했습니다.단백질 구조 예측 및 설계 작업을 단번에 해결하기 위해 통합 원자 확산 모델 아키텍처가 구축되었습니다. 항체 및 항원-항체 복합체 구조 예측 측면에서 GeoFlow V2는 탁월한 정확도와 속도로 유사 제품보다 앞서 있습니다. ByteDance에서 출시한 Seedance 1.0은 빠르고 효율적인 AI 비디오 생성을 위해 가변 자동 인코더와 확산 변환기를 결합한 기술 솔루션을 사용하는 차별화된 접근 방식을 취합니다. 이러한 속도 이점은 실시간 생성 및 대화형 애플리케이션에 새로운 지평을 열어주며, 이는 상용 애플리케이션 분야에서 폭넓은 전망을 제시합니다.
학계의 이러한 과학적 획기적 발전과 기업계의 혁신적인 관행은 원자 시스템 3차원 구조 생성 모델링 분야의 발전을 함께 촉진하고 있습니다. 기술의 지속적인 발전과 함께, 이 분야는 신소재 연구 개발 및 약물 설계 등 여러 측면에서 더욱 중요한 역할을 수행하여 세계적인 과학 문제 및 산업적 과제 해결에 강력한 기반을 제공할 것입니다.
참고문헌:
1.https://mp.weixin.qq.com/s/oF3-y7z8u1XpEtjd4q1u4w
2.https://mp.weixin.qq.com/s/tK0-