Enhancing Your Alignment via RM-LLM Game (2024)

aims to fine-tune the LLM response policy $\pi_{\theta}({\bm{y}}|{\bm{x}})$ with a group of human preference data ${\mathcal{D}}_{\text{P}}=\{({\bm{x}},{\bm{y}}^{w},{\bm{y}}^{l})\}$, so that the LLM can generate more satisfying responses to improve the human-model interaction quality. In each preference triplet $({\bm{x}},{\bm{y}}^{w},{\bm{y}}^{l})$, ${\bm{y}}^{w}\succ{\bm{y}}^{l}$ means response ${\bm{y}}^{w}$ is more “preferred” than ${\bm{y}}^{l}$ with respect to input ${\bm{x}}$. To align the LLM, a reward model (RM)(Christiano etal., 2017; Ouyang etal., 2022) $r_{\phi}({\bm{x}},{\bm{y}})$ is commonly utilized to score the quality of the LLM generated samples. RM learns human preferences ${\mathcal{D}}_{\text{P}}$ with a ranking loss(Bradley and Terry, 1952) ${\mathcal{L}}_{\text{rank}}(r_{\phi};{\mathcal{D}}_{\text{P}}):=$

where $\sigma(\cdot)$ is the Sigmoid function.For a response pair $({\bm{y}},\tilde{{\bm{y}}})$, the reward difference $r_{\phi}({\bm{x}},{\bm{y}})-r_{\phi}({\bm{x}},\tilde{{\bm{y}}})$ provides a preference probability :

With equation2, training RM with the Bradley-Terry ranking loss can be explained as the log-likelihood maximization of $Q_{\phi}$:

With a learned RM $r_{\phi}({\bm{x}},{\bm{y}})$, human preference alignment methods(Ouyang etal., 2022; Rafailov etal., 2023; Liu etal., 2023c) target on maximizing the reward expectation of generated responses:

where $\pi_{\text{ref}}({\bm{y}}|{\bm{x}})$ is a reference language model. $\text{KL}[\pi_{\theta}({\bm{y}}|{\bm{x}})\|\pi_{\text{ref}}({\bm{y}}|{\bm{x}})]$ prevents $\pi_{\theta}({\bm{y}}|{\bm{x}})$ from the degeneration of repeating a single response with the highest reward score, which also preserves the generation diversity. Since response samples $y$ are discrete, it is challenging to directly back-propagate from reward $r_{\phi}({\bm{x}},{\bm{y}})$ to policy $\pi_{\theta}({\bm{y}}|{\bm{x}})$. The typical solution to equation4 is reinforcement learning from human feedback (RLHF)(Ouyang etal., 2022), via the proximal policy optimization (PPO) algorithms(Schulman etal., 2017).

However, PPO suffers from implementation complexity and training instability(Yuan etal., 2023; Sun etal., 2023). Recent studiestry to avoid online reinforcement learning with offline schemes. DPO(Rafailov etal., 2023) finds a connection between the reward model and LLM’s optimal solution, then replaces the reward model with the likelihood ratio of $\pi_{\theta}$ and $\pi_{\text{ref}}$, as ${\mathcal{L}}_{\text{DPO}}(\pi_{\theta}):=$

Analogously, other methods consider human feedback learning from the perspective of contrastive learning. For example, RRHF(Yuan etal., 2023) propose a ranking loss as ${\mathcal{L}}_{\text{RRHF}}(\pi_{\theta}):=$

where ${\bm{y}}^{\text{best}}$ is the corresponding response to ${\bm{x}}$ with the highest reward, and the preference data ${\mathcal{D}}$ can be built from human annotation ${\mathcal{D}}_{\text{P}}$ or RM ranking results. Besides, rejection sampling (RJS)(Touvron etal., 2023b) (also called RAFT(Dong etal., 2023) and best-of-N(Stiennon etal., 2020)) directly fine-tunes LLM on ${\bm{y}}^{\text{best}}$ to further simplify the alignment process, ${\mathcal{L}}_{\text{RJS}}(\pi_{\theta}):=$

