Equivariant Efficient Joint Discrete and Continuous MeanFlow for Molecular Graph Generation

Equivariant Joint MeanFlow for Efficient Molecular Graph Generation

Introduction

This paper introduces Equivariant MeanFlow (EQUIMF), a generative modeling framework for molecular graphs that enforces $\mathrm{SE}(3)$-equivariance and unifies the synthesis of discrete topological structure and continuous 3D geometry. Standard approaches in graph generative modeling separately treat topology (nodes/edges) and geometry (3D coordinates), typically using iterative, decoupled denoising, resulting in physically inconsistent samples or slow inference. In contrast, EQUIMF simultaneously models both domains through synchronized MeanFlow dynamics and mutual conditioning, leveraging a backbone SE(3)-equivariant graph neural network (EGNN). The framework introduces a new discrete MeanFlow formulation and efficient joint training, enabling few-step, high-fidelity molecular generation with rigorous geometric and physical constraints.

Framework Architecture

The core architecture involves a shared EGNN-based encoder $\phi_{\theta_1}$ for cross-modal fusion, with two generation heads: a continuous MeanFlow head $\phi_{\theta_2}$ for coordinates and a discrete MeanFlow head $\phi_{\theta_3}$ for categorical graph attributes. Both heads evolve synchronously via a shared time bridge and jointly optimized loss, producing mutually conditioned updates where geometry and structure consistently inform one another.

Figure 1: Architecture of Equivariant MeanFlow. The framework processes both discrete graph structure and continuous 3D molecular geometry via a shared SE(3)-equivariant encoder. Discrete and continuous MeanFlow heads are optimized for mutually conditioned joint generation.

Mutual conditioning is realized by feeding the latent representation from the shared encoder to both heads, such that discrete topology predictions depend on geometric features and vice versa. This coupling is critical for maintaining chemical validity and realistic conformations, directly addressing the mode collapse and conformation-inconsistency issues of prior approaches.

MeanFlow Dynamics: Discrete and Continuous Domains

Discrete MeanFlow

The discrete component introduces a MeanFlow parameterization for categorical states (nodes and edges) based on a Continuous-Time Markov Chain (CTMC). Unlike previous discrete flow matching (DFM) models, which estimate instantaneous rate matrices and require fine-grained time steps, EQUIMF directly parameterizes the average rate matrix over a finite interval. This allows for efficient few-step sampling and accelerates training convergence without sacrificing fidelity.

Continuous MeanFlow

For the continuous (coordinate) domain, EQUIMF applies a conditional, symmetry-preserving MeanFlow. The model encodes the average velocity field connecting instances along the time bridge, contrasting with standard ODE/velocity-based flows which rely on local, instantaneous information. Utilizing the EGNN backbone ensures strict SE(3)-equivariance in both coordinate updates and neural parameterization, which is theoretically guaranteed by explicit proofs in the paper.

Theoretical Properties

The model's design enforces that (i) all discrete features are SE(3)-invariant; (ii) continuous MeanFlow velocity fields and (iii) discrete transition kernels are SE(3)-equivariant or invariant as appropriate. The universal equivariance ensures physical plausibility: generated molecules have consistent atomic types, bonding, and stable 3D conformations under any rigid transformation. Theoretical results, justified via explicit propositions and proofs, confirm that the coupled generation process is equivariant, a property not simultaneously achieved in previous frameworks.

Experimental Evaluation

Molecular Generation Quality

The framework is evaluated on QM9 (small organics) and GEOM-DRUG (large molecules) benchmarks. Key metrics include atom stability, molecule stability, chemical validity (RDKit), and uniqueness. EQUIMF attains atom stability of 98.9% and molecule stability of 93.0% on QM9, outperforming state-of-the-art equivariant diffusion and flow-matching methods on all axes. Furthermore, validity and uniqueness exhibit substantial gains, showing the model's capacity for both fidelity and diversity.

Property-Controlled Generation

On conditional molecule generation (property matching), EQUIMF achieves lower mean absolute errors than both EDM and EQUIFM baselines on key electronic and thermodynamic properties, indicating enhanced controllability and reduced bias in conditional scenarios.

Sampling Efficiency

EQUIMF enables efficient few-step denoising: at a 0.95 stability threshold, it requires half as many steps as the current SOTA baseline, confirming the practical advantage of MeanFlow-based interval modeling.

Figure 2: Comparison of molecular stability versus sampling steps. EQUIMF surpasses the baseline in convergence rate and final stability.

Ablation Analyses

Ablations demonstrate that SE(3)-equivariant inductive bias in the backbone is pivotal—removal leads to significant drops (up to 5% absolute) in molecule stability. Similarly, the mutual conditioning strategy is essential; decoupling geometry and structure produces chemically invalid or unstable samples, with the fully independent generation pipeline collapsing to as low as 28% molecule stability.

Implications and Future Directions

The primary practical implication of EQUIMF is the reliable, efficient synthesis of physically plausible, chemically valid molecular graphs, which is critical for generative design in chemistry and materials science. By bridging discrete and continuous modalities in a single equivariant framework, EQUIMF provides a scalable basis for large-scale molecule libraries, property-conditioned generation, and downstream tasks such as drug design and targeted discovery.

Theoretically, this work demonstrates that explicit mutual conditioning and interval-based MeanFlow dynamics offer concrete advantages over decoupled or instantaneous-velocity paradigms, both in sample quality and efficiency. Remarkably, the discrete MeanFlow parameterization points toward broader applications in generative modeling over general discrete structures, including text or program synthesis, if suitably extended.

Still, the current design balances between few-step efficiency and sample accuracy; future research will likely explore the feasibility of true one-step generation and direct optimal transport couplings in discrete-continuous domains. Extending equivariant MeanFlow to handle multi-modal distributions, large graphs, or active learning paradigms represents an open, impactful direction for generative modeling research.

Conclusion

EQUIMF establishes an efficient, theoretically principled, and empirically validated framework for molecular graph generation. Its main strengths are (i) joint, SE(3)-equivariant synthesis of discrete and continuous domains; (ii) substantial improvements in sample fidelity, validity, and uniqueness; and (iii) order-of-magnitude faster inference through few-step MeanFlow processes. This advances the state of generative modeling for structured, symmetry-rich data and provides a valuable methodology for both scientific and engineering applications.