기술 요약: SymDrift

문제 정의

분자와 같은 물리 시스템의 생성 모델링은 회전 ($\mathbb{R}^3$ 내) 과 동일 원자의 치환과 같은 전역 대칭성 하에서 불변인 분포를 학습해야 합니다. 공변적 (equivariant) 확산 및 흐름 매칭 모델은 이러한 불변성을 효과적으로 통합할 수 있지만, 일반적으로 비용이 많이 드는 다단계 샘플링 절차에 의존하므로 가상 약물 스크리닝과 같은 고처리량 응용 분야에는 계산적으로 불가능합니다.

최근 "드리프트 (drifting)" 모델들이 분포 매칭 문제로 생성을 공식화하여, 목표 분포와 일치하는 고정점을 학습하는 벡터장을 학습하는 효율적인 원샷 대안으로 등장했습니다. 그러나 저자들은 대칭 도메인에 드리프트 모델을 적용할 때 중요한 한계를 발견했습니다: 공변적 생성기는 일반적으로 대칭화된 목표 분포에서 얻은 것과 동일한 드리프트 장을 생성하지 않습니다.

구체적으로, 대칭군 $G$ 하에서 불변이 아닌 유한한 경험적 분포 $p$로 훈련할 때, 학습된 드리프트 장은 각 동치 클래스의 임의의 대표자에 편향됩니다. 이론적 분석 (정리 1) 은 공변적 생성기에서 유도된 집계된 벡터장 $\hat{V}_p(x)$가 완전히 대칭화된 분포가 유도하는 벡터장 $V_{p_G}(x)$와 수학적으로 구별됨을 보여줍니다. 이 불일치는 모델이 대칭화된 목표와 관련된 올바른 드리프트를 복구하는 것을 방해하여, 비용이 많이 드는 데이터 증강 (대칭화) 을 수행하지 않는 한 훈련 붕괴나 비최적 생성으로 이어집니다.

방법론: SymDrift

명시적 데이터 증강의 계산 비용을 들이지 않고 이 대칭성 특화 과제를 해결하기 위해, 저자들은 드리프트 장 자체를 대칭 인식적으로 만드는 SymDrift 프레임워크를 제안합니다. 이 접근법은 대칭화된 목표 분포를 근사하는 드리프트를 구축하기 위해 두 가지 상호 보완적인 전략을 도입합니다:

1. 명시적 정렬을 통한 대칭화 드리프트

이 전략은 좌표 공간에서 작동합니다. 생성된 샘플 $x$와 목표 샘플 $y$에 대해, 거리 $\|x - gy\|$를 최소화하는 최적의 군 변환 $g^*(x, y)$를 명시적으로 찾습니다.

• 메커니즘: 드리프트 장은 표준 쌍 $(x, y)$를 최적 정렬 쌍 $(x, g^*(x, y)y)$로 대체하여 구축됩니다.
• 구현: 회전 정렬은 Kabsch 알고리즘을 사용하여 효율적으로 계산됩니다. 그러나 동일 원자의 최적 치환을 결정하는 것은 여전히 계산적으로 어렵습니다 (계승 복잡도). 저자들은 반복적 접근법 (할당 알고리즘과 Kabsch 정렬을 번갈아 수행) 이 실용적인 근사치를 제공하지만, 지역 최소값과 높은 오버헤드에 시달릴 수 있다고 지적합니다.

2. G-불변 임베딩 공간에서의 드리프트

명시적 정렬의 계산 병목 현상을 우회하기 위해, 저자들은 고정된 $G$-불변 임베딩 공간에서 드리프트를 계산하는 것을 제안합니다.

• 임베딩 구성: 그들은 유형 조건부 쌍거리에 기반한 비모수적 임베딩 $\phi(x)$를 정의합니다. 분자 구성의 경우, 원자 쌍의 상호작용 유형 (예: C-C, C-H) 에 따라 쌍거리를 그룹화하고 각 그룹 내에서 정렬합니다.
• 특성: 이 임베딩은 전역 회전 ($SO(3)$) 과 동일 원자의 치환 ($\hat{S}_N$) 하에서 불변하도록 구성됩니다. 입력을 이 공간으로 매핑함으로써, 커널 유사도 측정은 물리 시스템의 대칭성을 본질적으로 존중하여 드리프트 계산 중 명시적 정렬 없이도 모호성을 제거합니다.
• 유연성: 이 프레임워크는 특정 임베딩에 국한되지 않으며, 모든 미분 가능한 불변 임베딩을 활용할 수 있습니다.

