GOAT: GFlowNet-guided Distribution Alignment

• GFlowNet-guided Distribution Alignment (GOAT) is a framework that formulates generation as trajectory flow optimization to steer models toward target distributions rather than solely maximizing endpoint rewards.
• It encompasses varied instantiations such as diffusion finetuning, concept erasure, and autoregressive TTS, each utilizing distinct objectives like residual gradients and trajectory balance for alignment.
• GOAT integrates rigorous theoretical guarantees, including optimal-transport interpretations, which underlie improved diversity preservation, safety editing, and accuracy in generative tasks.

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 (Liu et al., 2024). 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 (Kusumba et al., 2 Nov 2025). 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 (Liu et al., 21 Aug 2025). 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 (Maksimov et al., 4 Jun 2026).

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” (Liu et al., 2024). 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” (Kusumba et al., 2 Nov 2025). 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 (Liu et al., 2024).

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$ (Liu et al., 21 Aug 2025). In diffusion, states are noisy latents indexed by time; one formulation takes $c$0 with action $c$1, trajectory $c$2, forward policy given by the reverse diffusion conditional, and backward policy given by the fixed noising kernel $c$3 (Kusumba et al., 2 Nov 2025). 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 (Maksimov et al., 4 Jun 2026).

A canonical trajectory factorization is

$c$4

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

$c$5

with $c$6 (Liu et al., 21 Aug 2025). In the graph setting, normalized edge flows induce the policy directly,

$c$7

