# Adversarial Moment-Matching Distillation of Large Language Models

Chen Jia
SI-TECH Information Technology
<EMAIL_ADDRESS>

###### Abstract

Knowledge distillation (KD) has been shown to be highly effective in guiding a
student model with a larger teacher model and achieving practical benefits in
improving the computational and memory efficiency for large language models
(LLMs). State-of-the-art KD methods for LLMs mostly rely on minimizing
explicit distribution distance between teacher and student probability
predictions. Instead of optimizing these mandatory behaviour cloning
objectives, we explore an imitation learning strategy for KD of LLMs. In
particular, we minimize the imitation gap by matching the action-value moments
of the teacher’s behavior from both on- and off-policy perspectives. To
achieve this action-value moment-matching goal, we propose an adversarial
training algorithm to jointly estimate the moment-matching distance and
optimize the student policy to minimize it. Results from both task-agnostic
instruction-following experiments and task-specific experiments demonstrate
the effectiveness of our method and achieve new state-of-the-art performance.

##  1 Introduction

Large language models (LLMs) like GPT-4 [1] and LLaMA [32] have revolutionized
natural language processing, significantly enhancing the quality of text
generation across various tasks. This success is largely due to the extensive
scale of training data and the substantial increase in model parameters [17].
However, the high computational and memory requirements of these models
present significant challenges for practical deployment. To address these
issues, knowledge distillation (KD) [15] has emerged as a key technique. KD
involves transferring knowledge from a large, complex teacher model to a
smaller, more efficient student model, thereby maintaining high performance
while reducing resource demands. Most distillation methods for auto-regressive
text generation models, including LLMs, employ metrics of probability
distribution distance, such as Kullback-Leibler (KL) divergence [18] and
reverse KL divergence [13], aiming to align the token-level probability
distributions between the teacher and student models.

The distribution matching-based distillation methods can be viewed as behavior
cloning on a decision-making problem from the perspective of imitation
learning [21; 13; 2]. Based on this concept, early works based on the teacher-
generated outputs [18] or a supervised dataset [27] can be viewed as an off-
policy approach. Recent works further incorporate an on-policy approach,
training the student on its self-generated outputs [21], using KL-based
divergence [13; 2; 19] and total variation (TV) distance [35]. Accordingly,
such distribution matching-based methods face the sub-optimality problem. The
objective functions aimed at aligning the probability distributions between
the teacher and student models can be straightforward but cannot fully capture
the goal of distilling language knowledge. First, intuitively, the correct
output for an input can vary, and thus behavior cloning cannot capture the
full knowledge of a teacher. Besides, there is no standardized definition for
the quality of a generated output given an input, which makes it difficult to
define the objective of knowledge distillation. This imposes a significant
limitation on the generalization performance of the student model through
distillation.

Figure 1: The comparison between the distribution-matching-based distillation
and the action-value moment-matching distillation is outlined. $\pi_{\theta}$
and $\pi_{*}$ denote the student policy and the teacher policy, respectively.
For both on-policy (using student-generated outputs) and off-policy (using
teacher-generated outputs) perspectives, our approach optimizes moment-
matching of action-value functions ($Q$-functions) instead of minimizing the
distribution distance measured by $\mathcal{M}$ = KL, RKL, TV, etc.

To address the aforementioned issues, we employ a reinforcement learning (RL)
formulation for the auto-regressive text generation problem and utilize the
definition of imitation gap to describe the high-level goal of knowledge
distillation. Additionally, we address the imitation gap for KD by matching
moments of the action-value function, which reflects the quality of token-
level predictions for the entire output. In addressing the action-value
function, we adopt the approach of Swamy et al. [30], considering a two-player
minimax game between the language policy and the action-value functions,
aiming to minimize an upper bound of the moment-matching objective. For this
purpose, we introduce an adversarial training algorithm based on the policy
gradient to jointly optimize the on-/off-policy objectives. Figure 1
illustrates the overall approach.

We evaluate our approach on both the instruction-following dataset and three
task-specific datasets for text summarization, machine translation, and
commonsense reasoning. Results demonstrate that the proposed adversarial
moment-matching approach effectively optimizes the moment-matching distance of
the imitation gap and outperforms state-of-the-art KD methods and a range of
distribution-matching-based methods. The code and implementation is released
at https://github.com/jiachenwestlake/MMKD.

##  2 Related Work

Distillation of large language models. There has been an increasing interest
in knowledge distillation (KD) of auto-regressive LMs, especially concerning
large language models (LLMs) [37; 38]. This process effectively transfers
elicited knowledge from teacher LLMs to smaller student models, aiming to
compress the large size of neural network parameters and make LLMs more
efficient. Sequence-level KD (SeqKD) [18] is a variation of supervised fine-
tuning (SFT) in KD. It can be viewed as the simplest method for distillation
of black-box LLMs by fine-tuning the student model with teacher-generated
outputs. This method has been extensively used for LLMs and has achieved
success [31; 5]. In contrast, distillation of white-box LLMs can make full use
of internal information of the teacher model, such as logits [27; 35] and
hidden states [20], for distribution alignment, making it more effective and
efficient for KD. However, unlike previous work that explicitly clones the
distribution of teacher LLMs into student models, this work learns an
auxiliary $Q$-value function to guide KD.

Distillation via distribution matching. Most promising results in the
distillation of white-box LLMs are achieved by minimizing divergence between
the probability distributions of the teacher model and student models.
Kullback-Leibler (KL) divergence, reverse Kullback-Leibler (RKL) divergence,
and Jensen–Shannon (JS) divergence are three widely used KD objectives for
auto-regressive LMs [35; 13; 2; 19; 37]. Wen et al. [35] have shown the
equivalent formulations of sequence-level KL, RKL, JS divergences, and the
step-wise terms. Additionally, they also present the strong performance of
step-wise total variation (TV) distance for KD, which can upper bound the
sequence-level term. As a result, most recent works focus on on-policy
approaches for KD [2] and combine the real-time-generated outputs by students
(on-policy) with the real-time-generated outputs by teachers (or from
supervised datasets) (off-policy). Following this line, Gu et al. [13] further
propose a policy gradient-based method to address the high variance issues of
RKL-based methods, while Ko et al. [19] propose a more efficient and effective
method using a skew KL divergence loss and an adaptive off-policy approach. We
also focus on a combination of on-policy and off-policy objectives for KD, but
we introduce a more sophisticated moment-matching approach instead of directly
using the well-studied distribution-matching metrics such as KL, RKL, JS
divergences, and TV distance.

Distillation via reinforcement learning. In a common formulation of RL in text
generation [40; 23; 14], an auto-regressive model can be viewed as a language
policy, making decisions on the next token (action) based on the currently
generated sequence (state). From this perspective, KD corresponds to behavior
cloning in imitation learning [18; 6; 13; 2]. For imitation learning in text
generation, early works such as SeqGAN [40] and TextGAIL [36] utilize a
generative adversarial framework to balance between the reward model,
optimized by discriminating generated/real-word text, and the language policy,
optimized by policy gradient-based methods using the reward model. Existing
work on KD via imitation learning refers to ImitKD [21], which optimizes the
student policy by learning from demonstrations of the teacher model. RL-based
distillation can also be especially relevant for leveraging the feedback from
the teacher to train student models [3; 7], in which teacher models are used
to generate the feedback data for training a reward model. We build our method
upon a RL-based imitation learning framework. However, unlike previous work
[18; 13; 2], we propose an adversarial moment-matching approach to enhance
behavior cloning.

##  3 Method

###  3.1 Notations and Definitions

In this section, we consider the text generation task as a decision-making
process and give a corresponding reinforcement learning (RL) formulation.

Text generation. Given an input $\boldsymbol{x}$, the auto-regressive
generation task in our work aims to generate a sequence of tokens as the
output $(y_{1},\ldots,y_{T})$, where $y_{t}$ comes from a vocabulary
$\mathcal{V}$. For simplicity, we define
$\boldsymbol{y}=(y_{0},y_{1},\ldots,y_{T})$ as the full input-output sequence,
where $y_{0}=\boldsymbol{x}$. The generator is modeled by a conditional
probability distribution
$p_{\theta}(\boldsymbol{y}|\boldsymbol{x})=\Pi_{t=1}^{T}p_{\theta}(y_{t}|\boldsymbol{y}_{<t})$,
where $\boldsymbol{y}_{<t}$ denotes the prefix $(y_{0},y_{1},\ldots,y_{t-1})$.

