Property-Guided Molecular Generation and Optimization via Latent Flows

Alexander Arjun Lobo; Urvi Awasthi; Leonid Zhukov

cs.LGMar 27, 2026

Property-Guided Molecular Generation and Optimization via Latent Flows

Authors: Alexander Arjun Lobo, Urvi Awasthi, Leonid Zhukov

Organizations: BCG X AI Science Institute

Abstract

Molecular discovery is increasingly framed as an inverse design problem: identifying molecular structures that satisfy desired property profiles under feasibility constraints. While recent generative models provide continuous latent representations of chemical space, targeted optimization within these representations often leads to degraded validity, loss of structural fidelity, or unstable behavior. We introduce MoltenFlow, a modular framework that combines property-organized latent representations with flow-matching generative priors and gradient-based guidance. This formulation supports both conditioned generation and local optimization within a single latent-space framework. We show that guided latent flows enable efficient multi-objective molecular optimization under fixed oracle budgets with controllable trade-offs, while a learned flow prior improves unconditional generation quality.

Figures & tables

Figure 1: MoltenFlow overview. A VAE maps discrete molecular inputs $x$ to continuous latent representations $z$ and reconstructs molecules via a decoder. The latent space is organized through auxiliary property prediction: a differentiable property surrogate $f_{\psi}(z)$ induces an objective $\mathcal{J}(z;c)$ (e.g., target values or directional goals specified by a guidance vector $c$ ), and provides a guidance signal $g=\nabla_{z}\mathcal{J}(z_{t};c)$ during inference. In parallel, a flow-matching model learns a time-dependent vector field $v_{\omega}(z,t)$ that transports samples from a simple base distribution (random noise) to the empirical distribution of valid latent representations. The flow-based latent transport visualization is adapted from Sabour et al. (2025) , and the property-guidance schematic is inspired by ChemFlow ( Wei et al., 2024 ) .

Figure 4: Qualitative effect of guidance strength on optimization. As the guidance strength $\gamma$ increases, MoltenFlow produces progressively larger structural changes. Small $\gamma$ yields minimal edits dominated by the latent flow prior, intermediate $\gamma$ induces coherent improvements in QED with reasonable SA, and large $\gamma$ leads to aggressive modifications indicative of over-optimization.

Figure 5: Effect of guidance strength on optimization and distributional properties. Each panel shows a metric as the guidance strength

\gamma

increases (log scale). From left to right: hypervolume improvement (HVI), validity, Fréchet distance (FD-FP), and scaffold diversity. HVI increases sharply in an intermediate regime, indicating effective Pareto-front advancement, while excessive guidance leads to rising Fréchet distance and collapsing diversity, signaling over-optimization. Starred points denote the

\gamma

achieving the best value for each metric.

Panel of Figure 6 — Figure 6: Pareto-front advancement and guidance-strength effects. ( Left ) Optimized Pareto fronts obtained at different guidance strengths $\gamma$ . At intermediate $\gamma$ ( green ), MoltenFlow produces fronts that strictly dominate the original, expanding into the high-QED, low-SA region while retaining a broad set of trade-offs. ( Right ) Zoomed view of the high-QED/low-SA corner (top) and summary metrics across guidance strengths (bottom), illustrating the trade-off between hypervolume improvement, scaffold diversity, and distributional shift (see Appendix E.3.4 for metric definitions).

Panel of Figure 7 — Figure 7: UMAP of posterior-mean latent embeddings $\mu_{\phi}(x)$ for real molecules (fit set) and for generated molecules after re-encoding. Flow-prior samples align more tightly with the support of real embeddings compared to VAE-prior samples, consistent with improved validity and reduced FID-FP in Table 1 .

