🤖 machine learning

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

The paper introduces SE(3)-MeanFlow, a few-step generative framework that extends MeanFlow to Lie group geometry for protein backbone design, achieving high-quality generation with significantly fewer sampling steps than existing diffusion or flow-matching models by utilizing closed-form training targets and a specialized SE(3) alpha-Flow objective.

Original authors: Yikun Bai, Binghang Lu, Yikai Liu, Elaheh Akbari, Soheil Kolouri, Linxuan Wang, Ping He, Shuchan Wang, Ruqi Zhang, Guang Lin

Published 2026-07-31

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Yikun Bai, Binghang Lu, Yikai Liu, Elaheh Akbari, Soheil Kolouri, Linxuan Wang, Ping He, Shuchan Wang, Ruqi Zhang, Guang Lin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are an architect trying to build a new kind of living machine out of tiny, flexible Lego bricks. These machines are called proteins, and they are the workhorses of life, acting as everything from chemical factories to immune system soldiers. To build a new one from scratch, you don't just stack bricks randomly; you have to arrange them into a specific 3D shape, like a complex origami crane. If the shape is wrong, the machine won't work. For decades, scientists have been trying to teach computers to design these origami cranes, but there's a catch: the computer has to figure out how to twist and turn each brick in 3D space, which is a mathematically tricky dance.

The problem gets even harder when you want the computer to be fast. Current methods are like a slow-motion movie where the computer takes hundreds of tiny steps to figure out the final pose of the protein. It's accurate, but it's too slow to design millions of new proteins at once, which is what doctors and drug companies need to fight diseases quickly. The math behind this involves a special kind of geometry called "Lie groups," which sounds scary but is just a fancy way of describing how things rotate and move without breaking the rules of 3D space. Think of it as the difference between sliding a toy car on a flat table (easy) and rolling a ball around the surface of a globe (harder, because the surface curves).

This paper introduces a new trick called SE(3)-MeanFlow that acts like a "fast-forward" button for designing these protein shapes. Instead of taking hundreds of tiny, cautious steps to figure out the final pose, this new method learns to predict the average speed and direction of the entire journey in just one or a few giant leaps. It's like teaching a GPS not just to tell you the speed at your current location, but to calculate the entire route's average speed so you can jump straight to the destination. The authors show that by using this "average velocity" approach specifically tailored for the curved geometry of protein rotations, they can generate high-quality protein designs in as few as 10 or 20 steps, whereas older methods needed 100 or more. They didn't just guess this would work; they tested it on a massive dataset of real protein structures and found that their fast method produced designs just as strong and functional as the slow ones, proving that you don't need to take the scenic route to get a good result.

The Big Idea: From Slow Motion to Fast Forward

In the world of computer-generated protein design, the goal is to create new shapes that nature hasn't seen before. To do this, AI models act like a sculptor, starting with a blob of noise and slowly chipping away until a perfect protein backbone emerges. For a long time, the best sculptors used a technique called "diffusion" or "flow matching." Imagine trying to walk from your house to the park by taking 500 tiny, careful steps. You have to check your direction at every single step to make sure you don't wander off. This works well, but it takes forever if you need to visit a million parks.

The authors of this paper asked a simple question: Can we take bigger steps without getting lost?

They looked at a newer idea called MeanFlow, which was originally designed for flat, simple spaces (like a flat sheet of paper). The core idea of MeanFlow is to stop asking, "Which way should I go right now?" and start asking, "What was the average speed and direction I needed to get from point A to point B?" If you know the average speed for the whole trip, you can cover the distance in one giant stride instead of 500 tiny ones.

However, proteins don't live on a flat sheet of paper. They live on a curved, twisting surface where rotation is tricky. If you try to apply the flat-paper version of MeanFlow to proteins, the math breaks because the "average" of two rotations isn't just a simple average; it depends on the order you do them. This is where the paper's main breakthrough comes in.

The Secret Sauce: The Lie Group Shortcut

The authors realized that protein rotations belong to a special mathematical family called Lie groups (specifically $SE(3)$). They figured out a way to calculate the "average velocity" for these rotations without needing to do the heavy, slow math that usually slows things down.

