Title: Preference Alignment with Flow Matching
Venue: NeurIPS 2024
Accepted: yes
Keywords: Preference Alignment

Flow Matching
Abstract: We present Preference Flow Matching (PFM), a new framework for preference alignment that streamlines the integration of preferences into an arbitrary class of pre-trained models. Existing alignment 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 preference alignment objectives. Experimental results indicate the practical effectiveness of our method, offering a new direction in aligning a pre-trained model to preference. Our code is available at https://github.com/jadehaus/preference-flow-matching.

Preference Alignment with Flow Matching

Abstract

We present Preference Flow Matching (PFM), a new framework for preference alignment that streamlines the integration of preferences into an arbitrary class of pre-trained models. Existing alignment 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 preference alignment objectives. Experimental results indicate the practical effectiveness of our method, offering a new direction in aligning a pre-trained model to preference. Our code is available at https://github.com/jadehaus/preference-flow-matching.

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

̂

Sampling

y

-

y

+

SFT Model

π

ref

(

⋅

|

x

)

Final Model

π θ

(

⋅

|

x

)

Fine-tune

Train

Reward Model

r

ϕ

(

x

,

y

)

For training

y

-

y

+

Sampling

v

θ

(

t

,

y

|

x

)

Flow Matching Model

v

θ

(

t

,

y

|

x

)

Train

SFT Model

π

ref

(

⋅

|

x

)

y

-

y

+

Sampling

Flow Matching Model

v

θ

(

t

,

y

|

x

)

Solve

SFT Model

π

ref

(

⋅

|

x

)

RLHF Methods

Flow Matching (Training)

Flow Matching (Inference)

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., 2024b] 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 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.

2 Preliminaries

2.1 Reinforcement Learning from Human Feedback (RLHF)

Reinforcement learning from human feedback generally begins with obtaining a pre-trained reference policy π ref that can generate samples y ∼ π ref ( ·| 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 π ref . We then collect a dataset of N preference pairs D = { ( x i , y + i , y -i ) } N i =1 , where each x i denotes the context, and each y + i , y -i ∼ π ref ( ·| 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 × Y → R ≥ 0 , where X and Y are context space and response space, respectively. The goal of general RLHF is to recover an

optimal policy π ∗ such that

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 P ( y + > y -| x ) = σ ( r ∗ ( x, y + ) -r ∗ ( x, y -)) , where σ is the logistic function. Under this model, the general RLHF framework learns the reward model r ϕ ≈ r ∗ by minimizing the negative log-likelihood:

Once the reward model r ϕ is trained, we then use it to optimize for (1) to obtain π θ ≈ π ∗ 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., 2024b, Azar et al., 2023]. Direct Policy Optimization (DPO) [Rafailov et al., 2024b] 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 ˆ r θ ( x, y ) = β log( π θ ( y | x ) /π 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 R d with densities y -∼ p 0 and y + ∼ p 1 , possibly unknown (but able to sample). The flow matching considers the task of fitting a mapping f : R d → R d that transforms p 0 to p 1 , that is, if y -∼ p 0 , then f ( y -) ∼ p 1 . Inspired as in the motivation for the diffusion models, one can define a smooth time-varying vector field u : [0 , 1] × R d → R d that defines an ordinary differential equation (ODE),

where we use the notation u ( t, y ) interchanably with u t ( y ) . Denote the solution of the above ODE by ϕ ( t, y ) (or ϕ t ( y ) ) with initial condition ϕ 0 ( y ) = y . In other words, ϕ 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 ϕ 1 ( y ) where y ∼ p 0 . The integration map ϕ t induces a pushforward measure p t ≜ [ ϕ t ] ♯ ( p 0 ) , which is the density of points y ∼ p 0 transported via u from 0 to t .

To train the vector field v θ with a 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]:

where q ( z ) = π ( 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 an 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) preferences. In the first subsection, we illustrate our flow matching framework for learning preferences, 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 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:

In other words, we view that our collected data D is in fact generated from each of the marginal distributions y -∼ p 0 ( ·| x ) and y + ∼ p 1 ( ·| x ) obtained from 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 D as follows:

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 ) = N ( y | ty + +(1 -t ) y -, σ 2 ) . Once we obtain the vector field v θ , we can improve upon the generated negative samples y -∼ p 0 ( ·| x ) by solving (3) using an off-the-shelf numerical ODE solver [Runge, 1895, Kutta, 1901] to obtain samples f ( y -) ∼ p 1 . Specifically, we start from a sample y -with t = 0 , and "flow" along the ODE trajectory using v θ 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 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 notice 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 π ref . Instead, we simply start from y ∼ π ref as the starting point, and apply flow matching to obtain f ( y ) ≈ y + ∼ p 1 . We emphasize that PFM can still robustly generate positive samples, if we assume non-deterministic preferences i.e. , supp( p 1 ) ⊇ supp( p 0 ) . We also empirically find that using π 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 2a. We then pre-train a Gaussian mixture model to obtain samples as in Figure 2c. 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 ̂ r ϕ via (2). As can be seen in Figure 2b, the learned reward model overfits in the unseen region, which causes the RLHF method to diverge (Figure 2e). DPO also fails to learn the correct preference, as can be seen in Figure 2f. 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 2g). 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 2d and Figure 2g.)

