How to Leverage Demonstration Data in Alignment for Large Language Model? A Self-Imitation Learning Perspective

We formulate the objective of imitation learning as minimizing the reverse KL-divergence between $\pi_{\boldsymbol{\theta}}$ and the demonstration distribution  $\pi_{\rm{data}}$ Kostrikov et al. (2019); Fu et al. (2018):

	$\displaystyle\min_{{\boldsymbol{\theta}}}\ell_{\rm{GSIL}}(\boldsymbol{\theta})$	$\displaystyle={\mathrm{KL}}\big{(}\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})\|\pi_{\rm{data}}(\mathbf{y}\mid\mathbf{x})\big{)}$
		$\displaystyle=\mathbb{E}_{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}\Big{[}\log\frac{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\rm{data}}(\mathbf{y}\mid\mathbf{x})}\Big{]},$		(6)

where GSIL finds the model parameters by minimizing the reverse KL divergence, instead of optimizing the forward KL divergence in SFT as shown in Equation (2). In theory, while minimizing these two divergences theoretically leads to the same optimal solution $\pi_{\boldsymbol{\theta}^{*}}$, achieving this in practice requires full data coverage and infinite model expressive ability that are rarely met. Consequently, in practical settings, minimizing either KL divergence results in learned policies that exhibit different properties, as discussed in Murphy (2012). Specifically, forward KL ${\mathrm{KL}}(\pi_{\rm{data}}\|\pi_{\boldsymbol{\theta}})$ promotes mass-covering behavior, whereas reverse KL ${\mathrm{KL}}(\pi_{\boldsymbol{\theta}}\|\pi_{\rm{data}})$ encourages mode-seeking behavior Tajwar et al. (2024); Nachum et al. (2016); Agarwal et al. (2019); Xiao et al. (2021) as shown in Figure 1. Thus, forward KL encourages all responses in datasets to have equal probability, leading to an overestimation of the long tail of the target distribution, whereas reverse KL sharpens the probability mass on certain high-quality regions. Alignment commits to generating a certain subset of high-quality responses, which is achieved more effectively by minimizing the reverse KL, a shown by the recent work Tajwar et al. (2024).

In Section 4, we empirically demonstrate the results of optimizing these two divergences in practice and show the superiority of optimizing reverse KL divergence, especially on reasoning-heave downstream tasks such as math problem-solving, code generation, and logical reasoning.

However, performing mode-seeking is generally more challenging than mass-covering. Directly optimizing Equation (6) hardly leverages demonstration data effectively, especially since the data policy $\pi_{\rm{data}}$ is always unknown. In the RL literature, these challenges have been addressed through adversarial training Ho and Ermon (2016); Fu et al. (2018). These methods involve learning a reward function from demonstrations using complex and unstable adversarial training, which can be difficult to implement and adapt for LLM alignment.

In this paper, we propose a straightforward alternative that leverages demonstration data without necessitating the learning of a reward function via adversarial training. We observe that optimizing the objective (6) with respect to $\pi_{\boldsymbol{\theta}}$ requires the log density ratio $\log\frac{\pi_{\rm{data}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}$ between the data distribution and the current optimization policy. To circumvent this chicken-and-egg problem, we reformulate the imitation learning objective in Equation (6) into the following surrogate objective:

	$\displaystyle\max_{{\boldsymbol{\theta}}}\mathbb{E}_{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}\big{[}\log\frac{\pi_{\rm{data}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}-\log\frac{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}\big{]}=$
	$\displaystyle\mathbb{E}_{\pi_{\boldsymbol{\theta}}(\mathbf{y}|\mathbf{x})}\big{[}r(\mathbf{x},\mathbf{y})\big{]}-{\rm{KL}}\big{(}\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})\|\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})\big{)},$		(7)