so alignment can be phrased either in terms of policies or in terms of flows (Maksimov et al., 4 Jun 2026).

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 (Kusumba et al., 2 Nov 2025). 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 ($c$8-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

$c$9

which is the intended GOAT target distribution (Liu et al., 2024).

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,

$c^*$0

and minimizes

$c^*$1

The paper reports that a reward-based DB variant was explored but found inferior to TB for this asymmetric task (Kusumba et al., 2 Nov 2025).

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,

$c^*$2

with a linearly decayed reward temperature $c^*$3 (Liu et al., 21 Aug 2025).

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

$c^*$4

This establishes a distribution-alignment guarantee phrased as transport rather than reward matching (Maksimov et al., 4 Jun 2026).

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 $c^*$5-GFlowNet and residual $c^*$6-GFlowNet. The target distribution

$c^*$7

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 $c^*$8-DB losses, uses a forward-looking parameterization of residual flow scores, and adds a prior-preserving regularizer of the form $c^*$9 (Liu et al., 2024).

The reported empirical pattern is a reward–diversity–prior-preservation trade-off that differs from mode-seeking RL baselines. On Aesthetic, residual $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$0-DB achieves mean reward $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$1–$p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$2 with DreamSim diversity $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$3–$p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$4 and FID $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$5–$p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$6, whereas ReFL and DRaFT achieve high reward $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$7–$p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$8 but very low diversity $p_R(x)\propto p_{\theta_0}(x)\exp(\beta R(x))$9–$\nabla$0 and FID $\nabla$1–$\nabla$2k, which the paper interprets as mode collapse and forgetting. Similar Pareto behavior is reported on HPSv2 and ImageReward (Liu et al., 2024).

In concept erasure, GOAT is specialized to align an unsafe concept prompt $\nabla$3 to a safe anchor prompt $\nabla$4 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 $\nabla$5, then for every timestep $\nabla$6,

$\nabla$7

so the unsafe concept distribution is replaced by the safe anchor distribution (Kusumba et al., 2 Nov 2025).

The empirical results are reported on SD v1.4. Against non-adversarial baselines on UDAtk ASR, EraseFlow records $\nabla$8 versus DUO $\nabla$9, MACE $\tau_{c^*}$0, ESD $\tau_{c^*}$1, and UCE $\tau_{c^*}$2. As a plug-and-play component, EraseFlow + AdvUnlearn reaches $\tau_{c^*}$3 versus AdvUnlearn $\tau_{c^*}$4, and EraseFlow + SAFREE reaches $\tau_{c^*}$5 versus SAFREE $\tau_{c^*}$6. On NSFW datasets, EraseFlow reports I2P $\tau_{c^*}$7, Ring-a-Bell $\tau_{c^*}$8, and MMA-Diff $\tau_{c^*}$9. For prior preservation, FID is $t$0 and CLIP Score is $t$1. Reported cost is $t$2 minutes on one A100 GPU for SD v1.4, with memory $t$3 GB peak; the comparison figures are RACE $t$4 minutes and AdvUnlearn $t$5 minutes (Kusumba et al., 2 Nov 2025).

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 $t$6 yields $t$7 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 (Kusumba et al., 2 Nov 2025).

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 $t$8 and Spearman rank correlation coefficient $t$9 ($s_t=a_1\ldots a_t$0) (Liu et al., 21 Aug 2025). 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 $s_t=a_1\ldots a_t$1 with enhanced SubTB so that the marginal over complete token sequences satisfies $s_t=a_1\ldots a_t$2. The reward is based on backbone LM probabilities and sharpened by an inverse temperature $s_t=a_1\ldots a_t$3. To balance performance and suppress reward hacking, the reward temperature decays linearly from $s_t=a_1\ldots a_t$4 to $s_t=a_1\ldots a_t$5, and learning-rate optimization combines a $s_t=a_1\ldots a_t$6-step warm-up with cosine annealing over $s_t=a_1\ldots a_t$7 to a maximum learning rate of $s_t=a_1\ldots a_t$8 (Liu et al., 21 Aug 2025).

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 $s_t=a_1\ldots a_t$9 NVIDIA H100, evaluation uses $a^\top$0 V100, and the added training overhead comes from subtrajectory losses, which are potentially $a^\top$1 per sequence if fully enumerated (Liu et al., 21 Aug 2025).

The reported results emphasize hallucination reduction, uncertainty reduction, and preservation of speech quality. On test-hard with RMS, baseline CER is $a^\top$2, while GOAT models achieve $a^\top$3–$a^\top$4 in Chinese and $a^\top$5–$a^\top$6 in English/mix under RAS, which the paper summarizes as more than $a^\top$7 reduction. Across test-zh and test-en, CER/WER reduce substantially; one reported example is test-en WER dropping from $a^\top$8 for baseline RMS to approximately $a^\top$9–$\top$0 for GOAT RAS. Utterance Uncertainty Ratio reductions reach up to $\top$1, with examples including test-zh UUR $\top$2, test-en UUR $\top$3, and test-hard UUR $\top$4. UTMOS improves across configurations, while speaker similarity remains comparable to baseline (Liu et al., 21 Aug 2025).

The ablation results isolate the effect of the GOAT design. On test-hard, CER decreases from approximately $\top$5 under TB to approximately $\top$6–$\top$7 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 (Liu et al., 21 Aug 2025).

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

$\top$8

Under mass balance and fixed boundary flows, the paper proves $\top$9, and sampling trajectories from the learned policy recovers the corresponding optimal coupling (Maksimov et al., 4 Jun 2026).

The empirical validation includes hypergrid and permutation experiments. On a hypergrid with $c$00 and a moon-shaped source distribution, the learned sampler achieves $c$01 versus $c$02. On a Cayley graph with adjacent swaps, results include $c$03 versus $c$04 and $c$05 versus $c$06. For $c$07, exact OT is intractable, and the GFlowNet formulation is presented as a scalable approximation route (Maksimov et al., 4 Jun 2026).

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 (Kusumba et al., 2 Nov 2025). The TTS formulation uses intrinsic rewards derived from the backbone LM probabilities rather than external labels (Liu et al., 21 Aug 2025). By contrast, residual $c$08-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 (Liu et al., 2024).

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 (Kusumba et al., 2 Nov 2025). 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 (Kusumba et al., 2 Nov 2025). In diffusion reward finetuning, performance is sensitive to reward scale and quality, and reward-gradient computation through decoding adds cost (Liu et al., 2024). In TTS, stability depends on reward temperature decay and learning-rate optimization, and SubTB increases training complexity (Liu et al., 21 Aug 2025). In the OT setting, mis-specified graph topology or edge costs can distort the coupling, and poor optimization can yield suboptimal flows (Maksimov et al., 4 Jun 2026).

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.