Trajectory Matching: Discrete Flow Matching with Structured Actions

Niklas Kemper; Marten Lienen; David Lüdke; Stephan Günnemann

cs.LGOct 23, 2026

Trajectory Matching: Discrete Flow Matching with Structured Actions

Authors: Niklas Kemper, Marten Lienen, David Lüdke, Stephan Günnemann

Organizations: School of Computation, Information and Technology & Munich Data Science Institute, Technical University of Munich

Abstract

The natural units of structured discrete generation are often not independent single-component edits but richer actions: attaching a fragment to a molecule, inserting a node together with an incident edge, or growing a hierarchy by adding a child. Yet most existing discrete flow matching methods (DFM) are restricted to independent component-wise edits over fixed-size states. This restriction stems from endpoint conditioning: endpoints fix the source and target during training but not the action sequence between them, making tractable conditional rates hard to define for complex actions that cannot be factorized component-wise. We propose trajectory matching, which conditions the flow on a full sequence of actions connecting the endpoints. This yields principled, simulation-free training targets for user-defined action spaces, including actions that change state size, modify multiple components, or depend on earlier actions. We prove that the resulting objective recovers the marginal flow and that trajectory matching strictly subsumes factorized DFM, while additionally capturing actions with prerequisite structure such as hierarchy growth or fragment-based molecule generation. Empirically, on tree and molecular graph generation, trajectory matching with domain-aligned actions matches or surpasses factorized DFM baselines (lifting MOSES validity from 92.8% to 99.2%) and transfers without retraining to substructure extension and linking. Trajectory matching enlarges the design space of discrete flow matching by turning the action vocabulary into a modeling choice. An interactive explanation of trajectory matching is available at https://niklaskemper.com/blog/trajectory-matching/.

Figures & tables

Figure 1: Training step of trajectory matching. Conditioning on a full trajectory (rather than endpoints alone) yields tractable training targets for general action spaces (including multi-component, dependent, or size-changing actions). Pseudocode in Fig. 4 .

Figure 2: Four trajectory-matching instantiations on molecular graphs, BondEdit (factorized bond edits), FragGrow (fragment-based growth), AtomGrow (atom/bond growth from a seed), and MolEdit (noisy-initial editing), each shown with a sampled source $s_{I}$ , action set, example trajectory ${\tau}$ , and Rao-Blackwellized target rate $\bar{u}_{t}$ . Full definitions in App. E , more instantiations in Fig. 6

Table 3

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: FragmentGrow decompostion and dependency tree: Given the decomposition of $s_{F}$ into fragments and atom/bonds, there is one unique set of actions to generate $s_{F}$ from $s_{I}$ (marked in red), and the dependencies between these actions form a tree.

Figure 6: Examples of different existing methods that can be seen as instances of Trajectory Matching (text unmasking, text denoising, DeFog GraphGeneration Qin et al., 2025 and EditFlows Havasi et al., 2025 ) and our new instantiations of Trajectory Matching (AtomGrow, FragGrow, MolEdit, BondEdit, and GrowHierarchy). for each instantiation, we show an initial state

s_{I}\sim p_{\text{src}}

, the action space

\mathcal{A}

,a sample trajectory

\tau

from

s_{I}

to a final state

s_{F}\sim p_{\text{tgt}}

and the Rao-Blackwellized target rate

\bar{u}_{t}(s,a\mid s_{I},s_{F})

of an intermediate state

s

Objective	Degree $\downarrow$	Spectral $\downarrow$	Wavelet $\downarrow$	Validity (%) $\uparrow$
Decomposed	$0.0006\pm 0.0002$	$0.0072\pm 0.0006$	$0.0033\pm 0.0002$	$100.0\pm 0.0$
Direct Bregman	$0.0008\pm 0.0002$	$0.0085\pm 0.0009$	$0.0034\pm 0.0004$	$100.0\pm 0.0$

Appendix

Table 4: Objective ablation on Tree with GrowHierarchy . Entries report means and standard deviations over sampling seeds from a single trained model per objective.

Method	Trajectory sampling + target construction (ms)	Observed loader wait (ms)
AtomGrow	4.56	0.75
FragGrow	188.83	0.94
MolEdit	30.55	0.91

Appendix

Table 5: CPU preprocessing time and observed loader wait on QM9. Preprocessing overlaps with GPU training; loader wait is measured per training step and is not the total overhead relative to DeFoG.

Method	Training throughput (steps/s)	Sampling throughput (samples/s)
DeFoG	13.79	10.22
AtomGrow	11.89	7.96
FragGrow	12.01	7.63
MolEdit	12.45	7.69

Appendix

Table 6: End-to-end throughput on QM9 using one A100 GPU, averaged after warmup. Sampling uses 500 integration steps for every method.

Model	Class	Deg. $\downarrow$	Clus. $\downarrow$	Orbit $\downarrow$	Spec. $\downarrow$	Wavelet $\downarrow$	Ratio $\downarrow$	Valid $\uparrow$	Unique $\uparrow$	Novel $\uparrow$	V.U.N. $\uparrow$
		Tree Dataset