(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.)

1.75

200

1.00

Pretrained Model

1.00

Preferred Data

1.00

0.75

1.25

1.50

100

150

50

0.75

0.50

1.75

1.50

0.75

0.50

1.75

1.50

0.50

0.25

0

0.25

10

1.25

0.25

1.25

1.0

0.00

1.0

50

0.25

0.00

0.75

Mad

0.25

0.00

0.75

1.0

0.5

0.0

0.5

0.0

0.5

人

1.0

0.5

0.0

0.5

0.0

0.5

y

0.50

0.75

0.50

0.25

0.50

0.75

0.50

0.25

X

0.5

1.0

1.0

0.5

1.0

1.0

1.00

1.00 0.75 0.50 0.25 0.00 0.25 0.50 0.75 1.00

1.00

1.00 0.75 0.50 0.25 0.00 0.25 0.50 0.75 1.00

1.00

RLHF Fine-tuned Model

1.0

DPO Trained Model

1.00

Learned Vector Field

1.00

Learmed Vector Field

0.25

0.00

0.75

0.50

0.8

0.6

0.75

0.00

0.50

0.25

1.50

1.25

1.75

0.75

0.50

0.25

0.00

1.75

1.50

1.25

0.75

0.50

0.25

000

1.50

1.25

1.00

1.75

pae

0.25

0.25

0.75

0.25

0.75

0.25

0.75

0.50

0.2

0S'0-

0.50

0.25

0.50

0.50

0.25

0.50

0.50

0.75

0.0

0.75

0.00

0.75

0.75

0.25

1.00 +

1.00 0.75 0.50 0.25 0.00 0.25 0.50 0.75

1.00

1.00

1.00 0.75 0.50 0.25 0.00 0.25 0.50 0.75 1.00

1.00

050050520000520050-S000

1.00

0050050520000520050-500-00T

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 π ref in the first step) with the marginal distribution p 1 obtained in the latest iteration. This iterative process can be summarized as follows.

where we denote p ( n ) 0 and p ( n ) 1 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 P ( y + > y -| x ) , we can effectively "narrow" the sampling distribution towards the outputs with higher preference probability. See Figure 2h 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 π 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 π ∗ in (1) with β = 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:

where σ -1 ( ξ ) = log( ξ/ (1 -ξ )) 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 overfit the reward. Consider a case where the preferences are deterministic , i.e. , P ( y > y ′ ) = 1 , so that y is always preferred to y ′ . If we plug it into (11), we see that σ -1 ( P ( y > y ′ )) → + ∞ . Therefore, the solution π ∗ of (11) ends up overfitting for the preference likelihood, resulting in a weak or even null KL regularization, regardless of the size of β .

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 P ( y > y ′ ) ≈ 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 p 1 denote the marginal distribution of the positively-labeled samples y + . Then the marginal distribution p 1 obtained from the preference model P ( y > y ′ | x ) is an optimizer of the optimization problem

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 π ref towards the points where the preference probability P ( y > y ′ ) is increasing. However, unlike the RLHF or DPO objectives, the objective (12) is bounded even in the deterministic case P ( y > y ′ ) = 1 , making it robust to reward overfitting.

Interestingly, maximizing the objective (12) is equivalent to minimizing the KL distance between the policy π and the normalized preference score with a cross-entropy constraint:

Hence, the objective pushes the policy π to match the preference ˜ P , while encouraging π to align with the high-probability samples in π ref . Since the constraint restricts the policy to the high-probability regions of π ref where the preference labels are collected from, our method is less prone to reward overfitting in the out-of-distribution samples due to the distribution shift in π from π ref [Gao et al., 2023, Rafailov et al., 2024a]. Though we find this information-theoretic formulation interesting, we leave further analysis to future work.

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 β = 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 π ∗ in (11) with β → 0 .

Theorem 4.2 ( Convergence of Iterative Method ) . Assume π ref ∈ L 2 and P ( y > y ′ ) ∈ 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 E y -∼ π ref ( P ( y > y -)) is the largest, i.e.,

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 more beneficial than methods optimizing for explicit/implicit reward model? and Q3. Is PFM more beneficial than naïve add-on methods, e.g., separately training generative models to imitate preferred samples? To answer these questions, we validate our method in three domains: Conditional text and image generation, and offline reinforcement learning tasks.

