SE(3)-MeanFlow: Lie 그룹에서의 단 몇 단계로 단백질 백본 생성

SE(3)-MeanFlow: Few-Step Protein Backbone Generation on Lie Groups

단백질 백본의 생성 모델링은 특정 구조적 및 기능적 특성을 가진 새로운 단백질을 설계하는 데 유망합니다. 기존의 확산 및 플로우 매칭 모델은 SE(3)^N 상에서 고품질의 백본을 생성하지만, 추론 과정에서 수백 번의 네트워크 평가를 거치는 ODE를 수치적으로 계산해야 하며, 각 단계는 Lie 그룹 지수 함수 연산을 포함하므로 고처리량 설계에 병목 현상을 일으킵니다. 본 논문에서는 MeanFlow 프레임워크를 유클리드 공간에서 Lie 그룹 기하학인 단백질 프레임으로 확장한 SE(3)-MeanFlow라는 새로운 생성 모델을 제안합니다. Lie 대수 so(3)와 R^3 공간에서 직접 작동하며, 회전 및 이동에 대한 닫힌 형태의 평균 속도 식을 유도하여 시뮬레이션 없이 학습 목표를 설정합니다. 또한, Jacobian-vector 곱셈 연산을 회전 분기에서 제거하고 초기 단계로 사용되는 SE(3) alpha-Flow 목적 함수를 도입했습니다. 이후 학습은 안정화된 MeanFlow 손실 함수를 사용하여 진행되며, 이는 사전 훈련 및 수정 기반의 후처리 과정에 모두 사용됩니다. 단백질 백본 생성 실험에서 SE(3)-MeanFlow는 여러 배수의 샘플링 단계를 사용하는 플로우 매칭 모델과 동등하거나 더 나은 성능을 보이며, 특히 샘플링 단계가 적은 경우 더욱 우수한 성능을 나타냅니다. 이는 다양성에 약간의 손실이 있더라도 효율적인 설계가 가능하도록 합니다.

Generative modeling of protein backbones promises the de novo design of proteins with prescribed structural and functional properties. Existing diffusion and flow-matching models produce high-quality backbones on SE(3)^N, but inference requires numerically integrating an ODE over hundreds of network evaluations, each involving a Lie group exponential map - a bottleneck for high-throughput design campaigns. We introduce SE(3)-MeanFlow, a few-step generative framework that extends MeanFlow from Euclidean space to the Lie group geometry of protein frames. Working natively in the Lie algebra so(3) and in R^3, we derive closed-form average-velocity identities for rotations and translations, giving simulation-free training targets. We further introduce an SE(3) alpha-Flow objective that removes the Jacobian-vector product from the rotation branch and serves as a warm-up stage, after which training switches to a small-t stabilized MeanFlow loss that is used for the remainder of pretraining and for rectification-based post-training. In protein backbone generation, SE(3)-MeanFlow matches or exceeds flow-matching baselines that use several times more sampling steps, and its advantage widens in the few-step regime, where rectification lets it lead at every matched budget - at a modest cost in diversity.