Trajectory-Shaped Discrete Flow Matching

• TS-DFM is a framework that explicitly shapes the trajectory between source and target distributions using designed probability paths and time reparameterization.
• It unifies methods across diverse domains, including chemical transition-state prediction, autonomous driving, speech recognition, and text generation.
• By learning targeted velocity fields and employing efficient few-step solvers, TS-DFM achieves significant performance improvements in both accuracy and sampling efficiency.

Trajectory-Shaped Discrete Flow Matching (TS-DFM) denotes a class of flow-matching methods in which the path from a source distribution to a target distribution is explicitly shaped so that generation can proceed through structured intermediate states under a discrete or discretized solver. In the literature represented here, the label is used in two closely related senses. First, it is the explicit name of a transition-state prediction framework that transports an initial transition-state guess to a true transition-state geometry in molecular distance space (Luo et al., 21 Nov 2025). Second, it functions as a broader methodological description for discrete or few-step flow-matching systems whose trajectories are shaped by probability-path design, couplings, time reparameterization, in-the-loop perturbations, or training-only guidance (Monsefi et al., 8 May 2026, Gat et al., 2024). Across these uses, the central premise is that path design is not ancillary: it determines what velocity field is learned, what occupancies are visited during training and sampling, and how accurately the resulting process can be approximated with a small number of discrete steps.

1. Terminology and scope

In its narrowest sense, TS-DFM is the method introduced for chemical transition-state generation in distance-geometry space. That framework predicts transition-state geometries from reactant and product structures by learning a velocity field over pairwise distance matrices, using an optimal-transport conditional flow-matching construction and a purpose-built network, TSDVNet (Luo et al., 21 Nov 2025). In a broader sense, several recent works instantiate the same design logic without necessarily using the label: a flow is defined over a whole trajectory, sequence, or multimodal token array; the probability path is deliberately shaped; and inference is performed with a small, fixed number of discrete updates (Wang et al., 26 Sep 2025, Navon et al., 5 Oct 2025, Luo et al., 15 Oct 2025).

This broader usage is especially visible in discrete sequence modeling. Discrete Flow Matching formalizes generation on $\mathcal{D}=[d]^N$ by specifying a family of probability paths $p_t$, conditional token paths $p_t(x_i\mid x_0,x_1)$, and a probability velocity $u_t$ that defines a continuous-time Markov chain on the discrete state space (Gat et al., 2024). Under that formalism, “trajectory shaping” refers to the deliberate choice of path family, scheduler $\kappa_t$, intermediate distributions, or solver-time geometry. Drax makes this interpretation explicit at the design level by arguing that the training path should resemble likely intermediate inference errors in ASR rather than a direct random-noise-to-target transition (Navon et al., 5 Oct 2025). NExT-OMNI adopts metric-induced probability paths and kinetic-optimal velocities over multimodal token spaces, which is another direct instance of shaping the discrete transport trajectory itself (Luo et al., 15 Oct 2025).

The term therefore covers both a specific method family and a more general design doctrine. What unifies these usages is not a single architecture or domain, but the view that the trajectory between source and target is a primary modeling object.

2. Mathematical formulation

A standard continuous formulation begins with a source sample $x_0\sim p_0$, a data sample $x$, and a linear rectified path

$x_t = (1-t)x_0 + tx,\qquad t\in[0,1].$

Along that path, the target instantaneous velocity is constant, $x-x_0$, and the model learns a conditional vector field $v_\theta(t,x_t,c)$ by regressing to that target velocity. FlowDrive applies this construction directly to full future driving trajectories, and FlowTS applies the same rectified-flow logic to multivariate time series (Wang et al., 26 Sep 2025, Hu et al., 2024). In both cases, training uses flow-matching regression, while inference integrates the learned ODE with a small number of Euler steps.

Discrete Flow Matching replaces the continuous vector field with a probability velocity on a discrete state space. In the general formulation, a conditional token path is written as

$p_t$0