(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

β < <

1

(0.5397)

(h) Iterative PFM (

0.9996

)

3

345G

89

.1

345

6.7

4

0345

100

50

2

3

3

3

4

+

4

5

G

7

7

4

9

100

50

C

G

G

2

3

2.

5

6

6

100

50

C

山

2

3

3

3

Y

4.

3

6

7

7

7

g

99999999

100

50

000000

3

3

3

4

4

5

5

白

6

6

7

7

6

3

h

3

C

3

3

5

G

1

3

4

5

ta

7

150

3

G

？

了

150

4

5

150

Q

3

9

150

/

2.

3

y

b.

777

3

y

5

g

9

C

3

4

2

3

h

5

4

7

1

3

4.

5

200

4

5

7

g

200

2

3

h

5

7

2

200

2

3

4.

7

E

200

3

5

5

6

250

1

2

2.

3

3

51

4.

5

多

250

2.

3

3

4.

6

5

7

E

250

1

2

3

3

5.

6

G

7

7

9

250

2345

2.

3

4.

5

47

7

9

50

100

150

200

250

50

100

150

200

250

V

50

100

150

200

250

50

100

150

200

250

D

0.3

.4.

7

71

9.3

4.5

G

9

3.

5

6

7

8

9

100

OS

C

2

2.

3

3

3

4.

4.

5

6

7

7

g

9

100

50

Q

2

3

5

4.

4

7.

7

99977009

100

50

家

3

5

5

安

9

9

100

50

1

2

3

3

4.

4

5

S.

G

6

G

?

7

7

9

9

3

4

5

7

3

G

2

3

y

Los

7

200

150

f

3.

3

3

Y.

h

4

5

6

7

999

150

200

00.

0

4

3

3

4.

9

h

7

7.

200

150

6

7

200

150

Q

Q

Q

2

2.

3

3

3

n

4

5

6

G

7

7

7

3

4

7

8

3

4.

6

7

4

7

3

4.

7

9

250

34

9

250

1

3.

5

G

250

3

K

5

250

2.

3

4.

5

6

7

B

9

50

100

150

200

250

50

100

150

200

250

0

50

100

150

200

250

50

100

150

200

250

5.1 Conditional Image Generation

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 π ref and collect sample pairs from π ref ( ·| x ) conditioned on the digit labels x ∈ { 0 , · · · , 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 θ that transports y -to y + given a condition x . More experimental details are provided in the Appendix D.

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

5.2 Conditional Text Generation

Next, we adopt a controlled (positive) sentiment review generation task. As done in Rafailov et al. [2024b], to perform a controlled evaluation, we adopt the pre-trained sentiment classifier as the preference annotator. We train a preference flow on randomly selected pairs of movie reviews y + , y -from the IMDB dataset [Maas et al., 2011]. For our PFM framework to be applied to variable-length inputs, we employ a T5-based autoencoder to work with fixed-sized embeddings. We adopt GPT-2 SFT model on the IMDB dataset as a reference model π ref . We also compare our method with a RLHF (PPO) fine-tuned policy π PPO . See Appendix D for detailed experimental settings.

Below, we report the average preference score (from the classifier annotator) of 100 randomly generated review samples for each method. As shown in Table 1, PFM is able to improve the preference score of any baseline model to which it is attached. We emphasize here that our method requires relatively smaller training cost compared to the standard RLHF frameworks, even in the iterative settings. See Table 6 in Appendix D for the number of parameters that require training for each framework in tackling this task. PFM requires training a much smaller number of parameters (around 1.2%) while still achieving better performance. We also note here that instead of iteratively training the PFM module as described in Section 3.3, simply applying the same learned preference flow iteratively to the improved samples achieves the best performance. (See Appendix D.)

(a) SFT vs PFM

(b) PPO vs PFM

Distribution ofPreferenceScores of Test Samples

DistributionofPreferenceScores of TestSamples

20.0

F 0°96

92.5

Iref

Irer + PFM

Iref + PFM ×5

80

40

75

IpPO + PFM

IpPO + PFM ×5

17.5

35

Frequency

30

12.5

25

10.0

20

7.5

15

5.0

10

2.5

5

0.0

0

+

Interestingly, we observe that the distribution of the scores tends to shift more toward that of the preferred samples with an increasing number of PFM iterations. (See Figure 4.) This result aligns with our theoretical insights: the PFM framework learns to shift the source distribution (i.e., the distribution of the reference policy) toward the marginal distribution of the more preferred samples.

We also compute the win rate with GPT-4 evaluation. The results are summarized in Table 2. Both PFM and PPO fine-tuned policy excel the reference policy with win rates 100%. (First column of Table 2.) Furthermore, we observe that the iterative PFM with 5 iterations on the reference model outperforms the PPO fine-tuned policy. If PFM is added on top of the PPO fine-tuned policy, we observe near 100% win rates for both PPO + PFM and PPO + Iterative PFM.

5.3 Offline Reinforcement Learning

Finally, 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 trajectorybased preference alignment [Hejna et al., 2023, Kim et al., 2023]. We first randomly choose an starting state s 0 ∼ S , and sample two trajectories τ + and τ -with fixed length ℓ ≥ 2 from π ref :

Then, we obtain the preference τ + > τ -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 τ = ( a t , · · · , a t + ℓ ) from π ref ( ·| s t ) , and apply flow matching to obtain a better action sequence.

The baseline methods for comparing the performance of PFM include behavior cloning (BC), which we adopt as our pre-trained reference model π 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 + ∼ p 1 , aiming to replicate the marginal distribution of the "good" trajectories. Further experimental details are deferred to Appendix D.

Table 3 presents the outcomes from evaluations conducted on 12 offline datasets. Our findings indicate that PFM consistently demonstrates comparable or even superior performance with lower

(c)

(a)

(b)

(d)

(a)

(b)

(c)

(d)

under-estimated reward, stable

correct reward estimation, stable

over-estimated reward, fail

over-estimated reward, fail

100

Estimated Reward fy(x, e)

True Reward r(x, a)

80

60

40

20

0

75

Tmestep

100

125

150

175

200

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.4 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 5, 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 ∼ π ref has a near zero probability of being sampled from the true marginal p 0 , suggesting that PFM with prior as π 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 3). 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 BradleyTerry paradigm. In particular, they focus on finding an objective that is bounded even in a deterministic preference regime, by replacing the function σ -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 4.

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.

Despite these promising results, the current design of the PFM framework entails several challenges in the natural language processing (NLP) domain. The PFM framework, as currently designed, relies on the autoencoder to work with fixed-sized embeddings to handle variable-length texts. To scale our method to more complex NLP tasks, future research should explore ways to adapt the PFM framework to long-form texts.

References

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.

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.

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.

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.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.

Leo Gao, John Schulman, and Jacob Hilton. Scaling laws for reward model overoptimization. In International Conference on Machine Learning , pages 10835-10866. PMLR, 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.

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.

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.

Wilhelm Kutta. Beitrag zur näherungsweisen Integration totaler Differentialgleichungen . Teubner, 1901.

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.

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.

Sergey Levine, Peter Pastor, Alex Krizhevsky, Julian Ibarz, and Deirdre Quillen. Learning handeye coordination for robotic grasping with deep learning and large-scale data collection. The International journal of robotics research , 37(4-5):421-436, 2018.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022.

Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies , pages 142-150, 2011.

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.

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:2773027744, 2022.

Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434 , 2015.

Rafael Rafailov, Yaswanth Chittepu, Ryan Park, Harshit Sikchi, Joey Hejna, Bradley Knox, Chelsea Finn, and Scott Niekum. Scaling laws for reward model overoptimization in direct alignment algorithms. arXiv preprint arXiv:2406.02900 , 2024a.

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, 2024b.

Carl Runge. Über die numerische auflösung von differentialgleichungen. Mathematische Annalen , 46(2):167-178, 1895.

Daniel Shin, Anca D Dragan, and Daniel S Brown. Benchmarks and algorithms for offline preferencebased reward learning. arXiv preprint arXiv:2301.01392 , 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.

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.

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 klconstraint. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models , 2023.

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.

Weizhe Yuan, Richard Yuanzhe Pang, Kyunghyun Cho, Sainbayar Sukhbaatar, Jing Xu, and Jason Weston. Self-rewarding language models. arXiv preprint arXiv:2401.10020 , 2024.

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.

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.

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.

A Flow Matching Algorithm for Preference Alignment

Below, we outline our method in Algorithm 1. Although the vector field v θ is trained to transport the source distribution of less preferred samples p 0 to the target distribution, we use π 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 ϕ 1 ( y ) ∼ p 1 , even if the sample y is obtained from π ref instead of p 0 . In particular, we assume supp( p 1 ) ⊇ supp( p 0 ) , i.e. , for any sample y + obtained from the target distribution p 1 , the probability of y ∼ p 0 is non-zero. Then, for any sample y ∼ π 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 π ref instead of p 0 for inference.

Algorithm 1: PFM: Preference Flow Matching

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 π ref and p 0 as source distribution, respectively.

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 τ + , τ -∼ π ref ( ·| s t ) and chooses an action sequence with higher environmental reward.

As can be seen in Table 5, 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 π ref instead of p 0 . However, we observe that the flow matching model with π 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.

C Proof of Theorems

In this section, we provide proof to the theorems in the paper. Throughout this section, we write µ as the reference policy instead of π 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 D = { ( x, y + , y -) } is collected from a pre-trained reference policy µ , i.e., y + , y -∼ µ ( ·| x ) , where the preference y + > y -is labeled with probability 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 P ( y > y ′ | x ) is an optimizer of the optimization problem

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

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 (17):

where Z = ∫ µ ( y ) E y ′ ∼ µ ( P ( y > y ′ ) ) 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.

Since the Z is a constant, the KL-divergence objective is minimized at 0 if and only if the two distributions π and p 1 are identical, by the Gibbs' inequality. It follows that p 1 is the optimal solution of the optimization problemm (17), 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 π ref and the probability model P ( y > y ′ ) is generally valid in practical domains. For instance, assuming π ref being trained on L 2 loss as in many practical domains, π ref ∈ L 2 generally holds. Also, the preference model 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 µ ∈ L 2 , and P ( y > y ′ ) ∈ L 2 . Consider an iterative update of the marginal distribution p 1 as follows.

i.e., we iteratively compute the marginal distribution of more preferred samples from the update marginal distribution. Then, the limiting marginal distribution p ( ∞ ) 1 converges to the uniform distribution U of points y where the value E y ′ ∼ µ ( P ( y > y ′ )) is the largest, i.e.,

Proof. Let us denote by T the transformation applied at each step:

where Z p = ∫ p ( y ) ∫ p ( y ′ ) P ( y > y ′ ) dy ′ dy is the normalization constant. Then, the update rule in (25) can be written simply as p ( n ) 1 = T [ p ( n -1) 1 ] . We aim to show that this iterative procedure converges to a distribution that places uniform weight on the set of points y where E y ′ ∼ µ ( P ( y > y -)) is maximized. Given a probability distribution p , Let us define the function f p as follows.

Observe that 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 T is compact. Notice that the kernel K ( y, y ′ ) = p ( y ′ ) P ( y > y ′ ) is bounded, provided that P ( y > y ′ ) 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. ,

T can be shown to be a compact operator, from the square integrability assumptions.

Next, consider the behavior of T on any sequence of probability densities { p ( n ) } ∞ n =0 . 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 ′ ) and the compactness of T suggest that T [ p ∗ ] = p ∗ , establishing that p ∗ is a fixed point of T . To determine the nature of p ∗ , observe that

