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
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θ supplies node and connection probabilities for discrete flow matching, using refresh scheduling (b) and chemistry-informed identity scores (c). Constrained assembly 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 ↓
NSPDK ×103↓
Valid (%) ↑
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 ↓
NSPDK ↓
Novel (%) ↑
(×103)
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 ↓
Valid (%)
Unique (%)
Novel (%)
(MMD 2×103 )
↑
↑
↑
OMG-VAE
Native
30.498±0.458
57.03
97.32
100.00
Graph DiT
500
101.325±2.048
35.49
100.00
100.00
MIRRA
500
0.079±0.030
100.00
99.96
98.57
MIRRA
250
0.021±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 ±0.05 QED and ±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 to HT . (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,…,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 HT ; 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.
Setting
Valid n
FCD ↓
NSPDK ↓
Valid
Unique
Novel
ms/attempt
(×103)
(%)
(%)
(%)
↓
Common training data, one evaluation protocol
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.
Arm
Attempts
Rep 1
Rep 2
Rep 3
Rep 4
Median
per rep
(ms)
(ms)
(ms)
(ms)
(ms)
Matched 25 vs. 500 steps
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 ↓
Projected FCD ↓
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.
Setting
FCD ↓
NSPDK ↓
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
End-to-end timing and peak allocated memory
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.
Correction active at
FCD, median [range] ↓
NSPDK ↓
Valid
Unique
Novel
(×103)
(%)
(%)
(%)
RTX6000
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 ↓
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 ± SD
0.5349±0.0247
0.5393±0.0257
0.5267±0.0087
0.5265±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.
Model
Distinct novel
Hit rate
Distinct novel
Target-window
Valid
hits / second
(%)
hits / attempt
FCD ↓
(%)
QED, low (10th) quantile : target 0.529 ± 0.05
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.
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 develops computational methods to create chemical structures and conformations for discovery, optimization, and structure elucidation.
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× acceleration in sampling speed while maintaining high molecular quality.
Xinyuan Wei, Zian Li, Shaoheng Yan +2
Institute for Artificial Intelligence, Peking University · Yuanpei College, Peking University · School of Intelligence Science and Technology, Peking University +2
Text-to-molecule generation is typically formulated as a one-shot sequence generation problem, where a model directly maps target descriptions to molecular representations. However, molecular descriptions often contain informative structural constraints, and violating such constraints can change the molecular identity. This makes chemical verification and error correction important but underexplored. To fill this gap, we propose MolGVR, a chemistry-grounded Generator--Verifier--Refiner framework. The Generator infers structural evidence and generates candidate molecules. The Verifier addresses the lack of chemical validation by converting descriptions into chemical constraints and checking candidates against them. The Refiner addresses generation failures by revising candidates rejected by the Verifier. Experiments on ChEBI-20 and PCDes show that MolGVR improves exact-match performance. These results suggest that coupling generation with executable verification and feedback-guided refinement is an effective way to improve text-to-molecule generation.
Qian Tan, Xuanyu Zhu, Lei Jiang +3
University of Science and Technology of China · Shanghai Jiao Tong University · Shanghai Artificial Intelligence Laboratory
Modern graph generative models typically operate directly in the discrete graph space, explicitly generating node and edge variables, which can become costly as graphs grow. In this paper, we perform generation explicitly on latent representations of entire graphs obtained from a pretrained Variational Autoencoder with high reconstruction fidelity. The generated representations, obtained through flow matching, are then decoded only at the final step. Across molecular benchmarks of increasing size, our approach achieves strong validity and FCD while offering a favorable quality-efficiency trade-off compared with state-of-the-art explicit graph generative models. One of the main advantages of this formulation is that the graph representation only needs to be learned once, after which the same one can be reused across multiple generative objectives without retraining. We demonstrate generation guided by molecular properties and further introduce validity-aware generation though a classifier learned directly in latent space. All code will be made available upon acceptance.
Mathis Goupillon, Roman Bresson, Konstantinos Divriotis +1
Mohamed bin Zayed University of Artificial Intelligence