RL formulation. We formulate text generation as a sequential decision-making
process. At each time step $t\in\\{1,\ldots,T\\}$, the policy $\pi_{\theta}$
takes an action $y_{t}\in\mathcal{V}$ based on the current state
$\boldsymbol{y}_{<t}\in{\mathcal{Y}}$, transits to the next state
$\boldsymbol{y}_{<t+1}\in{\mathcal{Y}}$111In text generation, the state-
transition is commonly assumed to be deterministic [40; 23], i.e.,
$p(\boldsymbol{y}_{<t+1}|\boldsymbol{y}_{<t},y_{t})=1$. and receives a reward
$r(\boldsymbol{y}_{<t},y_{t})$ by a reward function
$r:{\mathcal{Y}}\times\mathcal{V}\rightarrow\mathbb{R}$. The policy
corresponds to the generation model
$\pi_{\theta}(y_{t}|\boldsymbol{y}_{<t})=p_{\theta}(y_{t}|\boldsymbol{y}_{<t})$.
A trajectory
$\tau\sim\pi_{\theta}=\\{\boldsymbol{y}_{<t},y_{t}\\}_{t=1,\ldots,T}$ refers
to a sequence of state-action pairs generated by first sampling a state
$y_{0}=\boldsymbol{x}$ from $p_{\boldsymbol{x}}$ and then repeatedly sampling
an action $y_{t}$ from $\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})$ and obtain
the next state $\boldsymbol{y}_{<t+1}$ for $T$ time steps. We also define our
value function and $Q$-value function as
$V^{\pi_{\theta}}(\boldsymbol{y}_{<t})=\mathbb{E}_{\tau\sim\pi_{\theta}|\boldsymbol{y}_{<t}}\sum_{t^{\prime}=t}^{T}r(\boldsymbol{y}_{<t^{\prime}},y_{t^{\prime}})$
and
$Q^{\pi_{\theta}}(\boldsymbol{y}_{<t},y_{t})=\mathbb{E}_{\tau\sim\pi_{\theta}|\boldsymbol{y}_{<t},y_{t}}\sum_{t^{\prime}=t}^{T}r(\boldsymbol{y}_{<t^{\prime}},y_{t^{\prime}})$.
We define the RL objective in our generation task to maximize the performance
$J(\pi_{\theta})=\mathbb{E}_{\tau\sim\pi_{\theta}}\left[\sum_{t=1}^{T}r(\boldsymbol{y}_{<t},y_{t})\right]$.

###  3.2 Knowledge Distillation as Moment-Matching Imitation Learning

Based on the RL formulation of auto-regressive generation, we can view the
goal of knowledge distillation at a high-level as to bridge the performance
gap between the teacher policy and the student policy.

######  Definition 1 (Imitation gap).

We define the imitation gap between the teacher policy and student policy as:

$\displaystyle
J(\pi_{*})-J(\pi_{\theta})=\mathbb{E}_{\tau\sim\pi_{*}}\left[\sum_{t=1}^{T}r(\boldsymbol{y}_{<t},y_{t})\right]-\mathbb{E}_{\tau^{\prime}\sim\pi_{\theta}}\left[\sum_{t=1}^{T}r(\boldsymbol{y}^{\prime}_{<t},y^{\prime}_{t})\right],$
(1)

where $\tau\sim\pi_{*}=\\{\boldsymbol{y}_{<t},y_{t}\\}_{t=1,\ldots,T}$ denotes
the trajectory of teacher policy and
$\tau^{\prime}\sim\pi_{\theta}=\\{\boldsymbol{y}^{\prime}_{<t},y^{\prime}_{t}\\}_{t=1,\ldots,T}$
denotes the trajectory of student policy.

From the perspective of imitation learning [30; 29], the objective of
distillation from the teacher policy $\pi_{*}$ to the student policy
$\pi_{\theta}$ can be represented as to minimize the imitation gap of Eq. (1)
w.r.t. the parameters of student policy $\theta$. An direct idea from Eq. (1)
is to use moment matching over the reward to optimize the imitation gap [30].
However, we actually care about the long-term reward, at each time step, we
should consider the accumulated reward in the future output rather than the
immediate reward to the fitness of previous tokens (prefix). To this end, we
can alternatively use the $Q$-value function (def. in Sec. 3.1) for each
timestep to represent the overall reward from the current timestep to the last
timestep. Similar to [30], we can apply the Performance Difference Lemma (PDL)
to expand the imitation gap in Eq. (1) into either off-policy or on-policy
expressions.

######  Proposition 1 (Off-policy bound of imitation gap).

Let $\mathcal{F}_{Q}$ denote the set of $Q$-value functions induced by
sampling actions from $\pi_{\theta}$, then we have:

$\displaystyle\begin{split}J(\pi_{*})-J(\pi_{\theta})\leq\sup_{f\in\mathcal{F}_{Q}}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\underbrace{\left[\sum_{t=1}^{T}\left(f(\boldsymbol{y}_{<t},y_{t})-\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\left[f(\boldsymbol{y}_{<t},y)\right]\right)\right]}_{\mathcal{U}^{\rm
off}(\tau,f,\theta):=}\end{split}$ (2)

In the following sections, we will use $\mathcal{U}^{\rm off}(\tau,f,\theta)$
to represent a sampled off-policy imitation gap with an trajectory
$\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}$ w.r.t. a $Q$-value
function $f$ and a teacher policy $\pi_{\theta}$.

###### Proof..

The derivation is mostly followed Swamy et al. [30]. See Appendix A.1 for the
complete derivation. ∎

The off-policy upper bound of the imitation gap in Proposition 1 only requires
a collected dataset of teacher-generated trajectories to be evaluated and
minimized.

######  Proposition 2 (On-policy bound of imitation gap).

Let $\mathcal{F}_{Q_{*}}$ denote the set of $Q$ functions induced by sampling
actions from $\pi_{*}$, then we have:

$\displaystyle\begin{split}J(\pi_{*})-J(\pi_{\theta})\leq\sup_{f\in\mathcal{F}_{Q_{*}}}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\underbrace{\left[\sum_{t=1}^{T}\left(\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{*}(\cdot|\boldsymbol{y}_{<t})}\left[f(\boldsymbol{y}_{<t},y)\right]-f(\boldsymbol{y}_{<t},y_{t})\right)\right]}_{\mathcal{U}^{\rm
on}(\tau,f):=}\end{split}$ (3)

In the following sections, we will use $\mathcal{U}^{\rm on}(\tau,f)$ to
represent a sampled on-policy imitation gap with an trajectory
$\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}$ w.r.t. a $Q$-value
function $f$ given the teacher policy $\pi_{*}$.

###### Proof..

The derivation is mostly followed Swamy et al. [30]. See Appendix A.2 for the
complete derivation. ∎

It is notable from Proposition 2 that the on-policy upper bound of the
imitation gap requires interactions with the teacher to tell us what action
they would take in any state visited by the student as well as on-policy
samples from the student’s current policy $\tau\sim\pi_{\theta}$.

In the remaining content of this section, we will exploring the relationship
between the proposed upper bounds of the imitation gap and the existing
distribution-matching objectives [35]. At the beginning, we draw a general
formulation of the state-of-the-art methods for distillation of LLMs [35; 13;
2; 19] that rely on distribution-matching between the student’s and teacher’s
predictions, through minimizing the step-wise probability distribution
distance between the teacher policy and student policy.

######  Definition 2 (Generalized step-wise distribution distance).

The off-policy and on-policy versions are defined as follows,

$\displaystyle d_{\rm\mathcal{M}}^{\rm off}(\pi_{\theta},\pi_{*}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\\!\\!\left[\sum_{t=1}^{T}{\mathcal{M}}(\pi_{*}(\cdot|\boldsymbol{y}_{<t}),\pi_{\theta}(\cdot|\boldsymbol{y}_{<t}))\right];$
(4) $\displaystyle d_{\mathcal{M}}^{\rm on}(\pi_{\theta},\pi_{*}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\\!\\!\left[\sum_{t=1}^{T}{\mathcal{M}}(\pi_{*}(\cdot|\boldsymbol{y}_{<t}),\pi_{\theta}(\cdot|\boldsymbol{y}_{<t}))\right],$
(5)

where $\mathcal{M}(\cdot,\cdot)$ denotes a distribution distance, consisting
of total variation (TV) distance [35] and Kullback-Leibler (KL)-based
divergence [13; 2]. Detailed definitions for these distances refer to Appendix
A.3.

It is notable from Wen et al. [35] that the sequence-level KL, RKL and JS
divergences can be equivalently represented as the step-wise terms, and the
sequence-level TV distance can be upper bounded by the step-wise terms, which
can be actually implemented by algorithms. To make a connection with the step-
wise distribution distance (Definition 2), we use the following definition.

######  Definition 3 (Distribution-matching formulation of moment-matching
bounds).

Based on Definition 2, we can re-formulate the on-policy and off-policy
moment-matching (MM) bounds (Proposition 2 and Proposition 1, respectively)
via step-wise distribution-matching, which can be defined as $d_{\mathcal{\rm
MM}}^{\rm on}(\pi_{\theta},\pi_{*})$ and $d_{\rm MM}^{\rm
off}(\pi_{\theta},\pi_{*})$ respectively, where the distance metric ${\rm
MM}(\cdot,\cdot)$ can be defined as follows,

$\displaystyle\begin{split}{\rm
MM}(\pi_{*}(\cdot|\boldsymbol{y}_{<t}),\pi_{\theta}(\cdot|\boldsymbol{y}_{<t}))&=\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{*}(\cdot|\boldsymbol{y}_{<t})}\left[f_{*}(\boldsymbol{y}_{<t},y)\right]-\mathop{\mathbb{missing}}{E}\limits_{y^{\prime}\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\left[f_{*}(\boldsymbol{y}_{<t},y^{\prime})\right],\\\
\text{\rm\bf Off-policy:}\
f_{*}&=\mathop{\arg\max}_{f\in\mathcal{F}_{Q_{*}}}{\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}}\left[\mathcal{U}^{\rm
off}(\tau,f,\theta)\right];\\\ \text{\rm\bf On-policy:}\
f_{*}&=\mathop{\arg\max}_{f\in\mathcal{F}_{Q}}{\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}}\left[\mathcal{U}^{\rm
on}(\tau,f)\right]\end{split}$ (6)