where ${\bm{y}}^{\text{best}}=\operatorname*{arg\,max}_{1\leq s\leq S}\{r_{\phi}({\bm$ is the sampled response with the highest reward score. Azar etal. (2023) extend the alignment objective into a more general form by replacing RM $r_{\phi}$ with the human preference probability $P({\bm{y}}\succ\tilde{{\bm{y}}}|{\bm{x}})$:

where $\Psi(\cdot)$ is a non-decreasing real-value function. This general alignment objective is called $\Psi$PO.

are a classical group of unsupervised machine learning approaches that can fit complicated real-data distributions in an adversarial learning scheme(Goodfellow etal., 2014). GANs use a discriminator $D(\cdot)$ and a generator $G(\cdot)$ to play a min-max game. The generator tries to cheat the discriminator with real-looking generated samples, while the discriminator aims to distinguish the true data and the samples:

where ${\bm{z}}$ is a random vector from prior $P_{{\bm{z}}}({\bm{z}})$ to induce the generation sample distribution. The objective equation8 has been theoretically justified as the Jensen–Shannon (JS) divergence between distributions of real data and samples(Goodfellow etal., 2014). Arjovsky etal. (2017) replace the JS divergence with the Wasserstein distance(Villani, 2009) and propose the Wasserstein GAN (WGAN):

where $\|f\|_{\text{L}}\leq K$ requires $f(\cdot)$ to be a $K$-Lipschitz continuous function. Wasserstein GANs have been recognized with higher training stability than the original GANs(Arjovsky etal., 2017).

In policy optimization of reinforcement learning, inspired by GANs, Ho and Ermon (2016) propose generative adversarial imitation learning (GAIL):

where ${\bm{a}}$ is the corresponding action based on the state ${\bm{s}}$, $D$ is a discriminator distinguishing difference between the learning policy $\pi_{\theta}$ and an expert policy $\pi_{\text{E}}$, and $\text{H}(\pi_{\theta})$ is the entropy of $\pi_{\theta}$.

For RM optimization of APO, we fix LLM $\pi_{\theta}({\bm{y}}|{\bm{x}})$ and update $r_{\phi}({\bm{x}},{\bm{y}})$. Note that in equation12 $\text{KL}[\pi_{\theta}({\bm{y}}|{\bm{x}})\|\pi_{\text{ref}}({\bm{y}}|{\bm{x}})]$ has no relation with $r_{\phi}$, so we can simplify the objective for RM updates:

The equation13 indicates that the APO-RM should enlarge the reward gap between golden answers and generated responses to challenge $\pi_{\theta}({\bm{y}}|{\bm{x}})$ for better generation quality. Note that equation13 has a similar form as WGANs in equation9, which can be intuitively explained as the calculation of the Wasserstein distance between distributions $P_{\theta}$ and $P_{\text{gold}}$. However, equation13 is not rigorously a Wasserstein distance because $r_{\phi}({\bm{x}},{\bm{y}})$ does not satisfy the Lipschitz continuity as described in Arjovsky etal. (2017).

To practically implement APO-RM training, we first collect a set of user queries $\{{\bm{x}}_{m}\}\sim P_{\mathcal{D}}({\bm{x}})$, then annotate each ${\bm{x}}_{m}$ with a golden response ${\bm{y}}^{\text{gold}}_{m}$, ${\mathcal{D}}_{\text{gold}}=\{({\bm{x}}_{m},{\bm{y}}_{m}^{\text{gold}})\}_{m=1$. Each $({\bm{x}}_{m},{\bm{y}}^{\text{gold}})$ can be regarded as a sample drawn from $P_{\text{gold}}({\bm{x}},{\bm{y}})$. Meanwhile, we generate ${\bm{y}}^{s}_{m}\sim\pi_{\theta}({\bm{y}}|{\bm{x}}_{m})$, so that $({\bm{x}}_{m},{\bm{y}}^{s}_{m})$ is a sample from distribution $P_{\theta}({\bm{x}},{\bm{y}})=P_{\mathcal{D}}({\bm{x}})\pi_{\theta}({\bm{y}}|{$. Denote ${\mathcal{D}}_{\text{sample}}=\{({\bm{x}}_{m},{\bm{y}}^{s}_{m})\}_{m=1}^{M}$. Combining ${\bm{y}}^{\text{gold}}$ and ${\bm{y}}^{s}$, we obtain an APO sample set ${\mathcal{D}}_{\text{APO}}=\{({\bm{x}}_{m},{\bm{y}}_{m}^{\text{gold}},{\bm{y}}$. Then the APO-RM objective in equation13 can be calculated:

Note that equation14 also enlarges the reward difference between pairs of responses as the Bradley-Terry (BT) loss (equation1)does. Hence, for training stability, we empirically use the BT loss to optimize equation14 instead, ${\mathcal{L}}_{\text{rank}}(r_{\phi};{\mathcal{D}}_{\text{APO}}):=$

With a Lagrange multiplier $\beta_{2}>0$, we convert the KL constraint in equation13 to a regularizer:

Note that $\text{KL}[P\|Q_{\phi}]=\mathbb{E}_{P}[\log P-\log Q_{\phi}]=-\text{H}(P)-$, where $\text{H}(P)$ is the entropy of ground-truth human preference $P({\bm{y}}\succ\tilde{{\bm{y}}}|{\bm{x}})$ as a constant for $r_{\phi}$ updating. As introduced in equation2, with a preference set ${\mathcal{D}}_{\text{P}}=\{({\bm{x}}_{n},{\bm{y}}^{w}_{n},{\bm{y}}^{l}_{n})\}$ representing samples of $P({\bm{y}}\succ\tilde{{\bm{y}}}|{\bm{x}})$, we have $-\mathbb{E}_{P}[\log Q_{\phi}]={\mathcal{L}}_{\text{rank}}(r_{\phi};{\mathcal{$. Then, the overall loss ${\mathcal{L}}_{\text{APO-RM}}(r_{\phi})$ is equivalent to:

The above APO-RM loss involves two datasets ${\mathcal{D}}_{\text{APO}}$ and ${\mathcal{D}}_{\text{P}}$. Since the golden responses consume much larger annotation resources than pair-wised response comparison, ${\mathcal{D}}_{\text{APO}}$ practically has a significantly smaller size than ${\mathcal{D}}_{\text{P}}$.In experiments, we find the re-weighting parameter $\beta$ requires to be larger to avoid over-fitting on the relatively smaller APO sample set ${\mathcal{D}}_{\text{APO}}$.We conduct more detailed ablation studies in the experimental Section4.