🤖 machine learning

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

This paper introduces Reward Transport, a novel flow matching framework that embeds controllable molecular properties directly into the noise-data coupling via optimal transport, enabling inference-time steering of generated molecules toward specific targets like logP or QED without requiring additional reward models or gradient guidance.

Original authors: Kehan Guo, Yili Shen, Yujun Zhou, Yue Huang, Chujie Gao, Shiyi Du, Xiangliang Zhang

Published 2026-07-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kehan Guo, Yili Shen, Yujun Zhou, Yue Huang, Chujie Gao, Shiyi Du, Xiangliang Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, magical factory that builds molecules (the tiny building blocks of medicines and materials). Usually, to tell this factory what kind of molecule to build, you have to give it a specific instruction manual for every single item, or you have to hire a super-smart inspector to check every item as it comes off the line and yell, "No, make that bigger!" or "No, make that stickier!" This is slow, expensive, and requires a lot of extra brainpower.

The researchers in this paper, Kehan Guo and his team, discovered a much sneakier, simpler way to control the factory. They realized that the "coupling"—the rule that decides which random noise (like static on an old TV) gets paired with which molecule during the factory's training—isn't just a boring technical detail. It's actually a control knob.

The Magic Trick: Sorting the Chaos

Think of the factory's training process like a dance. Normally, the factory pairs random dancers (noise) with random molecules (data) without any rhyme or reason. The authors say, "Wait a minute! What if we sort the dancers by how energetic they are, and sort the molecules by how 'greasy' (a property called logP) or how 'drug-like' (a property called QED) they are? Then, we pair the most energetic dancer with the greasiest molecule, the second-most energetic with the second-greasiest, and so on."

They call this Reward Transport.

By doing this sorting during training, they embed a secret map into the factory's brain. The factory learns that "High Energy Noise" always leads to "Greasy Molecules."

The One-Knob Control

Here is the best part: Once the factory is trained, you don't need an inspector, a reward model, or complex instructions. You just turn a single dial (a number called $s$ ) that controls how much "energy" you feed into the machine.

Turn the dial up? The factory spits out molecules that are greasier (higher logP).
Turn the dial down? It makes them less greasy.
Do the exact same thing for "drug-likeness" (QED)? The factory makes molecules that are more drug-like.

The paper shows that on a dataset of 224,568 molecules (ZINC-250K), turning this single knob changed the "greasiness" (logP) by a massive 137% (from 1.50 to 5.44) and the "drug-likeness" (QED) by a smaller but consistent amount, all while keeping the molecules 100% valid and unique.

What This Is NOT (The "Nope" List)

The authors are very careful to tell you what this trick is not:

It's not a "Size" cheat code. You might think, "Oh, they just made bigger molecules." But the paper proves this is wrong. When they turned the knob to make molecules greasier, the molecules got bigger (growing from 12 to 23 heavy atoms). But when they turned the same knob to make molecules more "drug-like," the molecules actually got smaller (shrinking from 23 to 16 atoms). If it were just a size trick, both would have grown or shrunk together. They didn't. The knob is smart enough to know which property it's controlling.
It's not a magic wand for any property. The paper explicitly notes that this trick works for properties that change smoothly with size (like greasiness). It failed to control "Synthetic Accessibility" (how hard a molecule is to build in a lab). Why? Because building difficulty depends on weird, specific shapes and patterns, not just a simple "bigger or smaller" scale. The single knob couldn't capture that complexity.
It doesn't work with every type of AI. The authors tested this on a specific type of AI setup (called $\hat{x}_1$ -prediction). They found that if you try to use this trick with a different, very common setup (called $\epsilon$ -prediction), the magic disappears. The signal gets too weak to control the factory.

How Sure Are They?

The team didn't just guess; they measured this rigorously.

They ran the experiment on two different massive datasets (ZINC-250K and GuacaMol) and got the same results.
They proved that the "knob" works by showing a perfect mathematical link (a rank correlation of 1.000) between the dial setting and the average property of the molecules produced.
They even showed that this control happens without any extra computer power during the final generation. It's a "free" feature built right into the training.