Since P ( y > y ′ ) is maximized uniformly over y when E y ′ ∼ µ ( P ( y > y ′ )) 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

Formally, recall that from Theorem 4.1, the updated distribution p ( n ) 1 is a solution to the following optimization problem:

Hence, we note that if any point y not in this set were to have a positive probability under p ∗ , then T [ p ∗ ] would not be equal to p ∗ , due to the strict maximization condition, contradicting the fixed-point property. Thus, p ( ∞ ) 1 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.

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. The time required varies by task, but as it only involves collecting preference datasets and learning the flow, it completes within a few hours.

D.1 Conditional Image Generation

Task We employ the MNIST [LeCun et al., 1998] 2 dataset to evaluate PFM on a conditional image generation task.

Preference Dataset We utilize a pre-trained DCGAN [Radford et al., 2015] generator as π ref and collect sample pairs from π 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 ( π ref ), a RLHF fine-tuned model and a DPO fine-tuned model of π 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 Conditional Text Generation

Task To evaluate PFM on the NLP domain, we adopt a controlled (positive) sentiment review generation task with the IMDB dataset [Maas et al., 2011].

Preference dataset As done in Rafailov et al. [2024b], to perform a controlled evaluation, we adopt the pre-trained sentiment classifier 3 as the preference annotator. The preference dataset is constructed from randomly selected pairs of moview reviews y + , y -from the IMDB dataset, where the preference is obtained from the classifier logit probability p (positive | y + ) > p (positive | y -) . We then train our PFM model on this preference dataset, to obtain the marginal distribution of the preferred (positive sentiment) review p 1 ( y + ) .

