Controlled Generation with Equivariant Variational Flow Matching

Dear reviewer uXdp,

Thank you for the detailed and thoughtful review. We appreciate your close reading and thank you for your constructive suggestions. Below we address the main points raised in your review and clarify several aspects that were not sufficiently emphasized in the original submission.

1. Contribution in Unconditional Generation (Table 2)

We agree that the motivation behind the unconditional generation section was not clearly explained. Our goal with Table 2 was not to introduce a new architecture, but to show that strong performance can be achieved via a simple, unified formulation under the VFM framework. By linearly interpolating between noise and data—whether continuous or discrete—and selecting an appropriate variational distribution (eg, Gaussian or categorical), one can train a single model with a shared objective across modalities. This avoids the need for specialized architectures or training regimes typically required in standard FM for mixed data types, and even generalizes to other distributions (eg, Poisson for neural activity). We have clarified this conceptual simplicity in the revision and added citations for baselines directly in the tables.

2. Comparison to D-Flow and Efficiency Tradeoff

While D-Flow performs better in Table 4, it does so by evaluating the forward process over many candidate initializations ($n \times K$ evaluations, eg, $10 \times 100$). In contrast, our method uses a single forward pass and a short fixed-point calibration for conditional guidance—requiring only $K$ forward evaluations and no additional gradient computations. We prioritize simplicity and scalability, whereas D-Flow trades compute for precision. Both are valid design choices, and we now highlight this tradeoff explicitly in the paper, including a discussion of guidance cost and number of function evaluations (NFE).

3. Equivariance Evaluation and Practical Utility

You are correct that we do not isolate the empirical effect of equivariance. While the theoretical formulation is a key contribution, we agree that deeper empirical validation would strengthen the work. We now clarify which models enforce which equivariant properties and will provide an ablation in the final version.

4. Generality Beyond Molecular Generation

While molecular generation is a natural testbed due to its symmetry constraints, the proposed methods are general. VFM provides a flexible framework for combining learned dynamics with structural inductive biases. We now emphasize this more clearly in the discussion and are actively exploring applications beyond molecules, including images (to be included in the final version, at least as a proof of concept).

5. Comparison to CG / CFG Methods

You raise a valid point. Our method assumes access to a classifier over $x_1$ rather than a time-dependent $p_t(y \mid x_t)$. While related, these regimes differ in flexibility and cost. CG/CFG typically requires backpropagating through a denoised prediction at every timestep, which is expensive and needs conditional training. Our fixed-point procedure operates post-hoc on the endpoint and does not require joint training. We now clarify this key distinction in the manuscript.

6. Theoretical Novelty of Equivariance Result

While standard flow matching methods can exhibit equivariant behavior under certain conditions, our contribution is to show that equivariance can be guaranteed through the variational distribution—without requiring the velocity field itself to be equivariant. This decoupling is novel and enables greater modularity: symmetry constraints can be enforced directly through the variational family, independent of the learned dynamics. This allows for clean, flexible design of symmetry-preserving generative models and simplifies implementation in practice. We now clarify this point more explicitly in the final version.

Once again, we thank the reviewer for their careful and constructive feedback. Your comments helped us significantly improve the clarity, presentation, and positioning of our contributions.

Dear reviewer 6iCA,

We sincerely thank the reviewer for their comprehensive assessment of our paper.

We appreciate your recognition that these contributions are "conceptually insightful and practically impactful" and that our work addresses an important gap in flow matching frameworks. We are glad to hear the reviewer is on the same page about the interesting avenues to explore regarding integrating 'classical' inference techniques with SOTA generative modeling.

We thank the reviewer for the clear summary and thoughtful comments. We appreciate the recognition of our core contributions, as well as the constructive suggestions that helped us clarify and strengthen the presentation. Below, we address the reviewer’s main points.

1. Inference-Time Control: Assumptions and Approximation Accuracy

We agree that our formulation suggests simplifying assumptions (eg, log-concavity of the classifier). These are now stated explicitly and discussed in the revised manuscript.

While VFM connected variational inference to unconditional flow matching, our goal is to extend this to conditional generation and introduce an equivariant formulation. We show that one can perform variational inference on $p_t(x_1 \mid x, y)$, and more importantly, use this perspective to rethink conditioning in generative models. A simple proof-of-concept demonstrates this idea.

