MIRRA: Fast, High-Fidelity Molecular Generation with Chemistry-Informed Discrete Flow Matching

Ge Sun; Juan J. de Pablo

cs.LGOct 12, 2026

MIRRA: Fast, High-Fidelity Molecular Generation with Chemistry-Informed Discrete Flow Matching

Authors: Ge Sun, Juan J. de Pablo

Organizations: Courant Institute of Mathematical Sciences, New York University, New York, NY 10012 · Tandon School of Engineering, New York University, Brooklyn, NY 11201 · Department of Physics, New York University, New York, NY 10003

Abstract

Fast, high-fidelity molecular generation requires accurate chemical choices and efficient sampling. Chemical building blocks reduce graph size, but make each node a choice among tens of thousands of related identities. We introduce MIRRA (Molecular Identity Representation, Refresh and Assembly), a discrete flow-matching generator that learns from chemical knowledge and evaluates identity probabilities only when needed, while keeping every identity in its vocabulary available. Chemical features share information across related blocks, and learned identity corrections distinguish them in molecular context. By deciding which positions will refresh before computing their probabilities, the sampler avoids unused predictions while preserving each sampling transition. Against the strongest baselines in each task, MIRRA reduces chemical and graph distribution errors on ZINC250k by 29-31% and 53-60%, respectively, with generation up to 8.6 times faster. Property targeting yields an estimated 7.8-20.8 times more distinct novel hits per second. On a held-out polymer test set, MIRRA achieves 100% periodic validity with known chain topology, versus 57.0% for a polymer variational autoencoder. These results demonstrate fast, high-fidelity generation with the full chemical vocabulary available throughout sampling.

Figures & tables

Figure 1: MIRRA combines chemistry-informed prediction, refresh scheduling and constrained assembly. (a) A 26-heavy-atom molecule becomes five building-block nodes. The predictor

p_{\theta}

supplies node and connection probabilities for discrete flow matching, using refresh scheduling (b) and chemistry-informed identity scores (c). Constrained assembly

\mathcal{D}

uses final identities and these probabilities to construct a tree over the blocks and pair compatible attachment sites, with limited identity changes when needed. (b) Independent Bernoulli refresh decisions select positions for vocabulary scoring using the full graph as context. A selected identity can remain unchanged. At 250 steps, expected scoring evaluations per position fall from 250 to 6.1. (c) Chemical-feature and identity-correction scores, plus a frequency offset, define probabilities over all 37,821 entries. The four chemical examples are vocabulary members, not a shortlist; intermediate states, matrices and probability bars are schematic.

(a) Fidelity and validity: all baselines
Model (sampling steps)	FCD $\downarrow$	NSPDK $\times 10^{3}$ $\downarrow$	Valid (%) $\uparrow$
Training reference	0.336	0.110	n/a
GDSS (1000)	14.656	19.500	97.01
DiGress (1000)	3.482	2.100	94.99
GruM (1000)	2.257	1.500	98.65
GraphBFN (500)	2.116	1.300	99.22

Table 1: High fidelity at few sampling steps on ZINC250k. Both MIRRA settings lead on FCD and NSPDK. Bold / underline mark best/second-best among compared models, using unrounded values.

Figure 2: Strong chemical and graph fidelity at low sampling cost. (a) FCD and NSPDK discrepancies; (b) FCD versus steps (reported number of function evaluations, NFE, for HDFM); (c) FCD versus time per attempted molecule. Lower is better. Quality for MIRRA and FragFM uses 10,000 attempts; MIRRA250 is a three-draw median. MIRRA and FragFM suffixes denote sampling steps. Hatching marks the training-data floor, measured from a fixed training sample. The 8.6-fold arrow is the matched MIRRA25–FragFM500 timing comparison; their times are medians of four runs (three for FragFM250/1000). See Appendix A .

Figure 3: Schedule the computation while preserving the choices. (a) The positions and entries scored by three implementations of one checkpoint. (b) FCD and NSPDK: circles are three-draw medians; ticks show every draw. (c) Median time over three separate 1,024-attempt runs and peak allocated tensor memory, with individual runs shown. Scheduling is 3.0 times faster than dense scoring and slightly faster than the tested subset. Quality follows the 10,000-attempt protocol of Figure 2 ; timings use the same GPU and batch size 128. Candidate restriction has no separate training or proposal correction, and variation reflects sampling and timing repetitions.

Output predictor	FCD $\downarrow$	NSPDK $\downarrow$	Novel (%) $\uparrow$
		$(\times 10^{3})$
Conventional learned predictor	0.1793	0.1223	98.05
Shuffled chemical features	0.1757	0.1547	97.55
Chemistry shuffled within matched groups	0.1655	0.1188	97.51
MIRRA: correctly assigned chemical features	0.1561	0.1063	97.81