Baseline We consider the GPT-2 SFT model on the IMDB dataset ( π ref ) and a RLHF fine-tuned model of π ref using PPO ( π PPO ) as baselines. We apply PFM to baseline models as an add-on module without fine-tuning. For the iterative PFM, we iteratively apply the initially obtained learned preference flow without collecting new data from the improved PFM model.

Embedding for variable-length input For our PFM framework to be applied to variable-length inputs, we employ a T5-based autoencoder 4 to obtain fixed-sized (1024-dimensional vector) embedding of the input texts, allowing us to work within the fixed-size latent space. Once the fixed-size embeddings z + and z -are obtained for each text sample y + and y -, we learn the conditional flow using PFM from z -to z + . During inference, we apply PFM to the latent embedding z of the given input text y , and decode the improved latent embedding using the T5 decoder. We adopt the same U-Net architecture used in our MNIST experiment, where we reshape the 1024-dimensional vector into a two-dimensional (32, 32) image tensor. As shown in Table 6, PFM requires a much smaller number of parameters (around 1.2%) than RLHF or DPO which involve fine-tuning LLMs.

Table 6: Parameters required for training for each method. PFM only requires 1.2% parameters to be trained compared to naive approaches (RLHF, DPO, etc.), and still achieves better performance in preference alignment.

Iterative PFM For our iterative PFM, we do not iteratively collect data and re-train the module. Instead, we simply iteratively apply the learned preference flow to the output samples. In particular, we apply the learned flow to the embedding z iteratively with the same flow module.

Prompt for GPT-4 Evaluation Below we include the prompts used to generate win rates.

reviews has a more positive sentiment?

