Riemannian MeanFlow

Motivation

• Known: Diffusion 및 flow model은 SE(3), simplex, torus 등 Riemannian manifold 상의 생성 모델링(단백질 backbone 생성, DNA 서열 설계 등)에서 지배적인 패러다임으로 자리잡았으며, geodesic을 따라 noise를 data로 옮기는 vector field를 학습한다.
• Gap: 기존 manifold 생성 모델은 ODE/SDE를 수치적으로 적분해야 하므로 추론 시 수십~수백 번의 neural network evaluation이 필요해 대규모 과학적 샘플링 워크플로우에서 계산 병목이 된다. Euclidean 공간에서는 consistency model이나 flow map 방법을 통한 few-step 생성 연구가 활발하지만, Riemannian manifold 설정에서는 이러한 접근이 충분히 탐구되지 않았다.
• Why: 단백질 backbone 생성이나 DNA 서열 설계와 같은 과학적 설계 파이프라인에서 생성 모델은 단발성 샘플러가 아니라 property-guided optimization 등 반복적 루프 내의 proposal 메커니즘으로 사용되므로, 각 proposal에 수백 번의 forward pass가 필요하면 탐색 범위가 제한된다. 따라서 manifold 상에서 few-step 생성이 가능해지면 과학적 탐색의 효율성과 확장성이 크게 개선된다.
• Approach: Euclidean flow map 항등식(Boffi et al., 2025)을 Riemannian manifold로 일반화하여 exponential/logarithmic map과 covariant derivative를 이용한 average velocity의 세 가지 동치 표현(Eulerian, Lagrangian, semigroup)을 유도하고, stop-gradient 배치를 통해 고차원 학습의 안정성과 메모리 효율을 확보한다.

Achievement

Figure 4. Performance vs. NFE. RMF variants outperform Fisher

1. Riemannian MeanFlow 항등식 및 목적함수 도출: Eulerian, Lagrangian, semigroup 세 가지 이론적으로 동치인 average velocity 특성화를 유도하고, stop-gradient 위치를 다르게 설정해 고차 미분(higher-order derivative) 계산을 완전히 제거함으로써 메모리 사용량을 거의 2배 절감하고 고차원 설정에서 학습 안정성을 향상시켰다.
2. 확장 가능한 파라미터화 제안: tangent vector 대신 manifold-valued endpoint를 직접 예측하는 x1-prediction을 제안하여 v-prediction과 동등하거나 더 나은 성능을 보이며 D=2048까지의 고차원에서도 잘 확장됨을 확인했고, semigroup objective와 x1-prediction의 조합이 고차원 manifold에서 가장 안정적인 학습 조합임을 규명했다.
3. 과학적 설계 응용 실증: DNA promoter 설계(simplex in R^1024x4)와 단백질 backbone 생성(SE(3)^N, N up to 128)에서 기존 state-of-the-art multi-step 모델과 comparable한 샘플 품질을 최대 10배 적은 function evaluation으로 달성했으며, reward look-ahead를 통한 reward-guided design을 manifold 상에서 최초로 시연했다.

How

Figure 2. Parameterization choices: One-step generation results

• Riemannian manifold M 상에서 flow map Φs,t와 average velocity us,t를 정의하고, logxs xt = (t-s)us,t(xs) 관계식을 s 또는 t에 대해 미분하여 Eulerian, Lagrangian 두 가지 항등식을 도출
• semigroup property(Φr,t∘Φs,r=Φs,t)로부터 세 번째 semigroup identity를 도출하여 세 항등식이 이론적으로 동치임을 증명
• 각 항등식에서 stop-gradient 연산자의 위치를 조정해 covariant derivative의 고차 항 계산을 회피하는 학습 목적함수를 설계
• 네트워크가 tangent vector(v-prediction) 대신 manifold-valued 종점(x1-prediction)을 직접 예측하도록 하는 파라미터화를 제안하고 SE(3), simplex 등에서 closed-form log/exp map 유도 활용
• promoter DNA 설계와 단백질 backbone 생성 벤치마크에서 NFE(number of function evaluations) 대비 성능을 FrameDiff, FrameFlow 등 기존 multi-step 방법과 비교
• reward look-ahead 기법을 통해 중간 시점에서 terminal state를 예측하여 reward-guided generation을 최소 추가 비용으로 수행

Originality

• Euclidean 공간에 국한되어 있던 flow map/average-velocity 기반 few-step 생성(MeanFlow 계열)을 Riemannian manifold로 최초로 체계적으로 일반화하고 세 가지 동치 항등식을 도출
• stop-gradient 배치를 달리하여 고차 미분(higher-order derivative) 평가를 완전히 제거하는 새로운 학습 목적함수 설계는 기존 Euclidean MeanFlow 및 동시대 독립 연구(Davis et al., 2025)와 차별화되는 지점
• manifold-valued endpoint를 직접 예측하는 x1-prediction 파라미터화를 제안하여 기존 과학적 아키텍처(단백질 구조 예측 네트워크 등)와의 호환성을 확보
• 기존 manifold consistency/flow-map 연구가 저차원 벤치마크(S2, SO(3), torsion angle)에 국한된 것과 달리, 고차원 실제 과학적 과제(DNA promoter 설계, 단백질 backbone 생성)에서 few-step 생성의 실용적 확장성을 최초로 실증
• manifold 상에서 reward look-ahead를 통한 reward-guided design을 최초로 시연

Limitation & Further Study

• 발췌된 본문에서는 Proposition 2.1(Eulerian RMF) 이후의 구체적 수식과 semigroup identity의 상세 유도, 그리고 안정화 기법의 세부 구현이 제시되지 않아 이론적 완전성 검증이 어려움
• 실험이 promoter DNA(simplex)와 단백질 backbone(SE(3)) 두 도메인에 한정되어 있어, 다른 manifold(예: SO(3) 단독, 구면, torus 등 다양한 기하 구조)나 더 큰 스케일의 실제 응용에 대한 일반화 가능성은 추가 검증이 필요
• x1-prediction과 semigroup objective 조합이 최선이라는 결론이 특정 실험 설정에 기반하므로, 다른 아키텍처나 데이터셋에서도 일관되게 성립하는지에 대한 후속 연구가 필요
• reward-guided design의 reward look-ahead 기법이 다양한 reward 함수 및 실제 과학적 최적화 루프에서 얼마나 강건하게 작동하는지에 대한 추가적 검증이 요구됨

Evaluation

총평: Riemannian manifold 상에서 few-step 생성을 위한 이론적 기반(세 가지 동치 항등식)과 실용적 파라미터화(x1-prediction)를 함께 제시하고, 실제 과학적 과제에서 최대 10배의 NFE 절감을 실증한 견실하고 시의적절한 연구이다. 이론적 엄밀성과 과학적 응용성을 모두 갖추어 manifold 생성 모델의 실용화에 중요한 기여를 할 것으로 평가된다.