---
title: 'GOAT: GFlowNet-guided Distribution Alignment'
url: https://www.emergentmind.com/topics/gflownet-guided-distribution-alignment-goat
type: topic
---

# GOAT: GFlowNet-guided Distribution Alignment

Searching arXiv for the cited GOAT-related papers to ground the article in current preprints.
GFlowNet-guided Distribution Alignment (GOAT) denotes a family of alignment procedures that cast generation as trajectory flow optimization and use GFlowNet training signals to steer a pretrained model toward a target distribution rather than performing pure reward maximization. In diffusion alignment, GOAT has been used to align a finetuned policy to $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$ while preserving the pretrained prior [2412.07775]. In concept erasure, it aligns the reverse diffusion trajectory distribution under an unsafe prompt $c$ to the trajectory distribution of a safe anchor prompt $c^*$ via a constant-reward Trajectory Balance objective [2511.00804]. In LM-based TTS, it reformulates autoregressive decoding as a trajectory flow problem and aligns sequence sampling to a sharpened intrinsic reward derived from backbone LM probabilities [2508.15442]. In non-acyclic GFlowNets with fixed initial flow, the same alignment viewpoint has an optimal-transport interpretation: the learned policy encodes an optimal coupling between source and target distributions on a graph [2606.06272].

## 1. Concept and representative instantiations

In the diffusion-alignment literature, GOAT is used to denote “the general strategy of aligning a pretrained diffusion model to a target distribution using GFlowNet training signals rather than pure reward maximization” [2412.07775]. A closely related formulation describes GOAT as learning “a stochastic policy over generation trajectories so that their probability is proportional to (or matches) a target, possibly unnormalized, ‘reward’ or flow density” [2511.00804]. Across applications, the common object is not a single terminal sample but a full generation trajectory, whether that trajectory is a reverse denoising path, an autoregressive token sequence, or a path in a directed graph.

Representative instantiations differ primarily in their alignment target, their objective, and the modality in which trajectories are defined.

| Setting | Alignment target | Objective family |
|---|---|---|
| Diffusion finetuning | $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$ | residual $\nabla$-DB |
| Concept erasure | safe anchor trajectory distribution $\tau_{c^*}$ | TB with constant reward |
| LM-based TTS | sharpened intrinsic sequence reward | enhanced SubTB |
| Non-acyclic graph alignment | target terminal distribution with fixed initial flow | minimum-flow GFlowNet / TB |

This usage indicates that GOAT is better understood as a distribution-alignment paradigm than as a single algorithm. Some instantiations rely on external reward gradients, some use intrinsic rewards, and some replace learned rewards with a constant on anchor trajectories, but all retain the GFlowNet principle of matching a target distribution through flow constraints rather than through endpoint-only reward maximization [2412.07775].

## 2. Shared trajectory formalism

GOAT formulations define generation as movement through a state space equipped with forward and backward transitions. In LM-based TTS, the state at step $t$ is the partial speech token sequence $s_t=a_1\ldots a_t$, the action is the next token, and the trajectory is the sequence from the empty prefix to the terminal state $a^\top$ ending in the termination token $\top$ [2508.15442]. In diffusion, states are noisy latents indexed by time; one formulation takes $s_t\equiv x_t\in\mathbb R^d$ with action $a_t:s_t\mapsto s_{t-1}$, trajectory $\tau=(x_T,x_{T-1},\dots,x_0)$, forward policy given by the reverse diffusion conditional, and backward policy given by the fixed noising kernel $q(x_t\mid x_{t-1})$ [2511.00804]. In non-acyclic graph GFlowNets, the environment is a finite directed graph with distinguished initial and sink nodes, and trajectory probabilities factorize over edge-local forward or backward policies [2606.06272].

A canonical trajectory factorization is
$$
P_F(\tau=(s_0\to\cdots\to s_n))=\prod_{i=0}^{n-1}P_F(s_{i+1}\mid s_i).
$$
The terminal-state distribution is induced by marginalizing over all trajectories ending at that terminal state. In the TTS formulation, the reward proportionality condition is written as
$$
R(x)=Z\sum_{\tau=(s_0\to\cdots\to s_n=x)}P_F(\tau), \qquad x\in\mathcal X,
$$
with $Z=\sum_{x\in\mathcal X}R(x)$ [2508.15442]. In the graph setting, normalized edge flows induce the policy directly,
$$
\pi(a\mid s)=\frac{\mathcal F(s\to a)}{\sum_{a'\in\mathrm{Out}(s)}\mathcal F(s\to a')},
$$
so alignment can be phrased either in terms of policies or in terms of flows [2606.06272].

