🤖 AI

NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation

본 논문은 순서 무관 학습 전략을 통해 원자 수준의 치환 불변성을 보장함으로써 QM9 및 GEOM-Drugs 데이터셋에서 3 차원 분자 생성 품질과 속도 측면에서 최첨단 성능을 달성하는 새로운 이웃 안내 자기회귀 집합 트랜스포머인 NEAT 를 소개합니다.

원저자: Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

게시일 2026-05-07

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

복잡한 3D 구조, 예를 들어 막대기 (결합) 로 연결된 다양한 색의 공 (원자) 으로 만든 분자 모델을 만드는 법을 로봇에게 가르치려 한다고 상상해 보세요. 목표는 로봇이 기존 것을 단순히 복사하는 것이 아니라, 처음부터 새로운 유효한 구조를 창조하도록 하는 것입니다.

이 논문은 NEAT(Neighborhood-Guided, Efficient, Autoregressive Set Transformer, 이웃 기반 효율적 자기회귀 집합 트랜스포머) 라는 새로운 로봇 건축가를 소개합니다. 작동 원리를 간단히 설명하면 다음과 같습니다:

문제: "줄서기" 혼란

대부분의 이전 로봇들은 학생들이 교실에 한 명씩 차례로 들어오듯, 분자를 엄격한 줄서기 방식으로 만들려고 했습니다. 그들은 "누가 먼저? 누가 두 번째?"를 결정해야 했습니다.

문제점: 분자에서 원자들은 실제로 "첫 번째"나 "두 번째" 자리가 없습니다. 그들은 3 차원 클러스터 안에서 동시에 존재합니다. 그들을 줄서기에 강제로 넣는 것은 결정들을 특정 순서로 나열하여 눈송이를 설명하려는 것과 같습니다. 이는 임의적이고 혼란스럽습니다. 로봇이 순서를 잘못 잡으면 분자가 깨집니다.
기존 방식: 이전 방법들은 "표준" (canonical) 순서를 만들어 이 문제를 해결하려 했지만, 이는 로봇에게 편향을 주었습니다. 로봇은 분자의 모양보다 목록의 순서를 더 잘 배우게 되었습니다.

해결책: "이웃" 전략

NEAT 는 게임을 바꿉니다. 줄 대신 분자를 군중으로 취급합니다.

작동 방식: 벽돌로 집을 하나씩 짓는다고 상상해 보세요. "목록에서 다음 벽돌은 무엇인가?"라고 묻는 대신, NEAT 는 당신이 지은 현재 벽을 보고 "이 벽 바로 옆에 서 있는 이웃들은 누구인가?"라고 묻습니다.
마법 같은 점: NEAT 는 현재 구조의 "경계"를 봅니다. 현재 구조에 닿아 있는 모든 원자는 다음에 추가될 유효한 후보임을 알고 있습니다. 순서는 중요하지 않으며, 오직 이웃만 중요합니다.
결과: 이로 인해 로봇은 "순열 불변 (permutation invariant)"이 됩니다. 분자를 왼쪽, 오른쪽, 혹은 거꾸로 보더라도 로봇은 같은 이웃을 보고 같은 유효한 구조를 만듭니다.

건축 과정: 한 번에 두 단계

NEAT 가 새로운 원자를 추가하기로 결정할 때, 두 가지 일을 동시에 수행합니다:

무엇인가? 원자의 종류를 추측합니다 (예: "이것은 탄소 공이어야 한다").
어디인가? Flow Matching이라는 특수한 수학 기법을 사용하여 그 원자가 3 차원 공간에서 정확히 어디에 떠 있어야 하는지 파악합니다. 이는 마치 강물이 나뭇잎 (원자) 을 무작위 시작점에서 분자 내의 완벽하고 안정적인 위치로 부드럽게 밀어내는 것과 같습니다.

NEAT 가 특별한 이유

이 논문은 NEAT 가 세 가지 이유로 주요 업그레이드라고 주장합니다:

빠릅니다: 다른 로봇들 (수시간 동안 돌덩이를 조각해 동상을 만드는 것과 같은) 이 수천 개의 노이즈가 많은 단계를 시뮬레이션할 필요가 없기 때문에, NEAT 는 분자를 훨씬 더 빠르게 만듭니다. 논문은 NEAT 가 경쟁사보다 5 배에서 16 배까지 빠르다고 말합니다.
유연합니다: 고정된 줄서기에 의존하지 않기 때문에 "스타터 키트 (부분 분자)"를 주고 작업을 마무리하도록 요청할 수 있습니다. 스타터 키트를 어떻게 전달하든 구조를 완성할 수 있습니다. 다른 로봇들은 종종 스타터 키트가 특정 "줄서기" 규칙과 일치하지 않으면 실패합니다.
정확합니다: 표준 테스트 (QM9 및 GEOM-Drugs 라는 데이터셋) 에서 NEAT 는 화학적으로 유효하고 고유한 분자를 만들어 현재 최고의 기존 방법들과 맞먹거나 능가합니다.

함정 (한계점)

저자들은 NEAT 가 아직 할 수 없는 것에 대해 솔직합니다:

크기 제한: 학습 중에 본 분자만큼 큰 분자만 만들 수 있습니다. 작은 분자만 학습했다면 거대한 분자를 만들어내는 외삽 (extrapolation) 은 할 수 없습니다.
결합 모호성: 이웃을 파악하기 위해 어떤 원자들이 연결되어 있는지 (결합) 를 알고 있어야 합니다. 결합이 불분명한 이상한 결정 구조를 주면 혼란스러워할 수 있습니다.
3D 규칙의 "마법" 부재: 3D 회전 (SE(3)-equivariance) 을 위한 내장된 물리 규칙은 없지만, 저자들은 학습 데이터를 무작위로 회전시키는 것만으로도 잘 작동한다는 것을 발견했습니다.

요약

간단히 말해, NEAT 는 컴퓨터에게 3D 분자를 설계하도록 가르치는 새로운 방법입니다. 로봇에게 엄격하고 인위적인 목록을 따르도록 강요하는 대신, 즉각적인 이웃을 보고 자연스럽게 만들도록 가르칩니다. 이는 더 빠르고, 더 유연하며, 현재 최첨단 방법만큼 정확합니다.

기술 요약: NEAT – 3D 분자 생성을 위한 이웃 유도형, 효율적, 자기회귀 집합 트랜스포머

문제 제기

생리활성 분자를 발견하기 위한 현재의 계산 워크플로우는 가상 스크리닝이나 de novo 생성에 크게 의존합니다. 가상 스크리닝은 화학 라이브러리의 유한한 크기와 증가하는 계산 비용에 의해 제한되는 반면, 생성 모델링은 유망한 대안을 제공합니다. 그러나 3D 분자 구조를 생성하는 것은 고유한 어려움을 수반합니다. 분자는 본질적으로 3 차원적이며, 그 특성은 종종 입체구조 (conformation) 에 의해 결정됩니다.

기존의 최첨단 3D 생성 모델들은 주로 확산 (diffusion) 및 흐름 매칭 (flow-matching) 기법을 활용합니다. 이러한 방법들은 종종 메시지 전달 업데이트에서 2 차 스케일링 문제를 겪으며, 가변적인 분자 크기에 대한 고유한 적응력이 부족합니다 (종종 초기화 시 결정된 고정된 원자 수를 요구함) 그리고 탈노이즈 (denoising) 중 동적인 원자 삽입을 관리하는 데 어려움을 겪습니다.

자기회귀 (Autoregressive) 모델은 가변적인 원자 수를 자연스럽게 지원하고 토큰당 단일 인코더 호출만 필요로 하는 유연한 대안을 제공합니다. 그러나 표준 자기회귀 트랜스포머는 토큰의 순차적 순서를 요구합니다. 분자 그래프에서 원자들은 고유한 순차적 순서를 가지고 있지 않습니다. 이전 연구들은 표준 원자 순서를 사용하여 이를 해결해 왔으나, 이러한 접근 방식들은 원자에 대해 치환 불변성 (permutation-invariant) 이 아닙니다. 이는 다음 토큰 예측을 특정 순서 규칙으로 편향시켜, 모델의 일반화 능력이나 최종 분자의 표준 순서와 일치하지 않을 수 있는 시작 원자 순서를 요구하는 접두어 완성 (prefix completion) 과 같은 작업을 수행하는 능력을 제한합니다. 반면, 치환 불변성을 보장하는 그래프 신경망 (GNN) 접근 방식들은 종종 화학적 유효성이 낮은 분자를 생성합니다.

방법론