Under Definition 3, we observe that the main difference between the moment-
matching bounds and other step-wise distribution distance, e.g., TV distance
and KL-based divergences in formulation comes from the optimal $Q$-value
function $f_{*}$, aiming to maximize the discrepancy of its expectations based
on $\pi_{*}(\cdot|\boldsymbol{y}_{<t})$ v.s.
$\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})$ for each step
$t\in\\{1,2,\ldots,T\\}$. To look deeper, we draw a connection between moment-
matching bounds and the step-wise TV distance using the following corollary.

######  Corollary 1.

Under a constrain on the class of $Q$-value functions:
$\mathcal{F}_{Q_{\theta}}=\mathcal{F}_{Q_{*}}=\\{f:\|f\|_{\infty}\leq 1\\}$,
the moment-matching bounds in Proposition 1 and Proposition 2 can be upper-
bounded by the step-wise TV distance in Definition 2, Formally, for the off-
policy and on-policy perspectives, we have:

$\displaystyle J(\pi_{*})-J(\pi_{\theta})\leq$
$\displaystyle\sup_{f:\|f\|_{\infty}\leq
1}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
off}(\tau,f,\theta)\right]\leq d_{\rm TV}^{\rm off}(\pi_{\theta},\pi_{*});$
(7) $\displaystyle J(\pi_{*})-J(\pi_{\theta})\leq$
$\displaystyle\sup_{f:\|f\|_{\infty}\leq
1}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
on}(\tau,f)\right]\leq d_{\rm TV}^{\rm on}(\pi_{\theta},\pi_{*})$ (8)

###### Proof..

See Appendix A.4 for the complete derivation. ∎

We can observe from Corollary 1 that minimizing the step-wise TV distance can
achieve sub-optimal results for the moment-matching bounds. Optimizing the
moment-matching bounds can potentially achieve better optimization results for
imitation learning objectives.

###  3.3 Adversarial Training Algorithm

Optimization objective. As shown in previous work [13; 2; 19] that
incorporating both the on-policy and off-policy distillation benefits
effectiveness and efficiency. We thus consider a training objective to jointly
minimize the on-policy upper bound of Proposition (2) and the off-policy upper
bound of Proposition (1). Both the moment-matching objectives can be optimized
by viewing the learning procedure as solving a game. More specifically, we
consider a two-player minimax game between the student policy and the
$Q$-value functions. To this end, we consider using parameter estimation
methods for the $Q$-value functions and assume that both $\mathcal{F}_{Q}$ and
$\mathcal{F}_{Q_{*}}$ can be represented using a parameterized class of
functions $\\{f_{\phi}\\}_{\phi\in\Phi}$, wherein we use $\phi_{1}$ and
$\phi_{2}$ to denote parameters of the off-policy and on-policy $Q$-value
functions, respectively. Thereby, the learning problem can be represented as
follows,

$\displaystyle\min_{\theta\in\Theta}\max_{\phi_{1},\phi_{2}\in\Phi}\underbrace{\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\left[\mathop{\mathbb{missing}}{E}\limits_{\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
off}(\tau,f_{\phi_{1}},\theta)\right]+\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]\right]}_{\mathcal{L}(\theta,\phi_{1},\phi_{2}):=}$
(9)

We denote the learning objective by $\mathcal{L}(\theta,\phi_{1},\phi_{2})$.
To minimize the objective of $\mathcal{L}(\theta,\phi_{1},\phi_{2})$ w.r.t the
policy parameters $\theta$, we use a policy gradient-based approach and derive
the policy gradient of $\mathcal{L}(\theta)$ w.r.t $\theta$ in Appendix A.5,
represented as follows,

$\displaystyle\begin{split}\nabla_{\theta}\mathcal{L}(\theta,\phi_{1},\phi_{2})&=\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\bigg{[}-\mathop{\mathbb{missing}}{E}\limits_{\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{G}_{\rm
off}(\tau,\theta)\right]+\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}^{\prime}_{0}=\boldsymbol{x}}\left[\mathcal{G}_{\rm
on}(\tau^{\prime},\theta)\right]\bigg{]}\\\ {\rm s.t.}\quad\mathcal{G}_{\rm
off}(\tau,\theta)&=\sum_{t=1}^{T}\mathop{\mathbb{missing}}{E}\limits_{y^{\prime}\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\nabla_{\theta}\log\pi_{\theta}(y^{\prime}|\boldsymbol{y}_{<t})f_{\phi_{1}}(\boldsymbol{y}_{<t},y^{\prime});\\\
\mathcal{G}_{\rm
on}(\tau^{\prime},\theta)&=\sum_{t=1}^{T}\nabla_{\theta}\log\pi_{\theta}(y^{\prime}_{t}|\boldsymbol{y}^{\prime}_{<t})\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\end{split}$ (10)

Besides, to maximize the objective of $\mathcal{L}(\theta,\phi_{1},\phi_{2})$
w.r.t. parameters of the on-policy $Q$-value function $\phi_{1}$ and
parameters of the off-policy $Q$-value function $\phi_{2}$, we use a
stochastic gradient ascent method and represent the corresponding gradients as
follows,

$\displaystyle\nabla_{\phi_{1}}\mathcal{L}(\theta,\phi_{1},\phi_{2})$
$\displaystyle=\mathop{\mathbb{missing}}{E}\limits_{\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\nabla_{\phi_{1}}\mathcal{U}^{\rm
off}(\tau,f_{\phi_{1}},\theta)\right]$ (11)
$\displaystyle\nabla_{\phi_{2}}\mathcal{L}(\theta,\phi_{1},\phi_{2})$
$\displaystyle=\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\nabla_{\phi_{2}}\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]$ (12)

Input: Dataset $\mathcal{D}_{\boldsymbol{x}\boldsymbol{y}}$ with inputs and
ground-truth outputs;

Teacher policy $\pi_{*}$, student policy $\pi_{\theta}$ with initial
parameters $\theta$ pretrained on
$\mathcal{D}_{\boldsymbol{x}\boldsymbol{y}}$, off-policy $Q$-value function
$f_{\phi_{1}}$, on-policy $Q$-value function $f_{\phi_{2}}$;

Learning rate $\eta$, batch size $M$, control factor $\alpha$

Output: A optimized student policy $\pi_{\theta}$

while  _$\theta$ has not converged_ do

for  _$k=0,1,2,\ldots,K$_ do

Sample an input $\boldsymbol{x}\sim\mathcal{D}_{\boldsymbol{x}}$ and generate
an output trajectory $\tau_{\rm
off}\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}$

$\phi_{1}\leftarrow\phi_{1}+\alpha\eta\nabla_{\phi_{1}}\mathcal{U}^{\rm
off}(\tau_{\rm off},f_{\phi_{1}},\theta)$ $\hfill\rhd$ maximize
$\mathcal{L}(\theta,\phi_{1},\phi_{2})$ by Eq. (11)

Sample an input $\boldsymbol{x}\sim\mathcal{D}_{\boldsymbol{x}}$ and generate
an output trajectory $\tau_{\rm
on}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}$

$\phi_{2}\leftarrow\phi_{2}+\alpha\eta\nabla_{\phi_{2}}\mathcal{U}^{\rm
on}(\tau_{\rm on},f_{\phi_{2}})$ $\hfill\rhd$ maximize
$\mathcal{L}(\theta,\phi_{1},\phi_{2})$ by Eq. (12)

end for

Sample an input $\boldsymbol{x}\sim\mathcal{D}_{\boldsymbol{x}}$ and generate
output trajectories $\tau_{\rm
off}\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}$ and $\tau_{\rm
on}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}$

$\theta\leftarrow\theta-\eta\left(\mathcal{G}_{\rm on}(\tau_{\rm
on},\theta)-\mathcal{G}_{\rm off}(\tau_{\rm off},\theta)\right)$ $\hfill\rhd$
minimize $\mathcal{L}(\theta,\phi_{1},\phi_{2})$ by Eq. (10)

end while

Algorithm 1 Adversarial training procedure

Training procedure. The training objective is to achieve a Nash equilibrium
between the parameters of student policy $\theta$ and the parameters of on-
policy and off-policy $Q$-value functions $\phi_{1},\phi_{2}$ for the
optimization problem of
$\min_{\phi}\max_{\phi_{1},\phi_{2}}\mathcal{L}(\theta,\phi_{1},\phi_{2})$ in
Eq. (9). To this end, we use an adversarial training strategy in Algorithm 1,
by starting from a student model fine-tuned on a dataset
$\mathcal{D}_{\boldsymbol{x}\boldsymbol{y}}$. In the training algorithm, we
iteratively maximize the objective w.r.t. the parameters of $Q$-value
functions $\phi_{1}$ $\phi_{2}$ and simultaneously minimize the objective
w.r.t. the parameters of student policy $\theta$. In each iteration of policy
updating, we first perform $K$ steps of gradient ascent w.r.t. the parameters
of $Q$-value functions $\phi_{1},\phi_{2}$, where the gradients in Eq. (11)
and Eq. (12) are approximately computed by sampling inputs from the dataset
and sampling the trajectories from the policies. Then, the parameters of
student policy $\theta$ are updated by gradient descent with the estimated
gradient of Eq. (11) with sampling inputs and policies.