The significance of this shared formalism is that credit assignment can be imposed along the trajectory rather than only at its endpoint. In concept erasure, this is explicit: the paper attributes prior limitations of concept erasure to “a myopic view of the denoising trajectories” and replaces endpoint alignment with alignment of the entire reverse diffusion trajectory [2511.00804]. This suggests that GOAT is centrally concerned with redistribution of probability mass across trajectories, not merely with rescoring final outputs.

## 3. Objective families and distributional guarantees

GOAT has been instantiated with several GFlowNet objectives. In diffusion finetuning, “Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets” formulates gradient-informed detailed balance ($\nabla$-DB) and its residual prior-preserving version. The residual guarantee states that if residual losses and the terminal condition are satisfied, then the finetuned marginal equals
$$
p_\theta(x_T)\propto R(x_T)^\beta\,p_{\theta_0}(x_T),
$$
which is the intended GOAT target distribution [2412.07775].

In concept erasure, “EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment” uses Trajectory Balance rather than detailed balance. The core instantiation sets a constant reward on anchor trajectories,
$$
R(\tau)=
\begin{cases}
\beta, & \tau\in\tau_{c^*} \\
0, & \text{otherwise},
\end{cases}
\qquad \beta>0,
$$
and minimizes
$$
\mathcal{L}_{c\leftarrow c^*}^{\mathrm{EraseFlow}}
=
\left(
\log Z_\phi
+
\sum_{t=1}^T \log p_\theta(x_{t-1}^*\mid x_t^*,t,c)
-
\log\beta
-
\sum_{t=1}^T \log q(x_t^*\mid x_{t-1}^*)
\right)^2.
$$
The paper reports that a reward-based DB variant was explored but found inferior to TB for this asymmetric task [2511.00804].

In LM-based TTS, GOAT trains with an enhanced Subtrajectory Balance objective rather than full-trajectory TB, because long autoregressive sequences are susceptible to “fragmentary collapse.” The target distribution is specified by a sharpened intrinsic reward derived from the backbone LM,
$$
R(\mathbf a_k\mid \mathbf q)=p_{\text{LM}}(\mathbf a_k\mid \mathbf q)^{\frac{1}{T}}
=
\left(\prod_{i=1}^k p_{\text{LM}}(a_i\mid a_{1:i-1},\mathbf q)\right)^{\frac{1}{T}},
$$
with a linearly decayed reward temperature $T_t$ [2508.15442].

In non-acyclic graph GFlowNets, the relevant objective is minimum total internal flow under fixed initial and terminal boundary conditions. The paper proves that this minimum-flow problem is equivalent to a Kantorovich OT problem, and that the learned trajectory distribution induces the optimal coupling
$$
\Pi^\star_{u,x}=\sum_{\tau:u\leadsto x}\mathcal P^\star(\tau).
$$
This establishes a distribution-alignment guarantee phrased as transport rather than reward matching [2606.06272].

A recurring point across these formulations is that objective choice is task-dependent. Diffusion reward finetuning emphasizes per-step score constraints, concept erasure emphasizes asymmetric trajectory redistribution, autoregressive TTS emphasizes subtrajectory consistency, and graph alignment emphasizes minimum-flow feasibility.

## 4. Diffusion-model alignment and concept erasure

In diffusion finetuning, GOAT appears in the form of $\nabla$-GFlowNet and residual $\nabla$-GFlowNet. The target distribution
$$
p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))
$$
is designed to emphasize high-reward samples while keeping the pretrained prior in support, thereby mitigating catastrophic forgetting and preserving diversity. The method trains the denoising policy and a residual flow-score network through residual $\nabla$-DB losses, uses a forward-looking parameterization of residual flow scores, and adds a prior-preserving regularizer of the form $\lambda\|\epsilon_\theta(x_t)-\epsilon_{\theta^\dagger}(x_t)\|^2$ [2412.07775].

The reported empirical pattern is a reward–diversity–prior-preservation trade-off that differs from mode-seeking RL baselines. On Aesthetic, residual $\nabla$-DB achieves mean reward $\approx 7.86$–$7.90$ with DreamSim diversity $\approx 29.2$–$29.7\times 10^{-2}$ and FID $\approx 317$–$318$, whereas ReFL and DRaFT achieve high reward $\approx 9.5$–$10.2$ but very low diversity $\approx 4$–$8\times 10^{-2}$ and FID $\approx 1.6$–$1.9$k, which the paper interprets as mode collapse and forgetting. Similar Pareto behavior is reported on HPSv2 and ImageReward [2412.07775].