훈련 목적 함수는 드리프트 장의 고정점을 정의하는 타겟과 생성기의 예측 사이의 거리를 최소화하며, 타겟을 고정된 참조로 취급하기 위해 stop-gradient 연산을 사용합니다.

주요 기여

1. SymDrift 프레임워크: 드리프트 장의 설계에 대칭성을 명시적으로 통합한, 고도로 대칭적인 분포의 원샷 생성 모델링을 위한 새로운 프레임워크.
2. 이론적 통찰: 공변적 생성기만으로는 드리프트 모델이 대칭화된 목표 분포를 복구하기에 부족함을 보여주었으며, 이는 드리프트 계산에서 명시적 대칭 처리가 필요함을 시사합니다.
3. 이중 전략: 대칭 인식을 달성하기 위한 두 가지 상호 보완적 방법의 개발: (i) 명시적 좌표 공간 정렬 및 (ii) 불변 임베딩 공간.
4. 효율성: 높은 정확도를 유지하면서 원샷 추론을 가능하게 함으로써 계산 오버헤드를 크게 감소시킴.

실험 결과

저자들은 분자 컨포머 생성 및 전이 상태 (TS) 생성이라는 두 가지 표준 벤치마크에서 SymDrift 를 평가했습니다.

• 분자 컨포머 생성 (GEOM-QM9):

  • SymDrift 는 정확도 지표 (AMR) 와 커버리지 (COV) 에서 원샷 추론 베이스라인 중 최첨단 성능을 달성했습니다.
  • GeoMol 및 AvgFlow 와 같은 기존 원샷 방법보다 우수한 성능을 보였습니다.
  • 추가 에너지 모델을 사용하는 EnFlow 보다 다양성은 약간 낮았지만, EnFlow 보다 3.5 배 빠릅니다.
  • 다단계 베이스라인 (예: 50 NFE 의 ET-Flow) 과 비교하여 SymDrift 는 추론 비용을 40 배 감소시키면서 동등한 정확도를 달성했습니다.

• 전이 상태 생성 (RDB7):

  • 더 큰 모델 변형 (18.7M 파라미터) 을 사용하여 SymDrift 는 25 단계 GoFlow 베이스라인과 비교 가능한 RMSD 및 DMAE 지표를 달성했습니다.
  • 원샷 추론은 다단계 GoFlow 베이스라인보다 87 배 빠릅니다.
  • 이는 속도가 중요한 대규모 반응 네트워크 탐색에 대한 방법의 적합성을 보여줍니다.

• 애블레이션 연구:

  • 실험은 대칭 처리 없이 공변적 아키텍처에 기본 좌표 드리프트를 적용하면 완전한 훈련 붕괴 (0% 커버리지) 로 이어짐을 확인했습니다.
  • 대칭성을 통합함에 따라 성능이 점진적으로 향상되었으며, 지역 최적점에 취약한 반복적 정렬 전략보다 불변 임베딩 공간이 가장 좋은 결과를 보였습니다.

중요성 및 주장

이 논문은 SymDrift 가 원샷 생성 모델의 효율성과 대칭 도메인에 필요한 물리적 정확성 사이의 간극을 성공적으로 메웠다고 주장합니다. 표준 드리프트 접근법 내재된 드리프트 장 불일치를 해결함으로써, SymDrift 는 다음과 같은 것을 가능하게 합니다:

• 고처리량 응용: 최대 40 배의 추론 시간 감소는 현재 반복적 샘플링이 병목 현상인 가상 약물 스크리닝 및 대규모 반응 네트워크 탐색과 같은 응용 분야에 이 방법을 유망하게 만듭니다.
• 훈련 효율성 대 추론 속도: 쌍거리 정렬 또는 임베딩 계산의 필요성으로 인해 훈련 비용이 더 많이 들지만, 이는 추론 속도의 한 자릿수 개선을 고려할 때 유리한 트레이드오프로 제시됩니다.
• 일반화 가능성: 이 프레임워크는 $G$-불변 분포가 관련된 모든 생성 작업에 적용 가능하며, 데이터 증강이나 복잡한 다단계 증류에 대한 원칙적인 대안을 제공합니다.

저자들은 제안된 불변 임베딩이 최대 표현력을 갖지 못한다는 점 (연결성 정보 제거) 과 표준 흐름 매칭에 비해 훈련 시간이 증가한다는 점을 지적하며 제한 사항에 대해 겸손하게 언급합니다. 그러나 그들은 고처리량 과학 응용 분야에서 추론 속도 향상이 이러한 비용을 정당화한다고 주장합니다.