Trajectory Matching

Consider building a molecule: it seems more natural to grow it by attaching whole fragments to a partial structure than by changing individual atom and bond features in a fixed-size noisy graph. Yet most discrete flow matching methods are restricted to such single-component edits on a fixed-size state. Trajectory matching extends discrete flow matching to structured actions that can change the size of a state, modify several components together, and depend on earlier actions.

The key idea is to condition training not only on a source and target state but on a full trajectory: a valid sequence of actions connecting the two endpoints. The trajectory fixes which actions happen and in what order; we randomize only when they occur. This gives us tractable training targets even for structured actions with dependencies. Once trained, the model generates samples as in standard discrete flow matching, without needing a trajectory. We’ll illustrate these ideas in this post through molecular generation and hierarchy generation.

To understand why trajectory conditioning helps, let’s first look at how standard discrete flow matching constructs its training targets.

Factorized Flows

Let us briefly recap the standard Discrete Flow Matching formulation to understand where its limitations come from.

Our goal is to transform an easy-to-sample source distribution p_{\mathrm{src}} into the data distribution p_{\mathrm{tgt}} over time t\in[0,1]. We do this through a continuous-time Markov chain: a process that moves between discrete states through random jumps. Writing S_t for its state at time t, the jump rate Q_t(s'\mid s) determines the probability of jumping from state s to s' over a short time interval:

We want to learn rates such that the probability path, the time-marginal state distribution, p_t(s)=\Pr[S_t=s] starts at p_0 = p_{\mathrm{src}} and ends at p_1 =p_{\mathrm{tgt}}.

Directly finding such rates is generally intractable. The central idea of flow matching is to construct simpler conditional processes connecting individual source and target states, s_I and s_F. To use this conditional process for training we need: i) easy sampling of intermediate states s_t, ii) easy computation of its conditional jump rates Q_t(s' \mid s_t, s_I, s_F). These conditional rates then provide supervision for an unconditional model Q^\theta_t(s' \mid s_t) that receives only the current state and time. Averaging this conditional supervision gives the marginal rates we want to learn.

However, constructing such a conditional process is still not straightforward. In continuous space, a straight line connecting s_I and s_F already works. In discrete space, there is no equally obvious interpolation: we need to connect the endpoints through random jumps while keeping intermediate states easy to sample and jump rates easy to compute. Moreover, the state space is typically combinatorially large, so we need to restrict the allowed jumps to keep the process tractable.

We can use that discrete states often consist of individual components: Tokens in a sequence or atom and bond features in a molecular graph. This lets us factorize the conditional process across components: each component simply switches from its source value to its target value at a random time, independently of the others. Sampling an intermediate state s_t is easy: we simply sample each component’s value independently. Additionally, the resulting process only jumps between states that differ in exactly one component, so we only need rates for these local edits between neighboring states.

This is the standard discrete flow matching construction which we call a factorized flow. It is easy to define and works well in many settings but makes three assumptions:

1. Fixed size: source, target, and intermediate states all have the same number of components.
2. Single-component edits: each jump changes only one component.
3. Independent conditional evolution: given the endpoints, components evolve independently of each other.

This makes clear why structured actions cannot be expressed by factorized flows: Consider attaching a whole ring fragment to an existing atom: this action changes the size of the state, affects multiple atoms and bonds at once, and must wait until the attachment atom exists—creating a dependency between actions. This single action therefore already breaks all three assumptions.

Trajectory Matching

How can we build discrete flows that support these structured actions? We need to solve two problems:

1. Defining the allowed jumps. Specifying a rate for every possible state-to-state transition is intractable. How can we restrict the transitions while still allowing complex jumps?
2. Constructing training targets. We can no longer factorize the conditional process across components. How can we construct a conditional process that can be used for training?

Action-based rates