In concept erasure, GOAT is specialized to align an unsafe concept prompt $c$ to a safe anchor prompt $c^*$ at the level of reverse denoising trajectories. The paper states a proposition: if TB holds with constant reward for anchor trajectories and the standard TB holds for the original model under $c^*$, then for every timestep $t$,
$$
p_{\theta^*}(x_{t-1}\mid x_t,t,c)=p_\theta(x_{t-1}\mid x_t,t,c^*)
\quad\Rightarrow\quad
p_{\theta^*}(x_0\mid c)=p_\theta(x_0\mid c^*),
$$
so the unsafe concept distribution is replaced by the safe anchor distribution [2511.00804].

The empirical results are reported on SD v1.4. Against non-adversarial baselines on UDAtk ASR, EraseFlow records $33.89$ versus DUO $64.40$, MACE $72.81$, ESD $78.81$, and UCE $87.28$. As a plug-and-play component, EraseFlow + AdvUnlearn reaches $1.42$ versus AdvUnlearn $16.94$, and EraseFlow + SAFREE reaches $24.57$ versus SAFREE $85.59$. On NSFW datasets, EraseFlow reports I2P $2.80$, Ring-a-Bell $0.00$, and MMA-Diff $0.60$. For prior preservation, FID is $17.93$ and CLIP Score is $25.67$. Reported cost is $\approx 2.8$ minutes on one A100 GPU for SD v1.4, with memory $\approx 42$ GB peak; the comparison figures are RACE $225$ minutes and AdvUnlearn $1440$ minutes [2511.00804].

The ablations also clarify what GOAT contributes in this setting. TB with reward improves markedly over DB with reward, and constant-reward TB gives the best overall NSFW performance. The paper further reports that $\log\beta\in[2,3]$ yields $\sim 96\%$ improvement in I2P, larger `STOP_SAMPLING` values improve stability by increasing anchor trajectory diversity, “semantically opposite” anchors outperform neutral or semantically close prompts, and mixing early and late timesteps outperforms only-early or uniformly random selections [2511.00804].

## 5. Autoregressive sequence alignment in LM-based TTS

In LM-based TTS, GOAT is a post-training framework for mitigating hallucinations in autoregressive decoding. The motivating empirical observation is an uncertainty analysis on SeedTTS-Eval test-hard with a CosyVoice2 backbone and stochastic multinomial sampling: utterance-level uncertainty is positively correlated with hallucination, with Pearson correlation coefficient $0.636$ and Spearman rank correlation coefficient $0.649$ ($p<1\mathrm E{-4}$) [2508.15442]. The uncertainty metric is predictive entropy at token-, word-, and utterance-level granularity.

The alignment mechanism uses intrinsic rewards rather than external reward models. GOAT trains a forward sampling policy $P_{\text{GFN}}$ with enhanced SubTB so that the marginal over complete token sequences satisfies $P_{\text{GFN}}^\top(\mathbf a^\top)\propto R(\mathbf a^\top)$. The reward is based on backbone LM probabilities and sharpened by an inverse temperature $T\in(0,1)$. To balance performance and suppress reward hacking, the reward temperature decays linearly from $1.0$ to $T_{\min}=0.825$, and learning-rate optimization combines a $20$-step warm-up with cosine annealing over $S_{\text{LRO}}\in\{1500,2500\}$ to a maximum learning rate of $1\mathrm e{-5}$ [2508.15442].

The practical pipeline is lightweight. The backbone is CosyVoice2, GOAT fine-tunes only lightweight adapters via LoRA, and as a post-training method it adds no inference-time components; generation uses the same decoder as the backbone, with no reranking or extra ASR feedback, so inference latency remains essentially unchanged. Training uses $4\times$ NVIDIA H100, evaluation uses $1\times$ V100, and the added training overhead comes from subtrajectory losses, which are potentially $O(n^2)$ per sequence if fully enumerated [2508.15442].