and the associated marginal probability velocity is obtained by averaging conditional velocities under the posterior $p_t$1 (Gat et al., 2024). The resulting process is a continuous-time Markov chain whose marginals satisfy the discrete continuity equation. The end-to-end DFM theory paper further studies this setting on $p_t$2, with factorized coordinate-wise rates $p_t$3, and proves that the total variation distance between generated and target distributions is controlled by the risk of the learned velocity field (Su et al., 26 Sep 2025).

The chemistry TS-DFM instantiates conditional flow matching in a more specialized way. It defines a deterministic initial transition-state guess in distance space,

$p_t$4

then transports it toward the true transition-state distance matrix $p_t$5 along a Gaussian optimal-transport path

$p_t$6

with target velocity $p_t$7 (Luo et al., 21 Nov 2025). Its training loss is the corresponding squared-error regression of the parametric velocity field $p_t$8 toward that constant displacement.

These formulations differ in state space and solver mechanics, but they share the same abstract structure: specify a path, derive its target velocity or rate field, and learn a neural surrogate that can be sampled with few discrete updates.

3. Probability-path design and time geometry

Probability-path design is the most direct form of trajectory shaping. In Drax, the standard two-way source-target mixture is replaced by a tri-mixture path

$p_t$9

where $p_t(x_i\mid x_0,x_1)$0 is an audio-conditioned middle distribution and $p_t(x_i\mid x_0,x_1)$1 peaks at $p_t(x_i\mid x_0,x_1)$2 under the scheduler $p_t(x_i\mid x_0,x_1)$3, $p_t(x_i\mid x_0,x_1)$4, $p_t(x_i\mid x_0,x_1)$5, $p_t(x_i\mid x_0,x_1)$6 (Navon et al., 5 Oct 2025). This construction is motivated by train-inference occupancy mismatch: acoustically plausible intermediate transcripts are closer to actual ASR decoding trajectories than direct source-target mixtures.

NExT-OMNI uses a different shaping mechanism. Instead of a convex source-target mixture, it defines a metric-induced probability path

$p_t(x_i\mid x_0,x_1)$7

with cosine distance in a shared token-embedding space and $p_t(x_i\mid x_0,x_1)$8, $p_t(x_i\mid x_0,x_1)$9, $u_t$0 (Luo et al., 15 Oct 2025). The associated kinetic-optimal velocity transports probability mass only toward tokens closer to the target under that metric. This yields a monotone, minimum-energy path through the discrete simplex.

Time geometry is another major shaping axis. FlowTS keeps the straight-line rectified path in data space but changes how training and sampling emphasize different regions of that path. It samples $u_t$1 from a Logit-Normal distribution during training, uses a shifted time schedule

$u_t$2

for unconditional inference, and uses $u_t$3 for conditional inference, with $u_t$4 shifting steps closer to the data end of the path (Hu et al., 2024). TR-CIE makes the same issue explicit for DFM samplers by introducing the schedule-based reparameterization

$u_t$5

which absorbs the factor $u_t$6 under standard factorized DFM rates and mitigates stiffness near terminal time (Fu et al., 23 Jun 2026).

A plausible synthesis is that TS-DFM is best understood as controlling not only where the path goes, but how time is spent along it. Middle distributions, metric-induced geodesic-like paths, and schedule-based time warps are all concrete realizations of that principle.

4. State representations, architectures, and in-the-loop shaping

Trajectory shaping is tightly coupled to representation choice. FlowDrive is explicitly trajectory-shaped: the state is an entire ego plan $u_t$7, with $u_t$8 waypoints and $u_t$9 features $\kappa_t$0, flattened to $\kappa_t$1 for flow matching (Wang et al., 26 Sep 2025). Its encoder builds scene tokens from neighbors, static objects, and lanes/routes; its decoder is a DiT-style transformer with adaptive LayerNorm-zero driven by $\kappa_t$2 and pooled context. Sampling then updates the whole trajectory at each flow step, rather than using autoregression over horizon positions.

The chemistry TS-DFM uses an entirely different state space: pairwise molecular distances. TSDVNet maintains atom representations $\kappa_t$3 and pair representations $\kappa_t$4, with two branches that process the current transition-state distances and the reactant/product context, then fuse them through distance-aware attention, mixing updates, and AlphaFold-like triangular updates before predicting a pairwise distance velocity field $\kappa_t$5 (Luo et al., 21 Nov 2025). This representation enforces rotation and translation invariance at the level of the flow state itself.

