Toward Better Geometric Representations for Molecule Generative Models

기술 요약: LENSEs – 비부드러운 구조 인코딩을 위한 잠재적 향상

문제 정의
기하학적 표현 조건부 생성 (GeoRCG) 은 3D 분자 생성의 효율적인 패러다임으로 부상하여, 분자 표현의 모델링과 3D 구조 생성을 분리했습니다. 이 프레임워크에서 사전 훈련된 분자 인코더 (예: UniMol, Frad) 는 잠재 공간을 정의하고, 표현 확산 모델 (RDM) 이 이 공간에서 분자 수준의 벡터를 샘플링하며, 생성기는 이 벡터를 조건으로 3D 구조를 재구성합니다.

그러나 GeoRCG 의 효과는 사전 훈련된 인코더가 제공하는 잠재 공간의 품질에 근본적으로 제한받습니다. 저자들은 두 가지 중요한 병목 현상을 식별했습니다:

1. 비부드러운 표현 다양체 (Non-smooth Representation Manifolds): 좌표 제거를 통한 속성 예측을 위해 주로 최적화된 사전 훈련된 인코더는 높은 Lipschitz 상수를 가진 표현을 생성합니다. 분자 구조의 작은 섭동이 잠재 공간에서 불균형적으로 큰 변화를 초래합니다. 이러한 비부드러움은 RDM 이 분포를 모델링하기 어렵게 만들어, 종종 수동으로 조정된 노이즈 주입과 같은 보상 메커니즘을 필요로 합니다.
2. 미활용된 다중 수준 의미론: 사전 훈련된 인코더는 얕은 층이 국소적 구조 모티프 (예: 기능기) 를 포착하고 깊은 층이 전역적 의미론을 인코딩하는 계층적 특징을 보입니다. 기존 GeoRCG 파이프라인은 일반적으로 최종 층 표현만 활용하여 이전 층에 보존된 풍부한 의미론 정보를 폐기합니다. 결과적으로 생성기는 이러한 의미론을 처음부터 다시 학습해야 하므로 최적화 난이도가 증가합니다.

방법론: LENSEs
이러한 문제를 해결하기 위해 저자들은 생성 훈련 중 조건부 잠재 공간을 능동적으로 정제하는 LENSEs(비부드러운 구조 인코딩을 위한 잠재적 향상) 를 제안합니다. 고정된 인코더에 수동적으로 의존하는 대신, LENSEs 는 세 가지 상호 보완적인 메커니즘을 도입합니다:

1. 학습 가능한 층 풀링을 갖춘 표현 헤드:
   인코더의 단일 결정론적 출력을 사용하는 대신, LENSEs 는 다중 층 특징을 집계합니다. 이는 다양한 인코더 깊이 ($H^{(l)}$) 의 표현에 가중치를 부여하는 학습 가능한 층 풀링 메커니즘을 사용합니다. 이러한 풀링된 특징은 투영 헤드를 통해 대각선 가우시안 사후 분포 ($\mu, \sigma$) 의 매개변수로 매핑됩니다. 잠재 코드 $z$는 재매개변수화 트릭을 사용하여 샘플링됩니다. 이 헤드는 생성기와 함께 훈련되어 의미론적으로 풍부하고 생성에 적합한 표현을 추출하며, 컴팩트한 잠재 공간을 보장하기 위해 KL-발산 손실로 정규화됩니다.

2. 분자 지각 손실 ($L_{PERC}$):
   컴퓨터 비전의 지각 손실에 영감을 받아, 이 손실은 제거된 (생성된) 분자의 고수준 의미론적 표현을 고정된 인코더의 특징 공간 내에서 깨끗한 실제 지상 진실 (ground-truth) 분자의 표현과 정렬시킵니다. 생성된 분자와 깨끗한 분자를 모두 인코더에 통과시켜 풀링된 그래프 수준 표현 간의 거리를 최소화함으로써, 생성기는 저수준 좌표 오류를 최소화하는 것뿐만 아니라 의미 있는 구조적 패턴 (예: 기능기) 을 보존하도록 장려됩니다. 실제 지상 진실 표현에는 스톱-그레이디언트 연산이 적용됩니다.