Table 2: Correct chemical features improve both chemical and graph fidelity on ZINC250k. Independently trained predictors share the input features, backbone, training budget and assembly rule, with comparable capacity. All rows use 250 sampling steps and the common final training update. Best values are bold; second-best values are underlined.

Figure 4: Assembly cannot replace accurate identity choices during generation. (a) A fixed-weight intervention includes or omits identity corrections during generation and assembly. (b) FCD on a logarithmic axis; (c) validity with an axis starting at 76%. Circles are medians and ticks show three sampling draws of the benchmark checkpoint. Assembly conditions within a generation arm share terminal states and edge probabilities. The full condition is exactly the MIRRA250 population in Figures 2 – 3 ; all use the same reference and evaluation. The roughly 95% FCD reduction compares full MIRRA with assembly-only correction.

Model	Steps	Structural distance $\downarrow$	Valid (%)	Unique (%)	Novel (%)
		(MMD ${}^{2}\times 10^{3}$ )	$\uparrow$	$\uparrow$	$\uparrow$
OMG-VAE	Native	$30.498\pm 0.458$	57.03	97.32	100.00
Graph DiT	500	$101.325\pm 2.048$	35.49	100.00	100.00
MIRRA	500	$\underline{0.079\pm 0.030}$	100.00	99.96	98.57
MIRRA	250	$\mathbf{0.021\pm 0.015}$	100.00	99.91	98.30

Table 3: Valid polymer repeat units with closer structural distributions. Evaluation uses 9,949 reserved periodic identities. MIRRA has 100% periodic validity at both step counts. Best/second-best values are bold/underlined; ties share rank.

Figure 5: Faster generation produces more useful property-targeted outputs. (a) Estimated distinct novel hits per second; (b) target-window FCD, including valid misses. Markers average three adapter runs sharing one backbone; ticks show each run. Rates combine 4,096-attempt yields with three separate 256-attempt timings; FCD uses 1,000 valid outputs and 1,000 target references. Tolerances are

\pm 0.05

QED and

\pm 0.5

logP.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Refresh scheduling reduces expected full-vocabulary evaluations per node from

T

H_{T}

. (a) At

T=250

, node refreshes per step against the schedule’s expectation

n/(T-k+1)

for

n

block nodes and one-based step index

k=1,\ldots,T

, on a logarithmic axis. Early steps refresh few positions and the final step refreshes every one. (b) The same at

T=25

. (c) Full-vocabulary evaluations per node. Dense scoring evaluates the vocabulary at every step; the schedule predicts the harmonic number

H_{T}

; the measured counts are total refreshes divided by the number of block nodes. Sampling uses 1,024 attempts at

T=250

and 128 at

T=25

from the benchmark checkpoint.

Common training data, one evaluation protocol
Setting	Valid $n$	FCD $\downarrow$	NSPDK $\downarrow$	Valid	Unique	Novel	ms/attempt
			$(\times 10^{3})$	(%)	(%)	(%)	$\downarrow$
MIRRA, 25 steps	9,927	0.3766	0.1795	99.27	99.98	97.69	15.7
MIRRA, 250 steps	9,980–9,985	0.3665	0.1537	99.81	99.97	97.61	23.3
MIRRA, 1000 steps	9,988	0.3824	0.1613	99.88	99.97	97.61	–
FragFM, 250 steps	9,878	0.5296	0.3782	98.78	99.97	99.22	81.5

Appendix

Table 4: ZINC250k results across sampling step counts. Both generators use the same training data. Bold marks the best value in each column, using unrounded values.

Matched 25 vs. 500 steps
Arm	Attempts	Rep 1	Rep 2	Rep 3	Rep 4	Median
	per rep	(ms)	(ms)	(ms)	(ms)	(ms)
MIRRA, 25 steps	4,096	26.31	15.77	15.61	15.04	15.69
FragFM, 500 steps	4,096	267.43	144.19	122.19	124.86	134.53
Matched 250 vs. 500 steps
MIRRA, 250 steps	4,096	23.40	22.90	23.24	24.48	23.32

Appendix

Table 5: Matched generation times. Individual timing repetitions and their medians.

Model	Steps	Median (s)	FCD $\downarrow$	Projected FCD $\downarrow$
MIRRA	25	17.09	0.1792	0.1814
MIRRA	250	27.83	0.1584	0.1585
Released FragFM	100	44.74	0.4220	0.3071
Released FragFM	250	78.66	0.3261	0.2228
Released FragFM	500	136.54	0.3356	0.2255

Appendix

Table 6: Fidelity and measured cost against the released FragFM checkpoint. Both measured MIRRA settings reach lower FCD and lower median generation time than the three tested released settings.