A third pattern is training-only in-the-loop shaping. In few-step text generation, TS-DFM replaces blind midpoint jumps in RK-4 distillation with guided navigation controlled by an energy compass $\kappa_t$6 implemented as a 6-layer Transformer encoder with no time conditioning (Monsefi et al., 8 May 2026). At each midpoint, the method first performs sequence-level candidate selection over $\kappa_t$7 CTMC jumps, default $\kappa_t$8, then a token-level refinement based on velocity alignment, and finally a safeguard that accepts the refinement only if the energy does not increase beyond $\kappa_t$9. All shaping is training-only; inference cost is unchanged.

FlowDrive provides a related but distinct mechanism at inference time. It injects a deterministic perturbation into the current trajectory inside the flow integration,

$x_0\sim p_0$0

then evaluates the vector field at the perturbed state and continues integration (Wang et al., 26 Sep 2025). The paper explicitly notes that this is not classifier guidance; it is a deterministic control perturbation inserted inside the flow so that the learned field can re-project candidates toward feasible trajectories.

Across these examples, TS-DFM is not reducible to one solver or one architecture. It is a coupling between state design, conditioning mechanism, and a specific strategy for steering intermediate states.

5. Domain-specific instantiations and empirical results

In chemical transition-state prediction, TS-DFM is the method name rather than an ex post interpretation. On Transition1X, TS-DFM improves over React-OT from RMSD mean 0.2436 to 0.1828 and DMAE mean 0.0934 to 0.0607, yielding “~30% improvement in both RMSD and DMAE.” It also reduces $x_0\sim p_0$1 from mean 0.4272 to 0.1614, improves saddle-point rate from 55.6% to 65.5%, and reduces force metrics from $x_0\sim p_0$2, $x_0\sim p_0$3 to $x_0\sim p_0$4, $x_0\sim p_0$5 (Luo et al., 21 Nov 2025). As an initialization for CI-NEB, TS-DFM yields the highest percentage of saddle points, 74.22%, and the fastest convergence, 19.7 s average. On RGD1, it improves RMSD and DMAE by at least 16% over React-OT on all test splits.

In autonomous driving, FlowDrive is not named TS-DFM, but it exemplifies the same design. Cluster-based sampling on trajectory patterns improves nuPlan Val14 (R) from 81.91 with no weighted sampling to 85.37, while scenario-based sampling drops performance to 80.08 (Wang et al., 26 Sep 2025). The model uses only 6–10 inference steps, with the best setting at 8 Euler steps and about 40 ms per pass for a single trajectory on an RTX 2080 Ti; the guided multi-candidate FlowDrive* reaches about 43 ms with batching for 30 candidates.

In speech recognition, Drax shows that a shaped discrete path can improve the accuracy-efficiency trade-off of non-autoregressive decoding. Its training-path ablation compares four path designs and reports that “the uniform source + audio middle (Drax’s design) achieves the best WER throughout training,” outperforming both the direct uniform-only path and an audio-conditioned source without middle (Navon et al., 5 Oct 2025). On English benchmarks, Drax (16/1) reports average WER 8.4 and RTFx 32.2, while Drax with Whisper rescoring at 8/16 reaches average WER 7.4 and RTFx 17.8.

In text generation, the few-step TS-DFM of energy-navigated distillation makes the trajectory itself the object of optimization. For a 170M-parameter model with uniform source, the 1,024-step teacher has GPT-2 perplexity 82.8, FS-DFM at 8 steps has 87.6, and TS-DFM at 8 steps reaches 56.1, which the paper summarizes as “32% lower perplexity than the 1,024-step teacher while being 128x faster” (Monsefi et al., 8 May 2026). With mask source, the 1,024-step teacher has 93.6, FS-DFM at 8 steps has 516.6, and TS-DFM at 8 steps with DFM initialization reaches 91.4. At 1.3B parameters with uniform source, the 1,024-step teacher has 57.5, FS-DFM at 8 steps has 81.5, and TS-DFM reaches 48.0.