##  4 Experiments

We consider task-agnostic instruction-following experiments and task-specific
experiments, including text summarization, machine translation, and
commonsense reasoning. We compare our approach with various KD baselines,
including: SFT, which fine-tunes the student model on the supervised dataset
$\mathcal{D}{\boldsymbol{x}\boldsymbol{y}}$; KD [15], which uses KL divergence
on the supervised dataset $\mathcal{D}{\boldsymbol{x}\boldsymbol{y}}$; SeqKD
[18], which applies SFT to the student model with teacher-generated outputs;
ImitKD [21], which uses KL divergence on the student-generated outputs;
MiniLLM [13], which uses RKL divergence with a policy gradient method; GKD
[2], which uses JS divergence with an on-policy method; and DistiLLM [19],
which uses an adaptive training method for off-policy optimization of a skew
KL divergence. Additionally, we focus on step-wise distance optimization for
KD and compare it with a range of well-known methods, including KL divergence,
RKL divergence, JS divergence, and TV distance, as discussed by Wen et al.
[35]. All the reported results are the average across three random seeds.

###  4.1 Task-Agnostic Distillation

Experimental Setup. We follow the previous works [13; 19] for the
implementation of the instruction-following experiment, aiming to evaluate the
distilled model’s ability to handle diverse tasks presented in the form of
instructions. We construct the training data from databricks-dolly-15k [8],
where we randomly select 15K samples for training and equally split 500
samples for validation and testing. We evaluate the trained model on five
instruction-following datasets: DollyEval, SelfInst [33], VicunaEval [5], S-NI
[34], and UnNI [16]. Following the previous works [13; 19], we also add the
OpenWebText [12] corpus, consisting of long-document plain text, for joint
training with a language modeling task. This has been shown to effectively
improve the performance of instruction tuning [13]. The evaluation metrics
include ROUGE-L [22] and GPT-4 feedback with the same prompts as in [19]. More
details on experimental setup refer to Appendix B.

Method | DollyEval | SelfInst | VicunaEval | S-NI | UnNI
---|---|---|---|---|---
GPT4 | R-L | GPT4 | R-L | GPT4 | R-L | R-L | R-L
OpenLLaMA2-7B (teacher) | 58.8±1.2 | 32.5±0.4 | 56.7±0.8 | 21.6±0.2 | 46.2±0.6 | 22.6±0.5 | 36.3±0.5 | 38.5±0.2
SFT (student) | 46.8±0.7 | 26.7±0.6 | 40.8±1.1 | 16.3±0.7 | 34.8±0.8 | 17.3±0.2 | 30.4±0.4 | 28.6±0.3
KD [15] | 43.9±0.8 | 22.4±0.4 | 43.5±0.5 | 17.4±0.5 | 33.7±0.3 | 16.4±0.2 | 29.3±0.6 | 23.4±0.3
SeqKD [18] | 50.2±0.6 | 26.2±0.4 | 46.8±0.3 | 15.8±0.5 | 38.8±1.2 | 18.0±0.6 | 29.7±0.3 | 27.8±0.1
ImitKD [21] | 53.7±1.6 | 25.3±0.3 | 45.0±0.7 | 18.4±0.4 | 41.7±1.2 | 19.1±0.2 | 33.1±0.7 | 28.7±0.5
MiniLLM [13] | 58.7±1.2 | 28.4±0.3 | 51.8±1.5 | 20.2±0.6 | 44.2±1.1 |  20.7±0.5 |  37.4±0.4 | 37.5±0.2
GKD [2] | 57.6±1.0 | 27.5±0.3 | 52.4±1.2 |  20.9±0.3 | 45.5±0.8 | 19.3±0.5 | 36.8±0.6 | 34.8±0.3
DistiLLM [19] |  59.2±1.2 |  29.5±0.2 |  53.4±1.0 | 20.8±0.7 |  46.3±0.9 | 20.4±0.3 | 37.2±0.1 |  38.2±0.1
ours | 59.8±0.8 | 30.7±0.4 | 54.2±1.2 | 21.7±0.5 | 47.8±0.7 | 21.4±0.4 | 38.7±0.4 | 39.1±0.3

Table 1: Comparison with state-of-the-art KD methods on the instruction-
following dataset using fine-tuned OpenLLaMA-7B as the teacher and fine-tuned
OpenLLaMA-3B as the student. We format the best, the second best and worse
than SFT results. The results based on GPT-2 are available in Appendix C.1.

Main results. Table 1 illustrates the instruction-following performances.
Compared with the SFT baseline, which indicates the student model without KD,
KD and SeqKD hardly improve the performances. This indicates that using only
supervised datasets or teacher-generated outputs does not benefit the KD of
large language models. In contrast, utilizing the student-generated outputs
with KL divergence [2], RKL divergence [13], and JS divergence [2] shows
effectiveness for KD in the instruction-following task. State-of-the-art
methods [13; 2; 19] tend to combine the student-generated outputs with the
teacher-generated output or supervised dataset to further improve the results
of KD. This shows that a mixture optimization of both on-policy and off-policy
objectives can effectively improve the KD performance of large language models
on the instruction-following task. In particular, we use an adversarial
moment-matching method and optimize both on-policy and off-policy objectives
for KD, thus achieving the best results on five test datasets with both GPT4
feedback and ROUGE-L evaluations.

###  4.2 Task-Specific Distillation

Experimental Setup. We evaluated the KD models on three tasks consisting of
text summarization, machine translation, and reasoning. For the text
summarization task, we follow Ko et al. [19] to conduct experiments on the
SAMSum [11] dataset. For the machine translation tasks, we follow Ko et al.
[19] to conduct experiments on the IWSLT’17 (en-de) [4] dataset. For the
commonsense reasoning task, we conduct experiments on the StrategyQA dataset
[10]. For all of the task-specific experiments, we use T5-XL [26] as the
teacher model and T5-Large/-Base/-Small as the student models. For the machine
translation experiments, we employ a multilingual pretrained model, mT5 [39],
to build the methods. For evaluation, we use ROUGE-L [22], BLEU [24], and
accuracy as the performance metrics on SAMSum, IWSLT’17 (en-de), and
StrategyQA, respectively. More details about the experimental setup refer to
Appendix B.

(a) DollyEval.

(b) SAMSum.

(c) IWSLT’17 (en-de).

(d) StrategyQA.

Figure 2: Performance of difference step-wise distribution distances.

Main results. Table 2 displays the performances on three task-specific
datasets. Since the original work of MiniLLM [13] does not consider these
tasks, we thus do not make comparisons with MiniLLM. The performance trend is
similar to the instruct-following results, revealing that KD of large language
models for specific tasks also benefits from the combination of on-policy
objectives with student-generated outputs and off-policy objectives with
teacher-generated outputs or supervised datasets. Additionally, we observe
that student models of different sizes all benefit from the KD methods to
improve performance. Overall, our approach achieves the best results on all
three task-specific datasets for student models of different sizes. This
demonstrates the effectiveness of an adversarial moment-matching approach for
KD of large language models on specific tasks.

Method | SAMSum | IWSLT’17 (en-de) | StrategyQA
---|---|---|---
T5-Small | T5-Base | T5-Large | T5-Small | T5-Base | T5-Large | T5-Small | T5-Base | T5-Large
T5-XL (teacher) | 52.5±0.4 | 35.2±0.2 | 64.5±0.8
SFT (student) | 40.6±0.2 | 47.3±0.3 | 49.8±0.2 | 21.5±0.1 | 30.1±0.0 | 33.7±0.1 | 52.4±0.5 | 57.5±0.8 | 60.7±0.8
KD [15] | 39.2±0.4 | 46.5±0.3 | 47.4±0.3 | 21.7±0.1 | 29.8±0.2 | 31.7±0.1 | 49.7±0.3 | 55.3±0.1 | 59.2±0.5
SeqKD [18] | 39.7±0.3 | 47.7±0.5 | 49.3±0.4 | 21.2±0.3 | 29.2±0.2 | 32.9±0.5 | 50.6±0.7 | 57.5±1.1 | 61.5±0.8
ImitKD [21] | 41.8±0.3 | 48.6±0.7 | 51.2±0.5 | 22.2±0.3 | 28.7±0.6 | 34.1±0.2 | 53.8±0.8 | 59.7±0.5 | 61.7±0.6
GKD [2] | 42.1±0.3 | 48.2±0.5 | 51.7±0.4 |  22.7±0.2 |  31.2±0.1 | 34.7±0.2 | 55.6±0.4 | 60.3±0.5 |  63.6±0.3
DistiLLM [19] |  42.6±0.2 |  49.4±0.6 |  52.1±0.4 | 22.5±0.1 | 30.8±0.2 |  35.5±0.1 |  56.3±0.3 |  61.2±0.7 | 62.8±0.2
ours | 43.7±0.4 | 50.4±0.3 | 52.7±0.3 | 23.7±0.1 | 32.4±0.3 | 36.0±0.2 | 58.2±0.4 | 62.9±0.3 | 65.3±0.7