where $r(\mathbf{x},\mathbf{y})\triangleq\log\frac{\pi_{\rm{data}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}$ can be viewed as an auxiliary reward function. Equations (6) and (7) are equivalent by adding and subtracting the same $\mathbb{E}_{\pi_{\boldsymbol{\theta}}(\mathbf{y}|\mathbf{x})}[\log{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})}]$ and $\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})$ can be the initial reference policy $\pi_{\rm{ref}}$ or the optimization policy in the last iteration used to sample the data. Interestingly, we find that even when only demonstration data is available, this objective takes a form similar to that used in the RLHF objective (3). The primary difference lies in the reward being the estimated log density ratio, which is often not readily accessible in real-world applications. The optimization of this objective, involving the density ratio $r(\mathbf{x},\mathbf{y})$, is not straightforward. We will demonstrate how to efficiently optimize it by effectively utilizing offline human demonstration data.

Before delving into the problem (7), we first describe how to calculate the auxiliary reward function in terms of the density ratio. In the tabular setting, we can directly compute $\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})$ and $\pi_{\rm{data}}(\mathbf{y}\mid\mathbf{x})$. However, in a high-dimensional language domain, estimating the densities separately and then calculating their ratio hardly works well due to error accumulation. A simple alternative is to estimate the log ratio via learning a classifier (discriminator) $s^{*}$ with logistic regression.

	$\displaystyle\min_{s}\ell_{\rm{DRE}}(s)=$	$\displaystyle-\mathbb{E}_{\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})}[\log\sigma(s(\mathbf{x},\mathbf{y}))]$		(8)
		$\displaystyle-\mathbb{E}_{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})}[\log(1-\sigma(s(\mathbf{x},\mathbf{y}))],$

where we view data samples as arising from data distribution over binary labels, where $\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})$ and $\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})$ are the densities of the class-conditional distribution. Thus, the log density ratio are related to the optimal classifier probabilities via following Bayes’ rule Bickel et al. (2009):

	$\displaystyle\log$	$\displaystyle\Big{(}\frac{\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})}\Big{)}^{\beta}=\log\frac{P(c=0){P}(c=1|\mathbf{x},\mathbf{y})}{P(c=1){P}(c=0|\mathbf{x},\mathbf{y})}$
		$\displaystyle=\log\Big{(}\frac{1}{\alpha}\frac{\sigma\big{(}s^{*}(\mathbf{x},\mathbf{y})\big{)}}{1-\sigma(s^{*}\big{(}\mathbf{x},\mathbf{y})\big{)}}\Big{)},$		(9)

where $\sigma(x)=\frac{1}{1+e^{-x}}$ is the sigmoid function that converts predicted scores into probabilities, and $\frac{P(c=1)}{P(c=0)}$ is the constant ratio between the priors of two classes. For simplicity, we heuristically introduce a hyperparameter $\alpha=\frac{P(c=1)}{P(c=0)}$ to denote the prior weight. Later, in our experiments, we will find it helpful to consider the class prior weight to the alignment process. $0<{\beta}<1$ introduced here is a power scaling parameter to control the trade-off between bias and variance, which interpolates between the uniform importance weights and the default weights Grover et al. (2019). While Equation(9) uses logistic regression for density ratio estimation, we can similarly derive expressions under arbitrary binary discrimination losses, resulting in different GSIL variants, as shown in Section 3.4.

The objective (7) with the discriminator for density ratio estimation can efficiently utilize the demonstration data. However, policy learning with the RL-style objective (7) is as challenging as in RLHF, and the computational costs for both density ratio estimation and policy learning are significantly higher than those of standard SFT. This makes them difficult to implement and use on large-scale problems, such as fine-tuning language models. We propose a simpler alternative that directly optimizes the imitation learning objective without needing RL training or a discriminator. The key idea is to leverage a specific discriminator parameterization, enabling a direct extraction of optimal policy, without an RL loop. Specifically, the optimal policy in (7) has a closed form as shown in Rafailov et al. (2024b):

$\displaystyle\pi^{*}(\mathbf{y}\mid\mathbf{x})=\frac{1}{Z(\mathbf{x})}\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})\exp\big{(}r(\mathbf{x},\mathbf{y})\big{)},$

(10)