D.3 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 π ref using behavior cloning (BC) for each offline dataset. We then collect segment pairs of the rollout trajectories from π 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 π ref , and a DPO fine-tuned model from the BC model. For DPO fine-tuned models, we search KL regularization coefficient β from 0.01 to 100 and adopt the best one. Additionally, we train a separate behavior cloning model to the collected preferred samples y + ∼ 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.

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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our claims made in abstract and introduction reflect the paper's contributions and scope. Theoretical and empirical results in Section 4 and Section 5 demonstrate our contributions. Furthermore, we address limitations in Section 7.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Please refer to Section 7.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Please refer to Section 4 and Appendix C.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide descriptions of our method in Section 3.1 and in Appendix A. Furthermore, we provide experimental details in Section 5 and Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We attach the code and related information in the supplementary material. Furthermore, we release the revised code at https://github.com/jadehaus/ preference-flow-matching .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide experimental settings and details in Section 5, Appendix D. We also release our code in supplementary materials and at https://github.com/ jadehaus/preference-flow-matching .

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Please refer to Section 5 and Appendix D. For RL tasks, we report mean episode returns with standard deviations for 5 different random seeds.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Please refer to Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: All authors read and followed the NeruIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: Please refer to Appendix E.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper poses no such risks since we do not release datasets or pre-trained models.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cited the original papers of the assets. URLs and licenses are mentioned in Appendix D.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide instructions to use our codes in README.md file which is included in the supplementary materials and at our repository https://github.com/ jadehaus/preference-flow-matching .

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel framework called Preference Flow Matching (PFM) for preference-based reinforcement learning (PbRL). The PFM approach aims to integrate human preferences into pre-trained models without the need for extensive fine-tuning, which is a common requirement in existing PbRL methods. The PFM framework utilizes flow matching techniques to learn directly from preference data, reducing the dependency on pre-trained model fine-tuning. It transforms less preferred data into more preferred outcomes by leveraging flow-based models, aligning model outputs with human preferences without explicit reward function estimation. The paper provides theoretical insights and empirical results that demonstrate the effectiveness of PFM in aligning pre-trained models with preferences and offers a new direction in PbRL. PFM is proposed as a solution to challenges such as scalability, inefficiency, and the need for model modifications, especially with black-box APIs like GPT-4.\n\nSTRENGTHS:\nInnovative Approach: PFM offers a fresh perspective on integrating human preferences into AI systems by using flow matching, which is a relatively under-explored method in preference alignment.\n\nTheoretical Foundation: The paper provides a solid theoretical basis for PFM, proving its alignment with standard PbRL objectives and robustness against overfitting in reward models.\n\nEmpirical Validation: Experimental results support the practical effectiveness of PFM, showing its ability to robustly align with preferences and achieve comparable performance to traditional RLHF methods and DPO.\n\nScalability and Efficiency: By eliminating the need for fine-tuning large pre-trained models, PFM addresses significant challenges related to scalability and computational efficiency.\n\nApplicability to Black-Box Models: PFM's ability to work as an add-on module for black-box models, like GPT-4, extends its applicability to scenarios where model access is limited.\n\nWEAKNESSES:\nExaggerated Motivation:  I believe the research motivation stated by the author in the abstract and introduction may be exaggerated. For example, the author claims that the method can be applied to black-box models like GPT-4, but there is no experimental validation. In fact, even for general language models, the author provides no experimental evidence to show that the proposed method is superior to DPO.\n\nLimited Domain Testing: The paper's experiments may not fully explore the potential of PFM across diverse domains, particularly in more complex tasks like general contextual generation. \n\nAssumption is too strong : The assumption is not general for RLHF scene. For example,  conditional flow is $y^+ - y^-$ (line 116), and the  probability path is subject to gaussian distribution,  which is not pratical for general AI model, such as GPT-4.\n\nPotential for Overfitting: Although PFM is designed to be robust against overfitting, the method might still be susceptible to overfitting in certain scenarios, particularly if the preference data is limited or biased.\n\nApplicability to NLP Tasks: The current design of PFM may not be directly applicable to natural language processing tasks due to challenges with variable-length data, which could be a significant limitation given the prevalence of NLP applications.\n\nQUESTIONS:\n1. I believe the assumptions in the paper are overly idealistic, as the method appears impractical for application to black-box models like GPT-4, and may even struggle with open-source models. If the authors could provide a specific plan on how to apply the method to open-source models (such as llama) and black-box models (such as GPT-4), including details of experiments and experimental results, I would consider raising the score. Otherwise, we suspect the research motivation in the paper has been exaggerated.\n\n2. How to apply the simple assumptions from lines 116-117 to LLM?\n\nLIMITATIONS:\nThe authors have made an effort to address the limitations and potential negative societal impacts of their work.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents Preference Flow Matching (PFM), which learns an ODE to transform less-preferred data into more-preferred ones. Due to PFM not explicitly or implicitly relying on reward functions, it avoids reward overfitting. Empirical evidence demonstrates that PFM outperforms DPO in image conditional generation and Offline RL tasks.\n\nSTRENGTHS:\n- The paper is well-written, with illustrative figures and toy examples that are intuitive and easy to understand.\n- The approach of using a flow model to transform less-preferred data into more preferred ones is straightforward and effective.\n\nWEAKNESSES:\n- The novelty of the paper appears somewhat limited, as PFM seems to use Conditional Flow Matching (CFM) to transform between less and more preferred data distributions. However, the paper does not thoroughly discuss the benefits of introducing CFM. Similar transformations can be achieved with Rectified Flow or Diffusion models, and work using Diffusion models (such as FTB) already exists.\n- The experimental validation in the paper is not sufficiently comprehensive, leading to less convincing results. For instance, in the context of Offline RL tasks, the following issues may exist:\n    - PbRL is commonly used to address tasks with difficult-to-define or sparse rewards. In D4RL, tasks like antmaze are more commonly used for better algorithm performance comparison.\n    - In MuJoCo tasks, PFM does not exhibit a significant improvement over DPO. Additionally, the standard deviation of the results appears to be quite large, making comparisons based solely on average means less reliable. Could this be due to limited random seeds? Typically, MuJoCo tasks have low standard deviations for having dense rewards.\n    - While the authors summarize many RLHF methods in the related works section, only DPO and vanilla RLHF are compared across the experiments, indicating a potential lack of baseline algorithms. For example, FTB uses a diffusion model to do similar tasks as PFM. Given the similarity between these two algorithms, why are they not compared in the study?\n\nQUESTIONS:\nPlease see Weaknesses.\n\nLIMITATIONS:\nPlease see Weaknesses.\n\nThe main areas for improvement in the paper are concentrated in the experimental section, as the current experimental results lack persuasiveness.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-28"}