In multimodal token generation, NExT-OMNI demonstrates that metric-shaped DFM can support unified any-to-any modeling. The paper reports that its final configuration reaches average score 45.6 in the paradigm-and-representation ablation, versus 41.4 for the autoregressive baseline, while also improving generation and retrieval metrics (Luo et al., 15 Oct 2025). In time series, FlowTS reports context FID scores of 0.019 and 0.011 on Stock and ETTh in the unconditional setting, solar forecasting MSE 213, and MuJoCo imputation MSE $x_0\sim p_0$6 in the conditional setting, while saturating in quality around $x_0\sim p_0$7 steps (Hu et al., 2024).

The empirical pattern is consistent across domains: when the path is shaped to reflect geometry, occupancy, feasibility, or sequence quality, few-step sampling improves disproportionately.

6. Theory, misconceptions, limitations, and open directions

A central theoretical result for DFM is that distributional error can be controlled by velocity-field error. In the end-to-end DFM analysis, for factorized velocities the total variation distance satisfies

$x_0\sim p_0$8

linking terminal distribution estimation directly to the learned velocity risk (Su et al., 26 Sep 2025). Drax complements this with an occupancy-based perspective: if $x_0\sim p_0$9 is the designed training path and $x$0 is the actual generation path under the learned velocity, then the occupancy mismatch is bounded by a time-weighted cumulative velocity error, and the generation risk obeys

$x$1

which is the formal motivation for shaping the training path toward inference-like intermediate states (Navon et al., 5 Oct 2025).

A different theoretical perspective comes from the sequential-data analysis of empirical flow matching. For Gaussian bridge paths between observed transitions, the optimal empirical field is a similarity-weighted mixture of per-transition instantaneous velocities, producing what the paper calls a “nonparametric, memory-augmented continuous-time dynamical system” (Lim et al., 9 Feb 2026). This implies a replay bias: with common path choices, flow matching may operate as similarity-weighted trajectory replay rather than as a fully abstract dynamical-law learner. For TS-DFM, that result makes path choice structurally important rather than merely heuristic.

Several misconceptions are addressed directly by the cited works. First, trajectory shaping is not synonymous with classifier guidance. FlowDrive’s moderated guidance is a deterministic perturbation inside the flow integration, not the gradient of an external log-probability or reward model (Wang et al., 26 Sep 2025). Second, training-time shaping need not appear at inference. Drax introduces the audio-conditioned middle distribution solely for training dynamics and explicitly removes it at test time, and the few-step text TS-DFM likewise performs all shaping during distillation while keeping inference unchanged (Navon et al., 5 Oct 2025, Monsefi et al., 8 May 2026). Third, better few-step performance does not necessarily require a larger student; the text TS-DFM argues that “the trajectory is the bottleneck, not the student” (Monsefi et al., 8 May 2026).

Limitations remain domain-specific but structurally similar. The chemistry TS-DFM still exhibits occasional large errors in torsion angles, orientation of distant substructures, and specific bond changes, and it is currently demonstrated on uncatalyzed organic reactions in gas phase (Luo et al., 21 Nov 2025). FlowDrive notes jerkiness in raw trajectories, dependence on HD maps, and the absence of differentiable safety constraints in training (Wang et al., 26 Sep 2025). Drax states that probability-path design for ASR remains largely unexplored and that its audio-shaped path is only one instantiation (Navon et al., 5 Oct 2025). TR-CIE depends on the standard factorized DFM rate parameterization $x$2, uses positivity clamping as a safeguard, and does not formally remove freezing-error terms from global KL bounds (Fu et al., 23 Jun 2026). NExT-OMNI notes that it is trained only at 7B / 2T tokens and that unified omnimodal training still carries substantial engineering complexity (Luo et al., 15 Oct 2025).

These limitations suggest a common research frontier. A plausible implication is that future TS-DFM systems will increasingly co-design path family, time geometry, state representation, and solver, rather than treating any of them as fixed infrastructure. The recent literature already points toward that direction through learned or metric-induced paths, history-aware samplers, training-only middle distributions, in-the-loop perturbations, and explicit energy-guided trajectory selection (Xing et al., 2023, Fu et al., 23 Jun 2026, Monsefi et al., 8 May 2026).