f-Trajectory Balance: a loss family that enables exploration in molecule generation, LLMs and generative models
Connect with us: Valence is constantly seeking talented individuals with diverse backgrounds and expertise to join our team. Explore open roles here.
Reinforcement learning-based fine-tuning has become the standard way that we post-train foundation models to achieve objectives that are more naturally described with a reward function. For example - if we want to teach LLMs to solve hard math or coding problems, it’s difficult to write up large sets of example traces of what the path to a solution should look like. But, we can finetune an LLM on a set of problems with a reward function that measures whether their solution is correct, and let reinforcement learning find its own way to discover reasoning traces that lead to “aha moments” (DeepSeek-AI [2024]). Similarly, if we want to generate molecules that strongly bind to a particular target, we can use SynFlowNets Cretu et al. [2024] which leverage reinforcement learning to sample from the distribution of strong binders. In all of these tasks, often the hardest problem to solve is the exploration problem: the model has to explore combinatorially large search spaces looking for high reward outcomes. And once it finds something with high reward, it faces an exploration / exploitation trade-off: do we keep trying something in the local neighbourhood of what we know to be good or, do we try something very different in the hopes that we might find a higher reward outcome?
Models don’t reason about this exploration / exploitation tradeoff explicitly - how their behaviour changes over the course of RL training is shaped by the optimisation process, which in turn is shaped by the loss functions that we optimise. In our recent ICML paper, f-Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data we explore how these loss functions can shape this behaviour.
An aside on math: I’m going to try to avoid the math-y details in this post and instead focus on the intuition so this block can be skipped while keeping the rest of the blog understandable. For the more technical reader: the key idea of the paper is that we can use reinforcement learning techniques on a big family of loss functions (anything of the loss form ℓ(logpθ(y)−logp⋆(y)) where ℓ is a strictly convex, differentiable function minimised at ℓ(0)) and we show that minimizing any of these losses corresponds to minimizing some f-divergence. We give a general loss that also works off-policy. f-divergences measure the difference between a model’s policy pθ, and some target distribution, p⋆(τ)∝p0(τ)exp(βR(τ)), over trajectories, τ. This target distribution is defined by the prior, p0(τ) (usually defined by the pretraining distribution) and a reward function, R(τ), that scores the outcomes of trajectories. f-divergences have well-studied properties around their mode covering vs mode seeking behaviour and this gives us a way of tuning exploration behaviour.
Utility functions and loss functions
We can understand the relationship between loss functions and exploration behaviour of a model through the classical economics literature on utility functions. If I offer you a choice between a guaranteed payout of (a) $50, or (b) a 50% chance of winning $120 and a 50% chance of getting nothing, different people will make different choices. Clearly, (b) has a higher expected value ($60) so if we were playing this game repeatedly, it would be the better choice. But, in a single interaction, there isn’t a correct answer. A risk-averse person might prefer the guaranteed payout. A risk-seeking person might prefer to take the gamble in the hopes of winning $120. Importantly, the two players don’t disagree about the facts of the world: they see the same outcomes and the same probabilities. We can interpret both behaviours as rational under different utility functions. A concave utility function makes the guaranteed option look more attractive; a convex utility function makes the risky option look more attractive. The utility function doesn’t change the world, but it changes the behaviour of the agent in the world.
Loss functions in reinforcement learning play a somewhat analogous role. These aren’t mathematically equivalent: a utility function transforms outcomes, while these loss functions score differences in distributions over trajectories. However they are similar in that they shape how we penalise and reward trajectories that lead to good and bad outcomes. If there are large penalties to a model putting some probability mass on trajectories that lead to low reward outcomes (strictly: low probability outcomes under p⋆(τ) but these probabilities have a monotonic dependence on rewards), then the model will explore less because it’ll optimize away from any policies that have a chance of putting probability mass on low reward outcomes.
This kind of trade-off is common in drug discovery: suppose you are running a pharma company trying to discover molecules that bind strongly to a target. The search space is enormous (up to 10⁶⁰ small molecules). Somewhere in that space there may be molecules that bind well, are not toxic, and have reasonable pharmacokinetics. But you don’t know where they are. Every molecule you try is a rollout. You synthesize it, test it, and get back a reward signal. Most molecules are bad. Some are moderately promising. A tiny number might be exceptional. Now suppose your team finds one moderately promising molecule. What should you do next? One strategy is to exploit: make local variants of this molecule, optimize around the scaffold, and put most of your experimental budget into the neighbourhood of something you already know works. Another strategy is to keep exploring: continue sampling very different regions of chemical space because the molecule you found may be the best one you have seen so far, but not the best one that exists. Both strategies are reasonable. Both aim to find high reward molecules. The difference is how quickly the decision maker either commits to the evidence they have already seen, or continues to seek out something better.
Similarly, at the beginning of a RL training run, the model has some prior distribution over possible trajectories. This prior comes from pre-training, or just the randomly initialized weights. Given a problem—e.g. a math problem in an LLM’s reasoning corpus—the model has many possible ways it might try to solve it: pattern matching, algebraic manipulation, decomposition into subproblems, brute-force search, proof-style reasoning, or just imitating solution formats it has seen before. Most of these rollouts will not get a reward. Occasionally, the model stumbles onto a reasoning pattern that works. At that point, the model faces the same kind of portfolio decision as the pharma company. Should it put more probability mass on the reasoning strategy that just worked? Or should it continue assigning probability mass to other reasoning strategies that currently look worse, but might turn out to generalize better across harder problems?
A priori, the model does not know whether better solutions exist: it only has partial information about the reward landscape from the rollouts it has sampled so far. So even though all of these objectives may have the same optimum in the limit, they can induce very different behaviour during training, because if we place low penalties on model selecting reasoning strategies (trajectories) that ultimately lead to low reward, then it is freer to explore more broadly, but at the cost of spending less time optimising known good solutions.
Different divergences penalize different mistakes. A mode-seeking divergence is relatively happy to concentrate probability mass on a small number of high reward trajectories once they have been discovered. This can be very efficient if the discovered strategy really is close to optimal. But it can also cause the model to over-commit to the first thing that works.
A mode-covering divergence behaves differently. It more strongly penalizes failing to assign probability mass to regions that might matter. This may look less efficient in the short run, but it can help the model avoid prematurely collapsing onto a narrow strategy.
This gives a slightly different way of thinking about the role of the loss function. While the reward function defines success (or more precisely, defines the distribution over which we want to sample), the divergence that we optimise implicitly defines the model’s attitude towards uncertainty on the way to success by determining which kinds of mistakes are expensive: wasting probability mass on trajectories that do not pay off, or failing to maintain probability mass on trajectories that might eventually pay off.
Off policy training
There is one more piece that matters here: whether the trajectories that the model is trained on have to come from the current version of the model? In on-policy RL, the answer is mostly yes: the model samples trajectories from its current policy, observes rewards, and updates itself based on those samples. Once the policy changes, the old samples do not represent the behaviour of the model we are now trying to train so they have to be discarded. This creates a tight coupling between search and learning.
The off-policy setting breaks this coupling. We can learn from trajectories generated by older versions of the model, by different models, or by different sampling strategies. In the drug discovery example, this is like learning not only from the molecules your current team would choose to synthesize, but from the entire archive of experiments run by previous teams: conservative local optimisation campaigns, exploratory screens over unusual scaffolds, and failed historical projects that still reveal something about the reward landscape.
This is useful because the trajectories produced by different policies encode different search behaviours. A mode-seeking policy may generate lots of local variants around known good solutions. A mode-covering policy may generate broader, more diverse rollouts. If our loss works off-policy, then we can use all of these rollouts as evidence. Search and learning no longer have to be the same process: we can use many different behaviours to generate experience, and a shared learner to distill what those behaviours discovered.
This is one of the reasons the f-Trajectory Balance family is attractive. On-policy, each loss inherits the exploration behaviour of its corresponding f-divergence. Off-policy, it remains a valid objective with the same global minimizer. So the loss function can both shape how agents search and allow us to reuse the data produced by other search behaviours. In large sparse-reward problems, this is especially valuable: we do not want to throw away rare successful trajectories just because they came from yesterday’s model, a more exploratory sampler, or a different part of the training system.
Tuning exploration during training
One final piece that I found exciting when we tried it: there’s a special class of f-divergences called α-divergence that have an explicit parameter α that allows us to tune between more mode seeking and mode covering behaviour. We can set it to some fixed value and get the desired outcome, but we found that we can also anneal it over the duration of training to start with more exploratory policies and then have more mode seeking behaviour later in training. There’s a lot more that needs to be done here in understanding how to tune this optimally, but even just linearly annealing it gave us the best performing models.
This post is part of “Inside Valence”, a series where you’ll get a behind-the-scenes look at our research, exploring new ways to predict, explain, and ultimately decode biology. If this resonates, consider subscribing!
 |
|