comments: SUMMARY:
This paper proposes an add-on method named Preference Flow Matching (PFM) to achieve preference alignment without learning a reward model. PFM can transfer $y$ generated from the original model to the preferred $y^+$ through a few flow iterations. Experiments demonstrate its superior performance compared with RLHF and DPO.

STRENGTHS:
-This paper is well-written. I really enjoy reading it.
- While flow matching is not novel, this paper proposes to adapt it to the preference alignment area and yields many advantages (e.g., simplicity and efficiency).
- Experiments on both image generation and offline RL demonstrate its comparable performance.

WEAKNESSES:
- During inference, PFM with prior as $π_{\rm ref}$ suffers from a shifted source distribution, which can affect the performance.
- Limited testing benchmarks. RLHF is getting a lot of attention these days, largely due to its success in the language model area. However, this paper does not include the language-based task because PFM is limited in its ability to generate variable text.

QUESTIONS:
- During inference on the D4RL tasks, how do you choose actions to interact with the environments? As the paper claims, PFM samples an action trajectory $\tau$ from $π_{\rm ref} (·|s_t)$, and then applies flow matching to obtain a better action sequence $\tau_{\rm new}$. How can you sample a trajectory $\tau$ without training a dynamic model? Do you apply all actions in $\tau_{\rm new}$ to interact with the environment, or select some actions from $\tau_{\rm new}$ to take?