Think of it like this: If you are spinning a top, calculating the average spin over a few seconds is hard if you try to measure it second-by-second. But if you have a special formula that tells you the "total spin" based on where the top started and where it ended, you can skip the middle part entirely. The authors derived a new formula that acts like this shortcut. They proved that by working in a specific mathematical space (the "Lie algebra"), they could get a clean, closed-form answer for the average velocity. This means the computer doesn't have to simulate the whole journey step-by-step to know the answer; it can just calculate the average directly.

What They Found: Speed Without Sacrifice

The team tested their new SE(3)-MeanFlow model against the best existing methods. They set up a challenge: generate protein backbones using different numbers of steps, from a slow 100 steps down to a lightning-fast 10 steps.

Here is what they discovered:

The Speed Boost: Their model could generate high-quality proteins in just 10 to 20 steps. In contrast, the older, slower methods needed 100 steps to get similar results. When the authors forced the old methods to take only 20 steps, the quality of the proteins dropped significantly.
Quality Stays High: Even with these huge jumps in speed, the proteins generated by SE(3)-MeanFlow were still "designable." This means that if you were to print them out and try to build them, they would fold into the correct shape. In fact, at 20 steps, their model had a success rate of 86.7%, while the next best competitor was only at 80.6%.
The Trade-off: There is a tiny cost. The proteins generated by the super-fast method were slightly less diverse (meaning they looked a bit more similar to each other) than the ones made by the slow, careful methods. However, the authors suggest this is a small price to pay for being able to generate millions of designs in the time it used to take to generate a few thousand.

Why This Matters

This isn't just a math trick; it's a practical tool for the future of medicine. Drug discovery often requires generating millions of candidate protein structures to find one that might fight a virus or a cancer cell. If a computer takes 100 steps to make one design, it might take days to find a good candidate. If it can do it in 10 steps, the same search could be done in hours.

The authors suggest that by using this "average velocity" approach, we can finally make high-throughput protein design a reality. They didn't just simulate this on a computer; they trained the model on a real dataset of over 3,600 proteins and tested it rigorously. While they note that their method is still being refined (especially regarding diversity), the results suggest that we are on the verge of being able to design life-saving proteins at a speed that matches the urgency of global health challenges.

In short, the paper shows that by understanding the geometry of rotation a little better, we can stop taking baby steps and start sprinting toward new medical breakthroughs.

Technical Summary: SE(3)-MeanFlow

Problem Statement

Generative modeling of protein backbones is critical for the de novo design of proteins with prescribed structural and functional properties. While existing diffusion and flow-matching models (e.g., FrameDiff, RFDiffusion, FoldFlow) have achieved high-quality generation on the space of protein frames, $SE(3)^N$ , they suffer from a significant inference bottleneck. Current methods require numerically integrating an Ordinary Differential Equation (ODE) over hundreds of steps (typically 100–500 network function evaluations, or NFE) to ensure accuracy. On the Lie group $SE(3)$, each step involves evaluating the network and applying the Lie group exponential map, making inference computationally expensive and unsuitable for high-throughput drug discovery pipelines that require millions of candidate structures.

While few-step generation has been explored in Euclidean space using Consistency Models or progressive distillation, these approaches often require pre-trained teachers or complex training curricula. MeanFlow, a recent method for Euclidean spaces, offers a self-contained alternative by modeling the average velocity of a flow trajectory rather than the instantaneous velocity, enabling one-step generation without distillation. However, naively extending MeanFlow to the curved manifold of $SE(3)$ is ill-posed because velocities at different points along a trajectory reside in different tangent spaces, necessitating parallel transport for averaging.

Methodology

The authors propose SE(3)-MeanFlow, a few-step generative framework specifically designed for the Lie group geometry of protein frames. The method leverages the decomposition $SE(3) \cong SO(3) \times \mathbb{R}^3$ to derive simulation-free training targets directly in the Lie algebra $\mathfrak{so}(3)$ and Euclidean space $\mathbb{R}^3$ , avoiding the need for parallel transport.

Core Theoretical Contributions