where $Z(\mathbf{x})=\sum_{\mathbf{y}}\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})\exp\left(r(\mathbf{x},\mathbf{y})\right)=\sum_{\mathbf{y}}\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})=1$, meaning that the optimal $\pi^{*}(\mathbf{y}|\mathbf{x})$ is forced to be self-normalized! This characteristic, determined by the reward definition in (7), is beneficial as it allows GSIL theoretically generalize to broader classes of loss functions beyond the pairwise BT preference model used in DPO Rafailov et al. (2024b) and SPIN Chen et al. (2024) (see Section 3.4). Combing Equations (9) and (10) with some simple algebra gives us:

	$\displaystyle\beta\log\frac{\pi^{*}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}=\log\left(\frac{1}{\alpha}\frac{\sigma(s^{*}(\mathbf{x},\mathbf{y}))}{1-\sigma(s^{*}(\mathbf{x},\mathbf{y}))}\right)$
	$\displaystyle\Rightarrow s^{*}(\mathbf{x},\mathbf{y})=\beta\log\frac{\pi^{*}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}+\log\alpha.$		(11)

With this reparameterization, we can express the density ratio estimation in terms of only the optimal policy $\pi^{*}$ and the sampling policy $\pi_{\boldsymbol{\theta}_{t}}$ as:

	$\displaystyle\ell^{*}_{\rm{DRE}}=-\mathbb{E}_{\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})}[\log\sigma(\beta\log\frac{\pi^{*}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}+\gamma)]-$
	$\displaystyle\mathbb{E}_{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})}[\log(1-\sigma(\beta\log\frac{\pi^{*}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}+\gamma))],$		(12)

where we define $\gamma=\log\alpha$ without loss of generality. Now, we have the probability of density ratio probability in terms of the optimal policy rather than the discriminator model, we can formulate the following maximum likelihood objective for a parameterized policy Rafailov et al. (2024b):

	$\displaystyle\ell_{\rm{GSIL}}(\boldsymbol{\theta})=-\mathbb{E}_{\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})}\big{[}\log\sigma(\beta\log\frac{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}+\gamma)\big{]}$
	$\displaystyle-\mathbb{E}_{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})}\big{[}\log(1-\sigma(\beta\log\frac{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}+\gamma))\big{]},$		(13)

where the gradient of this objective takes the form of the difference between two parts, one related to the demonstration data and the other related to data self-generated by the policy. Interestingly, this optimization also provides actionable and theoretical insights into a self-improvement pattern: we iteratively generate syntactic data from the model itself, improving the policy by contrasting these self-generated data with real demonstration data.

A central insight of this work is to frame imitation learning as a supervised binary classification between real demonstration data and self-generated syntactic data from the model itself. The discussion in Section 3.2 suggests that any binary losses for density ratio estimation can be used in GSIL. Existing density ratio estimation losses can be cast in the following general form Buja et al. (2005); Gneiting and Raftery (2007); Sugiyama et al. (2012):

	$\displaystyle\min_{s}\ell_{\rm{DRE}}(s)$	$\displaystyle=\mathbb{E}_{\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})}[\ell_{1}(s(\mathbf{x},\mathbf{y}))]$
		$\displaystyle+\mathbb{E}_{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})}[\ell_{-1}(s(\mathbf{x},\mathbf{y}))].$		(14)

Let $s^{*}$ be the optimal in the above estimation loss, following Section 3.2, the auxiliary reward in terms of density ratio can be written as follows:

$\displaystyle r(\mathbf{x},\mathbf{y})=\frac{1}{\beta}\log\left(\alpha\frac{\sigma\left(s^{*}(\mathbf{x},\mathbf{y})\right)}{1-\sigma\left(s^{*}(\mathbf{x},\mathbf{y})\right)}\right).$

(15)

As shown in (11), we can express the score function $s$ in terms of its corresponding optimal policy:

$\displaystyle s^{*}(\mathbf{x},\mathbf{y})=\beta\log\frac{\pi^{*}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}+\gamma,$

(16)

where $\gamma=\log\alpha$. With this alternative reparameterization, the general loss in (14) can be rewritten with respect to the parametrized policy $\pi_{\boldsymbol{\theta}}$ as:

	$\displaystyle\min_{\boldsymbol{\theta}}\ell_{\rm{GSIL}}(\boldsymbol{\theta})$	$\displaystyle=\mathbb{E}_{\pi_{\rm{data}}(\mathbf{y}|\mathbf{x})}[\ell_{1}(f_{\boldsymbol{\theta}}(\mathbf{x},\mathbf{y}))]$
		$\displaystyle+\mathbb{E}_{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}|\mathbf{x})}[\ell_{-1}(f_{\boldsymbol{\theta}}(\mathbf{x},\mathbf{y}))],$		(17)