The Bottom Line

This paper suggests that we don't always need complex, heavy-handed tools to control AI-generated molecules. Sometimes, the secret is just in how you pair up the random noise and the data during training. By sorting them like a deck of cards, you can turn a single number into a steering wheel that guides the AI to build exactly the kind of molecules you want—without needing a GPS or a map along the way. It's a clever, efficient way to navigate the chemical universe, provided you're steering toward properties that play nice with a single dial.

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

1. Problem Statement

In Flow Matching (FM) and related diffusion-based generative models, the coupling—the rule determining which noise vector is paired with which data point during training—is typically treated as a computational heuristic. Standard approaches use independent pairing for simplicity or minibatch optimal transport (OT) to straighten trajectories and reduce variance. In these paradigms, the coupling is optimized solely to facilitate training, not to encode specific structural properties into the learned flow field.

Consequently, achieving controllable generation (e.g., generating molecules with specific logP or QED scores) usually requires explicit mechanisms at inference time, such as:

Classifier-free guidance: Requiring a target property value as an input.
Gradient guidance: Computing gradients of a property predictor at every sampling step.
Reinforcement Learning (RL): Fine-tuning the generator against an oracle, often requiring oracle queries during sampling.

These methods introduce inference-time overhead, dependency on external reward models or oracles, or complex training pipelines. The paper asks whether the coupling itself can serve as an alignment interface, embedding controllable structure directly into the flow field during training so that control is available at inference without additional computation or oracle calls.

2. Methodology: Reward Transport

The authors propose Reward Transport, a method that redefines the noise-data coupling as a mechanism for injecting task structure into the flow field.

2.1 Property-Aligned Optimal Transport Coupling

Instead of pairing noise and data independently, Reward Transport constructs a monotone coupling based on a target molecular property $y(x)$ (e.g., logP, QED).

Scalar Keys: For each noise vector $z$ (a sequence of token embeddings), a scalar sorting key $s(z)$ is computed. The primary choice is the Euclidean norm of the mean-pooled noise: $s(z) = \|\bar{z}\|_2$ .
Rank-Based Pairing: Given a batch of noise vectors $\{z_n\}$ and molecules $\{x_n\}$ with properties $\{y_n\}$ , both sets are sorted by their respective keys ( $s$ and $y$ ). They are then paired rank-by-rank: the noise vector with the $k$ -th smallest $s$ is paired with the molecule having the $k$ -th smallest $y$ .
Training Objective: The flow matching model $v_\theta$ is trained to transport these aligned pairs. The loss combines a standard MSE on the endpoint prediction (in continuous embedding space) and a Cross-Entropy (CE) loss to ensure the decoded tokens form valid molecules.

2.2 Direction Conditioning

To ensure the model can access the noise-space coordinate $s$ during inference, a Direction MLP is introduced. This two-layer MLP maps the normalized scalar $s$ to a $d$ -dimensional embedding, which is added to the input of the transformer alongside time and positional embeddings. This allows the model to condition its velocity field on the specific noise coordinate used during the monotone pairing.

2.3 Inference Mechanism

At inference, the user specifies a scalar direction signal $s^*$ (a noise-space coordinate, not a target property value). The model generates molecules by integrating the flow field starting from noise vectors where $s(z) \geq s^*$ .

Theoretical Guarantee (Proposition 1): In the limit where the learned flow perfectly preserves the rank correspondence between $s$ and $y$ (coupling-preserving limit), conditioning on $s \geq s^*$ induces a distribution equivalent to one selection step of the Cross-Entropy Method (CEM). Specifically, it truncates the generated distribution to the top- $\epsilon$ fraction of the property distribution, where $\epsilon$ is determined by the threshold $s^*$ .
Zero Overhead: This control requires no oracle, no reward model, no gradient computation, and no rejection sampling at inference time.

3. Key Contributions