The reported results emphasize hallucination reduction, uncertainty reduction, and preservation of speech quality. On test-hard with RMS, baseline CER is $13.72\%$, while GOAT models achieve $6.61$–$6.53\%$ in Chinese and $6.56$–$6.76\%$ in English/mix under RAS, which the paper summarizes as more than $50\%$ reduction. Across test-zh and test-en, CER/WER reduce substantially; one reported example is test-en WER dropping from $7.10\%$ for baseline RMS to approximately $2.0$–$2.2\%$ for GOAT RAS. Utterance Uncertainty Ratio reductions reach up to $58\%$, with examples including test-zh UUR $0.42$, test-en UUR $0.58$, and test-hard UUR $0.39$. UTMOS improves across configurations, while speaker similarity remains comparable to baseline [2508.15442].

The ablation results isolate the effect of the GOAT design. On test-hard, CER decreases from approximately $11.7\%$ under TB to approximately $6.5$–$6.6\%$ under SubTB. Removing reward temperature decay degrades convergence and final performance; omitting learning-rate optimization induces reward hacking through premature sequence termination. The paper also notes that prosodic hallucinations such as pauses and silences do not consistently correlate with entropy, indicating that the current uncertainty proxy captures only a subset of hallucination phenomena [2508.15442].

## 6. Optimal-transport interpretation, misconceptions, and open problems

A distinctive theoretical development is the proof that a non-acyclic GFlowNet trained under a minimum-flow objective with fixed initial flow distribution implicitly solves a Kantorovich OT problem on a directed graph. Source nodes are the direct children of the initial node, target nodes are terminal states, and the transport cost is the graph-induced shortest-path cost
$$
c(u,x)=\min_{p\in\mathcal P(u\to x)}\sum_{e\in p}w_e.
$$
Under mass balance and fixed boundary flows, the paper proves $\mathrm{GFlow}^\star=\mathrm{OT}^\star$, and sampling trajectories from the learned policy recovers the corresponding optimal coupling [2606.06272].

The empirical validation includes hypergrid and permutation experiments. On a hypergrid with $H=10$ and a moon-shaped source distribution, the learned sampler achieves $\mathbb E[|\tau|]=4.352$ versus $\mathrm{OT}^\star=4.351$. On a Cayley graph with adjacent swaps, results include $n=4:\ \mathbb E[|\tau|]=0.557$ versus $\mathrm{OT}^\star\approx 0.567$ and $n=8:\ \mathbb E[|\tau|]=1.010$ versus $\mathrm{OT}^\star\approx 1.008$. For $n=20$, exact OT is intractable, and the GFlowNet formulation is presented as a scalable approximation route [2606.06272].

This theoretical perspective clarifies a common misconception: GOAT is not inherently tied to external reward models. The concept-erasure formulation replaces learned or handcrafted rewards with a constant on anchor trajectories, explicitly to avoid brittle and hackable reward critics [2511.00804]. The TTS formulation uses intrinsic rewards derived from the backbone LM probabilities rather than external labels [2508.15442]. By contrast, residual $\nabla$-GFlowNet in diffusion finetuning does depend on differentiable reward functions and their gradients, but does so within detailed-balance constraints designed to preserve diversity and the pretrained prior [2412.07775].

Another misconception is that GOAT is simply endpoint matching. The concept-erasure proposition is explicitly trajectory-level; the TTS use of SubTB is motivated by the need to learn from subsequences in long autoregressive chains; and the OT result shows that full path structure can encode a transport plan, not merely a terminal histogram [2511.00804]. A plausible implication is that GOAT is most distinctive when the geometry of the generation process matters and when local errors can propagate across long trajectories.

The current limitations are correspondingly diverse. In concept erasure, scaling to many visually similar concepts can cause interference and reduced retention, and gains on deterministic flows such as SDv3 and Flux are less pronounced than on diffusion models [2511.00804]. In diffusion reward finetuning, performance is sensitive to reward scale and quality, and reward-gradient computation through decoding adds cost [2412.07775]. In TTS, stability depends on reward temperature decay and learning-rate optimization, and SubTB increases training complexity [2508.15442]. In the OT setting, mis-specified graph topology or edge costs can distort the coupling, and poor optimization can yield suboptimal flows [2606.06272].

Taken together, these results place GOAT at the intersection of generative-model alignment, safety editing, sequence reliability, and structured transport. Its central claim is that trajectory-distribution alignment can be made explicit, optimized with GFlowNet objectives, and, in several settings, endowed with distributional guarantees that are stronger than endpoint-only reward maximization.

Source: https://www.emergentmind.com/topics/gflownet-guided-distribution-alignment-goat