기술 요약: 리만 평균 흐름 (Riemannian MeanFlow)

문제 제기

확산 및 흐름 모델은 리만 다양체 (Riemannian manifolds) 상의 생성 모델링에서 지배적인 패러다임이 되었으며, 단백질 백본 생성 및 DNA 서열 설계와 같은 작업에 성공적으로 적용되었습니다. 그러나 이러한 방법들은 일반적으로 추론 중 상미분 방정식 (ODE) 또는 확률 미분 방정식 (SDE) 을 수치적으로 적분하기 위해 수십에서 수백 개의 신경망 평가 (함수 평가, NFE) 를 필요로 합니다. 이 계산적 부담은 생성 모델이 속성 유도 최적화를 위한 반복적 루프 내에서 제안 메커니즘으로 작용하는 과학적 설계 파이프라인에서 특히 병목 현상을 만듭니다. 일관성 모델 (consistency models) 및 흐름 맵 학습과 같은 몇 단계 생성 기법들은 유클리드 환경에서 탐구되어 왔으나, 고차원 리만 다양체 환경에서는 여전히 충분히 탐구되지 않았으며 종종 불안정합니다.

방법론: 리만 평균 흐름 (RMF)

저자들은 **리만 평균 흐름 (Riemannian MeanFlow, RMF)**을 소개합니다. 이는 다양체 위에서 직접 흐름 맵을 학습하여 단 하나의 순방향 통과 (forward pass) 만으로도 고품질 생성을 가능하게 하는 프레임워크입니다. 이 방법론의 핵심은 임의의 시간 점 $s$와 $t$ 사이의 흐름 맵과 관련된 다양체 평균 속도 (manifold average velocity) 의 세 가지 동등한 특성을 유도하고 활용하는 데 있습니다.

1. 이론적 기반

이 프레임워크는 평균 속도 $u_{s,t}$를 $x_s$에서 $x_t$까지의 측지선 (geodesic) 을 따라 시간 $t-s$ 동안 유지될 경우 해당 점을 올바르게 이동시키는 일정한 속도로 정의합니다. 저자들은 학습 목적 함수의 기초가 되는 세 가지 항등식을 유도합니다:

• 오일러 항등식 (Eulerian Identity): 정의 관계를 소스 시간 $s$에 대해 미분하여 유도됩니다. 이는 평균 속도를 적분 곡선을 따른 속도장의 공변 미분 (covariant derivative) 및 로그 맵의 미분과 연관시킵니다.
• 라그랑주 항등식 (Lagrangian Identity): 타겟 시간 $t$에 대해 미분하여 유도됩니다. 이는 평균 속도를 순간 속도에 적용된 로그 맵의 미분 및 평균 속도의 편 시간 미분과 연관시킵니다.
• 반군 항등식 (Semigroup Identity): 흐름 맵의 반군 성질 ($\Phi_{r,t} \circ \Phi_{s,r} = \Phi_{s,t}$) 에서 유도된 대수적 조건입니다. 이는 $s$에서 $t$로의 직접 흐름과 중간 시간 $r$을 통한 흐름의 합성 사이의 일관성을 강제합니다.

2. 학습 목적 함수 및 안정화

저자들은 이러한 항등식을 회귀 목적 함수로 변환합니다. 중요한 설계 선택은 stop-gradient 연산자의 배치입니다. 이전 연구 (예: Generalized Flow Maps) 가 부분적으로 stop-gradient 를 적용한 것과 달리, RMF 는 전체 회귀 목표에 stop-gradient 를 적용합니다. 이 공식화는 다음과 같은 효과를 가집니다:

• 역전파 중 고차 미분 (예: 야코비안 - 벡터 곱을 통한 기울기) 평가의 필요성을 제거합니다.
• 메모리 사용량을 거의 2 배 감소시킵니다.
• 고차원 환경에서 학습 안정성을 향상시킵니다.

고차원에서의 최적화 어려움을 해결하기 위해, 논문은 몇 가지 안정화 기법을 제안합니다:

• 시간 샘플링: 오일러 및 반군 목적 함수를 위한 순서 있는 시간 쌍 ($s \le t$) 과 라그랑주 목적 함수를 위한 순서 없는 쌍.
• 적응형 손실 가중치: 미분 의존 항의 높은 분산을 처리하기 위해 회귀 목표의 크기에 기반하여 손실을 재가중치합니다.
• 시간 미분 제어: 네트워크 출력에서 시간 미분의 크기를 제한하기 위해 저주파 시간 임베딩을 사용합니다.

3. 매개변수화

