Reward Transport: Property Control in Flow Matching via Noise-Space Alignment

Abstract: The coupling in flow matching -- the rule pairing noise vectors with data points -- is typically treated as a computational choice. We show that this coupling can instead serve as an alignment interface: by matching noise and data according to a target molecular property, it embeds controllable structure directly into the learned flow field. Building on this view, we introduce Reward Transport, which uses optimal transport coupling at training time to align a scalar noise-space coordinate with molecular rewards; at inference, varying this coordinate steers the generated distribution without requiring an oracle, reward model, gradient guidance, or additional computation. In the coupling-preserving limit, thresholding this coordinate recovers the Cross-Entropy Method's truncated reward distribution, providing a principled, continuously adjustable distribution-level control knob. Empirically, on ZINC-250K and GuacaMol, sweeping the scalar induces monotone control of logP and consistent QED control over its operating range; most tellingly, the same knob produces opposite structural responses for different targets, growing molecules for logP but shrinking them for QED, which rules out a generic size bias. The interface is complementary to classifier-free guidance and conditional flow matching, while a negative result under epsilon-prediction diffusion clarifies where coupling-level alignment is structurally absent. Code: https://github.com/KehanGuo2/reward-transport

• The paper introduces Reward Transport, a monotone noise–property coupling that enables inference-time control of molecular logP and QED without reward models, gradients, conditioning inputs, or rejection sampling.
• The method achieves perfect group-level rank correlation on ZINC-250K, with logP spanning 1.50–5.44 and QED 0.336–0.477 while maintaining 100% validity and over 99% uniqueness.
• The paper identifies key limits and design requirements: warm-start fine-tuning, a direction MLP, unmasked padding loss, and x₁ or velocity prediction are important, while ε-prediction, SA control, and extreme QED extrapolation fail.

Reframing the coupling as an alignment interface

The paper's central claim is architectural rather than algorithmic: in flow matching, the coupling—the rule that pairs a noise vector $\mathbf{z}$ with a data point $\mathbf{x}$ during training—need not be treated as an implementation detail chosen for optimization convenience. Prior work has used independent pairing (Lipman et al., 2022) or minibatch optimal transport with Euclidean cost to straighten trajectories (Pooladian et al., 2023, Parmar et al., 2023). The authors instead argue that the coupling can carry task structure: if noise and data are paired according to a scalar molecular property, the learned flow field inherits a property-aligned organization in noise space, and steerability emerges without any inference-time oracle, reward model, gradient guidance, or conditioning input.

The instantiation, Reward Transport, is a monotone rearrangement: noise sequences are sorted by a scalar key $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$ (the $\ell_2$ norm of mean-pooled noise), molecules are sorted by target property $y$ (logP or QED), and the two are paired rank-by-rank. Because both keys are one-dimensional, this is the unique optimal transport plan under any convex cost—the classical monotone rearrangement [villani2003topics]. At inference, the user sets $s^*$ once, which is fed to the network through a small direction MLP alongside time and positional embeddings; sampling then integrates the velocity field as usual.

Theoretical content: implicit Cross-Entropy Method

The theoretical hinge is a proposition showing that, in the coupling-preserving limit (per-molecule Spearman correlation $\rho_{\text{per}} \to 1$ between $s(\mathbf{z})$ and the generated property), conditioning on $s \geq \tau$ yields exactly the data distribution truncated to the upper property quantile:

$$p_\theta(\mathbf{x} \mid s(\mathbf{z}) \geq \tau) \approx \frac{p_1(\mathbf{x})\,\mathbf{1}[y(\mathbf{x}) \geq y_\tau]}{Z(\tau)}, \qquad y_\tau = F_Y^{-1}(F_S(\tau)).$$

This is functionally one selection step of the Cross-Entropy Method [rubinstein1999cross] with elite fraction $\mathbf{x}$0, obtained in a single forward pass with no rejection sampling. The proof is a direct quantile argument; the paper is explicit that it relies on two conditions—sufficient model capacity, and the assumption that the 1D projection captures enough of the high-dimensional coupling structure—and that the realized $\mathbf{x}$1 (0.57 for logP, 0.22 for QED) measures the residual slack, making the proposition a distributional rather than deterministic statement. The gap between per-molecule $\mathbf{x}$2 and group-mean $\mathbf{x}$3 is presented as exactly this distributional approximation regime.

Empirical results on molecular generation

The method is implemented for variable-length SELFIES token generation with a 6-layer Transformer (~50M parameters, $\mathbf{x}$4-prediction, MSE plus cross-entropy loss), trained in two stages: a 120-epoch base model without coupling, followed by 5–10 epochs of OT-coupled fine-tuning. The headline results on ZINC-250K:

Metric	logP	QED
Property range across $\mathbf{x}$5	1.50 → 5.44 (+137%)	0.336 → 0.477
Group-mean $\mathbf{x}$6	1.000	1.000
Per-molecule $\mathbf{x}$7	0.57	0.22
Peak Cohen's $\mathbf{x}$8	+1.63	+0.36
Validity / uniqueness	100% / >99%	100% / >99%

