Geometry-Consistent Flow Distillation for Uncertainty-aware Protein Structure Generation

Diffusion and flow-matching models have achieved remarkable performance in protein structure generation and ensemble modeling, yet their practical use is often computationally constrained at scale by expensive multi-step sampling and costly iterative refinement. Recent work has explored applying distillation to protein generative models. However, directly transplanting the methods in image generation to 3D proteins overlooks fundamental geometric and physical constraints, most notably SE(3) gauge freedom and frame-dependent supervision, leading to single-step collapse with poor accuracy and structural quality. To address this challenge, we develop a geometry-consistent distillation(GCD) framework that makes one-step protein generation stable and well-posed. Our formulation captures local backbone correlations by explicitly coupling neighboring residues and further integrates an SE(3)-invariant corruption interface and a residue-local, frame-aligned distillation objective. For flow-matching teachers, we further derive the corresponding harmonic-prior score–velocity identity, connecting the teacher’s vector-field parameterization to score-based distillation. Empirically, naive Euclidean distillation collapses in the one-step setting, whereas GCD yields stable and competitive one-step students and transfers across multiple distillation paradigms. These results show that geometry-consistent supervision enables aggressive acceleration of protein generative models while retaining strong structural fidelity and substantial ensemble diversity.