Table 2: Comparison with the state-of-the-art KD methods on text
summarization, machine translation and commonsense reasoning datasets. We
report the ROUGE-L, BLEU and accuracy for SAMSum, IWSLT’17 (en-de) and
StrategyQA, respectively. We format the best, the second best and worse than
SFT results.

###  4.3 Analysis on Step-Wise Distance Optimization

Comparison with distribution matching. We make comparisons with different
step-wise distribution distances with a uniform formulation of Definition 2,
considering the on-policy, off-policy objectives as well as the joint form.
Results on four tasks are shown in Figure 2, more instruct-following results
are available in Appendix C.2. Compared with the KL divergence, RKL
divergence, JS divergence and total variation distance, the proposed moment-
matching distance achieves the best results under both the on-policy and off-
policy training objectives, which shows that the proposed moment-matching
approach is effective for KD of large language models. Besides, we observe
that using a joint objective of both on-policy and off-policy can further
significantly improve the performances. This shows that both on-policy and
off-policy moment-matching objectives contribute to the minimization of the
imitation gap and can thus benefit the KD of large language models.

(a) Training loss and $d^{\rm on}_{\rm MM}$, $d^{\rm off}_{\rm MM}$ against
training step.

(b) On-policy moment-matching distance $d^{\rm on}_{\rm MM}$ on the test sets.

(c) Off-policy moment-matching distance $d^{\rm on}_{\rm MM}$ on the test
sets.

Figure 3: Adversarial training procedure for optimizing the on-policy and off-
policy moment-matching distances $d^{\rm on}_{\rm MM}$, $d^{\rm off}_{\rm MM}$
on the instruction-following dataset.

Adversarial training procedure. We present the training loss and moment-
matching distance against the adversarial training steps. As depicted in
Figure 3 (a), the training loss initially increases within the first 0-1,000
steps, indicating that initially, the $Q$-value functions are stronger than
the policy in maximizing the loss function
$\mathcal{L}(\theta,\phi_{1},\phi_{2})$ in Eq. (9). Concurrently, the policy
gradient method contributes to minimizing the training loss, which eventually
converges to a much lower stable value. Additionally, both the on-policy and
off-policy moment-matching distances $d^{\rm on}{\rm MM}$ and $d^{\rm off}{\rm
MM}$ decrease and eventually reach a low value with only minor fluctuations.
For more results and details on experimental setups, please refer to Appendix
C.3.

Moment-matching distance optimization. We further illustrate the on-policy
moment-matching distance $d^{\rm on}{\rm MM}$ and the off-policy moment-
matching distance $d^{\rm off}{\rm MM}$ (defined in Definition 3) optimized by
different step-wise distances in Figure 3 (b) and (c), respectively.
Interestingly, we observe that the total variation (TV) distance obtains the
second-best results on average for both on-policy and off-policy distances.
This finding suggests a similarity between the formulations of TV distance and
moment-matching distances to some extent, as supported by the theoretical
result of Corollary 1. Across all instruction-following test sets, our
approach effectively optimizes both on-policy and off-policy moment-matching
distances more than other step-wise distribution distances used in KD,
including KL divergence, RKL divergence, JS divergence, and TV distance. This
observation also underscores the effectiveness of our policy gradient methods.
Extensive results on the task-specific datasets are available in Appendix C.4.

##  5 Conclusion

In this work, we investigated a moment-matching approach for knowledge
distillation of large language models. Specifically, we formulated knowledge
distillation from a perspective of imitation learning and derived both on-
policy and off-policy bounds for the imitation gap between the teacher model
and student model via moment-matching distance. Additionally, we proposed an
adversarial training algorithm to simultaneously estimate and minimize the
joint objective of on-policy and off-policy moment-matching distances. In
experiments, we evaluated the proposed algorithm on four instruction-following
datasets and three task-specific datasets, comparing it with a range of state-
of-the-art KD methods as well as four well-studied step-wise distribution
distances for KD of auto-regressive models. Results demonstrate that our
approach can effectively leverage the policy gradient method to optimize the
moment-matching distance and achieve the best results across all datasets.

Limitations and future work. The proposed adversarial training algorithm
requires additional computational steps for the inner-loop gradient ascent,
which may result in increased time complexity. Moreover, the proposed approach
necessitates auxiliary networks to build the $Q$-value functions, which may
incur additional memory costs. Therefore, in future work, we aim to enhance
the time and memory efficiency of our approach.

## References

  * [1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.  Gpt-4 technical report.  arXiv preprint arXiv:2303.08774, 2023.
  * [2] Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem.  On-policy distillation of language models: Learning from self-generated mistakes.  In The Twelfth International Conference on Learning Representations, 2024.
  * [3] Yuntao Bai, Saurav Kadavath, Sandipan Kundu, Amanda Askell, Jackson Kernion, Andy Jones, Anna Chen, Anna Goldie, Azalia Mirhoseini, Cameron McKinnon, et al.  Constitutional ai: Harmlessness from ai feedback.  arXiv preprint arXiv:2212.08073, 2022.
  * [4] Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stüker, Katsuitho Sudoh, Koichiro Yoshino, and Christian Federmann.  Overview of the iwslt 2017 evaluation campaign.  In Proceedings of the 14th International Workshop on Spoken Language Translation, pages 2–14, 2017.
  * [5] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E Gonzalez, et al.  Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality.  See https://vicuna. lmsys. org (accessed 14 April 2023), 2(3):6, 2023.
  * [6] Kamil Ciosek.  Imitation learning by reinforcement learning.  In International Conference on Learning Representations, 2021.
  * [7] Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun.  Ultrafeedback: Boosting language models with high-quality feedback.  arXiv preprint arXiv:2310.01377, 2023.
  * [8] Free Dolly.  Introducing the world’s first truly open instruction-tuned llm. databricks. com, 2023.
  * [9] Xinyang Geng and Hao Liu.  Openllama: An open reproduction of llama.  URL: https://github. com/openlm-research/open_llama, 2023.
  * [10] Mor Geva, Daniel Khashabi, Elad Segal, Tushar Khot, Dan Roth, and Jonathan Berant.  Did aristotle use a laptop? a question answering benchmark with implicit reasoning strategies.  Transactions of the Association for Computational Linguistics, 9:346–361, 2021.
  * [11] Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer.  Samsum corpus: A human-annotated dialogue dataset for abstractive summarization.  arXiv preprint arXiv:1911.12237, 2019.
  * [12] Aaron Gokaslan, Vanya Aaron, Ellie Pavlick, and Stefanie Tellex.  Openwebtext corpus.  2019\.
  * [13] Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang.  Minillm: Knowledge distillation of large language models.  In The Twelfth International Conference on Learning Representations, 2024.
  * [14] Yongchang Hao, Yuxin Liu, and Lili Mou.  Teacher forcing recovers reward functions for text generation.  Advances in Neural Information Processing Systems, 35:12594–12607, 2022.
  * [15] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  arXiv preprint arXiv:1503.02531, 2015.
  * [16] Or Honovich, Thomas Scialom, Omer Levy, and Timo Schick.  Unnatural instructions: Tuning language models with (almost) no human labor.  arXiv preprint arXiv:2212.09689, 2022.
  * [17] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei.  Scaling laws for neural language models.  arXiv preprint arXiv:2001.08361, 2020.
  * [18] Yoon Kim and Alexander M Rush.  Sequence-level knowledge distillation.  In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2016\.
  * [19] Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun.  Distillm: Towards streamlined distillation for large language models.  arXiv preprint arXiv:2402.03898, 2024.
  * [20] Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, and Tuo Zhao.  Less is more: Task-aware layer-wise distillation for language model compression.  In International Conference on Machine Learning, pages 20852–20867. PMLR, 2023.
  * [21] Alexander Lin, Jeremy Wohlwend, Howard Chen, and Tao Lei.  Autoregressive knowledge distillation through imitation learning.  In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 6121–6133, 2020.
  * [22] Chin-Yew Lin.  Rouge: A package for automatic evaluation of summaries.  In Text summarization branches out, pages 74–81, 2004.
  * [23] Richard Yuanzhe Pang and He He.  Text generation by learning from demonstrations.  In 9th International Conference on Learning Representations, ICLR 2021, 2021.
  * [24] Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu.  Bleu: a method for automatic evaluation of machine translation.  In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 311–318, 2002.
  * [25] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.  Language models are unsupervised multitask learners.  OpenAI blog, 1(8):9, 2019.
  * [26] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  Journal of machine learning research, 21(140):1–67, 2020.
  * [27] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter.  arXiv preprint arXiv:1910.01108, 2019.
  * [28] Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet.  On integral probability metrics,$\backslash$phi-divergences and binary classification.  arXiv preprint arXiv:0901.2698, 2009.
  * [29] Gokul Swamy, Sanjiban Choudhury, J Bagnell, and Steven Z Wu.  Sequence model imitation learning with unobserved contexts.  Advances in Neural Information Processing Systems, 35:17665–17676, 2022.
  * [30] Gokul Swamy, Sanjiban Choudhury, J Andrew Bagnell, and Steven Wu.  Of moments and matching: A game-theoretic framework for closing the imitation gap.  In International Conference on Machine Learning, pages 10022–10032. PMLR, 2021.
  * [31] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto.  Stanford alpaca: An instruction-following llama model, 2023.
  * [32] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971, 2023.
  * [33] Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannaneh Hajishirzi.  Self-instruct: Aligning language models with self-generated instructions.  arXiv preprint arXiv:2212.10560, 2022.
  * [34] Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Anjana Arunkumar, Arjun Ashok, Arut Selvan Dhanasekaran, Atharva Naik, David Stap, et al.  Super-naturalinstructions: Generalization via declarative instructions on 1600+ nlp tasks.  arXiv preprint arXiv:2204.07705, 2022.
  * [35] Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou.  f-divergence minimization for sequence-level knowledge distillation.  In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 10817–10834, 2023.
  * [36] Qingyang Wu, Lei Li, and Zhou Yu.  Textgail: Generative adversarial imitation learning for text generation.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 14067–14075, 2021.
  * [37] Taiqiang Wu, Chaofan Tao, Jiahao Wang, Zhe Zhao, and Ngai Wong.  Rethinking kullback-leibler divergence in knowledge distillation for large language models.  arXiv preprint arXiv:2404.02657, 2024.
  * [38] Xiaohan Xu, Ming Li, Chongyang Tao, Tao Shen, Reynold Cheng, Jinyang Li, Can Xu, Dacheng Tao, and Tianyi Zhou.  A survey on knowledge distillation of large language models.  arXiv preprint arXiv:2402.13116, 2024.
  * [39] Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami Al-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel.  mt5: A massively multilingual pre-trained text-to-text transformer.  arXiv preprint arXiv:2010.11934, 2020.
  * [40] Lantao Yu, Weinan Zhang, Jun Wang, and Yong Yu.  Seqgan: Sequence generative adversarial nets with policy gradient.  In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017.

##  Appendix A Proofs

###  A.1 Proof of Proposition 1

###### Proof..

Similar to the proof of Performance Difference Lemma (PDL), we have

$\displaystyle J(\pi_{*})-J(\pi_{\theta})$ $\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}r(\boldsymbol{y}_{<t},y_{t})\right]-\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\left[V^{\pi_{\theta}}(\boldsymbol{x})\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(r(\boldsymbol{y}_{<t},y_{t})+V^{\pi_{\theta}}(\boldsymbol{y}_{<t})-V^{\pi_{\theta}}(\boldsymbol{y}_{<t})\right)\right]-\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\left[V^{\pi_{\theta}}(\boldsymbol{x})\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(r(\boldsymbol{y}_{<t},y_{t})+V^{\pi_{\theta}}(\boldsymbol{y}_{<t+1})-V^{\pi_{\theta}}(\boldsymbol{y}_{<t})\right)\right]$
$\displaystyle\overset{(i)}{=}$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(Q^{\pi_{\theta}}(\boldsymbol{y}_{<t},y_{t})-V^{\pi_{\theta}}(\boldsymbol{y}_{<t})\right)\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(Q^{\pi_{\theta}}(\boldsymbol{y}_{<t},y_{t})-\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\left[Q^{\pi_{\theta}}(\boldsymbol{y}_{<t},y)\right]\right)\right]$
$\displaystyle\leq$
$\displaystyle\sup_{f\in\mathcal{F}_{Q}}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(f(\boldsymbol{y}_{<t},y_{t})-\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\left[f(\boldsymbol{y}_{<t},y)\right]\right)\right],$

