Interpreting The Optimization in RL Fine-Tuning within RLHF

Optimization Objective

$${\mathbb{E}}_{x\sim D,y\sim {\pi }_{\theta }(y|x)}[r_{\varphi }(x,y)]-\beta {\mathbb{D}}_{\text{KL}}[{\pi }_{\theta }(y|x)\Vert {\pi }_{\text{ref}}(y|x)]$$

• $\mathcal{D}$ : The distribution of the given data. It ensures that new prompts encountered by the model are consistent with the data it was trained on.

• ${\pi }_{\theta }$ : The LM (or policy) currently undergoing reinforcement learning. It’s typically initialized from a Supervised Fine-Tuned (SFT) model, often denoted as ${\pi }^{\text{SFT}}$.

• ${\pi }_{\text{ref}}$ : The reference policy, typically the ${\pi }^{\text{SFT}}$ model itself. It serves as a baseline to prevent the fine-tuned model from drifting too far from its original capabilities.

• $x\sim \mathcal{D}$ : The prompt $x$ follows the given data distribution. This implies that new input prompts won’t be drastically different from what the model has already encountered.

• $y\sim {\pi }_{\theta }(y|x)$ : The response $y$ is generated by the currently reinforced LM, ${\pi }_{\theta }$.

• $r_{\varphi }(x,y)$ : The reward model function, which scores how good a response $y$ is for a given prompt $x$. This model is trained to align with human preferences.

• ${\mathbb{D}}_{\text{KL}}[{\pi }_{\theta }(y|x)\Vert {\pi }_{\text{ref}}(y|x)]$ : The Kullback-Leibler (KL) divergence term. It measures how much the responses generated by the current RL-trained LM (${\pi }_{\theta }$) deviate from the responses of the reference LM (${\pi }_{\text{ref}}$). This term is crucial because if ${\pi }_{\theta }$ focuses too heavily on learning preferences (i.e., overfitting), it might lose its general response generation capabilities (e.g., factual accuracy).

• $\beta$ : A hyperparameter that controls the strength of the penalty applied when ${\pi }_{\theta }$ diverges from ${\pi }_{\text{ref}}$. A higher $\beta$ means a stronger penalty, keeping ${\pi }_{\theta }$ closer to ${\pi }_{\text{ref}}$.

Interpretation

The overall goal is to improve the average preference score of generated responses. During this process, a constraint is imposed to prevent the model from sacrificing general response quality by becoming overly focused on merely mimicking preference patterns. This ensures that the fine-tuned LM produces responses that are both preferred by humans and maintain high quality.

(Rafailov et al., 2024, pp. 3–4)

Reference

Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., & Finn, C. (2024). Direct Preference Optimization: Your Language Model Is Secretly a Reward Model. arXiv. https://doi.org/10.48550/arXiv.2305.18290