Train set	—	0.0001	0.0000	0.0000	0.0075	0.0030	1.0	100	100	0.0	0.0
GraphRNN ( You et al., 2018 )	Autoregressive	—	—	—	—	—	607.0	—	—	—	0.0
GRAN ( Liao et al., 2020 )	Autoregressive	0.1884	0.0080	0.0199	0.2751	0.3274	607.0	0.0	100	100	0.0
DiGress ( Vignac et al., 2023 )	Diffusion	0.0002	0.0000	0.0000	0.0113	0.0043	1.6	90.0	100	100	90.0
EDGE ( Chen et al., 2023 )	Diffusion	0.2678	0.0000	0.7357	0.2247	0.4230	850.7	0.0	7.5	100	0.0

Appendix

Table 7: Graph generation results on the Tree dataset.

Figure 7: Validity (left) and Ratio (right) on the Tree benchmark as a function of the number of Euler integration steps. GrowHierarchy preserves $100\%$ validity across three orders of magnitude of step budgets, while DeFoG’s validity collapses below $\sim\!10\%$ at $10$ steps. Average Ratio metric is comparable for both.

	Model	Val. $\uparrow$	R. Val. $\uparrow$	Uniq. $\uparrow$	FCD $\downarrow$
	Training set	99.3	99.5	99.2	0.03
Fixed-size	DiGress ( Vignac et al., 2023 )	99.0 ±0.0	–	96.2 ±0.1	–
	DisCo ( Xu et al., 2024 )	99.3 ±0.6	–	–	–
	Cometh ( Siraudin et al., 2024 )	99.6 ±0.1	–	96.8 ±0.2	0.25 ±0.01
	GBD ( Liu et al., 2024 )	99.8	–	–	0.093
	DeFoG ( Qin et al., 2025 )	99.3 ±0.0	99.4 ±0.1	96.3 ±0.3	0.12 ±0.00

Appendix

Table 8: Comparison of DFM and diffusion molecular graph generative models on QM9 (without hydrogens). Best results in bold , second best underlined .

Model	Validity $\uparrow$	Uniqueness $\uparrow$	Filters $\uparrow$	FCD $\downarrow$	SNN $\uparrow$	Scaffold $\uparrow$
GraphINVENT	96.4	99.8	95.0	1.22	0.54	12.7
FragGrow	99.2	99.9	99.3	1.21	0.55	13.7

Appendix

Table 9: Autoregressive comparison on MOSES. GraphINVENT results are taken from DiGress ( Vignac et al., 2023 ) . Results for all flow-matching/diffusion methods in Tab. 2 .

Figure 8: ZINC validity as a function of the number of Euler steps for the fragment-based FragGrow and the factorized BondEdit . FragGrow maintains substantially higher validity in the low-step regime (e.g., $78\%$ vs. $48\%$ at $5$ steps); both methods saturate at $\geq\!98\%$ for $\geq\!25$ steps.

Figure 9: Training dynamics on ZINC under AtomGrow : the Rao-Blackwellized target $\bar{u}_{t}$ converges faster and with lower epoch-to-epoch variance than the trajectory-conditioned target $u_{t}$ , while reaching the same final validity. The behavior matches the variance-reduction guarantee of Prop. B.3 .

Policy	Validity (%) $\uparrow$	Relaxed validity (%) $\uparrow$	Uniqueness (%) $\uparrow$	FCD $\downarrow$
Random	$98.8\pm 0.0$	$99.2\pm 0.0$	$97.0\pm 0.0$	$0.096\pm 0.00$
BFS	$98.9\pm 0.1$	$99.2\pm 0.1$	$96.6\pm 0.2$	$0.092\pm 0.00$
DFS	$98.8\pm 0.1$	$99.2\pm 0.1$	$96.8\pm 0.2$	$0.098\pm 0.00$

Appendix

Table 10: Effect of the trajectory distribution on AtomGrow on QM9, using single-trajectory training targets for all three policies. Results are means and standard deviations over five independent sampling runs from one trained model per policy.

Figure 10: Substructure extension on eleven seed substructures (rows) and four Euler start times (columns), reporting the valid-and-connected fraction of generated completions. AtomGrow (left) accommodates the task by setting $s_{I}$ to the seed; BondEdit (right) is adapted by sampling a target size and freezing the seed during generation. AtomGrow substantially outperforms the factorized baseline across all seeds and start times, including substructures absent from ZINC training molecules.

Panel of Figure 11 — Figure 11: Valid completions sampled by AtomGrow for two seed substructures at four Euler start times $t\in\{0.2,0.4,0.6,0.8\}$ (one row each). The seed (leftmost column) is preserved as a substructure of every completion. Lower $t$ leaves more of the trajectory remaining and produces larger, more elaborately decorated extensions; higher $t$ yields small decorations on the seed.

