# Prefernce Alignment with Flow Matching

Official Code base for the work [Prefernce Alignment with Flow Matching]() (PFM).

![](example/mnist/results/flow_matching.gif)

Given a pretrained (reference) model (or policy) $\pi_{\mathrm{ref}}$, and a pre-collected preference data $\mathcal{D} = \{(x_{i}, y_{i}^{+}, y_{i}^{-})\}_{i=1}^{N}$, where each $x_{i}$ denotes the context (or prompt), and each $y_{i}^{+}, y_{i}^{-} \sim \pi_{\mathrm{ref}}(\cdot | x_{i})$ are sampled from the reference policy, we learn the "*preference 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)$.

Our method serves as a simple add-on to any existing baseline model, and improves the quality of the model with respect to the provided preference.

|                                          | RLHF    | DPO     | $\Psi$-PO | CPL     | Ours    |
|------------------------------------------|---------|---------|-----------|---------|---------|
| Reward Model Free                        | ✘       | ✔       | ✔         | ✔       | ✔       |
| No Preference Assumptions                | ✘       | ✘       | ✔         | ✘       | ✔       |
| General Models and Tasks                 | ✔       | △       | △         | △       | ✔       |
| Preserves Pre-trained Model              | ✘       | ✘       | ✘         | ✘       | ✔       |


## Install
Download `PyTorch` of desired version according to the instructions in [the official PyTorch website](https://pytorch.org/). We recommend `PyTorch >= 2.2.2`.

Then, install relavant pacakges from `requirements.txt`. To train PFM and baselines, run `train.py`. Then, run `test.py` to evaluate the models.