논문은 평균 속도에 대한 세 가지 매개변수화 체계를 평가합니다:

• $v$-예측: 접선 벡터 (속도) 를 예측합니다.
• $x_t$-예측: 종점을 직접 예측합니다.
• $x_1$-예측: 현재 상태와 시간 간격을 기반으로 $t=1$에서의 궤적 종점을 예측합니다.
  저자들은 $x_1$-예측이 고차원 다양체에 가장 강력하다는 것을 발견했습니다. 이는 기존 과학적 아키텍처 (예: 단백질을 위한 SE(3) 프레임을 출력하는 아키텍처) 와 호환되며 $D=2048$까지의 차원으로 잘 확장되는 반면, $x_t$-예측은 실패하고 $v$-예측은 차원이 증가함에 따라 성능이 저하됩니다.

4. 보상 유도 추론

이 프레임워크는 모델 재학습 없이 보상 유도 생성을 지원합니다. 흐름 맵이 중간 단계로부터 최종 상태를 예측할 수 있는 능력을 활용하여, 저자들은 보안 선점 (reward look-ahead) 전략을 제안합니다. 현재 노이즈 상태 $x_t$에서의 보상 기울기를 사용하는 대신, 예측된 종점 $\Phi_{t,1}(x_t)$에서의 기울기를 사용합니다. 이 접근법은 생성물을 특정 목표 방향으로 유도하는 데 더 효과임이 입증되었습니다.

주요 결과

저자들은 RMF 를 두 가지 생물학적 작업에서 검증했습니다: 프로모터 DNA 설계 ($\mathbb{R}^{1024 \times 4}$의 심플렉스 상의 서열) 및 단백질 백본 생성 ($N$이 최대 128 인 $SE(3)^N$ 상).

• 효율성: RMF 는 최첨단 다단계 베이스라인 (예: Fisher FM, FrameFlow, FrameDiff) 과 비교 가능한 샘플 품질을 달성하면서 최대 10 배 적은 함수 평가를 요구합니다.
  • 단백질 백본 생성에서 RMF 는 5 단계에서 82% 의 설계 가능성 (designability) 을 달성하고 1 단계에서는 35% 를 달성하는 반면, 베이스라인은 이러한 소수 단계 영역에서 급격히 떨어지거나 완전히 붕괴됩니다.
  • DNA 설계에서 RMF 변형체는 단 1 단계만으로 Fisher FM 의 100 단계 성능과 일치합니다.
• 안정성: 제안된 안정화 기법들 (특히 반군 목적 함수와 $x_1$-예측의 조합) 은 다른 흐름 맵 방법들 (예: GFM) 이 수렴하지 않거나 고오류 샘플을 생성하는 고차원 공간에서 안정적인 학습을 가능하게 합니다.
• 보상 유도: $x_1$ 선점 전략을 사용한 보상 유도 추론은 단순한 유도나 유도 없는 생성에 비해 목표 조절 프로파일 (DNA) 및 2 차 구조 구성 (단백질) 과의 정렬을 크게 향상시킵니다.

중요성 및 주장

이 논문은 리만 다양체 상의 빠른 샘플링을 위한 원칙적인 이론적 기반과 실용적 지침을 확립했다고 주장합니다. 그 중요성은 다음과 같습니다:

1. 격차 해소: 몇 단계 생성 능력을 유클리드 공간에서 복잡하고 고차원인 리만 다양체로 확장하여 과학적 응용을 위한 중요한 단계를 제공합니다.
2. 확장성: 고차 미분을 피하는 항등식을 유도하고 $x_1$-예측을 활용함으로써, 기존 몇 단계 다양체 생성 모델이 접근할 수 없었던 차원 ($D \approx 2000$) 및 서열 길이에 확장 가능합니다.
3. 과학적 워크플로우의 효율성: 추론 비용의 감소 (최대 10 배) 는 신약 개발 또는 합성 생물학의 속성 유도 최적화와 같은 반복적 과학 설계 루프에서 제안 메커니즘으로서 생성 모델의 실용성을 높입니다.
4. 보상 선점: 흐름 맵이 중간 단계로부터 최종 상태를 예측함으로써 효율적인 보상 유도 설계를 가능하게 한다는 것을 보여주며, 이는 다양체 상에서 이전에 입증되지 않은 능력입니다.

저자들은 RMF 를 모든 생성 모델링을 대체하는 것이 아니라, 계산 비용이 제한 요인인 기하학적으로 제약된 과학적 영역에서 효율적이고 고품질의 생성을 가능하게 하는 전문화된 프레임워크로 위치시킵니다.