Rep.	Prior	Valid(%) $\uparrow$	Unique(%) $\uparrow$	Novel(%) $\uparrow$	Scaf. div. $\uparrow$	Avg KL $\downarrow$	FD-FP $\downarrow$
SMILES	VAE prior	5.76	98.44	99.47	0.944	0.299	13.23
SMILES	Flow prior (FM)	9.56	97.91	99.04	0.700	0.191	8.15
SELFIES	Flow prior (FM)	100.00	99.07	99.66	0.694	0.161	10.72
SMILES	VAE post. ( $\mu_{\phi}(x)$ ; ref.)	26.07	78.02	71.58	0.438	0.207	4.67

Table 1: Unconditional generation ablations (no surrogate guidance). A1 compares the VAE prior vs. flow-matching (FM) prior under SMILES. A2 compares SMILES vs. SELFIES under the FM prior. FD-FP is Fréchet distance in fingerprint embedding space; Avg KL is the mean KL divergence over RDKit descriptors.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Type	Range	Selected
$\gamma$ (guidance strength)	Log-uniform	$[10,50]$	36.07
$\sigma$ (noise scale)	Uniform	$[0.4,0.9]$	0.80
steps (integration steps)	Integer	$[10,50]$	12
$t_{\text{start}}$	Uniform	$[0.5,0.9]$	0.89
clip_norm	Categorical	$\{\varnothing,1,5,10\}$	5.0
normalize_gradient	Categorical	$\{\text{True},\text{False}\}$	True

Appendix

Table 2: Hyperparameter search space for MoltenFlow.

Panel of Figure 8 — Figure 8: Hyperparameter optimization diagnostics. Left: contour of hypervolume improvement as a function of guidance strength $\gamma$ and noise scale $\sigma$ . Center: optimization history across trials. Right: relative importance of hyperparameters for maximizing hypervolume improvement.

Figure 9: Representative Pareto fronts from best (top) and worst (bottom) runs for each method. Gray points: initial molecules. Colored points: final Pareto front. MoltenFlow achieves consistent improvements; BO methods show higher variance.

Figure 10: Runtime comparison across methods (

n=20

seeds). Error bars: 90% bootstrap CI. MoltenFlow is

\sim 8\times

faster than BO methods.

Experimental Setup
Parameter	Value
Oracle budget ( $B$ )	100
Initial set size	10
Batch size	1
Seeds per method	20
Initialization	Random sampling

Appendix

Table 3: Hyperparameters for budgeted optimization experiments.

$\gamma$	HVI $\uparrow$	HVI (%) $\uparrow$	HVI CI $\uparrow$	MSE $\downarrow$	Unique (%) $\uparrow$	Novel (%) $\uparrow$	Scaf. Div. $\uparrow$	FD-FP $\downarrow$	Avg. KL $\downarrow$
$0.001$	$0.023$	$0.39$	$0.023\pm 0.016$	$0.63$	$100$	$100$	$0.962$	$8.89$	$0.18$
$0.1$	$0.023$	$0.39$	$0.023\pm 0.016$	$0.64$	$100$	$100$	$0.963$	$8.91$	$0.18$
$1$	$0.023$	$0.39$	$0.023\pm 0.016$	$0.65$	$100$	$100$	$0.961$	$8.92$	$0.18$
$10$	$0.019$	$0.33$	$0.019\pm 0.015$	$0.77$	$100$	$100$	$0.956$	$9.00$	$0.18$
$\mathbf{92.4^{*}}$	$\mathbf{0.405}$	$\mathbf{6.96}$	—	$1.98$	$100$	$100$	$0.743$	$13.18$	$0.47$
$100$	$0.357$	$6.13$	$0.357\pm 0.077$	$2.02$	$100$	$100$	$0.702$	$13.46$	$0.53$

Appendix

Table 4: Complete gamma sweep results. HVI: hypervolume improvement. CI: mean

\pm

std across 3 seeds. Unique/Novel: fraction among valid molecules. All configurations achieve 100% validity except

\gamma=1000

(unique: 13%).