where $f_{\boldsymbol{\theta}}(\mathbf{x},\mathbf{y})\triangleq\beta\log\frac{\pi_{\boldsymbol{\theta}}(\mathbf{y}\mid\mathbf{x})}{\pi_{\boldsymbol{\theta}_{t}}(\mathbf{y}\mid\mathbf{x})}+\gamma$. Table 1 summarizes some notable density estimation methods developed over decades, including classification losses such as Hinge Cortes and Vapnik (1995), Brier Gneiting and Raftery (2007), Exponential Freund and Schapire (1995), and mean matching losses such as KLIEP and LSIF Sugiyama et al. (2012), each loss mapping into an alignment algorithm in our GSIL framework. Intuitively, these losses increase the likelihood of responses in the demonstration while promoting a decrease in the likelihood of synthetically self-generated data.

In this paper, we further investigate this implication in SPIN with fine-tuning on the demonstration data. Figure 2 shows that the likelihood of true responses in the demonstration counter-intuitively continues to decrease, although it remains higher than the likelihood of the generated response in SPIN. An undesirable consequence of this behavior is that the learned policy may increase the likelihood of unknown out-of-distribution responses Tajwar et al. (2024), instead of maximizing the likelihood of the chosen response, which is important in many practical applications of large language models, e.g., reasoning, coding, and mathematical problem solving, as shown in Pal et al. (2024); Yuan et al. (2024a) and our experiments.

Data. We evaluate GSIL on two widely used datasets for alignment: the UltraFeedback binarized dataset Tunstall et al. (2023) and the Anthropic-HH dataset Bai et al. (2022). Note that while these datasets provide paired chosen and rejected response, we only utilize the chosen response to form the demonstration dataset. The details of datasets are in Appendix A.1.

Table 2 compares the performance of GSIL against fine-tuning methods with demonstration data on UltraFeedback. As shown in the table, all variants of GSIL achieve remarkable improvements over SFT, particularly notable on challenging benchmarks. While SPIN can also enhance performance over the SFT model, GSIL, despite its simplicity, achieves the best overall performance on all benchmarks. These consistent and significant improvements highlight the robustness and effectiveness of GSIL. Notably, GSIL with logistic loss outperforms SPIN by 3.2 points in GSM8K (Math) and by 4.9 points in HumanEval (Code). We hypothesize that these improvements over SPIN can be attributed to the non-decreasing likelihood of real demonstrations in GSIL; as the likelihood of real samples decreases, it results in suboptimal performance, especially in mathematical reasoning and coding tasks where the chosen responses are very likely ground-truth answers. While almost all losses in our GSIL offer a significant improvement over SFT and SPIN, Logistic loss and Brier loss perform best in challenging tasks, making them worth considering as the initial attempts in practice. In Figure 3, we detail the model performances on MT-Bench with regard to different types of question.

We also compare the performance of GSIL, SPIN, and SFT on MT-Bench, which evaluate the models’ versatile conversational abilities across a diverse set of queries based on GPT-4. From the table, we observe that GSIL significantly boosts the performance, which demonstrates the effectiveness of GSIL on the instruction-following task.

To further evaluate the effectiveness of GSIL on safety alignment, we use the Anthropic-HH dataset, which contains 170k dialogues between a human and an automated assistant. Each transcript ends with a pair of responses generated by a large (although unknown) language model, along with a preference label denoting the human-preferred response in terms of harmlessness and helpfulness. Again, we only use the chosen responses as the demonstration data to train our policy. Figure 4 shows the win rates computed by GPT-4 over the chosen responses in the test set. Remarkably, GSIL aligns better with human preferences than the base model, SFT, and SPIN, achieving win rates of approximately 60% against the chosen responses. Additionally, we provide examples generated by both SPIN and GSIL in Tables 5 and 6 in Appendix D.1. These examples indicate that GSIL  shows strong promise in terms of aligning language models and can ensure that the generated responses are not only high-quality but also safe and harmless.