Geometry Informed Tokenization of Molecules for Language Model Generation

Essence

Geo2Seq 개요: 정규 표준화 라벨링 순서로 노드를 배열하고 각 노드에 원자 종류와 구면 좌표를 벡터 형태로 채워 수열로 변환

본 논문은 3D 분자 구조를 SE(3)-불변(invariant) 1D 이산 수열로 변환하는 Geo2Seq 토큰화 방법을 제안하여, 언어 모델(LM)이 3D 분자 생성 태스크를 효과적으로 수행할 수 있도록 한다. 기존의 확산 모델 기반 방법론보다 빠르면서도 더 나은 조건부 생성 성능을 달성한다.

Motivation

Known: 언어 모델은 자연언어처리(NLP)에서 뛰어난 성능을 보이며, SMILES나 SELFIES 같은 표현으로 2D 분자 생성에도 성공적으로 적용되었다.
Gap: 3D 분자 생성 문제에 대해 언어 모델을 직접 적용한 연구는 거의 없다. 기존 확산 모델 방식들은 수천 단계의 확산 스텝이 필요하여 생성 시간이 길다.
Why: 3D 기하학적 정보는 분자의 성질(예: 원자별 힘, 에너지)을 결정하는 데 결정적이므로, 2D 그래프만으로는 불충분하다. SE(3) 변환(회전, 병진)에 대해 불변인 표현이 필수적이다.
Approach: 분자의 3D 좌표를 SE(3)-불변 구면 좌표(spherical coordinates)로 변환하고, 정규 표준화 라벨링(canonical labeling)을 통해 순서를 결정하여 1D 이산 수열로 인코딩한다.

Achievement

편극율(Polarizability α) 조건으로 생성된 분자의 시각화

QM9 데이터셋에서 Geo2Seq + Mamba로 생성된 분자의 시각화

생성 품질 및 다양성: Geo2Seq을 다양한 언어 모델(Transformer, Mamba 등 SSM)과 결합하여 화학적으로 유효하고 다양한 3D 분자를 신뢰성 있게 생성할 수 있음을 입증
조건부 생성 성능: 조건부 생성(conditional generation) 태스크에서 강력한 확산 모델 기저선(EDM 등)을 큰 차이로 초과 달성
효율성: 확산 모델 대비 훨씬 빠른 생성 속도(자동회귀 디코딩)를 유지하면서 더 나은 성능 제공

How

Geo2Seq의 핵심 단계: 정규 표준화 라벨링, 구면 좌표 표현, 토큰화, 수열 생성

정규 표준화 라벨링(Canonical Labeling): 그래프 동형(isomorphism)을 제외한 정보 손실 없이 분자 그래프의 노드를 고정된 순서로 배열
구면 좌표 표현: 각 원자를 [z_i, d_i, θ_i, φ_i]로 표현 (z_i: 원자번호, d_i: 거리, θ_i, φ_i: 각도)
- 이전 원자로부터의 거리와 각도만 사용하므로 자동으로 SE(3)-불변 성질 획득
- 전역 기준 프레임(global equivariant frame)을 설정하여 회전과 평진에 대해 불변
토큰화 및 수열화: 각 원자의 정보를 연속 벡터로 인코딩하고, 양자화(quantization)를 통해 이산 토큰으로 변환
조건부 생성 확장: 양자 성질(quantum properties) 값을 수열에 포함시켜 LM의 장거리 문맥 의존성(long-context correlation) 활용
생성: 자동회귀 디코딩으로 토큰을 순차적으로 샘플링하고, Top-K 및 온도(temperature) 하이퍼파라미터로 다양성 제어

Originality

SE(3)-불변 토큰화 설계: 분자 기하학을 직접 1D 수열로 인코딩하면서 회전/병진 불변성을 보장하는 방법론의 참신함
모델-불가지론적(Model-agnostic) 접근: Geo2Seq이 입력 데이터에만 작동하므로, 다양한 LM 아키텍처(Transformer, SSM/Mamba 등)와 조합 가능
3D 분자 생성에 순수 언어 모델 적용: 기존 연구들이 GNN이나 기하학적 모듈을 병렬 또는 추가로 포함한 반면, 순수 LM 기반 접근의 가능성 입증
조건부 생성의 우수성: 장거리 의존성 학습에 우수한 LM 특성을 활용하여 양자 성질 조건부 생성에서 확산 모델 초과 달성

Limitation & Further Study

양자화 손실: 연속 좌표를 이산 토큰으로 변환하는 과정에서 정밀도 손실이 발생하며, 최적 양자화 레벨에 대한 상세한 분석 부족
수열 길이: 분자 크기가 커질수록 수열 길이가 증가하여 언어 모델의 컨텍스트 윈도우 제약에 직면할 수 있음
기하학적 제약 미실장: 물리적으로 타당한 분자 기하학(예: 결합 각도 범위, 입체 장애)에 대한 명시적 제약이 토큰화 단계에 포함되지 않음
정규 표준화의 계산 복잡도: 그래프 동형 판정의 NP-hardness 특성상, 매우 큰 분자에 대해 정규 표준화 계산이 병목이 될 수 있음
후속 연구 방향:
- 더 큰 분자 데이터셋(GEOM 등)에 대한 확장성 검증
- 에너지, 반응성 등 추가 양자 성질 조건부 생성 평가
- 다른 기하학적 표현(예: 상대 좌표) 비교 분석

Evaluation

총평: 본 논문은 언어 모델을 3D 분자 생성에 효과적으로 적용하기 위한 창의적인 토큰화 방법을 제안하며, SE(3)-불변성을 엄밀히 보장하면서도 조건부 생성에서 확산 모델을 초과하는 성능을 달성한다. 모델-불가지론적 설계로 확장성도 우수하나, 수치 정밀도와 계산 복잡도 측면의 실용적 제약에 대한 더 깊은 논의가 필요하다.