LIMITATIONS:
Limitations have been discussed in this paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a new framework for preference-based reinforcement learning, relying on Flow Matching. This framework utilizes flow-based models with optimal transport that interpolate between less preferred data and more preferred data, eliminating the need to estimate a reward function implicitly or explicitly. The authors provide a rigorous framework supported both by formal proofs and clear experimental evidence of the method's performance and demonstrate the effectiveness of iterative PFM, which iteratively narrows down the distribution towards preferred outputs.\n\nSTRENGTHS:\nThe paper notably excels along each of the dimensions: \n- originality: the paper presents a novel approach to preference alignment that hasn't been explored in the literature and shows a very promising potential for further exploration.\n- quality: the problem is well positioned, claims in the paper are clearly stated and the flow of argument circles back to support the claims throughout the paper. The authors do a good job of illustrating the nature of RLHF overfitting by starting with deterministic preferences, which are easy to illustrate and extend to non-deterministic cases. Then they prove that the proposed objective (12) for the marginal distribution $p1$ obtained from the preference model $\\mathbf{P}(y>y')$ is robust to overfitting even in the deterministic case.\n- clarity: the flow of argument is delivered with exceptional clarity. Notably, the authors provide both formal proofs of theorems and intuitive explanations, e.g. in lines 204-205 \" marginalization iteratively \"narrows\" down the distribution towards the outputs with higher preference\", which are then supported by the results in Table1, demonstrating performance improvements with lower variance across all baseline methods and supporting the claim in lines 36-37 \"By simply adding a preference flow matching module to black-box models, 37 PFM eliminates the need for fine-tuning the black-box model itself, providing a significant advantage.\"\n- significance: the contributions of the paper are significant yet the impact on more complex tasks is yet to be explored. In the current form, the method is not applicable to NLP tasks as LLMs, but flow matching per se is capable of working on multiple modalities and textual data. \n\n\n\n\nDemonstrate\n\nWEAKNESSES:\nThe paper is very strong, no notable weaknesses observed, suggestions for improvements are provided in the next section.\n\nQUESTIONS:\n- Consider an expansion of failure mode explanations for DPO from the paper https://arxiv.org/pdf/2404.10719, namely susceptibility to distribution shift, which can cause issues with generalizability in the regions without data support. They tie back to the author's critique of the failure modes of DPO, like overfitting.\n\n- line 210: \"Is PFM beneficial than methods optimizing for explicit/implicit reward model\" is this missing \"more beneficial\"? (same for Q3)\n\n- lines 126-129 highlight that the starting distribution $p_0$ is inaccessible during the inference step, and instead, they simply start from $y \\sim \\pi_{ref}$ as the starting point. This is an important implication for sampling/inference. Thus, a comment explaining why this works and whether it can be applied in any scenario under consideration is crucial. Consier is adding this to limitations if there are issues with applicability.\n\n- In the conclusion and limitations, the authors state limitations of applicability to NLP tasks: \" Future research should explore ways to adapt the PFM 314 framework for variable-length data, potentially through innovative alignment techniques or alternative 315 frameworks suited for text generation tasks.\" Diffusion Forcing: https://arxiv.org/abs/2407.01392 has been introduced recently to combine both autoregressive next token predictions with full-sequence diffusion. Consider this as a minor suggestion given the time constraints and timelines of the initial submission, the could be a useful source for future work.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations of their work throughout the paper and in the conclusions and limitations section.\n\nA minor suggestion:\n- The authors state the challenges of existing preference-based RL methods such as scalability, inefficiency, and the need for model modifications, especially with black-box, and later address them in the paper, except for scalability, which calls for stating more clear limitations or suggesting lines of future work to address this (extending the limitations in lines 309-310 would suffice).", "TITLE": "Official Review", "RECOMMENDATION": 9, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Preference Flow Matching (PFM), a novel framework for preference-based reinforcement learning that addresses key limitations of existing methods like Reinforcement Learning from Human Feedback (RLHF). PFM leverages flow matching techniques to directly learn from preference data. The core innovation of PFM lies in its formulation of preference alignment as a problem of learning a flow between marginal distributions of less preferred and more preferred samples, moving away from the implicit discriminative classification formulation of DPO into a discriminative/generative formulation. Methodologically, PFM employs a conditional flow matching objective to learn the vector field representing the preference flow. It also does not require the Bradley-Terry assumption. PFM is also compatible with 'reflow' like the original flow matching formulation. The authors derive the marginal distribution and prove it converges to a maximizing uniform. This all represents a clear, if straightforward, combination of (reward-free) preference learning and flow-matching providing a ready-made comparison with existing *-PO methods.\n\nSTRENGTHS:\nThe paper has multiple strengths. The combination of preference learning with flow matching is clear, if straightforward ; and its exposition is sound and easy to follow. The authors provide solid theoretical analysis, including proofs of optimality and convergence. This gives the method a strong mathematical grounding, which is often lacking in more heuristic approaches. In theory, PFM tackles important issues in existing methods, particularly the problem of reward overfitting and the challenges of aligning black-box models.\n\nWEAKNESSES:\n- Paper writing could be improved as several typos remain ('A careful reader might noticed' l.126, 'ordinary diffusional (sic) equation' l.89...). Some proofreading is in order !\n\n- Most importantly, the empirical contribution is not overly compelling. It is surprising that on MNIST the DPO results are that poor, both in absolute terms and relatively (worse than RLHF). IPO is not considered, DPO is probably taken offline rather than online (which generally tends to improve results), and most importantly, the beta parameter is not iterated upon/optimized (if it is, then this needs to be made clear, if possible in the main text, but it is not even in Appendix D.1). It is well known that the beta parameter is crucial for the performance of those methods and one advantage of very small scale experiments such as MNIST is that iterating is inexpensive. Tuning the baseline to the maximum would lend credence to these experiments; besides, it is done in the other suite of experiments.\n\n- Similarly, the offline RL experiments are small scale (Gym-MuJoCo. This time the beta parameter is optimized, and on all expert domains, the PFM performance is not massively different from DPO - in fact, the expert average for DPO fine-tuned is stronger than that of PFM (even though PFM provides variance reduction). While we know that DPO has tendencies to overfit and reward-hack, this is noted by the authors themselves who ask in section 5.3 'is learning a flow truly beneficial' ? Finding this detailed and balanced discussion in the paper was a nice display of intellectual fortitude and very welcome. Nonetheless, the empirical benefits of the proposed approach so far are quite ambiguous. The paper would truly benefit from showcasing one indisputable and conclusive empirical setting in favour of PFM.\n\n- Finally, the paper doesn't deeply explore the computational costs of PFM compared to other methods, which could be an important practical consideration. In particular, is reflow worth it ?\n\nQUESTIONS:\n- What stopped the authors (other than lack of time, infrastructure, compute) from trying for small-scale language modelling experiments, since this is the domain that originally stemmed the derivation of DPO and IPO ? In reference to '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', it would technically be possible to model distributions over fixed length sequences, along with padding and/or an EOS token. One great advantage of working in the text domain is the ability to display sample completions from finetuned models in a qualitative study. Also, testing with actual human preferences would actually be better.\n\n- Even if additional experiments are out of scope of this current review period, which other empirical designs can the authors think of that would better showcase the benefits of PFM compared to *-PO ?\n\nLIMITATIONS:\nLimitations are noted in the 'weaknesses' section below and mostly acknowledged in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-23"}