Two observations carry most of the evidential weight. First, the same scalar knob produces opposite structural responses for the two targets: sweeping $\mathbf{x}$9 grows logP molecules from 12 to 23 heavy atoms but shrinks QED molecules from 23 to 16. This rules out a generic size bias and supports the claim that the coupling writes target-specific structure into the flow. Second, the effect replicates on GuacaMol (1.6M molecules) with identical hyperparameters, where logP control is stronger ($s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$0 at $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$1, attributed to the broader empirical property range giving the sorted coupling more room to stretch), while QED control is monotone only on $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$2 and degrades at extreme $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$3—consistent with QED's bounded geometry. Seed-level replication preserves rank correlations ($s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$4 logP, $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$5 QED) while absolute shifts shift modestly, so the paper treats $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$6 as the reproducible statistic and $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$7 as recipe-sensitive.

A same-backbone Conditional Flow Matching baseline also achieves $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$8 with matched validity and diversity, but attains a wider logP span at extremes (1.79–7.99 vs. 1.50–5.50) at the cost of a steeper FCD climb off-manifold. The paper is candid that this is not a superiority claim: the two occupy different points on a steerability–interpretability trade-off, with conditioning offering explicit property targets and coupling offering a property-free percentile scalar with distribution-truncation semantics. Relative to RL and guidance-based methods, however, Reward Transport is the only paradigm surveyed with $s(\mathbf{z}) = \lVert \bar{\mathbf{z}} \rVert_2$9 inference overhead and no test-time oracle.

Failure modes and component analysis

The ablations establish that every component is load-bearing, and surface two previously unreported failure modes for flow models over discrete, variable-length sequences:

• Implicit signal erasure: removing the direction MLP reverses the sign of the effect, because Pre-LayerNorm transformers normalize activations and erase the magnitude information that a norm-based coupling encodes into noise. The direction MLP converts $\ell_2$0 into a directional activation pattern that survives normalization.
• The length confound: masking the MSE loss at padding positions causes the model to route the coupling signal into sequence length ($\ell_2$1) instead of the target property, since longer sequences trivially have higher noise norms. Computing the loss at padding positions forces regression of the raw noise vector there, anchoring global noise magnitude in the hidden state and closing the shortcut.

Removing the OT coupling entirely yields $\ell_2$2, confirming the coupling—not the direction conditioning—is the active ingredient. Notably, the specific 1D key is not privileged: $\ell_2$3, $\ell_2$4, a random linear projection, and a single coordinate all achieve group $\ell_2$5 with per-molecule $\ell_2$6 in $\ell_2$7, and $\ell_2$8 is not even the best performer. The monotone rearrangement itself carries the signal.

Training dynamics add a practical constraint: from-scratch training exhibits a transient signal window (emerging at epoch 3, erased by epoch 10 as the model overwrites the coupling), whereas warm-start fine-tuning preserves the signal throughout with ~60% higher peak effect size. Stronger base models yield higher absolute QED but smaller marginal gains, exactly as the CEM-truncation semantics predict.

Scope limits and negative results

The paper delineates where the interface exists rather than claiming generality. Three boundaries stand out:

Prediction target dependence. Under $\ell_2$9-prediction (and affine-equivalent velocity prediction), the Bayes-optimal network target depends on the coupled pair, so the coupling-induced gradient reaches the direction pathway. Under $y$0-prediction, the target is fresh noise independent of the coupling, and the coupling-induced gradient is attenuated by the forward-process signal-to-noise ratio $y$1; the appendix derives this analytically and reports no measurable coupling effect on QM9 with EDM. The failure is architectural and, per the authors, not recoverable by tuning.

One-dimensionality. The interface is a scalar, distribution-level knob: it truncates by property rather than assigning per-sample values, and it cannot control topological properties. SA score control fails ($y$2) because SA depends on discrete fragment features that a size-correlated scalar cannot distinguish; multi-property control would require a higher-dimensional coupling key whose inference semantics no longer reduce to 1D monotone rearrangement.

Empirical scope. Claims rest on two SELFIES datasets, one property at a time, and a single ~50M-parameter backbone. Unconditional QED (0.477 on ZINC) sits below the dataset mean, so absolute property levels are base-model-bounded even though steerability is not; FCD values (26–34) are elevated relative to literature baselines and attributed to the small backbone. QED control on GuacaMol degrades at extreme $y$3, and the extrapolation behavior is property-dependent: logP extrapolates monotonically to $y$4 (mean logP 14.4, 50 atoms, 100% novelty), while QED peaks near $y$5 and reverses.

Conclusion

The paper's contribution is a reframing with a working instantiation: the flow matching coupling, chosen to align noise-space coordinates with a scalar reward, functions as a distribution-level control channel that requires no oracle, reward model, gradient, or conditioning input at inference, and whose coupling-preserving limit recovers one Cross-Entropy Method selection step. The empirical evidence—monotone group-level control on two datasets, opposite structural responses ruling out size bias, and component ablations isolating the coupling as the active ingredient—is internally consistent with the theory, and the negative results under $y$6-prediction and on SA score are presented as principled scope conditions rather than incidental failures. The open questions are specific: whether higher-dimensional coupling keys can extend the interface to topological or multi-property control, whether the two-stage training requirement generalizes beyond this backbone, and whether coupling and conditioning can be composed in a single model with two orthogonal inference-time knobs.