where $(i)$ follows from Bellman equation. ∎

###  A.2 Proof of Proposition 2

###### Proof..

Similar to the proof of Proposition 1, we have

$\displaystyle J(\pi_{*})-J(\pi_{\theta})$ $\displaystyle=$
$\displaystyle-\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}r(\boldsymbol{y}_{<t},y_{t})\right]+\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\left[V^{\pi_{*}}(\boldsymbol{x})\right]$ $\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(V^{\pi_{*}}(\boldsymbol{y}_{<t})-\left(r(\boldsymbol{y}_{<t},y_{t})+V^{\pi_{*}}(\boldsymbol{y}_{<t})\right)\right)\right]+\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\left[V^{\pi_{*}}(\boldsymbol{x})\right]$ $\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(V^{\pi_{*}}(\boldsymbol{y}_{<t})-\left(r(\boldsymbol{y}_{<t},y_{t})+V^{\pi_{*}}(\boldsymbol{y}_{<t+1})\right)\right)\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(V^{\pi_{*}}(\boldsymbol{y}_{<t})-Q^{\pi_{*}}(\boldsymbol{y}_{<t},y_{t})\right)\right]$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{*}(\cdot|\boldsymbol{y}_{<t})}\left[Q^{\pi_{*}}(\boldsymbol{y}_{<t},y)\right]-Q^{\pi_{*}}(\boldsymbol{y}_{<t+1},y_{t})\right)\right]$
$\displaystyle\leq$
$\displaystyle\sup_{f\in{\mathcal{F}_{Q_{*}}}}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{*}(\cdot|\boldsymbol{y}_{<t})}\left[f(\boldsymbol{y}_{<t},y)\right]-f(\boldsymbol{y}_{<t},y_{t})\right)\right]$

∎

###  A.3 Existing Step-Wise Distribution Distance for KD

######  Definition 4 (Step-wise distribution distances for distillation).

Following Wen et al. [35], we define four groups of well-studied probability
distribution distances as follows,

  * •

Total variation (TV) distance. [35] The off-policy and on-policy step-wise TV
distance can be defined as follows,

$\displaystyle d_{\rm TV}^{\rm off}(\pi_{\theta},\pi_{*}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\\!\\!\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\left|\pi_{*}(y|\boldsymbol{y}_{<t})-\pi_{\theta}(y|\boldsymbol{y}_{<t})\right|\right];$
(13) $\displaystyle d_{\rm TV}^{\rm on}(\pi_{\theta},\pi_{*}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\\!\\!\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\left|\pi_{*}(y|\boldsymbol{y}_{<t})-\pi_{\theta}(y|\boldsymbol{y}_{<t})\right|\right]$
(14)

  * •

Kullback–Leibler (KL) divergence [35] The off-policy and on-policy step-wise
KL divergence between the teacher policy $\pi_{*}$ and the student policy
$\pi_{\theta}$ can be defined as follows,

$\displaystyle d_{\rm KL}^{\rm off}(\pi_{*},\pi_{\theta}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\pi_{*}(y|\boldsymbol{y}_{<t})\log\frac{\pi_{*}(y|\boldsymbol{y}_{<t})}{\pi_{\theta}(y|\boldsymbol{y}_{<t})}\right];$
(15) $\displaystyle d_{\rm KL}^{\rm on}(\pi_{*},\pi_{\theta}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\pi_{*}(y|\boldsymbol{y}_{<t})\log\frac{\pi_{*}(y|\boldsymbol{y}_{<t})}{\pi_{\theta}(y|\boldsymbol{y}_{<t})}\right]$
(16)

  * •

Reverse Kullback–Leibler (RKL) divergence [35] The off-policy and on-policy
step-wise RKL divergence between the teacher policy $\pi_{*}$ and the student
policy $\pi_{\theta}$ can be defined as follows,

$\displaystyle d_{\rm RKL}^{\rm off}(\pi_{\theta},\pi_{*}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\pi_{\theta}(y|\boldsymbol{y}_{<t})\log\frac{\pi_{\theta}(y|\boldsymbol{y}_{<t})}{\pi_{*}(y|\boldsymbol{y}_{<t})}\right];$
(17) $\displaystyle d_{\rm RKL}^{\rm on}(\pi_{\theta},\pi_{*}):=$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\pi_{\theta}(y|\boldsymbol{y}_{<t})\log\frac{\pi_{\theta}(y|\boldsymbol{y}_{<t})}{\pi_{*}(y|\boldsymbol{y}_{<t})}\right]$
(18)

  * •

Jenson–Shannon (JS) divergence [35] The JS divergence between the teacher
policy $\pi_{*}$ and the student policy $\pi_{\theta}$ can be defined based on
the KL divergence and RKL divergence as follows,

$\displaystyle d_{\rm JS}^{\rm off}(\pi_{\theta},\pi_{*}):=$
$\displaystyle\frac{1}{2}d_{\rm KL}^{\rm
off}(\pi_{*},\frac{\pi_{\theta}+\pi_{*}}{2})+\frac{1}{2}d_{\rm RKL}^{\rm
off}(\pi_{\theta},\frac{\pi_{\theta}+\pi_{*}}{2});$ (19) $\displaystyle d_{\rm
JS}^{\rm on}(\pi_{\theta},\pi_{*}):=$ $\displaystyle\frac{1}{2}d_{\rm KL}^{\rm
on}(\pi_{*},\frac{\pi_{\theta}+\pi_{*}}{2})+\frac{1}{2}d_{\rm RKL}^{\rm
on}(\pi_{\theta},\frac{\pi_{\theta}+\pi_{*}}{2})$ (20)

###  A.4 Proof of Corollary 1

