Technical Summary: f-Trajectory Balance

Problem Statement

The alignment and fine-tuning of generative models, including Large Language Models (LLMs) and Generative Flow Networks (GFlowNets), typically rely on minimizing the discrepancy between a model's distribution and a target distribution. In Reinforcement Learning (RL) fine-tuning, this is often framed as maximizing expected reward subject to a Kullback-Leibler (KL) divergence penalty.

A significant challenge in this domain is the trade-off between mode-seeking and mode-covering behaviors. Standard objectives, such as the reverse KL divergence (often implemented via Trajectory Balance or squared KL loss), tend to be mode-seeking, causing models to collapse onto a few high-probability modes. This is suboptimal for tasks requiring exploration, such as drug discovery or diverse candidate generation, where "mode-covering" (capturing the full diversity of the target distribution) is preferred.

Furthermore, existing methods often struggle with off-policy validity. While estimators like REINFORCE or PPO are common, they suffer from high variance. Recent approaches using the squared KL loss ($KL_{sq}$) offer low-variance gradients on-policy and validity off-policy, but they are inherently tied to the properties of the reverse KL divergence, limiting the ability to control the exploration-exploitation trade-off.

Methodology

The authors propose f-Trajectory Balance, a family of surrogate loss functions that generalizes the squared KL loss to the entire family of f-divergences.

Theoretical Foundation

The core theoretical contribution is establishing a one-to-one correspondence between translation-invariant loss functions on log-probabilities and f-divergences.

1. From Divergence to Loss: For any convex function $f$ defining an f-divergence $D_f(p_\theta \| p^*)$, the authors derive a loss function $L_f$ defined on the difference of log-probabilities $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$:
   $$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$$
   This loss has two critical properties:

   • On-Policy Gradient Equivalence: When samples are drawn from the model policy ($\mu = p_\theta$), the expected auto-differentiated gradients of $L_f$ match the gradients of the corresponding f-divergence $D_f$.
   • Off-Policy Validity: When evaluated on off-policy data (sampled from a behavior distribution $\mu$), the loss remains a valid objective with the same global minimizer ($p_\theta = p^*$), provided the support of the target is contained in the support of the behavior policy.

2. From Loss to Divergence: Conversely, any strictly convex, translation-invariant loss $\ell$ on log-probabilities corresponds to minimizing a specific f-divergence. The paper proves these mappings are inverses of each other.

Practical Implementations

To address practical challenges in training, the authors introduce several extensions:

• DevGrad Loss (Batch-wise Normalization): To handle intractable partition functions (common in GFlowNets and unnormalized targets), the authors generalize the VarGrad estimator. Instead of using batch variance, they use the batch-wise generalized deviation of the log-probability differences. This involves estimating the normalization constant $\log Z$ by minimizing the loss over the batch, effectively centering the score function coefficients and reducing variance.
• Tempered Loss: To prevent numerical overflow when dealing with small temperature parameters $\beta$ (where the target distribution is $p^* \propto \pi_{ref} \exp(r/\beta)$), the authors propose a tempered loss. This scales the loss by $1/\beta$ and applies it to the tempered distribution, ensuring gradients remain bounded and stable.
• f-Trajectory Balance for GFlowNets: The framework is applied to GFlowNets by enforcing trajectory balance constraints using the $L_f$ loss. This generalizes the standard Trajectory Balance (which corresponds to reverse KL) to arbitrary f-divergences.

Key Contributions

1. Generalized Loss Family: Derivation of a general family of translation-invariant surrogate losses ($L_f$) whose on-policy gradients match those of the corresponding f-divergence, while retaining off-policy validity.
2. Theoretical Equivalence: Proof of the inverse relationship between translation-invariant log-probability losses and f-divergences, allowing for the systematic design of new loss functions with desired properties (e.g., mode-covering).
3. Extension to Unnormalized Targets: Generalization of the VarGrad technique to arbitrary f-divergences via the DevGrad loss, addressing the intractable partition function problem in RL fine-tuning.
4. Control of Exploration: Demonstration that the choice of f-divergence (parameterized by $\alpha$ in $\alpha$-divergences) directly controls the mode-seeking vs. mode-covering behavior. Lower $\alpha$ values (e.g., Forward KL, Hellinger) promote mode-covering, while higher values (e.g., Pearson $\chi^2$) promote mode-seeking.

Experimental Results

The authors evaluate the loss family across four distinct tasks:

1. Synthetic Grid (GFlowNets): On a 2D grid task designed to test mode coverage, losses with lower $\alpha$ (Forward KL, Hellinger) successfully discovered all four modes of the target distribution significantly faster and more reliably than standard Trajectory Balance (Reverse KL). Conversely, Pearson $\chi^2$ ($\alpha=2$) exhibited strong mode-seeking behavior, getting stuck in the first mode found.
2. Molecule Discovery (SynFlowNet): Applied to the SynFlowNet framework for generating synthetically accessible molecules. The authors found that tuning $\alpha$ allowed control over the diversity of generated molecules. An annealing schedule (starting with lower $\alpha$ for exploration and increasing to $\alpha=1$ for exploitation) yielded the best trade-off, generating higher-reward molecules with greater diversity compared to standard Trajectory Balance.
3. Diffusion Models: In conditional sampling of MNIST digits (odd vs. even), standard Trajectory Balance failed to sample evenly across the target modes (oversampling specific digits like 0 and 6). f-Trajectory Balance losses, particularly those with lower $\alpha$, achieved more uniform sampling across the target modes.
4. Asynchronous LLM Tuning: The losses were applied to fine-tune LLMs (Qwen and OLMo families) on math problems (GSM8k, MATH) with a 50-step asynchronous delay. The f-divergence losses demonstrated stable training and consistent entropy-reward trade-offs, whereas standard PPO implementations exhibited instability under the same asynchronous conditions.

Significance and Claims

The paper claims that the effectiveness of the squared KL loss is not an isolated phenomenon but a specific instance of a broader family of losses. By extending the property of "on-policy gradient matching" and "off-policy validity" to the entire family of f-divergences, the authors provide a unified framework for tuning generative models.

The primary significance lies in the ability to explicitly control the exploration-exploitation trade-off simply by selecting a different surrogate loss function (or tuning an $\alpha$ parameter), without changing the underlying training algorithm. This allows practitioners to choose between mode-seeking (for reward maximization) and mode-covering (for diversity and exploration) behaviors while maintaining the stability and off-policy compatibility required for modern, large-scale asynchronous training of LLMs and GFlowNets. The work suggests that this framework is particularly valuable for applications like drug discovery and exploratory agents where covering the space of high-reward candidates is critical.