3. 표현 정렬 (REPA) 손실 ($L_{REPA}$):
   생성기 은닉 상태가 자연스럽게 인코더 특징과 정렬된다는 관찰에 영감을 받아, 이 손실은 생성기의 노드 수준 특징을 사전 훈련된 인코더의 특징과 명시적으로 정렬합니다. 학습 가능한 프로젝터가 생성기 특징을 인코더의 특징 공간으로 매핑하고, 대응하는 노드 간의 코사인 유사도를 최대화합니다. 이는 사전 훈련과 생성 간의 의미론적 격차를 줄여 훈련을 가속화하고 화학적으로 의미 있는 지식 획득을 개선합니다.

총 훈련 목적 함수는 기본 생성 손실 ($L_{gen}$) 과 KL 손실, 지각 손실, REPA 손실을 결합합니다.

주요 결과
저자들은 GEOM-DRUG 및 QM9 데이터셋에서 LENSEs 를 평가하여 GeoRCG, CanonFlow, SemlaFlow 를 포함한 최첨단 베이스라인과 비교했습니다.

• 생성 품질: 까다로운 GEOM-DRUG 데이터셋에서 LENSEs 는 최첨단 성능을 달성했습니다. 97.28% 유효성과 98.51% 분자 안정성을 기록하여 이전 최고 성능인 GeoRCG(유효성 95.3%, 안정성 97.20%) 와 CanonFlow(유효성 95.9%) 를 능가했습니다. 또한 가장 낮은 에너지 및 변형 지표를 달성하여 우수한 물리적 타당성을 나타냈습니다.
• 표현 부드러움: Lipschitz 상수 분석은 정제된 잠재 공간이 원시 인코더 공간보다 훨씬 더 부드러움을 보여줍니다. GEOM-DRUG 에서 경험적 Lipschitz 상수는 4.6 배 감소하여 표현 다양체가 확산 모델링에 더 적합함을 확인했습니다.
• 표현 풍부함: 표현의 유효 차원은 원시 인코더의 65.32 에서 LENSEs 의 88.42 로 증가하여 차원 붕괴가 감소하고 더 풍부하며 구조화된 잠재 공간이 됨을 나타냅니다.
• 하류 활용도: 고정된 인코더와 표현 헤드를 사용하여 QM9 속성에 대한 하류 회귀 작업에 사용될 때, LENSEs 는 모든 속성에서 원시 인코더를 능가하여 생성 훈련 과정 자체가 분자 표현을 정제하기 위한 효과적인 사전 훈련 패러다임으로 작용함을 입증했습니다.
• 샘플링 효율성: LENSEs 는 베이스라인에 비해 더 적은 생성 단계로 높은 유효성과 안정성을 달성하여 향상된 샘플링 효율성을 보여줍니다.

의의 및 주장
이 논문은 고품질 3D 분자 생성이 생성기 아키텍처뿐만 아니라 조건부 표현의 품질에 결정적으로 의존한다고 주장합니다. LENSEs 는 사전 훈련된 인코더의 잠재 공간이 더 부드러우며, 의미론적으로 더 유익하고, 생성 목표와 더 잘 정렬되도록 체계적으로 개선될 수 있음을 보여줍니다.

저자들은 그들의 작업이 생성 품질과 표현 품질 사이의 양방향 관계를 확립한다고 주장합니다. 구체적으로, 정렬 목적을 가진 생성적 훈련은 분자 인코더를 위한 잠재적인 사전 훈련 패러다임으로 작용할 수 있으며, 표현을 더 예측 가능하고 하류 작업에 유용하도록 정제할 수 있다고 제안합니다. 이 작업은 유용한 분자 표현이 본질적으로 생성에 친화적이지 않으며, 생성 모델링에서 그 잠재력을 완전히 발휘하기 위해 명시적인 정제가 필요함을 강조합니다.