Coupling as Alignment Interface: The paper reframes the noise-data coupling in flow matching from a training heuristic to a distribution-level alignment interface. It demonstrates that choosing the coupling based on a target property organizes the learned flow field such that noise-space coordinates correspond to property values.
Single-Knob Molecular Steering: The method achieves monotone control over molecular properties (logP and QED) using a single scalar noise coordinate. Crucially, the same knob produces opposite structural responses for different targets: increasing $s$ generates larger molecules for logP (lipophilicity) but smaller molecules for QED (drug-likeness). This rules out a generic size bias and confirms the coupling writes target-specific structure into the flow.
Reward Transport for Discrete Sequences: The authors implement this for variable-length SELFIES token sequences, identifying and addressing two specific failure modes in flow models generating discrete data:
- Implicit Signal Erasure: In Pre-LayerNorm architectures, magnitude information (encoded in $\|\bar{z}\|$ ) is erased. The Direction MLP resolves this by injecting a learned direction pattern.
- Length Confound: Without unmasked MSE loss on padding tokens, the model learns a shortcut where high $s$ simply increases sequence length to satisfy the norm constraint. Unmasked MSE forces the model to regress the raw noise at padding positions, anchoring the magnitude signal and preventing this shortcut.

4. Experimental Results

The method was evaluated on ZINC-250K and GuacaMol datasets.

Control Performance:
- logP: On ZINC-250K, sweeping $s$ from -3 to +7 shifted the mean logP from 1.50 to 5.44 (+137%), with a Spearman correlation $\rho=1.000$ between $s$ and the mean property. The effect size (Cohen's $d$ ) reached 1.63.
- QED: The same knob shifted QED monotonically (e.g., +0.059 at $s=5$ on ZINC), with $\rho=1.000$ .
- Structural Specificity: For logP, increasing $s$ increased the number of heavy atoms (12 $\to$ 23). For QED, increasing $s$ decreased the number of heavy atoms (23 $\to$ 16).
Generalization: Results replicated on GuacaMol (1.6M molecules), showing even stronger effect sizes for logP ( $d=2.31$ ) and consistent monotonicity for QED within the valid range.
Comparison to Baselines:
- Compared to Conditional Flow Matching (where property values are explicit inputs), Reward Transport achieves similar monotonicity ( $\rho=1.000$ ) but operates via distribution truncation rather than explicit target hitting.
- Compared to $\epsilon$ -prediction diffusion, the method shows that coupling-level alignment is structurally absent in $\epsilon$ -prediction because the Bayes-optimal target is independent of the coupling, attenuating the gradient signal.
Training Dynamics: Rank preservation (the correlation between $s$ and generated $y$ ) emerges rapidly during fine-tuning (within one epoch) and plateaus below 1.0 (e.g., 0.57 for logP), consistent with the theoretical "distributional approximation" gap described in Proposition 1.

5. Significance and Claims

The paper claims that Reward Transport establishes the coupling in flow matching as a distinct, distribution-level control channel that can carry alignment structure on its own.

Efficiency: It eliminates the need for oracle calls, reward models, or iterative optimization at inference time, reducing computational cost significantly compared to RL or guidance-based methods.
Mechanism: It demonstrates that the "steerability" of generative models need not rely on explicit conditioning inputs; it can be baked into the model's geometry via the training coupling.
Limitations and Scope: The authors are modest about the scope. The method provides a 1D distributional knob (truncating the distribution) rather than exact per-sample value assignment. It works well for scalar properties correlated with global molecular characteristics (like size or polarity) but fails for topological properties like Synthetic Accessibility (SA) where a 1D scalar key cannot distinguish complex structures of similar size. Furthermore, the method is shown to be incompatible with $\epsilon$ -prediction targets, highlighting that the interface depends on the specific prediction target of the network.

In summary, the work shifts the paradigm of controllable generation from "guiding the trajectory at inference" to "aligning the training distribution," offering a principled, zero-overhead mechanism for property control in flow matching.