###### Proof..

We first derive an upper bound for the on-policy moment-matching bound of Eq.
(3). Set $\mathcal{F}_{Q_{*}}=\\{f:\|f\|_{\infty}\leq 1\\}$, Eq. (3) can be
bounded as follows,

$\displaystyle\sup_{f:\|f\|_{\infty}\leq
1}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
on}(\tau,f)\right]$ $\displaystyle=$ $\displaystyle\sup_{f:\|f\|_{\infty}\leq
1}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\left(\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{*}(\cdot|\boldsymbol{y}_{<t})}\\!\\!\left[f(\boldsymbol{y}_{<t},y)\right]-\mathop{\mathbb{missing}}{E}\limits_{y^{\prime}\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t}))}\left[f(\boldsymbol{y}_{<t},y^{\prime})\right]\right)\right]$
$\displaystyle\overset{(i)}{\leq}$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\sup_{f:\|f\|_{\infty}\leq
1}\left(\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{*}(\cdot|\boldsymbol{y}_{<t})}\\!\\!\left[f(\boldsymbol{y}_{<t},y)\right]-\mathop{\mathbb{missing}}{E}\limits_{y^{\prime}\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t}))}\left[f(\boldsymbol{y}_{<t},y^{\prime})\right]\right)\right]$
$\displaystyle\overset{(ii)}{=}$
$\displaystyle\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\left|\pi_{*}(y|\boldsymbol{y}_{<t})-\pi_{\theta}(y|\boldsymbol{y}_{<t})\right|\right]\overset{\text{Def.
\ref{def:stepdist}}}{=}d^{\rm on}_{\rm TV}(\pi_{\theta},\pi_{*}),$

where $(i)$ follows from Jensen’s inequality and $(ii)$ follows from [28].

Similarly, we can bound the off-policy version of Eq. (2) as follows,

$\displaystyle\sup_{f:\|f\|_{\infty}\leq
1}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
off}(\tau,f,\theta)\right]$
$\displaystyle\leq\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}\atop\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\left|\pi_{*}(y|\boldsymbol{y}_{<t})-\pi_{\theta}(y|\boldsymbol{y}_{<t})\right|\right]$
$\displaystyle\overset{\text{Def. \ref{def:stepdist}}}{=}d^{\rm off}_{\rm
TV}(\pi_{\theta},\pi_{*})$

∎

###  A.5 Derivation of Policy Gradient

###### Proof..

Based on the definition of training objective in Eq. (9), we have

$\displaystyle\begin{split}\nabla_{\theta}\mathcal{L}(\theta,\phi_{1},\phi_{2})&=\nabla_{\theta}\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\left[\mathop{\mathbb{missing}}{E}\limits_{\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
off}(\tau,f_{\phi_{1}};\theta)\right]+\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]\right]\\\
&=\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\left[\mathop{\mathbb{missing}}{E}\limits_{\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\nabla_{\theta}\mathcal{U}^{\rm
off}(\tau,f_{\phi_{1}};\theta)\right]+\nabla_{\theta}\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]\right]\\\ \end{split}$ (21)

Based on the definition of $\mathcal{U}^{\rm off}(\tau,f_{\phi_{1}},\theta)$
in Eq. (2), we can compute the gradient $\nabla_{\theta}\mathcal{U}^{\rm
off}(\tau,f_{\phi_{1}},\theta)$ for any $\tau\in{\mathcal{Y}}$ and
$\phi_{1}\in\Phi$ as follows,

$\displaystyle\begin{split}&\nabla_{\theta}\mathcal{U}^{\rm
off}(\tau,f_{\phi_{1}},\theta)\\\
=&-\sum_{t=1}^{T}\nabla_{\theta}\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\left[f_{\phi_{1}}(\boldsymbol{y}_{<t},y)\right]\\\
=&-\sum_{t=1}^{T}\sum_{y\in\mathcal{V}}\pi_{\theta}(y|\boldsymbol{y}_{<t})\nabla_{\theta}\log\pi_{\theta}(y|\boldsymbol{y}_{<t})f_{\phi_{1}}(\boldsymbol{y}_{<t},y)\\\
=&-\sum_{t=1}^{T}\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\nabla_{\theta}\log\pi_{\theta}(y|\boldsymbol{y}_{<t})f_{\phi_{1}}(\boldsymbol{y}_{<t},y)\end{split}$
(22)

then, based on the definition of $\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})$ in Eq. (3), we have

$\displaystyle\begin{split}&\nabla_{\theta}\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]\\\
=&\int\nabla_{\theta}\left[p_{\theta}(\tau|\boldsymbol{y}_{0}=\boldsymbol{x})\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]d\tau^{\prime}\\\
=&\int\nabla_{\theta}p_{\theta}(\tau^{\prime}|\boldsymbol{y}_{0}=\boldsymbol{x})\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})d\tau^{\prime}\\\ =&\int
p_{\theta}(\tau^{\prime}|\boldsymbol{y}_{0}=\boldsymbol{x})\nabla_{\theta}\log
p_{\theta}(\tau^{\prime}|\boldsymbol{y}_{0}=\boldsymbol{x})\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})d\tau^{\prime}\\\
=&\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\nabla_{\theta}\log
p_{\theta}(\tau^{\prime}|\boldsymbol{y}_{0}=\boldsymbol{x})\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]\\\
\overset{(i)}{=}&\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\nabla_{\theta}\log\pi_{\theta}(y^{\prime}_{t}|\boldsymbol{y}^{\prime}_{<t})\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right],\end{split}$ (23)

where $(i)$ uses the fact that the state-transition is deterministic, and thus
$\nabla_{\theta}\log
p_{\theta}(\tau^{\prime}|\boldsymbol{y}^{\prime}_{0}=\boldsymbol{x})=\sum_{t=1}^{T}\nabla_{\theta}\log\pi_{\theta}(y^{\prime}_{t}|\boldsymbol{y}^{\prime}_{<t})$
for any
$\tau^{\prime}=\\{\boldsymbol{y}^{\prime}_{<t},y^{\prime}_{t}\\}_{t=1,\ldots,T}$.

Combining Eq. (21) with Eq. (22) and Eq. (23), we have

$\displaystyle\nabla_{\theta}\mathcal{L}(\theta,\phi_{1},\phi_{2})=\mathop{\mathbb{missing}}{E}\limits_{\boldsymbol{x}\sim
p_{\boldsymbol{x}}}\Bigg{[}$
$\displaystyle-\mathop{\mathbb{missing}}{E}\limits_{\tau\sim\pi_{*}|\boldsymbol{y}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\mathop{\mathbb{missing}}{E}\limits_{y\sim\pi_{\theta}(\cdot|\boldsymbol{y}_{<t})}\nabla_{\theta}\log\pi_{\theta}(y|\boldsymbol{y}_{<t})f_{\phi_{1}}(\boldsymbol{y}_{<t},y)\right]$
$\displaystyle+\mathop{\mathbb{missing}}{E}\limits_{\tau^{\prime}\sim\pi_{\theta}|\boldsymbol{y}^{\prime}_{0}=\boldsymbol{x}}\left[\sum_{t=1}^{T}\nabla_{\theta}\log\pi_{\theta}(y^{\prime}_{t}|\boldsymbol{y}^{\prime}_{<t})\mathcal{U}^{\rm
on}(\tau^{\prime},f_{\phi_{2}})\right]\Bigg{]}$

∎

##  Appendix B Experimental Setup

We use NVIDIA A40 GPUs with 40GB RAM to conduct all the experiments.

###  B.1 Instruction-Following Experiments

Base models. We conduct experiments on both GPT-2 [25] and OpenLLaMA [9]. For
the GPT-2 experiments, we use GPT-2 XL222https://huggingface.co/openai-
community/gpt2-xl with 1.5B parameters to construct the teacher policy and
GPT-2333https://huggingface.co/openai-community/gpt2 with 117M parameters to
construct the student policy. For the $Q$-value function, we use a GPT-2 model
for feature extraction and add a linear layer with the size of 768 x 1 to
output the logits of step $y_{t}$ as $f(\boldsymbol{y}_{<t},y_{t})$. For the
OpenLLaMA experiments, we use OpenLLaMA-7B444https://huggingface.co/openlm-
research/open_llama_7b with 6.7B parameters to construct the teacher policy
and OpenLLaMA-3B555https://huggingface.co/openlm-research/open_llama_7b with
2.7B parameters to construct the student model. We construct the $Q$-value
function with the OpenLLaMA-3B model and a linear layer with the size of 3,200
x 1.

Training details. We fine-tune the OpenLLaMA-7B teacher model and the
OpenLLaMA-3B student models on the corresponding supervised dataset with
10,000 steps. The GPT-2 teacher and student models use the fine-tuned
checkpoints by Gu et al. [13]. For the implementation of compared baselines,
we use the code by Ko et al. [19] and re-run the results. The optimization
protocol for KD training largely follows the previous work [13, 19]. In
particular, we search for the learning rates among a finite set for each
experiment to obtain the best result. The batch size for each experiments is
seleted to make full use of the 40GB RAM of an A40 GPU. To handle the
adversarial training, we choose the number of adversarial steps $K=5$ and the
adversarial control factor $\alpha=0.1$ based on the development experiments.
The hyperparameters for training are listed in Table 3.

