기술 요약: 디퍼모픽 최적화 (Diffeomorphic Optimization)

문제 정의

단백질 설계와 같은 도메인에서의 머신러닝 데이터는 고차원 주변 공간(ambient space) 내에 임베딩된 저차원 매니폴드(manifold) 상에 존재하는 경우가 많습니다. 이러한 매니폴드는 일반적으로 암시적이며 알려져 있지 않습니다. 미분 가능한 목적 함수(예: 단백질 안정성, 결합 친화도)를 주변 공간에서 직접 최적화하는 것은 도전적인 과제입니다. 왜냐하면 손실 함수 지형(loss landscape)이 고차원이고, 울퉁불퉁하며, 비볼록(non-convex)하기 때문입니다. 매니폴드의 기하학적 구조를 인지하지 못하는 표준 경사 하강법은 매니폴드에서 벗어나 분포 외(out-of-distribution) 솔루션을 생성하는 문제를 빈번히 일으킵니다. 가이던스(조건부 밀도를 근사함), 보상 미세 조정(재학습 필요), 또는 최적 제어(궤적을 매니폴드 밖으로 끌어낼 수 있는 제어를 주입함)와 같이 기존의 방법들은 계산 비용의 과다, 처리 불가능성, 또는 매니플드 제약 조건을 엄격하게 유지하지 못하는 등의 한계를 가집니다.

방법론

본 논문은 확산(diffusion) 및 플로우 기반 생성 모델의 디퍼모픽(diffeomorphic, 미분 가능하고 역변환 가능한) 성질을 활용하여 데이터 매니폴드 상에서 엄격하게 수행되는 디퍼모픽 최적화 프레임워크를 제 제안합니다.

핵심 개념

확산 및 플로우 모델은 다루기 쉬운 기저 공간 $Z$(일반적으로 가우시안과 같은 단순한 단봉형 분포)로부터 데이터 매니폴드 $X$로 가는 디퍼모픽 사상(map) $g: Z \to X$를 학습합니다. $X$에서 직접 손실 $L: X \to \mathbb{R}$을 최적화하는 대신, 이 방법은 문제를 기저 공간으로 재매개변수화합니다:
$$\tilde{L}(z) = L(g(z))$$
경사 하강법은 기저 변수 $z$에 대해 수행됩니다:
$$z^{(i+1)} = z^{(i)} - \lambda \nabla_z L(g(z^{(i)}))$$
최종 최적화된 점 $x^{(n)} = g(z^{(n)})$은 구성상 데이터 매니폴드 위에 존재함이 보장됩니다.

이론적 기초

저자들은 미분 기하학을 사용하여, 기저 공간에서의 경사 하강법이 $O(\lambda^2)$ 보정 항을 제외하면 데이터 매니폴드 상에서의 **리만 경사 하강법(Riemannian gradient descent)**과 동일함을 입증합니다. 디퍼모픽 사상은 매니폴드의 기하학적 구조를 암시적으로 학습하며, 이는 데이터 공간의 복잡하고 다봉적인 영역을 기저 공간의 더 부드러운 영역으로 효과적으로 재매개변수화합니다. 이를 통해 더 쉬운 모드 전환(mode-switching)과 더 부드러운 최적화 지형을 구현할 수 있습니다.

단백질을 위한 기술적 구현

단백질 구조는 흔히 행렬 리 군(matrix Lie groups), 구체적으로 백본 프레임을 위한 $SE(3)$와 회전을 위한$SO(3)$를 사용하여 표현됩니다. ODE 솔버를 통한 역전파를 가능하게 하기 위해, 본 논문은 두 가지 방법을 도출했습니다:

1. 재용도 변경된 오토그라드(Repurposed Autograd): 표준 행렬 미분을 래핑하여 $SO(3)$상의 리만 경사도를 계산하는 방법입니다. 이는$\nabla L(R) = 2 [\frac{dL}{dR}R^\top]_A$라는 관계식(여기서$[\cdot]_A$는 반대칭 성분을 의미함)을 활용합니다. 이를 통해 기존의 오토그라드 엔진과 원활하게 통합될 수 있습니다.
2. 아드조인트 상태 방법(Adjoint State Method): 행렬 리 군에 대한 아드조인트 상태 방법의 일반화입니다. 이 방법은 경사도 계산을 종단 값 문제(terminal value problem)로 정식화하여, 아드조인트 상태 $A_\tau$에 대한 역시간 ODE를 해결합니다. 이는 메모리 효율적($O(L)$)이지만, 오토그라드 방식에 비해 이산화 오차(discretization error)를 주의 깊게 다뤄야 합니다.