저자들은 표준 원자 순서에 의존하지 않고 원자 치환 불변성을 존중하는 3D 분자 생성을 위한 최초의 자기회귀 트랜스포머 모델인 NEAT(Neighborhood-Guided, Efficient, Autoregressive Set Transformer) 를 제안합니다.

핵심 아키텍처

NEAT 는 분자 그래프를 시퀀스가 아닌 원자의 집합으로 취급합니다. 아키텍처는 세 가지 주요 구성 요소로 이루어져 있습니다:

데이터 증강 모듈 (이웃 유도): 순차적 위치 인코딩 대신, 모델은 분자 성장을 모방하는 샘플링 전략을 사용합니다. 분자 그래프가 주어지면 연결된 부분 그래프 (소스 집합) 가 샘플링됩니다. 타겟 집합은 소스 집합의 "경계 노드"(1 홉 이웃) 로 정의됩니다. 이는 원자가 추가된 순서에 관계없이 다음 단계에 대한 유효한 토큰 연속체 집합을 제공합니다.
집합 트랜스포머 트렁크: 소스 집합은 집합 트랜스포머 (양방향 어텐션과 위치 임베딩이 없는 NanoGPT/GPT-2 기반) 를 사용하여 인코딩됩니다. 원자 유형은 선형적으로 임베딩되고, 위치는 푸리에 특징 (Fourier features) 을 사용하여 인코딩됩니다. 가법 풀링 (Additive pooling) 은 이러한 노드 표현들을 집합 수준의 표현 ( $z_n$ ) 으로 집계합니다. 이 설계는 소스 집합 내 원자 순서에 대한 치환 불변성을 보장합니다.
흐름 매칭 헤드: 집합 수준의 표현 $z_n$ $z_{n}$ 은 다음 토큰을 예측하는 데 사용되며, 원자 유형과 위치로 분해됩니다: $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$ $p (a_{n + 1}, x_{n + 1} ∣ z_{n}) = p (a_{n + 1} ∣ z_{n}) \cdot p (x_{n + 1} ∣ a_{n + 1}, z_{n})$ .
- 원자 유형: 선형 레이어와 소프트맥스를 통해 범주형 분포로 모델링됩니다.
- 원자 위치: 조건부 흐름 매칭 (CFM) 을 통해 모델링됩니다. 모델은 무작위 초기 위치 분포를 실제 원자 위치로 운반하는 벡터장을 학습합니다.

학습 전략

이웃 유도: 학습 중 모델은 단일 다음 토큰이 아닌 모든 유효한 경계 노드의 집합에 의해 감독됩니다. 이는 표준화 (canonicalization) 의 필요성을 제거합니다.
교사 강제 (Teacher Forcing): 원자 유형 예측과 위치 복원은 분리됩니다. 모델은 유형을 예측하고, 흐름 매칭 헤드는 해당 유형에 조건부인 위치를 학습합니다.
증강: 오차 누적에 대한 견고성을 향상시키기 위해, 학습 중 소스 집합 원자의 일부에 가우시안 노이즈가 주입되고 위치는 노이즈가 있는 소스 집합에 대해 재중앙화됩니다.
손실 함수: 총 손실은 원자 유형의 교차 엔트로피 손실과 위치의 흐름 매칭 손실의 합입니다. 소스 및 타겟 분포를 결합하기 위해 선형 할당 문제가 해결되어 운송 비용을 최소화하고 경로 교차를 방지합니다.

추론

생성은 무작위 초기 원자로 시작합니다. 각 단계에서 현재 원자 집합이 인코딩되고, 다음 원자 유형이 예측되며, 무작위 초기 위치가 샘플링됩니다. 그런 다음 흐름 매칭 벡터장이 오일러 (Euler) 또는 오일러 - 마루야마 (Euler-Maruyama) 방법을 사용하여 적분되어 최종 3D 좌표를 결정합니다. 이 과정은 정지 토큰이 예측될 때까지 반복됩니다.

주요 기여