Figure 12: Substructure linking on six seed pairs (rows) and four Euler start times (columns), reporting the valid-and-connected fraction of generated molecules. AtomGrow (left) sets $s_{I}$ to the union of both seeds and uses both as roots of the action DAG; BondEdit (right) is adapted by sampling a target size, freezing both seeds, and requiring inserted bonds to link them. AtomGrow substantially outperforms the factorized baseline across all pairs and start times.

Panel of Figure 13 — Figure 13: Valid, connected linkings sampled by AtomGrow for two seed pairs at four Euler start times $t\in\{0.2,0.4,0.6,0.8\}$ (one row each). Both seed substructures (leftmost column) appear in every generated molecule, joined by a linker whose complexity is controlled by $t$ .

Discrete Flow Matching

Discrete flow matching learns probability flows that transform simple categorical noise into structured discrete data.

See all 34 papers

May 8, 2026

Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation

Discrete flow matching generates text by iteratively transforming noise tokens into coherent language, but may require hundreds of forward passes. Distillation uses the multi-step trajectory to train a student to reproduce the process in a few steps. When the student underperforms, the usual explanation is insufficient capacity. We argue the opposite: the trajectory is the bottleneck, not the student. Each training trajectory is built through a chain of blind stochastic jumps with no evaluation of sequence quality; a single bad decision at an early midpoint propagates through subsequent steps, yet the student must imitate the result. Trajectory-Shaped Discrete Flow Matching (TS-DFM) replaces these blind jumps with guided navigation: a lightweight energy compass evaluates candidate continuations at each midpoint, selecting the most coherent. All shaping is training-only; inference cost is unchanged. On 170M-parameter language modeling, the shaped student at 8 steps achieves 32% lower perplexity than the 1,024-step teacher while being 128x faster, with gains consistent across source distributions and three evaluators of increasing scale. TS-DFM achieves the best perplexity of any discrete-generation baseline we compare against, including methods trained on 6x more data or using 5x larger models.

Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade +3

1The Ohio State University · 2Apple

Flow Matching Diffusion Language Model Inference Language Model Distillation Flow-Based Generative Modeling Teacher-Student Learning Discrete Flow Matching

Ask PDF

Oct 4, 2026

Discrete Action Matching: Learning Stochastic Dynamics from Samples via State Graphs

Learning population dynamics from unpaired temporal marginals is an ill-posed inverse problem that requires structural assumptions on the underlying dynamics. We introduce

\textit{Discrete Action Matching}

(DAM), a finite-state counterpart of Action Matching based on discrete Wasserstein geometry. For a prescribed marginal path and transport geometry, we derive an action-minimization objective for its canonical minimum-kinetic-energy current. Our key observation is that the density dependence of the discrete action reduces to neighboring density ratios. Along an empirical interpolation of the snapshots, DAM first estimates these ratios and then learns an action potential. The learned fields also define a graph-supported Markov sampler. Experiments on controlled synthetic dynamics and real mouse gastrulation data evaluate marginal reconstruction and interpolation. Additional experiments approximate numerical surface-transport paths from paired samples.

Mikhail Persiianov, Alexander Korotin

Applied AI Institute Moscow, Russia

Latent Dynamics Modeling Optimal Transport

Ask PDF

May 9, 2026

Discrete Flow Matching: Convergence Guarantees Under Minimal Assumptions

Flow Matching has recently emerged as a popular class of generative models for simulating a target distribution

μ_1

from samples drawn from a source distribution

μ_0

. This framework relies on a fixed coupling between

μ_0

and

μ_1

, and on a deterministic or stochastic bridge to define an interpolating process between the two distributions. The time marginals of this process can then be approximately sampled by estimating the transition rates, or more generally the generator, of its Markovian projection. This framework has recently been extended to the case of discrete source and target distributions, under the name Discrete Flow Matching (DFM). However, theoretical guarantees for such models remain scarce. In this paper, we study two DFM models on

\mathbb{Z}_m^d = \{0,\ldots,m-1\}^d

, sampled through time discretization, and derive non-asymptotic associated bounds for both of them. In contrast to previous work, we establish non-asymptotic bounds in Kullback--Leibler divergence for the early-stopped version of the target distribution. We also derive explicit convergence guarantees in total variation distance with respect to the true target distribution. Importantly, these bounds rely only on an approximation error assumption, relaxing standard score assumptions used in earlier works, while also yielding improved dependence on the vocabulary size

m

and the dimension

d

Le-Tuyet-Nhi Pham, Giovanni Conforti, Zhenjie Ren +1

École Polytechnique, CMAP, IP Paris Route de Saclay, 91120 Palaiseau, France · Università degli Studi di Padova, Via Trieste, 63, 35131 Padova, Italia · LaMME, Université Évry Paris-Saclay Boulevard François Mitterrand 91000, France

Flow Matching Generative Modeling Flow-Based Generative Modeling Discrete Flow Matching

Ask PDF

Abstract

Discrete Flow Matching

Explore similar work

Trajectory as the Teacher: Few-Step Discrete Flow Matching via Energy-Navigated Distillation

Discrete Action Matching: Learning Stochastic Dynamics from Samples via State Graphs

Discrete Flow Matching: Convergence Guarantees Under Minimal Assumptions