Figure 7: Consecutive MIRRA outputs beside training molecules. (a) The first 12 outputs of one 250-step run of the benchmark checkpoint, shown in generation order without filtering, deduplication or manual selection. All 12 are chemically valid and connected. (b) 12 molecules drawn uniformly at random, without replacement, from the training molecules. The strip beneath reports descriptor means over the first 48 outputs of the same run and 48 uniformly drawn training molecules.

End-to-end timing and peak allocated memory
Setting	FCD $\downarrow$	NSPDK $\downarrow$	Valid (%)	Unique (%)	Novel (%)
Dense, full vocabulary	0.3680	0.1543	99.81	99.97	97.62
Refresh, full vocabulary	0.3665	0.1537	99.81	99.97	97.61
Refresh, sampled subset	0.9469	0.8780	99.76	99.67	98.14
Setting	ms/attempt, median [range]		Peak allocation (GiB)
Dense, full vocabulary	131.56 [125.52, 191.77]		1.364

Appendix

Table 7: Exact refresh scheduling preserves full-support quality; candidate restriction changes it. One benchmark checkpoint at 250 steps.

Training pair	Aligned FCD	Shuffled FCD	Reduction (%)
A	0.1561	0.1757	11.19
B	0.1543	0.1727	10.63
C	0.1551	0.1841	15.74

Appendix

Table 8: Chemical alignment improves fidelity in three independently trained pairs. Primary 250-step, prior-assembly comparison.

RTX6000
Correction active at	FCD, median [range] $\downarrow$	NSPDK $\downarrow$	Valid	Unique	Novel
		$(\times 10^{3})$	(%)	(%)	(%)
Neither	7.4210 [7.4134, 7.4297]	16.8574	80.03	96.87	98.59
Assembly only	7.6384 [7.5242, 7.6598]	13.9544	88.10	89.52	98.54
Generation only	0.3675 [0.3538, 0.3726]	0.1600	99.68	99.97	97.65
Generation + assembly	0.3665 [0.3514, 0.3745]	0.1537	99.81	99.97	97.61

Appendix

Table 9: Corrections during generation dominate the fidelity improvement on both devices. The assembly effect is smaller and metric-dependent.

Correction	Valid / attempted	Validity (%)	FCD $\downarrow$
Neither	3,313 / 4,096	80.88	8.2378
Group only	4,047 / 4,096	98.80	4.0742
Identity only	3,266 / 4,096	79.74	4.1570
Both	4,090 / 4,096	99.85	0.6264

Appendix

Table 10: Both within-group and between-group corrections contribute. Group and identity corrections are varied during both generation and assembly in a fixed trained model.

Head	Run	250 steps		1000 steps
		Selected	Final	Selected	Final
Aligned	A	0.5282	0.5443	0.5169	0.5245
Aligned	B	0.5622	0.5622	0.5299	0.5299
Aligned	C	0.5142	0.5115	0.5333	0.5251
Mean $\pm$ SD		$0.5349\pm 0.0247$	$0.5393\pm 0.0257$	$0.5267\pm 0.0087$	$0.5265\pm 0.0030$
Shuffled	A	0.5626	0.5504	0.5401	0.5443

Appendix

Table 11: MOSES: three training runs and both checkpoint policies. FCD against TestSF; smaller is better.

Method	Property	Candidates	Selected
MIRRA	QED, logP	1; 2	2
Graph DiT	QED, logP	1; 2	2
FragFM	QED	(20, 1); (1042, 52.09)	(1042, 52.09)
FragFM	logP	(20, 1); (9.814, 0.491)	(20, 1)

Appendix

Table 12: Guidance candidates and validation-selected strengths. FragFM pairs list graph and bag strengths; graph-only guidance uses the graph entry.

QED, low (10th) quantile : target 0.529 $\pm$ 0.05
Model	Distinct novel	Hit rate	Distinct novel	Target-window	Valid
	hits / second	(%)	hits / attempt	FCD $\downarrow$	(%)
MIRRA	2.719	18.12	0.1739	5.156	99.71
FragFM, graph guidance	0.236	34.06	0.3386	5.330	90.92
FragFM, graph + bag guidance	0.225	35.45	0.3503	5.346	90.21
Graph DiT	0.119	9.69	0.0969	11.614	43.65

Appendix

Table 13: Complete property-targeted generation results. Distinct novel hits per second, hit rate, distinct novel hits per attempt, target-window FCD and validity for every target, model and guidance setting.

Figure 8: Hit rate and throughput behind Figure 5 a. (a) Target hits per attempted molecule, as a percentage of all 4,096 attempts. Both FragFM settings reach a higher hit rate than MIRRA at every target. (b) Attempts per second on a logarithmic axis. MIRRA completes 12.3–13.0 times as many attempts per second as the fastest baseline at each target, and 21.3–25.3 times as many as either FragFM setting. The rate in Figure 5 a is each run’s distinct novel fraction of attempts multiplied by its throughput. MIRRA markers average three property-adaptation runs sharing one frozen backbone; rates and throughputs are averaged separately.

