PAD: Personalized Alignment at Decoding-time

We define the standard text generation process of large language models (LLMs) with prompt $\mathbf{x}$ and response $\mathbf{y}$ as token-level Markov Decision Process (MDP). MDP is denoted as a tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},\mathcal{P},R,\rho)$, where the state space $\mathcal{S}$ consists of the prompt and all tokens generated so far (i.e., $\mathbf{s}_{t}=(\mathbf{x},\mathbf{y}_{1:t-1})$). The action space $\mathcal{A}$ is the tokens from the vocabulary (i.e., $\mathbf{a}_{t}=\mathbf{y}_{t}$). $\mathcal{P}$ is the transition kernel, which is deterministic that given state $\mathbf{s}_{t}=(\mathbf{x},\mathbf{y}_{1:t-1})$ and action $\mathbf{a}_{t}=\mathbf{y}_{t}$, the next state is $\mathbf{s}_{t+1}=(\mathbf{s}_{t},\mathbf{a}_{t})=(\mathbf{x},\mathbf{y}_{1:t})$. R : $S\times A\to\mathbb{R}$ represents the token-wise reward. $\rho_{0}$ is the initial state distribution of $\mathbf{s}_{0}$, i.e., $\mathbf{x}$. A (Markov) policy in MDPs $\pi:S\to\Delta(A)$ is a mapping from state to a distribution over actions.

Traditional RLHF approaches model the policy by a parametric distribution $\pi_{\theta}(\mathbf{a}|\mathbf{s})$ defined by a LLM over a finite set of actions. These approaches first learn a reward function from human feedback on prompt and response pairs $(\mathbf{x},\mathbf{y}^{w},\mathbf{y}^{l})$. The reward function is modeled as a contextual bandit using Bradley-Terry preference model:

$$p^{*}(\mathbf{y}^{w}\succeq\mathbf{y}^{l})=\frac{\exp R(\mathbf{x},\mathbf{y}^{w})}{\exp R(\mathbf{x},\mathbf{y}^{w})+\exp R(\mathbf{x},\mathbf{y}^{l})},$$

(1)

where $\mathbf{y}^{w}$ and $\mathbf{y}^{l}$ denote the preferred and not-preferred completions for the prompt $\mathbf{x}$. $p^{*}(\mathbf{y}^{w}\succeq\mathbf{y}^{l})$ denotes the probability that $\mathbf{y}^{w}$ is preferred to $\mathbf{y}^{l}$. Assuming access to a static dataset of comparisons $D=\{\mathbf{x}(i),\mathbf{y}(i)_{w},\mathbf{y}(i)_{l}\}^{N}_{i=1}$ sampled from $p^{*}$, we can parametrize a reward model $R(\mathbf{x},\mathbf{y})$ and estimate the parameters via maximum likelihood. Framing the problem as a binary classification we have the negative log-likelihood loss:

$$L_{R}(R,D)=-E_{(\mathbf{x},\mathbf{y}_{w},\mathbf{y}_{l})\sim D}[\text{log}\sigma(R(\mathbf{x},\mathbf{y}_{w})-R(\mathbf{x},\mathbf{y}_{l}))]$$

(2)

Then, the policy model (i.e., the LLM) $\pi_{\theta}$ is optimized with a gradient-based method like PPO using the following KL-constrained RL objective:

$$\max_{\pi_{\theta}}\mathbb{E}_{a_{t}\sim\pi_{\theta}(\cdot|\mathbf{s}_{t})}\left[\sum_{t=0}^{T}(R(\mathbf{s}_{t},\mathbf{a}_{t})-\beta\mathcal{D}_{KL}(\pi_{\theta}(\cdot|\mathbf{s}_{t}),\pi_{\text{ref}}(\cdot|\mathbf{s}_{t}))\right]$$

(3)

where $\pi_{\text{ref}}$ is a reference policy, often the language model resulting from supervised finetuning, from which the learned policy should not significantly deviate. In the optimization of the policy, only generating the EOS token carries a reward as output by the reward model which is combined with KL penalty, while for all other tokens in the vocabulary, only the KL component is non-zero. Additionally, it is worth mentioning that we omit the supervised finetuning (SFT) stage, which is often considered as part of the RLHF pipeline. This is due to that the SFT stage is not directly relevant to the focus of this paper.

Unlike classical RLHF, DPO, as derived in [rafailov2024direct], stays entirely within the contextual bandits setting entirely and also uses the bandit-based preference model. To circumvent the need for an RL algorithm, DPO uses the well-known closed form solution to the KL-contextual bandit version of the RL problem posed in Eq. 3 (Ziebart et al., 2008; Levine, 2018):

$$\pi^{*}(\mathbf{y}|\mathbf{x})=\frac{1}{Z(\mathbf{x})}\pi_{\text{ref}}(\mathbf{y}|\mathbf{x})e^{R(\mathbf{x},\mathbf{y})/\beta},$$

(4)

where $\pi^{*}$ is the optimal policy and $Z(x)$ is the partition function that normalizes it. DPO rearranges this equation to solve for reward as:

$$R(\mathbf{x},\mathbf{y})=\beta\text{log}\frac{\pi^{*}(\mathbf{y}|\mathbf{x})}{\pi_{\text{ref}}(\mathbf{y}|\mathbf{x})}-Z(\mathbf{x}).$$

(5)

Substituting this relationship into the standard binary cross-entropy loss function used for reward modeling (Eq. 2) yields the DPO loss equation as the partition function $Z(x)$ cancels from the Bradley Terry model:

$$L_{\text{DPO}}(\pi,D)=-\mathbb{E}_{(\mathbf{x},\mathbf{y}_{w},\mathbf{y}_{l})\sim D}[\text{log}\sigma(\text{log}\frac{\pi^{*}(\mathbf{y}_{w}|\mathbf{x})}{\pi_{\text{ref}}(\mathbf{y}_{w}|\mathbf{x})}-\text{log}\frac{\pi^{*}(\mathbf{y}_{l}|\mathbf{x})}{\pi_{\text{ref}}(\mathbf{y}_{l}|\mathbf{x})})].$$

(6)

We first rewrite Eq. 2 as entropy-regularized:

$$\max_{\pi_{\theta}}\mathbb{E}_{a_{t}\sim\pi_{\theta}(\cdot|s_{t})}\left[\sum_{t=0}^{T}\left(r(s_{t},a_{t})+\beta\log\pi_{\text{ref}}(a_{t}|s_{t})\right)+\beta\mathcal{H}(\pi_{\theta})\Big{|}s_{0}\sim\rho(s_{0})\right]$$

(7)

The relationship between future returns and the current timestep is captured by the Bellman Equation which are satisifed by any valid Q-function. We write this below for the optimal policy $\pi^{*}$ under the reward $R$ with a KL divergence penalty:

$$Q^{*}(\mathbf{s}_{t},\mathbf{a}_{t})=r(s_{t},a_{t})+\beta\log\pi_{\text{ref}}(a_{t}|s_{t})+V^{*}(\mathbf{s}_{t+1}),$$

(8)

where the optimal value function $V^{*}$ is a function of $Q^{*}$:

$$V^{*}(\mathbf{s}_{t})=\beta\text{log}\int_{A}e^{Q^{*}(\mathbf{s}_{t},\mathbf{a})/\beta}d\mathbf{a}.$$

(9)

Following [rafailov2024r], the main idea is first inverting the Bellman Equation to represent Q function:

$$\sum_{t=0}^{T-1}R(\mathbf{s}_{t},\mathbf{a}_{t})=V^{*}(\mathbf{s}_{0})+\sum_{t=0}^{T-1}\beta\text{log}\frac{\pi^{*}(\mathbf{a}_{t}|\mathbf{s}_{t})}{\pi_{\text{ref}}(\mathbf{a}_{t}|\mathbf{s}_{t})},$$

(10)

Then, the sum of rewards (Q function) in terms of the optimal policy can be directly substituted into the preference model in Eq. 1:

$$p_{\pi^{*}}(\mathbf{y}^{w}\succeq\mathbf{y}^{l})=\sigma(\sum_{t=0}^{N-1}\beta\text{log}\frac{\pi^{*}(\mathbf{a}_{t}^{W}|\mathbf{s}_{t}^{W})}{\pi_{\text{ref}}(\mathbf{a}_{t}^{W}|\mathbf{s}_{t}^{W})}-\sum_{t=0}^{M-1}\beta\text{log}\frac{\pi^{*}(\mathbf{a}_{t}^{l}|\mathbf{s}_{t}^{l})}{\pi_{\text{ref}}(\mathbf{a}_{t}^{l}|\mathbf{s}_{t}^{l})}).$$

(11)