$\gamma$	MolWt	MolLogP	HBD	HBA	TPSA	RingCount	RotBonds
$0.001$	$0.22$	$0.13$	$0.02$	$0.16$	$0.14$	$0.12$	$0.48$
$0.1$	$0.22$	$0.12$	$0.02$	$0.16$	$0.14$	$0.12$	$0.46$
$1$	$0.22$	$0.13$	$0.02$	$0.15$	$0.15$	$0.12$	$0.46$
$10$	$0.24$	$0.13$	$0.02$	$0.17$	$0.14$	$0.11$	$0.42$
$\mathbf{92.4^{*}}$	$0.68$	$0.22$	$0.03$	$0.76$	$0.67$	$0.27$	$0.68$
$100$	$0.76$	$0.23$	$0.05$	$0.86$	$0.77$	$0.32$	$0.71$

Appendix

Table 5: Individual descriptor KL divergences

D_{\text{KL}}(P_{\text{gen}}\|P_{\text{train}})

\downarrow

across guidance strengths.

Data
Parameter	Value
Dataset	ZINC250K
Representation	SELFIES
Max sequence length	128
Vocabulary size	111
Train/test split	81/9/10

Appendix

Table 6: Hyperparameters for ZINC optimization experiments.

Panel of Figure 11 — Figure 11: UMAP of posterior-mean latent embeddings $\mu_{\phi}(x)$ for real molecules (fit set) and generated molecules after re-encoding. Each UMAP is fit on the real embeddings for the corresponding model run.

Explore similar work

Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design
May 28, 2026Sven Gutjahr, Riccardo De Santi, Luca Schaufelberger +2Constrained Generative Modeling Molecular Optimization
Department of Computer Science, ETH Zurich · ETH AI Center · NCCR Catalysis, Switzerland +1
Latent Flow Matching for Molecular Graph Generation
Oct 5, 2026Mathis Goupillon, Roman Bresson, Konstantinos Divriotis +1Flow Matching Molecular Graph Generation
Mohamed bin Zayed University of Artificial Intelligence
Molecular representation shapes the balance between target fidelity and exploration in flow based polymer generation
Sep 10, 2026Tianren ZhangMolecular Generation Inverse Design
Department of Materials Science and Engineering, University of Delaware, Newark, Delaware 19716, United States

May 28, 2026J/K move · Enter open · S save

Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design

Sven Gutjahr, Riccardo De Santi, Luca Schaufelberger, Kjell Jorner+1

Department of Computer Science, ETH Zurich · ETH AI Center · NCCR Catalysis, Switzerland+1

Constrained Generative Modeling Molecular Optimization Generative Modeling Flow-Based Generative Modeling

Open paper AskPDF arXiv

Adapting generative foundation models, in particular diffusion and flow models, to optimize given reward functions (e.g., binding affinity) while satisfying constraints (e.g., molecular synthesizability) is fundamental for their adoption in real-world scientific discovery applications such as molecular design or protein engineering. While recent works have introduced scalable methods for reward-guided fine-tuning of such models via reinforcement learning and control schemes, it remains an open problem how to algorithmically trade-off reward maximization and constraint satisfaction in a reliable and predictable manner. Motivated by this challenge, we first present a rigorous framework for Constrained Generative Optimization, which brings an optimization viewpoint to the introduced adaptation problem and retrieves the relevant task of constrained generation as a sub-case. Then, we introduce Constrained Flow Optimization (CFO), an algorithm that automatically and provably balances reward maximization and constraint satisfaction by reducing the original problem to sequential fine-tuning via established, scalable methods. We provide convergence guarantees for constrained generative optimization and constrained generation via CFO. Ultimately, we present an experimental evaluation of CFO on both synthetic, yet illustrative, settings, and a molecular design task. Across these evaluations, CFO achieves consistent increases in reward while ensuring high constraint satisfaction, showcasing its practical utility for constrained generative optimization.

Abstract

Explore similar work

Constrained Flow Optimization via Sequential Fine Tuning for Molecular Design

Latent Flow Matching for Molecular Graph Generation

Molecular representation shapes the balance between target fidelity and exploration in flow based polymer generation