Hyperparameter | GPT-2 | OpenLLaMA
---|---|---
max. training steps | 10,000 | 10,000
adv. step $K$ | 5 | 5
batch size (per GPU) | {8, 16, 32} | {4, 8}
dropout rate | 0.1 | 0.1
adv. control factor ($\alpha$) | 0.1 | 0.1
learning rate (lr) | {5e-5, 1e-4,5e-4} | {5e-6, 1e-5, 5e-5}
warmup steps | 1,000 | 500
weight decay | 1e-2 | 1e-2
max length | 512 | 512
sampling top-p | 1.0 | 1.0
sampling temperature | 1.0 | 1.0
evaluation | Greedy Sampling | Greedy Sampling
#GPUs | 2 | 4

Table 3: Training hyperparameters for instruction-following experiments.

###  B.2 Task-Specific Experiments

Base models. For the text summarization and commonsense reasoning experiments,
we use T5-XL666https://huggingface.co/google/t5-v1_1-xl with 2.8B parameters
to construct the teacher policy and construct the student policy with
T5-Large777https://huggingface.co/google/t5-v1_1-large (770M parameters),
T5-Base888https://huggingface.co/google/t5-v1_1-base (220M parameters) and
T5-Small999https://huggingface.co/google/t5-v1_1-small (60M parameters). For
the machine translation experiments, we use mT5-XL [39] to construct the
teacher policy and use mT5-Large/-Base/-Small to construct the student policy.
The corresponding $Q$-value functions are constructed using the student model
and a linear layer of size 1024x1, 768x1 and 512x1 for the large, base and
small models, respectively.

Training details We initialize the corresponding teacher and student models
using 10,000-step-fine-tuning checkpoints on the SAMSum dataset, 80,000-step-
fine-tuning checkpoints on the IWSLT’17 (en-de) dataset and 3,000-step-fine-
tuning checkpoints on the StrategyQA dataset. We largely follow Ko et al. [19]
to set the hyperparameters for training. In particular, we search for the
learning rate from a preset range to obtain the best result for each baseline
and our method. The batch size is selected to make full use of the RAM of
GPUs. We use a relatively larger maximum number of training steps for IWSLT’17
(en-de) experiments to satisfy sufficient convergences for the machine
translation task. We use beam search for the evaluation on the IWSLT’17 (en-
de) dataset.

Hyperparameter | SAMSum | IWSLT’17 (en-de) | StrategyQA
---|---|---|---
max. training steps | 10,000 | 80,000 | 3,000
adv. step $K$ | 5 | 2 | 5
batch size (per GPU) | {16, 32, 64} | {16, 32, 64} | {16, 32, 64}
dropout rate | 0.0 | 0.3 | 0.1
adv. control factor ($\alpha$) | 0.1 | 0.1 | 0.1
learning rate (lr) | {5e-5, 1e-4,5e-4} | {1e-4,5e-4,1e-3} | {1e-4,5e-4,1e-3}
warmup steps | 1,000 | 4,000 | 300
weight decay | 1e-2 | 1e-4 | 1e-2
max. length | 1024 | 512 | 1,024
sampling top-p | 1.0 | 1.0 | 1.0
sampling temperature | 1.0 | 1.0 | 1.0
evaluation | Greedy sampling | Beam search | Greedy sampling
#GPUs | 2 | 4 | 1

Table 4: Training hyperparameters on three task-specific datasets.

##  Appendix C Additional Results

###  C.1 Results Based on GPT-2

In addition to the experimental results based on OpenLLaMA for instruction-
following tasks, we also conduct experiments based on GPT-2. Results are
illustrated in Table 5. Compared with current state-of-the-art KD approaches,
our method achieves the best results on five datasets with both GPT-4 feedback
and ROUGE-L evaluations.

Method | DollyEval | SelfInst | VicunaEval | S-NI | UnNI
---|---|---|---|---|---
GPT4 | R-L | GPT4 | R-L | GPT4 | R-L | R-L | R-L
GPT-2 XL (teacher) | 45.5±0.7 | 28.2±0.8 | 34.7±1.6 | 14.3±0.2 | 32.7±1.6 | 16.2±0.3 | 27.6±0.3 | 32.2±0.3
SFT (student) | 29.8±1.2 | 23.4±0.2 | 20.2±0.7 | 10.3±0.5 | 17.8±0.9 | 14.6±0.4 | 16.1±0.3 | 18.2±0.6
KD [15] | 29.5±0.8 | 23.8±0.3 | 18.0±1.0 | 12.3±0.2 | 17.2±0.7 | 15.2±0.4 | 20.8±0.5 | 22.5±0.3
SeqKD [18] | 29.8±0.5 | 24.2±0.2 | 18.2±0.8 | 11.6±0.4 | 18.2±0.7 | 15.5±0.3 | 15.5±0.6 | 20.1±0.1
ImitKD [21] | 26.4±0.6 | 22.7±0.5 | 18.2±0.5 | 11.5±0.4 | 18.6±0.4 | 14.5±0.3 | 18.2±0.3 | 21.8±0.4
MiniLLM [13] | 30.2±1.2 | 24.3±0.3 | 20.5±0.3 |  13.2±0.3 | 20.5±0.7 | 18.5±0.3 | 22.7±0.3 | 23.5±0.2
GKD [2] | 29.2±0.6 | 23.6±0.2 | 20.7±0.5 | 12.7±0.2 | 20.2±0.6 | 17.7±0.2 | 25.1±0.3 | 25.9±0.1
DistiLLM [19] |  31.2±0.4 |  25.2±0.4 |  21.7±0.5 | 12.5±0.3 |  22.5±1.2 |  19.2±0.5 |  27.7±0.2 |  27.6±0.4
ours | 31.7±0.5 | 26.1±0.3 | 22.7±0.5 | 14.2±0.3 | 23.6±0.8 | 20.5±0.2 | 28.6±0.2 | 29.9±0.5

Table 5: Comparison with state-of-the-art KD methods on the instruction-
following dataset using fine-tuned GPT-2 XL (1.5B) as the teacher and fine-
tuned GPT-2 (0.1B) as the student. We format the best, the second best and
worse than SFT results.

###  C.2 Comparisons on Step-Wise Distribution Distance

Figure 4 and Figure 5 illustrate performance comparison with well-studied
step-wise distribution distance, including KL, RKL, JS divergences and TV
distances. Results show that the optimization of proposed moment-matching
objectives outperforms other step-wise distribution distances via either on-
policy distillation and off-policy distillation. Besides, jointly using on-
policy and off-policy moment-matching further improves the performances and
achieves the best results on five instruction-following datasets with KD from
the OpenLLaMA-7B to OpenLLaMA-3B model, and achieves the best results on three
task-specific datasets with KD from the (m)T5-XL to (m)T5-Base model.

(a) DollyEval.

(b) SelfInst.

(c) VicunaEval.

(d) S-NI.

(e) UnNI.

Figure 4: Performance of difference step-wise distribution distances on five
instruction-following datasets based on OpenLLaMA-3B $\rightarrow$
OpenLLaMA-7B.

(a) SAMSum.

(b) IWSLT’17 (en-de).

(c) StrategyQA.

Figure 5: Performance of difference step-wise distribution distances on three
task-specific datasets based on (m)T5-XL $\rightarrow$ (m)T5-Base.

###  C.3 Adversarial Training Procedure

Figure 6 illustrates the training loss and on-/off-policy moment-matching
distances against the training steps on the instruction-following dataset and
three task-specific datasets. We can observe that the training losses on four
datasets have a similar trend, increasing at the beginning and then converging
to a relatively lower level. The trend of loss function aligns with the
characteristics of adversarial training with gradient descent ascent. In
contrast, both the on-policy moment-matching distance $d^{\rm on}_{\rm MM}$
and the off-policy moment-matching distance $d^{\rm off}_{\rm MM}$ reduce as
the number of training steps increases, which shows the effectiveness of our
adversarial training approach for moment-matching.

(a) Instruct-following.

(b) SAMSum.

(c) IWSLT’17 (en-de).

(d) StrategyQA.

Figure 6: Training loss and $d^{\rm on}_{\rm MM}$, $d^{\rm off}_{\rm MM}$
against training step on four datasets.

###  C.4 Moment-Matching via Distribution Matching

We investigate how the distribution-matching methods via KL, RKL, JS
divergences or TV distance can optimize the moment-matching distance in Figure
7 and Figure 8. Results show that the proposed adversarial training algorithm
are more effective in minimizing the moment-matching distance than the
distribution-matching methods.

(a) On-policy moment-matching distance.

(b) Off-policy moment-matching distance.

Figure 7: Moment-matching via distribution-matching on the instruction-
following dataset.

(a) On-policy moment-matching distance on SAMSum.

(b) On-policy moment-matching distance on IWSLT’17 (en-de).

(c) On-policy moment-matching distance on StrategyQA.

(d) Off-policy moment-matching distance on SAMSum.

(e) Off-policy moment-matching distance on IWSLT’17 (en-de).

(f) Off-policy moment-matching distance on StrategyQA.

Figure 8: Moment-matching via distribution-matching on three task-specific
datasets.