A general state-to-state rate Q_t(s'\mid s) could, in principle, represent complex jumps like a fragment insertion. However, specifying and learning rates for every possible state-to-state transition is not possible in practice. Instead, we describe the allowed transitions through a user-defined vocabulary of actions \mathcal{A}. An action maps the current state to a new state, s'=a(s). For example, an action might attach a particular ring fragment to atom 5 or add a new child to the root of a hierarchy. Instead of specifying a rate directly for next states, we define rates u_t(a \mid s) over all available actions a\in\mathcal A(s) at state s. The probability of a jump from s to s' then depends on the rate of taking an action that leads to this state:

Choosing the action vocabulary lets us restrict jumps to meaningful domain-specific edits, giving the model a useful inductive bias. For example, we design actions for a hierarchy-generation task that guarantee validity by construction. Additionally, actions also give us a description of transitions that generalize well across states: the same fragment-attachment action can apply to different intermediate molecules, producing a different next state in each case.

Single-component edits are a special case: the actions simply set one component to a particular value. Factorized flows therefore fit naturally into this action-based description.

The action vocabulary specifies which jumps the model can make. We still need a conditional process that provides tractable training targets for their rates.

Trajectory-conditioned training

We now need a conditional process connecting s_I to s_F through these actions, with intermediate states that are easy to sample and rates that are easy to compute. We can no longer rely on the factorized component-wise construction. Instead, we face two difficulties:

Which actions should be taken? In factorized flows, the remaining edits are clear from the current state and the target: every component that differs from its target value still needs to switch. With structured actions, different sets of actions may connect the same endpoints. For example, if the vocabulary allows both, we could build the same structure atom by atom or insert a whole fragment.

In which order can they occur? In factorized flows, the required component edits can occur in any order. Structured actions may have prerequisites: an attachment atom must exist before a fragment can be attached to it, a child can only be added to an existing parent. A valid conditional process must respect these dependencies.

The endpoints therefore leave us with choices about both which actions to take and how to order them. Defining a conditional process that accounts for these choices while remaining easy to sample and supervise seems very hard.

Defining Trajectories

Our central idea is to condition during training not only on the endpoints, but also on a complete sequence of valid actions connecting them: a trajectory.

The trajectory fixes which actions happen and in what order. All that remains is to specify when each action happens.

In factorized flows with a linear schedule, each required component edit happens at an independent, uniformly sampled time in [0,1]. We keep these random event times, but assign them according to the trajectory: sample n uniform times, sort them, and execute action a_i at the i-th event. At time t, the number of events that have occurred is therefore K_t\sim\mathrm{Binomial}(n,t). The state is simply the one reached after that many actions:

Sampling an intermediate state is therefore easy: draw K_t and select s_{K_t}. The conditional jump rate also has a simple interpretation. At state s_i, only the next action in the trajectory has positive rate:

As in factorized flows, the rate increases as the remaining time shrinks. Here, it also scales with the number of remaining actions: more actions left means the next event needs to happen sooner. The order is enforced by allowing only the next action to occur. Once we reach s_n, all rates are zero.

Given a trajectory, we thus have both ingredients we need for training: we can easily sample an intermediate state and compute its conditional action rates directly. But where do the trajectories come from?

Sampling trajectories

As we have already seen, many different valid trajectories may connect the two endpoints s_I and s_F. We model this using a distribution over trajectories:

This distribution specifies which ways of connecting the endpoints we see, and how often we see them during training. It therefore introduces an inductive bias about how objects should be constructed. Our main requirement is that we can efficiently sample valid trajectories from it. How we do this depends on the domain and the action vocabulary.

For example, consider growing a hierarchy by attaching children to existing nodes. Starting from the root of a target hierarchy, we can repeatedly choose a node whose parent has already been added and attach it. Different choices produce different valid construction orders; the rule for choosing among available nodes defines q.

Importantly, we do not need a dataset of observed trajectories but directly construct them from the endpoints.

Training

Training now works much like standard discrete flow matching, except that we sample not only endpoints but also a trajectory connecting them. We then sample a time and an intermediate state along that trajectory. The model receives only the sampled state and time—not the trajectory or target endpoint—and predicts action rates. The trajectory supplies the supervision: it tells us what rates to train those predictions against.

Training remains simulation-free because we can sample intermediate states directly from the trajectory. As in standard flow matching, the conditional supervision gives us the marginal rate we want to learn: the training objective has the same expected gradients as directly matching the marginal rate, and its optimum recovers that rate (when the model is sufficiently expressive).

Sampling

At generation time, we simulate the learned jump process using, for example, Euler sampling. We start from a sample of the source distribution and, at each step, predict action rates from the current state and time. Together with the step size, these rates determine the probability of applying each available action. No trajectory or target endpoint is needed.

Notice that while each trajectory-conditioned target assigns positive rate to only one next action, the learned marginal rate averages over trajectories and can assign positive rates to several available actions. A single Euler step can therefore apply multiple compatible actions, so trajectory conditioning does not force generation to proceed one action at a time.

Empirical results

Let’s see practical use cases of trajectory matching: To generate a tree, we can start from a root and repeatedly attach children. These actions guarantee validity by construction as they cannot introduce a cycle or disconnect the tree.

This pays off especially when using few sampling steps, where the factorized baseline often produces invalid trees.

For molecular graph generation, we can showcase the flexibility of our approach by using different action spaces that range from factorized single-component edits to insertion of full fragments:

In our experiments, fragment insertion works particularly well. Building with larger molecular pieces gives the model a useful inductive bias for generating valid molecules.

The choice of actions also shapes what we can do with the trained model. AtomGrow learns to grow molecules from partial structures, so we can directly use it for substructure extension and linking simply by changing its starting state. Without retraining, it substantially outperforms the adapted factorized baseline on both tasks.

We’ve now seen the practical potential of flow matching with structured actions. If you’d like to explore the theory further, the next section shows precisely how trajectory matching recovers factorized flows and extends them to actions with dependencies.

Trajectory Matching Subsumes Factorized Flows

Trajectory matching trains on a single next action in a trajectory, while factorized flows supervise on all independent component edits. How are these two constructions connected? The key is to look at what happens when we average over different action orderings.

Rao–Blackwellized Rates

So far, we have used one sampled trajectory to construct each training target. That trajectory selects a single next action. But other trajectories connecting the same endpoints may pass through the same intermediate state and select a different action. Can we directly supervise on all possible next actions?

To weight each trajectory correctly, we define its posterior probability given the current state, time, and endpoints:

Averaging the trajectory-conditioned action rates under this posterior gives the Rao–Blackwellized rate:

When this average is efficient to compute, we can use it directly as the training target. It preserves the expected gradient and optimum while reducing gradient variance. It also makes the connection to factorized flows explicit, as we’ll see next.

Recovering Factorized Flows

Let's see how to recover factorized flows using trajectory matching. To instantiate trajectory matching, we have to answer two questions: which actions should we use, and what trajectory distribution should we choose?

The action vocabulary is straightforward: single-component edits, each setting one component to a particular value.

The trajectory distribution also admits an intuitive form. Given the endpoints, the required actions are fixed: each component that differs from its target value must be set to that value. We denote this action set by \alpha(s_I,s_F). These actions affect different components, so they are independent and can occur in any order. A natural choice is therefore to sample uniformly over all orderings of this set.

Now suppose the remaining action set \alpha(s,s_F) at state s contains m actions. Each is equally likely to come next, with probability \frac{1}{m}, and receives rate \frac{m}{1-t} when it does. Averaging gives the Rao–Blackwellized rates:

These are exactly the rates used to train factorized flows with a linear schedule! The intermediate-state distribution also matches: each required component edit has independently occurred by time t with probability t.

With Rao–Blackwellized targets, we therefore recover the factorized flow matching objective exactly. Single-trajectory targets give the same expected gradients and minimizers, though with potentially higher gradient variance. Factorized flows are thus a special case of trajectory matching.

Beyond Independent Actions

Let's see what happens in trajectory matching if we go beyond independent actions. A natural next step is to keep the required action set fixed, but introduce tree-like dependencies. Both GrowHierarchy and FragGrow are examples of this.

We represent these dependencies as a forest over actions: each node is an action, and an edge means that the parent action must happen before the child. Actions with no directed path between them are required to be independent, so their relative order does not matter.

As before, we choose a trajectory distribution that is uniform over all valid action orderings, now the topological orderings of the dependency forest.

At an intermediate state s, consider the forest of remaining actions, \alpha(s,s_F). Only its roots are currently available: every other action still has an unmet prerequisite. Additionally, an available action is more likely to come next when more remaining actions depend on it, i.e., the size \operatorname{desc}(a) of the subtree rooted at a (including a itself) is large. The Rao–Blackwellized rates capture exactly this:

This connects two familiar extremes. With no dependencies, the forest has no edges: every remaining action is available and receives equal rate, recovering supervision on all actions as in factorized flows. With chain-like dependencies, only the next action receives positive target rate, resembling supervision in fixed-order autoregressive generation. Trajectory matching with general tree-like dependencies lies between these extremes, allowing some actions to occur in either order while enforcing prerequisites for others.

Our paper explores even more constructions, including more general dependencies and trajectory distributions in which different sequences connecting the same endpoints can use different sets of actions.

Trajectory matching brings these choices into a common framework: which actions to use, and how they can be ordered. Independent component edits recover factorized flows, while structured actions, potentially with prerequisites, extend the same training principle to richer generation processes.

This has been my favorite project of my PhD so far, and I would be very happy to hear from you if you have questions, thoughts, or feedback on the work or on this (my first!) blog post.