Our aim is not a full Bayesian treatment, but to show that minimal inference-time updates—without retraining and with negligible overhead—enable effective post-hoc control. Using only the posterior mean is sufficient: in endpoint-based models (eg, flows, diffusions), the expected conditional velocity depends only on this value (see VFM). This eliminates the need for higher-order terms, keeping the method scalable while still improving empirical performance. As such, issues like log-concavity are less critical for our setting—we propose a simple integration of inference with generative modeling and focus not on recovering the full posterior, but on obtaining a reliable estimate of the posterior mean.

Broadly, this work opens a direction within VFM: enabling modular inference-time control by integrating classical inference tools with learned approximations. We clarified this vision in the revision, and hope it inspires further work in this efficient design space.

2. Clarification of Notation and Presentation

We thank the reviewer for catching unclear notation and formatting issues. We confirm that $\mu_t(x)$ denotes the mean of $q_t(x_1 \mid x)$ and have made this explicit. We revised the explanation around Eq. 15 to clarify its relation to VFM: the fixed-point update approximates the mean of $p(x_1 \mid x, y) \propto p(x_1 \mid x)p(y \mid x_1)$.

We also corrected the bolding in Table 3 and improved formatting throughout.

3. Modeling Details of $q_t(x_1 \mid x)$

Yes, we use a mean-field Gaussian for continuous variables and a categorical distribution for discrete ones, as in the original VFM. We added a brief clarification in the revised text.

Once again, we thank the reviewer for the insightful feedback, which helped refine both the clarity and precision of our work.

Dear reviewer TqUv,

Thank you for the detailed and thoughtful review and for engaging deeply with both the theoretical and empirical aspects of our work. Below we respond to the key points and how they will be addressed in the revised manuscript.

1. Unified Objective and Variational Distributions

We fully agree: while VFM provides a unified objective, it still requires variational distributions per modality. However, our contribution is to show that these choices can be modular and plugged directly into the same objective, avoiding the need to redesign the loss or architecture per data type—a common challenge in standard FM pipelines. Though not appropriately explored here, this even naturally extends to other distributions (eg, Poisson for neural data) by learning the conditional rate and optimizing the NLL. We clarified this core motivation in the revision.

2. Fairness of Table 3 Comparisons

While our method is flexible enough to generate any subset of molecular features (e.g., just positions or positions plus atom types), in all experiments, we matched the generation scope of each baseline for a fair comparison. For example, if a baseline did not model bond structures or formal charges, we excluded those as well. We now clearly state this in the text and annotate Table 3 to indicate which components each model generates.

3. Simplifying Assumptions and Classifier Guidance

We now explicitly state and discuss simplifying assumptions (eg, classifier log-concavity) in the revised manuscript.

This work extends VFM with an inference-based view of conditional generation. We perform variational inference on $p_t(x_1 \mid x, y)$ in a simple proof-of-concept setup that yields meaningful improvements. Due to the linearity of the conditional velocity (as in flows/diffusions), matching only the posterior mean suffices (see VFM), and a full Bayesian treatment is unnecessary. Our main contribution is showing that this inference view enables scalable post-hoc control by combining classical inference tools with learned approximations. For more detail, we refer to our response to reviewer k6L2. This also addresses the reviewer's question about the normalization constant in Eq. 13—your interpretation is correct.

Furthermore, while related to classifier guidance (CG/CFG) in diffusion, our approach differs significantly in cost and flexibility. CG typically requires backpropagating through a denoised prediction at each timestep, which can be expensive and must be trained conditionally. In contrast, our fixed-point method operates post-hoc at the endpoint, with no joint training. Finally - even though is it true that true the Tweedie transform a similar effect can be obtained, the approach would arguably be rather noisy and therefore hard to learn. We now clarify this key distinction in the paper.

4. Author questions / Minor comments:

• Missing citations: We agree and will cite [4], [5], and other related work on conditional generation and classifier guidance. These citations have been added in Section 3.1.
• Missing Experimental Details (Table 2): We have included full model architecture and training details in the appendix, and clarified in Table 2 whether atom type modeling is included per experiment.
• Table Formatting and Metric Presentation: We corrected the inconsistent bolding and formatting in Tables 1–3. Bolding is now applied uniformly (best per column, including ties), and we added global method rankings to aid interpretability across datasets.
• Table 1 Repetition: We included Table 1 to demonstrate that G-VFM reproduces the results of standard VFM models as a consistency check. That said, we agree that this space could be better used to highlight new contributions, so we have moved Table 1 to the appendix.
• Typo: We have corrected “Fischer flow” to “Fisher flow” (line 422).

We again thank the reviewer for their constructive feedback. Your comments helped significantly improve the clarity, fairness, and completeness of the revised manuscript.