치환 불변 자기회귀 생성: NEAT 는 이웃 유도와 집합 트랜스포머를 사용한 새로운 학습 전략을 도입하여 원자 순서에 불변인 3D 분자 생성을 달성함으로써, 자기회귀 모델에서 표준 순서의 한계를 극복합니다.
효율성: 입력을 집합으로 취급하고 토큰당 단일 인코더 호출을 사용함으로써, NEAT 는 확산 기반 메시지 전달의 2 차 스케일링과 일부 흐름 매칭 베이스라인에서 발견되는 원자 생성 단계당 여러 통합 단계의 필요성을 피합니다.
접두어 완성 능력: 순서 중립적인 NEAT 의 특성은 접두어가 특정 표준 순서와 일치할 필요가 없이 고품질의 접두어 완성 (주어진 스캐폴드 주변의 분자 생성) 을 가능하게 하며, 이는 QUETZAL 과 같은 시퀀스 기반 모델이 어려움을 겪는 작업입니다.
최첨단 성능: 이 모델은 자기회귀 및 확산/흐름 매칭 베이스라인 모두에 비해 표준 벤치마크 (QM9 및 GEOM-Drugs) 에서 경쟁력 있거나 우수한 결과를 달성합니다.

실험 결과

이 모델은 QM9(소형 유기 분자) 및 GEOM-Drugs(약물 유사 분자) 데이터셋에서 평가되었습니다.

생성 품질: QM9 에서 NEAT 는 xyz2mol 유효성 및 유일성 측면에서 이전 자기회귀 모델 (G-SchNet, G-SphereNet, Symphony, Geo2Seq) 보다 우수한 성능을 보였습니다. xyz2mol 유효성 95.2% 및 유일성 91.6% 를 달성하여 조회 (lookup) 유효성 및 유일성에서 2 위를 기록했습니다. GEOM 에서 NEAT 는 자기회귀 모델 중 가장 높은 조회 유효성 (99.7%) 과 두 번째로 좋은 xyz2mol 유효성 (93.9%) 및 유일성 (91.6%) 을 달성하여 FlowMol3 와 같은 최첨단 흐름 매칭 모델의 성능과 일치했습니다.
접두어 완성: 100 개의 다양한 접두어에서 1,000 개의 분자를 생성하는 작업에서 NEAT 는 QUETZAL 을 크게 능가했습니다. NEAT 는 97.8% 의 조회 유효성과 92.7% 의 xyz2mol 유효성을 달성한 반면, QUETZAL 은 각각 27.8% 와 28.1% 에 그쳤습니다. 이는 NEAT 가 원자 순서 규칙과 독립적으로 구조적 특징을 일반화할 수 있음을 보여줍니다.
런타임 효율성: NEAT 는 상당한 속도 이점을 보여줍니다. SemlaFlow 보다 약 5 배 빠르고 FlowMol3 보다 16 배 빠릅니다. 최악의 시나리오에서 NEAT 의 복잡도는 $O(BN^3)$ (여기서 $N$ 은 원자 수) 인 반면, 확산/흐름 모델은 $O(BTN^2)$ (여기서 $T$ 는 통합 단계) 로 스케일링됩니다. 약물 유사 분자의 경우 $N$ 이 일반적으로 $T$ 보다 작기 때문에 NEAT 는 실용적인 속도 이점을 제공합니다.

중요성 및 주장

이 논문은 NEAT 가 조건부 3D 분자 생성을 포함한 복잡한 약물 설계 작업에서 트랜스포머 아키텍처를 활용하기 위한 강력한 기반을 제공한다고 주장합니다. 표준 원자 순서에 대한 의존성을 제거함으로써, 이 모델은 이전에는 자기회귀 접근 방식에서 어려웠던 스캐폴드 기반 생성 (접두어 완성) 과 같은 유연한 작업을 가능하게 합니다.

저자들은 현재 아키텍처가 명시적으로 SE(3)-공변 (equivariant) 은 아니지만, 학습 중 무작위 회전 증강을 사용하면 이를 효과적으로 보완하여 예측에서 오직 미미한 변동만 발생한다고 지적합니다. 이 연구는 NEAT 를 SE(3)-공변 트랜스포머와 통합하는 것이 향후 연구의 유망한 방향임을 시사합니다.

저자들은 모델이 고품질의 3D 구조를 생성하지만, 분자 생성은 광범위한 검증, 합성 및 테스트가 필요한 복잡한 파이프라인의 한 단계에 불과하다고 강조합니다. 또한, 이웃 유도를 위한 엣지 정보에 대한 의존성 (무기 결정에 대한 적용 가능성을 제한할 수 있음) 과 학습 중 관찰된 최대 원자 수를 넘어선 외삽의 부재와 같은 한계점도 인정합니다.