$SE(3)$의 경우, 경사 하강법은 회전과 이동 부분으로 분리되어, 이동에는 표준 유클리드 업데이트를, 회전에는 리 군 지수 사상(exponential map) 업데이트를 적용할 수 있게 합니다.

주요 기여

1. 디퍼모픽 최적화 프레임워크: 생성 모델을 재학습하거나 보조 네트워크를 사용하지 않고도 임의의 미분 가능한 비용 함수를 데이터 매니폴드 상에서 최적화하는 새로운 방법.
2. 이론적 동등성: 기저 공간의 경사 하강법이 매니폴드 상의 리만 경사 하강법과 동등함을 증명하여, 궤적이 매니폴드 상에 머물도록 보장함.
3. 리 군 역전파(Lie Group Backpropagation): $SO(3)$및$SE(3)$ 상의 ODE 솔버를 통한 효율적인, 오토그라드 호환 가능한 역전파 방법(일반화된 아드조인트 상태 방법 포함) 도출.
4. 구현: 이러한 경사도 계산을 위한 효율적인 PyTorch 구현 제공.

실험 결과

저자들은 최첨단 생성 모델(FrameFlow, DiffDock, AlphaFlow)을 사용하여 여러 단백질 설계 작업에 대해 이 방법을 평가했습니다:

• 이차 구조 타겟팅 (FrameFlow): 디퍼모픽 최적화는 타겟 라마찬드란 영역(베타 시트) 내의 잔기가 **91.3%**에 달한 반면, 튜닝된 가이던스는 **63.3%**를 기록했습니다. 이 방법은 매니폴드에 머물면서 모드를 전환하고 백본 구조를 최적화하는 데 성공했습니다.
• 펩타이드 결합 친화도 (DiffDock): 이 방법은 OC-Flow 프레임워크(최적 제어를 사용하는 방식)보다 펩타이드 결합 친화도 최적화 성능이 뛰어났으며, 실행 시간 측면에서 2배 더 빠른 결과를 얻었습니다.
• Rosetta 에너지 완화 (AlphaFlow): Rosetta 에너지를 최소화할 때, 디퍼모픽 최적화 후 표준 Rosetta Relax를 수행하면 PDB 테스트 세트 전반에서 표준 Rosetta Relax 단독 사용 시보다 훨씬 낮은 에너지(수천 단위 차이)를 나타냈습니다. 이러한 개선은 베이스라인의 계산 예산을 늘리더라도 일관되게 나타났습니다.
• 단백질-리간드 도킹: 동일한 계산 예산 하에서 iid 샘플링 대비 Vina 도킹 점수를 유의미하게 개선했습니다.

의의 및 주장

본 논문은 디퍼모식 최적화를 브루트 포스(brute-force) 방식의 샘플링 및 랭킹(현재 단백질 설계의 병목 현상인 실험실 용량 문제로 인해 표준으로 사용됨)에 대한 표적화된 대안으로 제시합니다. 강화 학습이나 보상 미세 조정과 달리, 이 방법은 재학습을 요구하지 않습니다. 가이던스와 달리, 계산이 불가능한 조건부 밀도를 근사하거나 하이퍼파라미터가 많이 필요한 스케줄에 의존하지 않습니다.

저자들은 생성 모델이 학습한 암시적 기하학을 활용함으로써, 디퍼모픽 최적화가 구조 완화, 이차 구조 수정, 결합 친화도 최적화와 같은 작업에 대해 더 높은 품질의 샘플을 생성하는 더 효율적이고 견고한 방법을 제공한다고 주장합니다. 이 방법은 미분 가능한 보상을 가진 모든 동결된(frozen) 사전 학습된 플로우 또는 확산 모델과 함께 사용할 수 있는, 실무자를 위한 즉시 사용 가능한 도구로 제시됩니다.