Target	Hit rate (%), two-run mean			Distinct novel hits/s, real labels
	Real	Shuffled	Uncond.	Run A	Run B	Run C
QED, low	17.91	8.01	8.29	2.540	2.748	2.869
QED, median	31.56	28.50	29.36	4.956	4.828	4.830
QED, high	53.14	26.09	25.67	8.503	8.471	8.317
logP, low	21.36	11.01	10.41	3.494	3.452	3.610
logP, median	31.46	27.98	26.40	4.914	4.765	4.673

Appendix

Table 14: Condition-label controls and individual adaptation rates. Hit-rate controls average two matched adapter runs; distinct novel hit rates are reported for each of three adapter runs.

Output features	Selected update	Selected loss	Final loss
Aligned	192,000	2.4617	2.4632
Permuted	198,000	2.5347	2.5365

Appendix

Table 15: Chemically aligned output sharing lowers polymer prediction loss. Object-mean endpoint cross-entropy in nats.

Draw	Aligned	Permuted	Difference	Novelty A (%)	Novelty P (%)
1	2.133	2.750	0.62	98.63	97.27
2	5.900	6.749	0.85	98.54	97.61
3	8.946	7.810	$-1.14$	98.76	97.97
4	5.193	8.755	3.56	98.27	97.41
Mean	5.543	6.516	0.97	98.55	97.56

Appendix

Table 16: Aligned output sharing lowers polymer structural MMD 2 in three of four additional 250-step draws. Paired draws of 4,096 attempts per arm against the full validation reference.

Checkpoint policy	Draw	Aligned head	Conventional head
Validation-selected EMA	1	3.243	11.36
	2	0.5184	9.745
	Mean	1.881	10.55
Final-update EMA	1	3.163	13.84
	2	1.081	11.19
	Mean	2.122	12.51

Appendix

Table 17: The shared output head achieves lower structural discrepancy than a capacity-matched conventional head under both checkpoint policies. Paired 250-step draws of 4,096 attempts each; unbiased structural MMD 2 in units of

10^{-5}

Molecular Generation

Molecular generation develops computational methods to create chemical structures and conformations for discovery, optimization, and structure elucidation.

See all 231 papers

FlashMol: High-Quality Molecule Generation in as Few as Four Steps
May 7, 2026Xinyuan Wei, Zian Li, Shaoheng Yan +2Diffusion Distillation Few-Step Diffusion Sampling
Institute for Artificial Intelligence, Peking University · Yuanpei College, Peking University · School of Intelligence Science and Technology, Peking University +2
MolGVR: A Chemistry-Grounded Framework for Text-to-Molecule Generation
Jul 31, 2026Qian Tan, Xuanyu Zhu, Lei Jiang +3Molecular Generation Text Generation
University of Science and Technology of China · Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory
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

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

FlashMol: High-Quality Molecule Generation in as Few as Four Steps

Xinyuan Wei, Zian Li, Shaoheng Yan, Cai Zhou+1

Institute for Artificial Intelligence, Peking University · Yuanpei College, Peking University · School of Intelligence Science and Technology, Peking University+2

Diffusion Distillation Few-Step Diffusion Sampling 3D Molecular Generation Molecular Generation Distribution Matching Distillation Conformer

Open paper AskPDF arXiv

Generating chemically valid 3D molecular conformations is critical for computational drug discovery. Classical diffusion-based models like GeoLDM perform well but require hundreds of steps, making large-scale in silico screening impractical. Recent efforts on few-step molecular generation have accelerated this process to 12-50 steps, but they often largely sacrifice sample stability. In this work, we present FlashMol, an ultra-fast molecule generative model producing high-quality molecular conformations in as few as 4 steps. To achieve this, we adapt distribution matching distillation (DMD) - a reverse KL-divergence minimization objective - to the molecular domain for effective distillation. Considering the local minimization behavior of DMD, we respace the molecule generation timesteps, providing the generator with much better initialization and enables effective distillation. Additionally, to mitigate the mode-seeking behavior of DMD and improve diversity, we further regularize it with a Jensen-Shannon divergence term, which incorporates the mean-seeking behavior of the forward KL divergence. Extensive experiments on QM9 and GEOM-DRUG datasets demonstrate that FlashMol matches and even surpasses the original 1000-step teacher, achieving up to 250

\times

acceleration in sampling speed while maintaining high molecular quality.

Abstract

Molecular Generation

Explore similar work

FlashMol: High-Quality Molecule Generation in as Few as Four Steps

MolGVR: A Chemistry-Grounded Framework for Text-to-Molecule Generation

Latent Flow Matching for Molecular Graph Generation