Integration-Based MeanFlow Identity: The authors define the average velocity over an interval $[s, t]$ $[s, t]$ via a time-ordered exponential. By differentiating this definition, they derive a closed-form identity linking the average velocity, its trajectory derivative, and the instantaneous velocity.
- For rotations, this identity involves the right Jacobian $J$ of $SO(3)$, which accounts for the non-commutative nature of the group.
- The resulting training objective allows the model to learn the average velocity $u(s, t, x_t)$ using only the instantaneous velocity as supervision, without requiring the full trajectory integration during training.
Closed-Form Targets: Unlike general Riemannian MeanFlow approaches that rely on parallel transport, SE(3)-MeanFlow derives closed-form average-velocity identities in $\mathfrak{so}(3)$ and $\mathbb{R}^3$ . This eliminates the computational cost and numerical instability associated with parallel transport.

Practical Training Framework

To address numerical instability, particularly when $t \to 0$ (where $1/t$ factors in the loss can amplify errors), the authors introduce a stable training curriculum:

$\alpha$ -Flow Warm-up: A JVP-free (Jacobian-Vector Product free) objective based on a two-segment construction is used as a warm-up stage. This objective replaces the differential target with a consistency target that avoids differentiating through the trajectory derivative, stabilizing the initial training phase.
Small- $t$ Stabilized MeanFlow: After the warm-up, training switches to a differential MeanFlow loss. The authors implement numerical stabilization techniques, including auxiliary variables to absorb $1/t$ factors and a small- $t$ reparameterization, ensuring robust training even with aggressive step reduction.
Rectification (Post-Training): A rectification stage (self-reflow) is applied to straighten the transport paths further, improving performance in the few-step regime.

Architecture

The model utilizes the ReQFlow backbone (an Invariant Point Attention network) with minimal modifications:

A shared two-time embedding for inputs $(s, t)$ .
Per-block AdaLN-Zero gates to condition the rigid updates on the time interval.
The model size is approximately 16.8M parameters, significantly smaller than concurrent Riemannian MeanFlow baselines (437M parameters).

Key Contributions

Theoretical Derivation: The paper provides a rigorous derivation of the MeanFlow identity on $SE(3)$, establishing closed-form average-velocity objectives in the Lie algebra without parallel transport.
Stable Training Algorithms: The introduction of the $\alpha$ -Flow warm-up and small- $t$ stabilization techniques makes few-step generation reliable for protein backbones, addressing the numerical fragility of differential targets on manifolds.
Performance in Few-Step Regime: The method achieves state-of-the-art designability with significantly fewer sampling steps compared to flow-matching baselines.

Experimental Results

The method was evaluated on the SCOPe dataset (3,673 protein backbones, lengths 60–128) against baselines including FrameFlow, QFlow, ReQFlow, and Riemannian MeanFlow (RMF).

Designability: SE(3)-MeanFlow matches or exceeds the designability (fraction of backbones with scRMSD < 2Å) of baselines that use 5–10 times more sampling steps.
- At 20 steps, SE(3)-MeanFlow achieves a designable fraction of 0.867, outperforming RMF (0.806) and QFlow (0.778).
- At 10 steps, the rectified version (RecSE3MF) maintains a designable fraction of 0.894, significantly higher than ReQFlow (0.837).
Efficiency: The method reduces the sampling budget from hundreds of steps to as few as 10–20 steps while maintaining high structural quality.
Diversity and Novelty: There is a modest trade-off in diversity (higher pairwise TM-scores) compared to coverage-oriented baselines, but the method preserves stable secondary structure statistics (helix/strand fractions) and C $\alpha$ validity across different step counts.
Ablation Studies: Controlled experiments on a synthetic $SO(3)^2$ benchmark confirm that the average-velocity parameterization is the primary driver of few-step success, rather than network architecture or coupling strategies.

Significance and Claims

The paper claims that SE(3)-MeanFlow advances the designability–efficiency frontier in protein generation. By exploiting the specific algebraic structure of $SE(3)$ as a Lie group, the authors demonstrate that it is possible to derive simulation-free training targets that enable high-quality generation with a fraction of the computational cost required by existing methods.

The authors emphasize that their approach is self-contained, requiring no pre-trained teacher or complex distillation, and is scalable (using a smaller model than concurrent Riemannian methods). While they acknowledge a modest reduction in diversity compared to some baselines, they argue that the ability to generate designable structures in 10–20 steps makes the method particularly valuable for high-throughput applications where millions of candidates must be screened. The work suggests that future directions include extending the method to conditional generation and motif scaffolding, and pushing toward true one-step generation.