# Preference Alignment with Flow Matching

Minu Kim1 Yongsik Lee1∗ Sehyeok Kang1 Jihwan Oh1
Song Chong1† Seyoung Yun1†
1KAIST AI
{minu.kim, dldydtlr93, kangsehyeok0329, ericoh929,
songchong<EMAIL_ADDRESS>Equal contribution. †Corresponding authors.

###### Abstract

We present Preference Flow Matching (PFM), a new framework for preference-
based reinforcement learning (PbRL) that streamlines the integration of
preferences into an arbitrary class of pre-trained models. Existing PbRL
methods require fine-tuning pre-trained models, which presents challenges such
as scalability, inefficiency, and the need for model modifications, especially
with black-box APIs like GPT-4. In contrast, PFM utilizes flow matching
techniques to directly learn from preference data, thereby reducing the
dependency on extensive fine-tuning of pre-trained models. By leveraging flow-
based models, PFM transforms less preferred data into preferred outcomes, and
effectively aligns model outputs with human preferences without relying on
explicit or implicit reward function estimation, thus avoiding common issues
like overfitting in reward models. We provide theoretical insights that
support our method’s alignment with standard PbRL objectives. Experimental
results indicate the practical effectiveness of our method, offering a new
direction in aligning a pre-trained model to preference.

##  1 Introduction

Preference-based reinforcement learning (PbRL) has emerged as a groundbreaking
approach with significant contributions to performance improvement (Akrour et
al., 2011; Wilson et al., 2012), particularly in the realm of artificial
intelligence where understanding and incorporating human preferences are
crucial. Unlike traditional reinforcement learning, which struggles due to the
absence of explicit reward functions or the infeasibility of defining
comprehensive environmental rewards, PbRL leverages a variety of feedback
forms from humans to guide the learning process. This class of PbRL method is
often referred to as reinforcement learning from human feedback (RLHF)
(Ziegler et al., 2019; Levine et al., 2018; Ouyang et al., 2022).

Despite their effectiveness, these methods necessitate fine-tuning pre-trained
models to align with user preferences, introducing several challenges such as
scalability, accessibility, inefficiency, and the need for model
modifications. For instance, with black-box APIs like GPT-4 (OpenAI et al.,
2024), customization based on user preferences is constrained due to
restricted access to the underlying model. Moreover, even if fine-tuning were
feasible, the large model size results in inefficient training and high
resource consumption. Aligning black-box models with user preferences remains
an under-explored area in research, despite its critical importance and
growing demand.

In this line of research, we propose Preference Flow Matching (PFM), which
redefines the integration of human preferences by directly learning a
preference flow from the less preferred data to the more preferred ones. This
direct modeling of preference flows allows our system to better characterize
and replicate the marginal distribution of the favored outcomes. We adopt a
novel flow matching framework (Lipman et al., 2022), which is a simple,
intuitive, yet relatively under-explored method for preference alignment. By
simply adding a preference flow matching module to black-box models, PFM
eliminates the need for fine-tuning the black-box model itself, providing a
significant advantage.

Additionally, our method offers a highly robust approach for preference
alignment, by circumventing the need for explicit or implicit reward function
estimation. In typical RLHF scenarios, a model is initially trained to
approximate a reward function based on human preferences. This reward model is
then used to guide the policy learning process, aiming to align agent
behaviors more closely with human preferences. However, this approach can
introduce complexities and potential biases in translating human preferences
into numerical rewards. In particular, learned reward models can often overfit
the ground truth preference model, especially in the finite data regime (Azar
et al., 2023). Recent advancements such as Direct Preference Optimization(DPO)
(Rafailov et al., 2024) address the complexities of RLHF by eliminating the
need for reward learning. However, these methods still inherently optimize for
the reward model, and hence they are also susceptible to reward overfitting.
In contrast, PFM directly learns preference flow, thereby removing the need
for any reward model assumptions and resolving the challenges associated with
the reward model learning.

We prove both theoretically and empirically that our method is able to learn
an object that is similar to the standard RLHF objectives, while being robust
to the preference overfitting observed in traditional RLHF pipelines. We also
demonstrate how we can further improve the quality of alignment via iterative
flow matching, with theoretical guarantees. Experimentally, we find that while
typical RLHF methods and DPO suffer from preference overfitting, our method
can robustly align with preference and still achieve comparable performances.

Figure 1: Illustration of our PFM framework. In the typical RLHF scenarios
(left), we first sample preference data from the supervised fine-tuned (SFT)
reference model. A reward model is learned from the collected dataset, either
implicitly (as in DPO) or explicitly. The reward model is then used to fine-
tune the reference policy to obtain the final model. Our method directly
learns the preference flow from the collected preference data, where the flow
is represented as a vector field $v_{\theta}$ (middle). For inference, we
again sample a point from the reference policy, and improve the quality of
alignment by using the trained flow matching model, without the need of fine-
tuning the existing reference model (right).

##  2 Preliminaries

###  2.1 Reinforcement Learning from Human Feedback (RLHF)

Reinforcement learning from human feedback generally begins with obtaining a
pre-trained reference policy $\pi_{\mathrm{ref}}$ that can generate samples
$y\sim\pi_{\mathrm{ref}}(\cdot|x)$ given a context $x$. For example, a context
$x$ could be a text prompt given by a user, and the sample $y$ could represent
an appropriate text response generated by the reference policy
$\pi_{\mathrm{ref}}$. We then collect a dataset of $N$ preference pairs
$\mathcal{D}=\\{(x_{i},y_{i}^{+},y_{i}^{-})\\}_{i=1}^{N}$, where each $x_{i}$
denotes the context, and each
$y_{i}^{+},y_{i}^{-}\sim\pi_{\mathrm{ref}}(\cdot|x_{i})$ are generated
responses to $x_{i}$ and marked as good or bad samples, respectively. Here, we
assume that the preference $y_{i}^{+}>y_{i}^{-}$ is generated from a ground-
truth reward $r^{*}:X\times Y\to\mathbb{R}_{\geq 0}$, where $X$ and $Y$ are
context space and response space, respectively. The goal of general RLHF is to
recover an optimal policy $\pi^{*}$ such that

$\pi^{*}=\underset{\pi}{\mathrm{argmax}}\
\mathbb{E}_{x}\left(\mathbb{E}_{y\sim\pi(\cdot|x)}\Big{(}r^{*}(x,y)\Big{)}-\beta\mathbb{D}_{\mathrm{KL}}\Big{(}\pi(\cdot|x)\Big{\|}\pi_{\mathrm{ref}}(\cdot|x)\Big{)}\right).$
(1)

RLHF pipelines generally require reward learning. One of the most popular
choices of the reward model is the Bradley-Terry model (Bradley and Terry,
1952), which assumes that the preference $y^{+}>y^{-}$ is generated from the
probability $\mathbb{P}(y^{+}>y^{-}|x)=\sigma(r^{*}(x,y^{+})-r^{*}(x,y^{-}))$,
where $\sigma$ is the logistic function. Under this model, the general RLHF
framework learns the reward model $r_{\phi}\approx r^{*}$ by minimizing the
negative log-likelihood:

$\mathcal{L}_{R}(\phi;\mathcal{D}):=-\mathbb{E}_{(x,y^{+},y^{-})\sim\mathcal{D}}\Big{(}\log\sigma\big{(}r_{\phi}(x,y^{+})-r_{\phi}(x,y^{-})\big{)}\Big{)}.$
(2)

Once the reward model $r_{\phi}$ is trained, we then use it to optimize for
(1) to obtain $\pi_{\theta}\approx\pi^{*}$ using standard reinforcement
learning algorithms.

There is also a class of _reward-free_ methods that eliminates the need of
reward learning phase (Rafailov et al., 2024; Azar et al., 2023). Direct
Policy Optimization (DPO) (Rafailov et al., 2024) is a representative reward-
free method that optimizes for (1) directly without learning a reward model.
Although being a reward-free method, DPO implicitly optimizes for the reward
function as in (2), by replacing
$\hat{r}_{\theta}(x,y)=\beta\log(\pi_{\theta}(y|x)/\pi_{\mathrm{ref}}(y|x))$
as the implicit reward estimate.

###  2.2 Flow Matching

Flow matching is a class of generative model, where given a prior distribution
$p_{0}$, we aim to model a target distribution $p_{1}$ from $p_{0}$. A key
difference of the flow matching to the other generative models is that the
prior $p_{0}$ can be an _arbitrary_ distribution, (diffusion, for example,
starts from a Gaussian prior $p_{0}$) and that the flow matching algorithm
learns to modify the prior distribution $p_{0}$ to the target distribution
$p_{1}$ with a neural network.

Throughout, we consider a pair of data distributions over $\mathbb{R}^{d}$
with densities $y^{-}\sim p_{0}$ and $y^{+}\sim p_{1}$, possibly unknown (but
able to sample). The flow matching considers the task of fitting a mapping
$f:\mathbb{R}^{d}\to\mathbb{R}^{d}$ that transforms $p_{0}$ to $p_{1}$, that
is, if $y^{-}\sim p_{0}$, then $f(y^{-})\sim p_{1}$. Inspired as in the
motivation for the diffusion models, one can define a smooth time-varying
vector field $u:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ that defines an
ordinary diffusional equation (ODE),

$dy=u(t,y)dt$ (3)

where we use the notation $u(t,y)$ interchanably with $u_{t}(y)$. Denote the
solution of the above ODE by $\phi(t,y)$ (or $\phi_{t}(y)$) with initial
condition $\phi_{0}(y)=y$. In other words, $\phi_{t}(y)$ is the point $y$
transported along the vector field $u$ from time $0$ to $t$. In order to
obtain samples from the target distribution $p_{1}$, we simply compute
$\phi_{1}(y)$ where $y\sim p_{0}$. The integration map $\phi_{t}$ induces a
pushforward measure $p_{t}\triangleq[\phi_{t}]_{\sharp}(p_{0})$, which is the
density of points $y\sim p_{0}$ transported via $u$ from $0$ to $t$.

To train the vector field $v_{\theta}$ with neural network that mimics the
vector field $u$ of our interest, we can solve for the conditional flow
matching objective, as proposed by Lipman et al. (2022):

$\mathcal{L}_{\mathrm{CFM}}(\theta)\triangleq\mathbb{E}_{t\sim[0,1],z\sim
q(\cdot),y\sim p_{t}(\cdot|z)}\left(\|v_{\theta}(t,y)-u_{t}(y|z)\|^{2}\right)$
(4)

where $q(z)=\pi(y^{-},y^{+})$ is some coupled distribution of samples
$y^{-},y^{+}$ and $u_{t}(y|z)=y^{+}-y^{-}$ is a straight path from a source
sample to a target sample. The conditional distribution $q(z)$ can be chosen
to be a independent coupling of source and target distribution
$q(z)=p_{0}(y^{-})p_{1}(y^{+})$ (Lipman et al., 2022), or the 2-Wasserstein
optimal transport plan as proposed by Tong et al. (2023).

##  3 Preference Flow Matching

In this section, we describe how we can use flow matching to learn (human)
preference. In the first subsection, we illustrate our flow matching framework
for learning preference, and compare it with typical RLHF pipelines. Then in
Section 3.2, we demonstrate our method in a simple 2-dimensional toy
experiment. Finally in Section 3.3, we provide an extension of our framework
that can iteratively improve the performance, with theoretical guarantees.

###  3.1 Flow Matching for Preference Learning

Instead of trying to optimize for the unknown reward $r^{*}$ or the preference
probability model $\mathbb{P}(y^{+}>y^{-}|x)$, we simply learn a flow from the
marginal distribution of less preferred data $p_{0}(y^{-}|x)$ to the marginal
distribution of more preferred data $p_{1}(y^{+}|x)$ by leveraging what is
explicitly characterized in the collected preference data:

$\displaystyle p_{0}(y^{-}|x)$
$\displaystyle\propto\pi_{\mathrm{ref}}(y^{-}|x)\int\pi_{\mathrm{ref}}(y|x)\mathbb{P}(y>y^{-}|x)dy$
(5) $\displaystyle p_{1}(y^{+}|x)$
$\displaystyle\propto\pi_{\mathrm{ref}}(y^{+}|x)\int\pi_{\mathrm{ref}}(y|x)\mathbb{P}(y^{+}>y|x)dy$
(6) $\displaystyle=\pi_{\mathrm{ref}}(y^{+}|x)\
\mathbb{E}_{y\sim\pi_{\mathrm{ref}}(\cdot|x)}\bigg{(}\mathbb{P}(y^{+}>y|x)\bigg{)}$
(7)

In other words, we view that our collected data $\mathcal{D}$ is in fact
generated from each of the marginal distributions $y^{-}\sim p_{0}(\cdot|x)$
and $y^{+}\sim p_{1}(\cdot|x)$ obtained from $\mathbb{P}(y^{+}>y^{-}|x)$,
respectively. Hence, following the conventions in the literature, (Tong et
al., 2023) we define the flow matching objective for preference dataset
$\mathcal{D}$ as follows:

$\mathcal{L}(\theta)=\mathbb{E}_{t\sim[0,1],z\sim\mathcal{D},y\sim
p_{t}(\cdot|z)}\left(\|v_{\theta}(t,y|x)-u_{t}(y|z)\|^{2}\right)$ (8)

where we define the condition $z=(x,y^{+},y^{-})$, conditional flow
$u_{t}(y|z)=y^{+}-y^{-}$, and the probability path
$p_{t}(y|z)=\mathcal{N}(y|ty^{+}+(1-t)y^{-},\ \sigma^{2})$. Once we obtain the
vector field $v_{\theta}$, we can improve upon the generated negative samples
$y^{-}\sim p_{0}(\cdot|x)$ by solving (3) using an off-the-shelf numerical ODE
solver (Runge, 1895; Kutta, 1901) to obtain samples $f(y^{-})\sim p_{1}$.
Specifically, we start from a sample $y^{-}$ with $t=0$, and "flow" along the
ODE trajectory using $v_{\theta}$ until $t=1$, to arrive at the target
$y^{+}$. Detailed algorithm can be found in Algorithm 1. Notably, generating
improved samples can be done without the need of fine-tuning the existing
model, since we learn a separate vector field that transports negative samples
from $p_{0}$ to $p_{1}$. Furthermore, we did not require any assumption for
the probability model $\mathbb{P}(y^{+}>y^{-}|x)$, so our method can extend to
general scenarios that do not adopt the Bradley-Terry model. Our method is
outlined in Figure 1.

A careful reader might noticed that for inference, we require negative samples
$y^{-}$ from the marginal distribution $p_{0}$, to obtain aligned samples
$y^{+}$. However, this $p_{0}$ is inaccessible during inference step, as we
must first acquire preference label $y^{+}>y^{-}$ for samples generated from
$\pi_{\mathrm{ref}}$. Instead, we simply start from $y\sim\pi_{\mathrm{ref}}$
as the starting point, and apply flow matching to obtain $f(y)\approx
y^{+}\sim p_{1}$. We emphasize that PFM can still robustly generate positive
samples, if we assume non-deterministic preferences i.e.,
$\mathrm{supp}(p_{1})\supseteq\mathrm{supp}(p_{0})$. We also empirically find
that using $\pi_{\mathrm{ref}}$ instead of $p_{0}$ as the source distribution
can produce comparable results in practical scenarios. Further details can be
found in Appendix B.

###  3.2 Illustrative Example: Preference Generated from 8-Gaussians Density

Here, we demonstrate how our method learns to improve generated samples to
better align with the preference, in a simple 2-dimensional toy experiment. We
consider a ground truth reward function generated from an 8-Gaussians density
as illustrated in Figure 2(a). We then pre-train a Gaussian mixture model to
obtain samples as in Figure 2(c). The pairwise preference labels are then
generated using the ground truth 8-Gaussians reward function, as done in many
existing preference-based reinforcement learning (PbRL) settings (Christiano
et al., 2017; Ibarz et al., 2018; Shin et al., 2023).

Once preference data are collected, we first learn a reward model
$\widehat{r}_{\phi}$ via (2). As can be seen in Figure 2(b), the learned
reward model overfits in the unseen region, which causes the RLHF method to
diverge (Figure 2(e)). DPO also fails to learn the correct preference, as can
be seen in Figure 2(f). We note here that DPO is also subjective to the reward
overfitting since DPO also implicitly learns to optimize for the reward using
the Bradley-Terry model (2) (Xu et al., 2024; Azar et al., 2023).

However, PFM is free of such reward overfitting issues, as we do not optimize
for the reward function using the Bradley-Terry model. Unlike other RLHF
methods, our model robustly learns to align with the preference from the
provided dataset (Figure 2(g)). Notably, our method does not try to overfit
beyond the unseen region, since the learned target distribution from the flow
matching model tries to mimic the distribution $p_{1}(y^{+})$ of collected
preferred samples. (Compare Figure 2(d) and Figure 2(g).)

(a) Ground-truth reward

(b) Learned reward model

(c) Pre-trained model

(d) Preferred samples

(e) RLHF

(f) DPO

(g) PFM

(h) PFM (5 iter.)

Figure 2: Comparison of RLHF, DPO, and PFM on a 2-dimensional toy experiment.
We generate preference labels from a ground truth reward in (a) and a pre-
trained Gaussian reference policy (c). Both the RLHF (e) and DPO (f) methods
struggle to align with the prefernce, due to overfitted reward model (b), even
with the presence of KL regularizer ($\beta=1$). PFM is able to mimic the
distribution of the positively-labeled samples (d), and therefore achieves the
highest performance (g). Repeating PFM iteratively to the marginal samples can
further improve the alignment with the preference (h).

###  3.3 Improving Alignment with Iterative Flow Matching

As done in iterative variants of DPO (Xiong et al., 2023; Yuan et al., 2024),
we can also further improve the quality of alignment with iterative flow
matching. Specifically, upon obtaining a marginal distribution $p_{1}$ by
applying flow matching, we again collect a new preference data $y^{-},y^{+}$
from the obtained marginal distribution $p_{1}$ in (6). We repeat this process
iteratively, by replacing the source distribution (which is
$\pi_{\mathrm{ref}}$ in the first step) with the marginal distribution $p_{1}$
obtained in the latest iteration. This iterative process can be summarized as
follows.

$\displaystyle p_{0}^{(n)}(y^{-}|x)$ $\displaystyle\propto
p_{1}^{(n-1)}(y^{-}|x)\int p_{1}^{(n-1)}(y|x)\mathbb{P}(y>y^{-}|x)dy$ (9)
$\displaystyle p_{1}^{(n)}(y^{+}|x)$ $\displaystyle\propto
p_{1}^{(n-1)}(y^{+}|x)\int p_{1}^{(n-1)}(y|x)\mathbb{P}(y^{+}>y|x)dy,\quad
p_{1}^{(0)}=\pi_{\mathrm{ref}},$ (10)

where we denote $p_{0}^{(n)}$ and $p_{1}^{(n)}$ to be the source and target
distribution of the flow matching model at the $n$-th iteration, respectively.
By repeatedly marginalizing the distribution with respect to the preference
$\mathbb{P}(y^{+}>y^{-}|x)$, we can effectively "narrow" the sampling
distribution towards the outputs with higher preference probability. See
Figure 2(h) for the results of the iterative method in our toy experiment.
Note that even during this iterative approach, we leave the parameters of the
pre-trained model $\pi_{\mathrm{ref}}$ untouched, and only require sampling
from this model throughout the whole process. Later in Section 4, we formally
prove that the iterative method allows us to obtain a distribution class that
maximizes the ground truth expected preference, and hence yields an optimal
policy $\pi^{*}$ in (1) with $\beta=0$. See Theorem 4.2.

##  4 Theoretical Analysis of Preference Flow

In this section, we theoretically analyze why PFM framework can effectively
learn to align with the preference. Interestingly, learning to generate
samples from the marginal distribution $p_{1}$ in (6) optimizes an objective
that is similar to the goal of general RLHF in (1). Following the formulation
provided by Azar et al. (2023), one can observe that the objective (1) is
equivalent to the below form:

$\pi^{*}=\underset{\pi}{\mathrm{argmax}}\
\mathbb{E}_{y\sim\pi}\bigg{(}\mathbb{E}_{y^{\prime}\sim\pi_{\mathrm{ref}}}\Big{(}\sigma^{-1}(\mathbb{P}(y>y^{\prime}))\Big{)}\bigg{)}-\beta\mathbb{D}_{\mathrm{KL}}\bigg{(}\pi\Big{\|}\pi_{\mathrm{ref}}\bigg{)}$
(11)

where $\sigma^{-1}(\xi)=\log(\xi/(1-\xi))$ is the logit function, and we drop
the conditional dependence of $x$ for simplicity. Note DPO also optimizes for
the same objective as in (11).

Let us take a step back, and characterize the failure modes of the RLHF and
DPO frameworks, by figuring out when these methods overfits the reward
function. Consider a simple example where the preferences are _deterministic_
, i.e., $\mathbb{P}(y>y^{\prime})=1$, so that $y$ is always preferred to
$y^{\prime}$. If we plug it into (11), we see that the term diverges to
$\sigma^{-1}(\mathbb{P}(y>y^{\prime}))\to+\infty$. Therefore, the solution
$\pi^{*}$ of (11) ends up overfitting for the preference likelihood, resulting
in a weak or even null KL regularization, regardless of the size of $\beta$.

Although in the case where the preference is not deterministic, this
phenomenon can still be pronounced in the finite data regime (Azar et al.,
2023). Even if the true preference is strictly less than 1, we may have access
to only a few data samples to estimate $\mathbb{P}(y>y^{\prime})\to 1$. This
means that overfitting can be a critical issue in general, especially if the
action space $Y$ or the context space $X$ is extremely large, as in the case
of aligning large language models to human preferences.

In contrast, the PFM framework learns to generate a marginal distribution
$p_{1}$. One can show that this marginal is a solution for the optimization
problem that is similar to the objective (11).

######  Theorem 4.1 (Characterization of the Marginal).

Let $\mathcal{D}=\\{(x,y^{+},y^{-})\\}$ be a preference dataset collected from
a pre-trained reference policy $\pi_{\mathrm{ref}}$, where the preference
$y^{+}>y^{-}$ is labeled with probability $\mathbb{P}(y^{+}>y^{-}|x)$. Let
$p_{1}$ denote the marginal distribution of the positively-labeled samples
$y^{+}$. Then the marginal distribution $p_{1}$ obtained from the preference
model $\mathbb{P}(y>y^{\prime}|x)$ is an optimizer of the optimization problem

$p_{1}=\underset{\pi}{\mathrm{argmax}}\
\mathbb{E}_{y\sim\pi}\bigg{(}\log\mathbb{E}_{y^{\prime}\sim\pi_{\mathrm{ref}}}\Big{(}\mathbb{P}(y>y^{\prime})\Big{)}\bigg{)}-\mathbb{D}_{\mathrm{KL}}\bigg{(}\pi\Big{\|}\pi_{\mathrm{ref}}\bigg{)}.$
(12)

We defer the proof to the Appendix C. Similar to the RLHF and DPO objective
(11), the solution $p_{1}$ of (12) drives the original distribution
$\pi_{\mathrm{ref}}$ towards the points where the preference probability
$\mathbb{P}(y>y^{\prime})$ is increasing. However, unlike the RLHF or DPO
objectives, the objective (12) is bounded even in the deterministic case
$\mathbb{P}(y>y^{\prime})=1$, making it robust to reward overfitting.

Despite its robustness, one may notice that the objective (12) is less
flexible compared to the original objective, due to the fixed regularization
constant with $\beta=1$. Below, we show that if we apply the iterative
algorithm provided in Section 3.3, one can further reduce the KL
regularization strength and obtain an optimal policy $\pi^{*}$ in (11) with
$\beta\to 0$.

######  Theorem 4.2 (Convergence of Iterative Method).

Assume $\pi_{\mathrm{ref}}\in L^{2}$ and $\mathbb{P}(y>y^{\prime})\in L^{2}$.
Consider an iterative update of the marginal distribution $p_{1}$ in (10).
Then, the iteration converges to the uniform distribution of points $y$ where
the value
$\mathbb{E}_{y^{-}\sim\pi_{\mathrm{ref}}}\left(\mathbb{P}(y>y^{-})\right)$ is
the largest, i.e.,

$p_{1}^{(\infty)}\to U\left(\left\\{y:y\in\underset{y}{\mathrm{argmax}}\
\mathbb{E}_{y^{-}\sim\pi_{\mathrm{ref}}}\left(\mathbb{P}(y>y^{-})\right)\right\\}\right),$
(13)

where $U$ stands for uniform distribution, and we drop the conditional
dependence of $x$ for simplicity.

We defer the proof to the Appendix C. Intuitively, the proof follows from the
fact that the marginalization iteratively "narrows" down the distribution
towards the outputs with higher preference. We note here that the $L^{2}$
assumptions are generally valid in practical domains. See Appendix C.

##  5 Experimental Results

In this section, we conduct experiments to address the following questions:
_Q1. Can PFM align generated samples from the black-box model with preference
and achieve comparable results in practical tasks?_ _Q2. Is PFM beneficial
than methods optimizing for explicit/implicit reward model?_ and _Q3. Is PFM
beneficial than näive add-on methods, e.g., separately training generative
models to imitate preferred samples?_ To answer these questions, we validate
our method in two distinct domains: Conditional image generation task and
offline reinforcement learning tasks.

###  5.1 Conditional Image Generation

(a) Pretrained (0.8389)

(b) Rejected $y^{-}$ (0.3951)

(c) Preferred $y^{+}$ (0.9976)

(d) PFM (0.9841)

(e) RLHF (0.9171)

(f) DPO (0.8936)

(g) DPO $\beta<\\!\\!<1$ (0.5397)

(h) Iterative PFM (0.9996)

Figure 3: Comparison of RLHF, DPO, and PFM on a conditional MNIST image
generation task. Numbers represent the preference score. PFM (d) demonstrates
superior sample quality and preference alignment compared to RLHF (e) and DPO
(f), where DPO collapses with a small size of $\beta$ (g). The iterative PFM
with only two iterations (h) results in almost perfectly aligning with the
preferences.

We first evaluate PFM on a conditional image generation task using the MNIST
dataset (LeCun et al., 1998). Specifically, we utilize a pre-trained DCGAN
(Radford et al., 2015) generator as $\pi_{\mathrm{ref}}$ and collect sample
pairs from $\pi_{\mathrm{ref}}(\cdot|x)$ conditioned on the digit labels
$x\in\\{0,\cdots,9\\}$. To construct preference datasets, we assign
preferences to sample pairs according to the softmax probabilities of the
labels from a LeNet (LeCun et al., 1998). Then, we learn a PFM flow
$v_{\theta}$ that transports $y^{-}$ to $y^{+}$ given a condition $x$. More
experimental details are provided in the Appendix D.

Figure 3(a) illustrates the generated samples from $\pi_{\mathrm{ref}}$, and
the rejected and preferred images are depicted in Figure 3(b) and Figure 3(c),
respectively, where the values in parenthesis are the measured preference
score. As shown in Figure 3(d), PFM achieves higher preference alignment and
better sample quality than RLHF (Figure 3(e)) and DPO (Figure 3(f)) without
fine-tuning $\pi_{\mathrm{ref}}$. Moreover, PFM achieves nearly perfect
alignment with the preferences after only two iterations (Figure 3(h)),
demonstrating the effectiveness of iterative PFM.

###  5.2 Offline Reinforcement Learning

Next, we employ the D4RL (Fu et al., 2020) benchmark to assess the performance
of PFM in reinforcement learning tasks. Following the prior works on the PbRL
literature, we adopt trajectory-based preference alignment (Hejna et al.,
2023; Kim et al., 2023). We first randomly choose an starting state $s_{0}\sim
S$, and sample two trajectories $\tau^{+}$ and $\tau^{-}$ with fixed length
$\ell\geq 2$ from $\pi_{\mathrm{ref}}$:

$\displaystyle\tau^{+}$
$\displaystyle:=(a_{0},a_{1},\cdots,a_{\ell})\sim\pi_{\mathrm{ref}}(\cdot|s_{0})$
(14) $\displaystyle\tau^{-}$
$\displaystyle:=(a_{0}^{\prime},a_{1}^{\prime},\cdots,a_{\ell}^{\prime})\sim\pi_{\mathrm{ref}}(\cdot|s_{0}).$
(15)

Then, we obtain the preference $\tau^{+}>\tau^{-}$ given the starting state
$s_{0}$ using a scripted teacher approach that has also been widely adopted in
the PbRL settings (Lee et al., 2021; Kim et al., 2023), which prioritizes
trajectories with higher rewards based on the ground truth reward. For
inference at a given state $s_{t}$, we again sample an action trajectory
$\tau=(a_{t},\cdots,a_{t+\ell})$ from $\pi_{\mathrm{ref}}(\cdot|s_{t})$, and
apply flow matching to obtain a better action sequence.

The baseline methods for comparing the performance of PFM includes behavior
cloning (BC), which we adopt as our pre-trained reference model
$\pi_{\mathrm{ref}}$, and a DPO fine-tuned model from the BC model.
Additionally, we train a separate behavior cloning model to the collected
preferred samples $y^{+}\sim p_{1}$, aiming to replicate the marginal
distribution of the "good" trajectories. Further experimental details are
deferred to Appendix D.

Table 1: Normalized results on MuJoCo datasets. Mean and standard deviation from 5 seeds are reported. | Pretrained (BC) | DPO Fine-tuned | PFM (Ours) | Marginal BC
---|---|---|---|---
ant-random-v2 | 1.003 $\pm$ 0.065 | 1.000 $\pm$ 0.103 |  1.004 $\pm$ 0.174 | 0.726 $\pm$ 6.166
ant-medium-v2 | 3.466 $\pm$ 28.57 | 3.671 $\pm$ 18.52 | 3.742 $\pm$ 3.287 |  3.866 $\pm$ 2.084
ant-expert-v2 | 4.966 $\pm$ 32.01 |  5.350 $\pm$ 5.002 | 5.234 $\pm$ 3.578 | 3.850 $\pm$ 46.20
hopper-random-v2 | 0.083 $\pm$ 0.261 | 0.085 $\pm$ 0.256 |  0.230 $\pm$ 0.086 | 0.158 $\pm$ 4.589
hopper-medium-v2 | 1.699 $\pm$ 5.173 | 1.720 $\pm$ 4.039 |  1.892 $\pm$ 2.696 | 1.296 $\pm$ 1.740
hopper-expert-v2 | 3.601 $\pm$ 0.493 | 3.609 $\pm$ 0.951 |  3.615 $\pm$ 0.790 | 1.034 $\pm$ 0.106
halfcheetah-random-v2 | -0.001 $\pm$ 0.020 | 0.000 $\pm$ 0.036 |  0.000 $\pm$ 0.012 | -0.006 $\pm$ 0.083
halfcheetah-medium-v2 | 4.806 $\pm$ 3.501 | 4.927 $\pm$ 2.659 |  5.119 $\pm$ 3.453 | 4.536 $\pm$ 4.998
halfcheetah-expert-v2 | 11.02 $\pm$ 4.868 |  11.16 $\pm$ 2.999 | 10.90 $\pm$ 3.244 | 0.312 $\pm$ 9.799
walker2d-random-v2 | 0.069 $\pm$ 0.152 | 0.065 $\pm$ 0.114 | 0.083 $\pm$ 0.190 |  0.114 $\pm$ 0.558
walker2d-medium-v2 | 2.772 $\pm$ 26.36 |  3.401 $\pm$ 17.50 | 3.334 $\pm$ 22.94 | 2.999 $\pm$ 18.26
walker2d-expert-v2 |  4.988 $\pm$ 0.568 | 4.977 $\pm$ 0.406 | 4.976 $\pm$ 0.306 | 0.747 $\pm$ 0.789
Random Average | 0.289 | 0.287 | 0.329 | 0.248
Medium Average | 3.186 | 3.430 | 3.522 | 3.174
Expert Average | 6.143 | 6.273 | 6.181 | 1.486
D4RL Average | 3.206 | 3.330 | 3.344 | 1.636
Figure 4: Analysis of a sample episode of a DPO fine-tuned model on the MuJoCo
ant environment. DPO fine-tuned model often over-estimates the reward due to
reward overfitting (e.g., $t=196$). This can cause the policy to choose
problematic actions. Here, the implicit reward estimation is
$\hat{r}_{\theta}(s,a)=\beta\log(\pi_{\theta}(a|s)/\pi_{\mathrm{ref}}(a|s))$.

Table 1 presents the outcomes from evaluations conducted on 12 offline
datasets. Our findings indicate that PFM consistently demonstrates comparable
or even superior performance with lower variance across all baseline methods.
Notably, our method excels particularly in datasets generated from suboptimal
behavioral policies, achieving better performance. Furthermore, PFM manages to
match the performance on expert datasets even in the absence of a reward
model, underscoring its robustness and effectiveness. This demonstrates that
PFM can effectively align black-box models with preference through flow
matching, without the need to fine-tune the pre-trained model.

###  5.3 Is Learning a Flow Truly Beneficial?

In this remaining section, we focus on answering the remaining questions, _Q2_
and _Q3_. We first investigate why PFM can be advantageous over previous
methods with explicit/implicit reward modeling. As can be seen in Figure 4,
DPO, like typical RLHF approaches, is also prone to reward overfitting, and
may cause the agent to fail. This is because if the preference estimate is
close to 0 or 1, these methods may end up overoptimizing to the exploding
reward model (Ziegler et al., 2019; Gao et al., 2023). PFM, on the other hand,
is inherently robust to such over-estimation, as we adopt a completely
different optimization framework that does not require a reward proxy. (See
Theorem 4.1.)

On the other hand, we observe less performance gain in the expert datasets.
This is a possible failure mode of PFM, where the generated samples are
already near-optimal. In such regime, an arbitrary source
$y\sim\pi_{\mathrm{ref}}$ has a near zero probability of being sampled from
the true marginal $p_{0}$, suggesting that PFM with prior as
$\pi_{\mathrm{ref}}$ might suffer from a shifted source distribution. We
verify this experimentally on walker2d, where PFM struggles the most. By
adopting a true marginal $p_{0}$ as the source, PFM with prior $p_{0}$ can
achieve the highest performance among all baselines. This performance is
evident even on the expert dataset, matching our theoretical analysis. See
Appendix B.

Next, we compare PFM to an alternative approach that simply tries to
approximate the marginal distribution directly from the positive samples.
Intuitively, training a generative model from scratch that replicates the
marginal $p_{1}$ is as computationally costly as training the original
reference model. Experimentally, we observe that PFM achieves better
performance than training a behavior cloning model (Marginal BC) to replicate
the distribution of the preferred samples (Table 1). However, it is also worth
mentioning that the Marginal BC model does occasionally yield the best
results, suggesting the potential of using a marginal distribution for
preference alignment.

##  6 Related Works

Contrastive Preference Learning (CPL) (Hejna et al., 2023) is a class of
reward-free methods that utilizes contrastive learning techniques to align
model outputs with the preferences observed in the dataset. By leveraging
contrastive loss, CPL encourages the model to distinguish between preferred
and less preferred outcomes effectively. Flow-to-Better (FTB) (Zhang et al.,
2023) innovatively uses a diffusion model to transition from less preferred
data to more preferred data, similar to the flow-based approach in our work.
However, FTB mainly focuses on data augmentation, where they used the trained
diffusion model to generate more data samples for behavior cloning. Despite
their strengths, both works rely on the Bradley-Terry model to implicitly
learn the reward function.

Identity Preference Optimization (IPO) (Azar et al., 2023) builds upon the
foundation laid by DPO, extending the framework to accommodate a broader range
of preference models beyond the Bradley-Terry paradigm. In particular, they
focus on finding an objective that is bounded even in a deterministic
preference regime, by replacing the function $\sigma^{-1}$ in (11) with an
identity function $I(x)=x$. This effectively mitigates the reward overfitting
problem observed in DPO and standard RLHF methods.

Our method distinguishes itself from these approaches by not requiring the
Bradley-Terry assumption nor the fine-tuning of pre-trained models. This
eliminates the risk of reward overfitting associated with the Bradley-Terry
model and reduces the computational cost significantly. By avoiding the need
for fine-tuning, our method offers a more efficient and scalable solution for
integrating human preferences into reinforcement learning systems. This makes
our approach particularly suitable for scenarios where computational resources
are limited or where quick adaptation to human feedback is essential. The
comparison of these related works is summarized in Table 2.

Table 2: Comparison of our method to other works. | RLHF | DPO | IPO | CPL | FTB | PFM (Ours)
---|---|---|---|---|---|---
Reward Model Free | ✗ | ✓ | ✓ | ✓ | ✓ | ✓
No Reward Assumptions (e.g. BT) | ✗ | ✗ | ✓ | ✗ | ✗ | ✓
Applicable to Black-Box Models | ✗ | ✗ | ✗ | ✗ | ✓ | ✓

##  7 Conclusion and Limitations

In conclusion, this research introduces Preference Flow Matching (PFM), a
novel add-on approach that offers a practical, efficient, and scalable
solution for integrating human preferences. This research highlights the
potential of flow matching as a powerful tool for preference alignment and
opens new avenues for further exploration and development in the field of
RLHF. The ability to align black-box models with human preferences without
extensive model modifications marks a critical step forward, with broad
implications for the deployment and usability of AI systems in real-world
applications.

Our theoretical and empirical analyses demonstrate that PFM achieves alignment
performance comparable to standard RLHF methods while being more resilient to
preference overfitting. The iterative flow matching technique further enhances
alignment quality, by continually refining the preference alignment without
modifying the underlying pre-trained model parameters. This iterative
refinement, grounded in the flow matching framework, ensures that the sampling
distribution increasingly favors higher preference probabilities, leading to a
robust and optimal policy.

Despite these promising results, the current design of the PFM framework
entails several challenges and limitations. In this work, we mainly focused on
classical PbRL settings (Kim et al., 2023; Lee et al., 2021; Zhu et al.,
2024), where we adopt the scripted teacher as a preference model. Although
theoretically generalizable, our experiments did not cover human preference
scenarios, which we leave as a future work. Further experiments on more
complex tasks like general contextual image generation are also desirable as
future directions.

Another notable limitation is its applicability to the natural language
processing (NLP) domain. The PFM framework, as currently designed, cannot be
directly applied to NLP tasks because the source and target texts may have
different lengths. Future research should explore ways to adapt the PFM
framework for variable-length data, potentially through innovative alignment
techniques or alternative frameworks suited for text generation tasks.

## Acknowledgements

We thank Taehyeon Kim at KAIST for pointing out the strengths of our work and
providing motivations of this work. We would also like to appreciate Sihyeon
Kim and Yongjin Yang at KAIST for providing experimental details and
suggesting relevant settings for PbRL tasks. Finally, we thank Junghyun Lee at
KAIST for revising the details of the proofs of theorems.

## References

  * Akrour et al. (2011) Riad Akrour, Marc Schoenauer, and Michele Sebag.  Preference-based policy learning.  In _Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2011, Athens, Greece, September 5-9, 2011. Proceedings, Part I 11_ , pages 12–27. Springer, 2011.
  * Azar et al. (2023) Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos.  A general theoretical paradigm to understand learning from human preferences.  _arXiv preprint arXiv:2310.12036_ , 2023.
  * Bradley and Terry (1952) Ralph Allan Bradley and Milton E Terry.  Rank analysis of incomplete block designs: I. the method of paired comparisons.  _Biometrika_ , 39(3/4):324–345, 1952.
  * Christiano et al. (2017) Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei.  Deep reinforcement learning from human preferences.  _Advances in neural information processing systems_ , 30, 2017.
  * Fu et al. (2020) Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning, 2020.
  * Gao et al. (2023) Leo Gao, John Schulman, and Jacob Hilton.  Scaling laws for reward model overoptimization.  In _International Conference on Machine Learning_ , pages 10835–10866. PMLR, 2023.
  * Hejna et al. (2023) Joey Hejna, Rafael Rafailov, Harshit Sikchi, Chelsea Finn, Scott Niekum, W Bradley Knox, and Dorsa Sadigh.  Contrastive prefence learning: Learning from human feedback without rl.  _arXiv preprint arXiv:2310.13639_ , 2023.
  * Ibarz et al. (2018) Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei.  Reward learning from human preferences and demonstrations in atari.  _Advances in neural information processing systems_ , 31, 2018.
  * Kim et al. (2023) Changyeon Kim, Jongjin Park, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee.  Preference transformer: Modeling human preferences using transformers for rl.  _arXiv preprint arXiv:2303.00957_ , 2023.
  * Kutta (1901) Wilhelm Kutta.  _Beitrag zur näherungsweisen Integration totaler Differentialgleichungen_.  Teubner, 1901.
  * LeCun et al. (1998) Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Lee et al. (2021) Kimin Lee, Laura Smith, and Pieter Abbeel.  Pebble: Feedback-efficient interactive reinforcement learning via relabeling experience and unsupervised pre-training.  _arXiv preprint arXiv:2106.05091_ , 2021.
  * Levine et al. (2018) Sergey Levine, Peter Pastor, Alex Krizhevsky, Julian Ibarz, and Deirdre Quillen.  Learning hand-eye coordination for robotic grasping with deep learning and large-scale data collection.  _The International journal of robotics research_ , 37(4-5):421–436, 2018.
  * Lipman et al. (2022) Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le.  Flow matching for generative modeling.  _arXiv preprint arXiv:2210.02747_ , 2022.
  * OpenAI et al. (2024) OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shapiro, Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brockman, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gibson, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Gontijo-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hallacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey, Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin, Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, Łukasz Kaiser, Ali Kamali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Łukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lee, Jan Leike, Jade Leung, Daniel Levy, Chak Ming Li, Rachel Lim, Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning, Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reiichiro Nakano, Rajeev Nayak, Arvind Neelakantan, Richard Ngo, Hyeonwoo Noh, Long Ouyang, Cullen O’Keefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr H. Pong, Tolly Powell, Alethea Power, Boris Power, Elizabeth Proehl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross, Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Selsam, Kyla Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalie Staudacher, Felipe Petroski Such, Natalie Summers, Ilya Sutskever, Jie Tang, Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Preston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vijayvergiya, Chelsea Voss, Carroll Wainwright, Justin Jay Wang, Alvin Wang, Ben Wang, Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng, Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Workman, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph.  Gpt-4 technical report, 2024.
  * Ouyang et al. (2022) Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al.  Training language models to follow instructions with human feedback.  _Advances in neural information processing systems_ , 35:27730–27744, 2022.
  * Radford et al. (2015) Alec Radford, Luke Metz, and Soumith Chintala.  Unsupervised representation learning with deep convolutional generative adversarial networks.  _arXiv preprint arXiv:1511.06434_ , 2015.
  * Rafailov et al. (2024) Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn.  Direct preference optimization: Your language model is secretly a reward model.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Runge (1895) Carl Runge.  Über die numerische auflösung von differentialgleichungen.  _Mathematische Annalen_ , 46(2):167–178, 1895.
  * Shin et al. (2023) Daniel Shin, Anca D Dragan, and Daniel S Brown.  Benchmarks and algorithms for offline preference-based reward learning.  _arXiv preprint arXiv:2301.01392_ , 2023.
  * Tong et al. (2023) Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio.  Improving and generalizing flow-based generative models with minibatch optimal transport.  _arXiv preprint arXiv:2302.00482_ , 2023.
  * Wilson et al. (2012) Aaron Wilson, Alan Fern, and Prasad Tadepalli.  A bayesian approach for policy learning from trajectory preference queries.  _Advances in neural information processing systems_ , 25, 2012.
  * Xiong et al. (2023) Wei Xiong, Hanze Dong, Chenlu Ye, Ziqi Wang, Han Zhong, Heng Ji, Nan Jiang, and Tong Zhang.  Iterative preference learning from human feedback: Bridging theory and practice for rlhf under kl-constraint.  In _ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models_ , 2023.
  * Xu et al. (2024) Shusheng Xu, Wei Fu, Jiaxuan Gao, Wenjie Ye, Weilin Liu, Zhiyu Mei, Guangju Wang, Chao Yu, and Yi Wu.  Is dpo superior to ppo for llm alignment? a comprehensive study.  _arXiv preprint arXiv:2404.10719_ , 2024.
  * Yuan et al. (2024) Weizhe Yuan, Richard Yuanzhe Pang, Kyunghyun Cho, Sainbayar Sukhbaatar, Jing Xu, and Jason Weston.  Self-rewarding language models.  _arXiv preprint arXiv:2401.10020_ , 2024.
  * Zhang et al. (2023) Zhilong Zhang, Yihao Sun, Junyin Ye, Tian-Shuo Liu, Jiaji Zhang, and Yang Yu.  Flow to better: Offline preference-based reinforcement learning via preferred trajectory generation.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Zhu et al. (2024) Tianchen Zhu, Yue Qiu, Haoyi Zhou, and Jianxin Li.  Decoding global preferences: Temporal and cooperative dependency modeling in multi-agent preference-based reinforcement learning.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 38, pages 17202–17210, 2024.
  * Ziegler et al. (2019) Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving.  Fine-tuning language models from human preferences.  _arXiv preprint arXiv:1909.08593_ , 2019.

##  Appendix A Flow Matching Algorithm for Preference Alignment

Below, we outline our method in Algorithm 1. Although the vector field
$v_{\theta}$ is trained to transport the source distribution of less preferred
samples $p_{0}$ to the target distribution, we use $\pi_{\mathrm{ref}}$ for
inference instead of $p_{0}$ due to inaccessibility to the preference model.

Note that with the non-deterministic preference assumption, PFM is still
theoretically guaranteed to obtain $\phi_{1}(y)\sim p_{1}$, even if the sample
$y$ is obtained from $\pi_{\mathrm{ref}}$ instead of $p_{0}$. In particular,
we assume $\mathrm{supp}(p_{1})\supseteq\mathrm{supp}(p_{0})$, i.e., for any
sample $y^{+}$ obtained from the target distribution $p_{1}$, the probability
of $y\sim p_{0}$ is non-zero. Then, for any sample $y\sim\pi_{\mathrm{ref}}$
sampled from the reference model, we can guarantee that this sample has non-
zero probability of being sampled as less preferred, indicating that there is
a learned flow from this point to a better sample, with non-zero probability.
In the next section, we provide empirical evidence that further justifies the
use of $\pi_{\mathrm{ref}}$ instead of $p_{0}$ for inference.

1 # Training

2repeat

3 Sample $z=(x,y^{+}.y^{-})\sim\mathcal{D}$

4 $u_{t}(y|z)\leftarrow y^{+}-y^{-}$

5 $p_{t}(y|z)\leftarrow\mathcal{N}(y|ty^{+}+(1-t)y^{-},\sigma^{2})$

6 Sample $t\sim[0,1]$ and $y\sim p_{t}(\cdot|z)$

7 $\theta\leftarrow\theta-\eta\nabla\mathcal{L}(\theta)$ # from (8)

8until  _$v_{\theta}$ converges_;

9# Inference

10Sample $y\sim\pi_{\mathrm{ref}}(\cdot|x)$

11Solve ODE (3) with initial condition $\phi_{0}(y)=y$ using $v_{\theta}$

12return $\phi_{1}(y)$

Algorithm 1 PFM: Preference Flow Matching

##  Appendix B Empirical Evidence for Using Reference Policy Instead of the
True Marginal During Inference

We conduct experiments on the D4RL walker2d environment to compare the
performance of the flow matching method using $\pi_{\mathrm{ref}}$ and $p_{0}$
as source distribution, respectively.

Table 3: Normalized results on MuJoCo datasets. Mean and standard deviation from 5 seeds are reported. | Pretrained (BC) | PFM from $\pi_{\mathrm{ref}}$ | PFM from $p_{0}$ | Planning
---|---|---|---|---
walker2d-random-v2 | 0.069 $\pm$ 0.152 | 0.083 $\pm$ 0.190 |  0.095 $\pm$ 1.460 | 0.090 $\pm$ 0.190
walker2d-medium-v2 | 2.772 $\pm$ 26.35 | 3.334 $\pm$ 22.94 |  3.701 $\pm$ 2.607 | 3.503 $\pm$ 11.57
walker2d-expert-v2 | 4.988 $\pm$ 0.568 | 4.976 $\pm$ 0.306 |  5.014 $\pm$ 0.538 | 5.009 $\pm$ 0.308
Average | 2.610 | 2.798 | 2.937 | 2.867

To replicate the actual marginal distribution of less preferred samples
$p_{0}$, we sample two trajectories from a given state $s_{t}$, and use the
ground truth reward function to select the less preferred data (with lower
reward). Note that this use of the environment reward is restricted in
practice, as it assumes access to the preference model during inference. We
also compare another baseline model, namely a planning model, that searches
among the two generated trajectories
$\tau^{+},\tau^{-}\sim\pi_{\mathrm{ref}}(\cdot|s_{t})$ and chooses an action
sequence with higher environmental reward.

As can be seen in Table 3, the flow matching model with the source
distribution matched to the actual marginal distribution $p_{0}$ (as in the
training process) achieves the highest performance among all baselines,
including a flow matching model that uses $\pi_{\mathrm{ref}}$ instead of
$p_{0}$. However, we observe that the flow matching model with
$\pi_{\mathrm{ref}}$ as the source distribution yields comparable results to
the model using $p_{0}$ as source. It is also worth noting that the flow
matching model (with true $p_{0}$) achieves better performance even compared
to the planning method, that explicitly uses a trajectory with higher reward.
This suggests that using a flow matching for preference alignment can provide
better alignment than simply conducting an exhaustive search.

##  Appendix C Proof of Theorems

In this section, we provide proof to the theorems in the paper. Throughout
this section, we write $\mu$ as the reference policy instead of
$\pi_{\mathrm{ref}}$ for the ease of writing.

First, we prove Theorem 4.1, where we rewrite the theorem statement for
completeness.

######  Theorem C.1 (Characterization of the Marginal).

Suppose that a preference dataset $\mathcal{D}=\\{(x,y^{+},y^{-})\\}$ is
collected from a pre-trained reference policy $\mu$, i.e.,
$y^{+},y^{-}\sim\mu(\cdot|x)$, where the preference $y^{+}>y^{-}$ is labeled
with probability $\mathbb{P}(y^{+}>y^{-}|x)$. Let $p_{1}$ denote the marginal
distribution of the positively-labeled samples $y^{+}$. Then the marginal
distribution $p_{1}$ obtained from the preference model
$\mathbb{P}(y>y^{\prime}|x)$ is an optimizer of the optimization problem

$p_{1}=\underset{\pi}{\mathrm{argmax}}\
\mathbb{E}_{y\sim\pi}\bigg{(}\log\mathbb{E}_{y^{\prime}\sim\mu}\Big{(}\mathbb{P}(y>y^{\prime})\Big{)}\bigg{)}-\mathbb{D}_{\mathrm{KL}}\bigg{(}\pi\Big{\|}\mu\bigg{)}.$
(16)

###### Proof.

For simplicity, we drop the conditional dependence of $x$. Notice that the
general RLHF optimization objective (11) attains a unique optimal solution

$\pi^{*}(y)\propto\mu(y)\
\exp\left(\frac{1}{\beta}\mathbb{E}_{y^{\prime}\sim\mu}\Big{(}\sigma^{-1}(\mathbb{P}(y>y^{\prime}))\Big{)}\right)$
(17)

As it is a widely known result, we defer the proof to prior works, e.g., Azar
et al. [2023]. Intuitively, comparing the above solution with the target
distribution $p_{1}$ as rewritten in (7) allows us to formulate a similar
objective to (11) which $p_{1}$ optimizes. Formally, we begin with considering
the optimization objective in (16):

$\displaystyle\underset{\pi}{\mathrm{argmax}}$ $\displaystyle\
\mathbb{E}_{y\sim\pi}\bigg{(}\log\mathbb{E}_{y^{\prime}\sim\mu}\Big{(}\mathbb{P}(y>y^{\prime})\Big{)}\bigg{)}-\mathbb{D}_{\mathrm{KL}}(\pi\|\mu)$
(18) $\displaystyle=\underset{\pi}{\mathrm{argmax}}\
\mathbb{E}_{y\sim\pi}\left(\log\mathbb{E}_{y^{\prime}\sim\mu}\Big{(}\mathbb{P}(y>y^{\prime})\Big{)}-\log\frac{\pi(y)}{\mu(y)}\right)$
(19) $\displaystyle=\underset{\pi}{\mathrm{argmin}}\
\mathbb{E}_{y\sim\pi}\left(\log\frac{\pi(y)}{\mu(y)}-\log\mathbb{E}_{y^{\prime}\sim\mu}\Big{(}\mathbb{P}(y>y^{\prime})\Big{)}\right)$
(20) $\displaystyle=\underset{\pi}{\mathrm{argmin}}\
\left(\log\frac{\pi(y)}{\mu(y)\mathbb{E}_{y^{\prime}\sim\mu}\big{(}\mathbb{P}(y>y^{\prime})\big{)}}\right)$
(21) $\displaystyle=\underset{\pi}{\mathrm{argmin}}\
\left(\log\frac{\pi(y)}{\frac{1}{Z}\mu(y)\mathbb{E}_{y^{\prime}\sim\mu}\big{(}\mathbb{P}(y>y^{\prime})\big{)}}-\log
Z\right)$ (22)

where
$Z=\int\mu(y)\mathbb{E}_{y^{\prime}\sim\mu}\big{(}\mathbb{P}(y>y^{\prime})\big{)}dy$
is the normalization constant. Notice that the distribution in the denominator
is precisely the marginal distribution $p_{1}$ of our interest, which is
indeed a valid distribution. Hence, the last optimization objective can be
reorganized as follows.

$\underset{\pi}{\mathrm{argmin}}\ \mathbb{D}_{\mathrm{KL}}(\pi\|p_{1})-\log
Z.$ (23)

Since the $Z$ is a constant, the KL-divergence objective is minimized at 0 if
and only if the two distributions $\pi$ and $p_{1}$ are identical, by the
Gibbs’ inequality. It follows that $p_{1}$ is the optimal solution of the
optimization problemm (16), which completes the proof. ∎

Next, we aim to prove Theorem 4.2. Before stating and proving the theorem, we
note here that the $L^{2}$ assumptions of the pre-trained model
$\pi_{\mathrm{ref}}$ and the probability model $\mathbb{P}(y>y^{\prime})$ is
generally valid in practical domains. For instance, assuming
$\pi_{\mathrm{ref}}$ being trained on $L_{2}$ loss as in many practical
domains, $\pi_{\mathrm{ref}}\in L^{2}$ generally holds. Also, the preference
model $\mathbb{P}$ is usually provided in a finite support, or modeled as
Bradley-Terry model, which generally behaves well.

######  Theorem C.2 (Convergence of Iterative Method).

Assume that $\mu\in L^{2}$, and $\mathbb{P}(y>y^{\prime})\in L^{2}$. Consider
an iterative update of the marginal distribution $p_{1}$ as follows.

$p_{1}^{(n)}(y)\propto p_{1}^{(n-1)}(y)\int
p_{1}^{(n-1)}(y^{\prime})\mathbb{P}(y>y^{\prime})dy^{\prime},\quad
p_{1}^{(0)}=\mu,$ (24)

i.e., we iteratively compute the marginal distribution of more preferred
samples from the update marginal distribution. Then, the limiting marginal
distribution $p_{1}^{(\infty)}$ converges to the uniform distribution $U$ of
points $y$ where the value
$\mathbb{E}_{y^{\prime}\sim\mu}\left(\mathbb{P}(y>y^{\prime})\right)$ is the
largest, i.e.,

$p_{1}^{(\infty)}\to U\left(\underset{y}{\mathrm{argmax}}\
\mathbb{E}_{y^{\prime}\sim\mu}\left(\mathbb{P}(y>y^{\prime})\right)\right).$
(25)

###### Proof.

Let us denote by $\mathcal{T}$ the transformation applied at each step:

$\mathcal{T}[p](y)=\frac{1}{Z_{p}}p(y)\int
p(y^{\prime})\mathbb{P}(y>y^{\prime})dy^{\prime},$ (26)

where $Z_{p}=\int p(y)\int p(y^{\prime})\mathbb{P}(y>y^{\prime})dy^{\prime}dy$
is the normalization constant. Then, the update rule in (24) can be written
simply as $p_{1}^{(n)}=\mathcal{T}[p_{1}^{(n-1)}]$. We aim to show that this
iterative procedure converges to a distribution that places uniform weight on
the set of points $y$ where
$\mathbb{E}_{y^{\prime}\sim\mu}\left(\mathbb{P}(y>y^{-})\right)$ is maximized.
Given a probability distribution $p$, Let us define the function $f_{p}$ as
follows.

$f_{p}(y)\triangleq\mathbb{E}_{y^{\prime}\sim
p}\bigg{(}\mathbb{P}(y>y^{\prime})\bigg{)}=\int
p(y^{\prime})\mathbb{P}(y>y^{\prime})dy^{\prime}.$ (27)

Observe that $\mathcal{T}[p](y)$ increases $p(y)$ proportionally to
$f_{p}(y)$. Consequently, the regions where $f_{p}(y)$ is higher will see an
amplification in probability mass relative to regions where $f_{p}(y)$ is
lower. This amplification occurs iteratively, with higher $f_{p}(y)$ regions
increasingly dominating the distribution.

Formally, we claim that the fixed-point iterator $\mathcal{T}$ is compact.
Notice that the kernel $K(y,y^{\prime})=p(y^{\prime})\mathbb{P}(y>y^{\prime})$
is bounded, provided that $\mathbb{P}(y>y^{\prime})$ is bounded by 1, and that
$p$ is a probability density, which integrates to 1. By Schur’s test and the
properties of Hilbert-Schmidt operators, i.e.,

$\iint|K(y,y^{\prime})|^{2}dydy^{\prime}<\infty,$ (28)

$\mathcal{T}$ can be shown to be a compact operator, from the square
integrability assumptions.

Next, consider the behavior of $\mathcal{T}$ on any sequence of probability
densities $\\{p^{(n)}\\}_{n=0}^{\infty}$. By the properties of compact
operators in the space of continuous functions, any sequence $\\{p^{(n)}\\}$
has a convergent subsequence in the weak topology. Let $p^{*}$ be the limit of
any convergent subsequence. The uniform boundedness of $K(y,y^{\prime})$ and
the compactness of $\mathcal{T}$ suggest that $\mathcal{T}[p^{*}]=p^{*}$,
establishing that $p^{*}$ is a fixed point of $\mathcal{T}$. To determine the
nature of $p^{*}$, observe that

$\int p^{*}(y)\int p^{*}(y^{\prime})\mathbb{P}(y>y^{\prime})dy^{\prime}dy=1.$
(29)

Since $\mathbb{P}(y>y^{\prime})$ is maximized uniformly over $y$ when
$\mathbb{E}_{y^{\prime}\sim\mu}(\mathbb{P}(y>y^{\prime}))$ is maximized,
$p^{*}$ must concentrate its mass on the set where this expectation is
maximized. Therefore, $p^{*}$ converges to a uniform distribution over the set

$\arg\max_{y}\mathbb{E}_{y^{\prime}\sim\mu}(\mathbb{P}(y>y^{\prime})).$ (30)

Formally, recall that from Theorem 4.1, the updated distribution $p_{1}^{(n)}$
is a solution to the following optimization problem:

$p_{1}^{(n)}=\underset{\pi}{\mathrm{argmax}}\
\Psi(\pi):=\mathbb{E}_{y\sim\pi}\bigg{(}\log\mathbb{E}_{y^{\prime}\sim
p_{1}^{(n-1)}}\Big{(}\mathbb{P}(y>y^{\prime})\Big{)}\bigg{)}-\mathbb{D}_{\mathrm{KL}}\bigg{(}\pi\Big{\|}p_{1}^{(n-1)}\bigg{)}.$
(31)

Hence, we note that if any point $y$ not in this set were to have a positive
probability under $p^{*}$, then $\mathcal{T}[p^{*}]$ would not be equal to
$p^{*}$, due to the strict maximization condition, contradicting the fixed-
point property. Thus, $p_{1}^{(\infty)}$ converges to the uniform distribution
over the optimal set, as stated. We note here that because the space of
probability densities is closed under the topology induced by the function
space norm, $p^{*}$ should also a probability density, and ultimately the
unique minima. ∎

##  Appendix D Experimental Details

In this section, we describe our experimental details.

Resource All experiments were conducted on a single Nvidia Titan RTX GPU and a
single i9-10850K CPU core for each run.

###  D.1 Conditional Image Generation

Task We employ the MNIST [LeCun et al., 1998] dataset to evaluate PFM on a
conditional image generation task (http://yann.lecun.com/exdb/mnist/).

Preference Dataset We utilize a pre-trained DCGAN [Radford et al., 2015]
generator as $\pi_{\mathrm{ref}}$ and collect sample pairs from
$\pi_{\mathrm{ref}}$ conditioned on the digit labels as contexts. To construct
preference datasets, we assign preferences to sample pairs according to the
softmax probabilities of the labels from a LeNet [LeCun et al., 1998].

Baseline We consider the pre-trained DCGAN generator ($\pi_{\mathrm{ref}}$), a
RLHF fine-tuned model and a DPO fine-tuned model of $\pi_{\mathrm{ref}}$ as
baselines. All methods are trained until convergence, and we report the
normalized episodes returns with the standard deviation from 5 different
random seeds.

###  D.2 Offline Reinforcement Learning

Task To assess the performance of PFM in reinforcement learning tasks, we
employ the D4RL [Fu et al., 2020] benchmark from https://github.com/Farama-
Foundation/D4RL where the datasets and code are licensed under the Creative
Commons Attribution 4.0 License (CC BY) and the Apache 2.0 License,
respectively. We consider four different tasks (ant, hopper, halfcheetah,
walker2d) from Gym-Mujoco domain with three different levels of offline
dataset quality (random, medium, expert) for each task.

Preference dataset We first pre-train $\pi_{\mathrm{ref}}$ using behavior
cloning (BC) for each offline dataset. We then collect segment pairs of the
rollout trajectories from $\pi_{\mathrm{ref}}$, with each pair starting from
the same state as a context. To construct preference datasets, we utilize a
scripted teacher which prioritizes trajectories with higher rewards based on
the ground truth rewards provided by the environment. This approach has been
widely adopted in the PbRL settings [Kim et al., 2023, Lee et al., 2021, Zhu
et al., 2024]. The preference datasets consist of 1,000 pairs of preferred and
rejected segments and their context for each offline dataset, with the segment
length 10.

Baseline The baseline methods for comparing the performance of PFM includes
behavior cloning (BC), which we adopt as our pretrained reference model
$\pi_{\mathrm{ref}}$, and a DPO fine-tuned model from the BC model. For DPO
fine-tuned models, we search KL regularization coefficient $\beta$ from 0.01
to 100 and adopt the best one. Additionally, we train a separate behavior
cloning model to the collected preferred samples $y^{+}\sim p_{1}$, aiming to
replicate the marginal distribution of the "good" trajectories. All methods
are trained until convergence, and we report the normalized episodes returns
with the standard deviation from 5 different random seeds.

##  Appendix E Broader Impacts

As an add-on module, PFM can be seamlessly integrated into various real-world
AI applications, such as generative models and continuous control systems,
without the need to modify the underlying application models. This integration
enables the applications to deliver personalized results that align with
individual user preferences. However, since PFM utilizes preference data, it
raises potential privacy concerns similar to those found in typical PbRL
methods. These concerns can be mitigated by ensuring that access to user
preferences is granted only with explicit user consent.