Title: Adversarial Moment-Matching Distillation of Large Language Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Knowledge distillation

Large language models

Imitation learning

Moment-matching

Adversarial training
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 metrics measuring the divergence between teacher and student probability predictions. Instead of optimizing these mandatory 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 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.

Adversarial Moment-Matching Distillation of Large Language Models

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 behavior 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 momentmatching 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 [35] 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 [19]. However, the high computational and memory requirements of these models present significant challenges for practical deployment. To address these issues, knowledge distillation (KD) [16] 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 [20] and reverse KL divergence [14], 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 decisionmaking problem from the perspective of imitation learning [24, 14, 2]. Based on this concept, early works based on the teacher-generated outputs [20] or a supervised dataset [30] can be viewed as an offpolicy approach. Recent works further incorporate an on-policy approach, training the student on its self-generated outputs [24], using KL-based divergence [14, 2, 21] and total variation (TV) distance [39]. 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

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

'

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

∗

ℳ

~GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

(GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

, GLYPH(cmap:d835)

GLYPH(cmap:d835)+1)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

(GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

, GLYPH(cmap:d835)

'

)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)+1

GLYPH(cmap:d835)

'

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

∗

ℳ

~GLYPH(cmap:d835)

∗

GLYPH(cmap:d835)

GLYPH(cmap:d835)∗

(GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

, GLYPH(cmap:d835)

GLYPH(cmap:d835)+1)

GLYPH(cmap:d835)

GLYPH(cmap:d835)∗

(GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

, GLYPH(cmap:d835)

'

)

(a) On-policy distribution-matching distillation.

(b) Off-policy distribution-matching distillation.

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

GLYPH(cmap:d835)

∗

GLYPH(cmap:d835)

∗

GLYPH(cmap:d835)

∗

GLYPH(cmap:d835)

∗

GLYPH(cmap:d835)'

GLYPH(cmap:d835)

GLYPH(cmap:d835)'

GLYPH(cmap:d835)-1

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)-1

GLYPH(cmap:d835)'

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)'

GLYPH(cmap:d835)-1

GLYPH(cmap:d835)

GLYPH(cmap:d835)

...

GLYPH(cmap:d835)

GLYPH(cmap:d835)-1

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)+1

GLYPH(cmap:d835)

GLYPH(cmap:d835)+1

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)+1

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)+1

GLYPH(cmap:d835)

≤GLYPH(cmap:d835)+1

GLYPH(cmap:d835)

GLYPH(cmap:d835)+1

GLYPH(cmap:d835)

GLYPH(cmap:d835)+1

(c) On-policy

Q

-value moment-match. distillation (

ours

).

(d) Off-policy

Q

-value moment-match. distillation (

ours

).

imposes a significant limitation on the generalization performance of the student model through distillation.

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. [33], 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.

Theoretically, we compare the moment-matching objective with other distribution-matching measurements such as step-wise TV distance and analyze the convergence rate of our algorithm to an ϵ -accurate stationary point for optimization. Empirically, 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 momentmatching 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 are 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) [41, 42]. 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. Sequencelevel KD (SeqKD) [20] 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 teachergenerated outputs. This method has been extensively used for LLMs and has achieved success [34, 6]. In contrast, distillation of white-box LLMs can make full use of internal information of the teacher model, such as logits [30, 39] and hidden states [23], 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 [39, 14, 2, 21, 41]. Wen et al. [39] 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 realtime-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. [14] further propose a policy gradient-based method to address the high variance issues of RKL-based methods while Ko et al. [21] 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 [44, 26, 15], 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 [20, 7, 14, 2]. For imitation learning in text generation, early works such as SeqGAN [44] and TextGAIL [40] 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 [24], 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 [4, 9], in which teacher models are used to generate the feedback data for training a reward model. We build our method upon an RL-based imitation learning framework. However, unlike previous work [20, 14, 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 x , the auto-regressive generation task in our work aims to generate a sequence of tokens as the output ( y 1 , . . . , y T ) , where y t comes from a vocabulary V . For simplicity, we define y = ( y 0 , y 1 , . . . , y T ) as the full input-output sequence, where y 0 = x denotes the input. The generator is modeled by a conditional probability distribution p θ ( y | x ) = Π T -1 t =0 p θ ( y t +1 | y ≤ t ) , where y ≤ t denotes the prefix ( y 0 , y 1 , . . . , y t ) , t ∈ { 0 , 1 , . . . , T -1 } .

RL formulation. We model text generation as a finite-horizon, time-independent Markov decision process. At each time step t ∈ { 0 , . . . , T -1 } , the policy π θ takes an action ( t ): y t +1 ∈ V based on the current state ( t ): y ≤ t ∈ Y , transits to the next state ( t + 1 ): y ≤ t +1 ∈ Y and receives a reward ( t ): r ( y ≤ t , y t +1 ) by a reward function r : Y × V → R . The policy corresponds to the generation model π θ ( y t +1 | y ≤ t ) = p θ ( y t +1 | y ≤ t ) . We focus on a (conditional) trajectory { y 1 , y ≤ 1 , y 2 , . . . , y ≤ T -1 , y T } =: τ ∼ π θ | x which refers to a sequence of state-action pairs generated by given an initial state y 0 = x ∼ p x and then repeatedly sampling an action y t +1 ∼ π θ ( ·| y ≤ t ) and obtain the next state y ≤ t +1 ∼ T ( ·| y ≤ t , y t +1 ) 1 for T time steps. In such case, the probability of a (conditional) trajectory is formally represented as p ( τ | x , π θ ) = Π T -1 t =0 T ( y ≤ t +1 | y ≤ t , y t +1 ) π θ ( y t +1 | y ≤ t ) . We also define our value function and Q -value function as V π θ ( y ≤ t ) = E τ ( t ) ∼ π θ | y ≤ t [ ∑ T -1 t ′ = t γ t ′ -t r ( y ≤ t ′ , y t ′ +1 ) ] and Q π θ ( y ≤ t , y t +1 ) = E τ ( t ) ∼ π θ | y ≤ t ,y t +1 [ ∑ T -1 t ′ = t γ t ′ -t r ( y ≤ t ′ , y t ′ +1 ) ] , where γ ∈ (0 , 1) denotes the discounting factor. We define the RL objective in our generation task to maximize the performance J ( π θ ) = E x ∼ p x E τ ∼ π θ | x [ ∑ T -1 t =0 γ t r ( y ≤ t , y t +1 ) ] .

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:

From the perspective of imitation learning [33, 32], the objective of distillation from the teacher policy π ∗ to the student policy π θ can be represented as to minimize the imitation gap of Eq. (1) w.r.t. the parameters of student policy θ . A direct idea from Eq. (1) is to use moment matching over the reward to optimize the imitation gap [33]. 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 §3.1) for each timestep to represent the overall reward from the current timestep to the last timestep. Similar to [33], we can apply the Performance Difference Lemma (PDL) [18, 3, 33] to expand the imitation gap in Eq. (1) into either off-policy or on-policy expressions.

Proposition 1 ( Off-policy bound of imitation gap [33]). Let F Q denote the set of Q -value functions induced by sampling actions from π θ , then we have:

In the following sections, we will use L off ( π θ , f ) to represent the off-policy moment-matching objective of mitation learning for KD.

The off-policy moment-matching objective 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 [33]). Let F Q ∗ denote the set of Q -value functions induced by sampling actions from π ∗ , then we have:

In the following sections, we will use L on ( π θ , f ) to represent the on-policy moment-matching objective of an imitation learning for KD.

Proof. See Appendix A.1 and Appendix A.2 for the complete derivations of Proposition 1 and Proposition 2, respectively.

It is notable from Proposition 2 that the on-policy moment-matching objective 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 τ ∼ π θ | x .

In the remaining content of this section, we will explore the relationship between the momentmatching objectives and the existing distribution-matching objectives [39]. At the beginning, we draw a general formulation of the state-of-the-art methods for distillation of LLMs [39, 14, 2, 21] 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,

where M ( · , · ) denotes a distribution distance, consisting of total variation (TV) distance [39] and Kullback-Leibler (KL)-based divergence [14, 2]. Detailed definitions for these distances refer to Appendix A.3. For simplicity, we directly replace M with TV, KL, RKL, etc in the following sections.

It is notable from Wen et al. [39] 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 objectives). Based on Definition 2, we can re-formulate the off-policy and on-policy moment-matching (MM) objectives (Proposition 1 and Proposition 2, respectively) via step-wise distribution-matching, which can be defined as d off MM ( π θ , π ∗ ) and d on MM ( π θ , π ∗ ) respectively, where the distance metric MM( · , · ) can be defined as follows,

where L off ( π θ , f ) and L on ( π θ , f ) denote the off-policy and on-policy moment-matching objectives, which are defined in Proposition 1 and Proposition 2, respectively.

Under Definition 3, we observe that the main difference between the moment-matching objectives and other step-wise distribution distance, e.g., TV distance and KL-based divergences in formulation comes from the optimal Q -value function f off(on) ∗ , aiming to maximize the discrepancy of its expectations based on π ∗ ( ·| y ≤ t ) v.s. π θ ( ·| y ≤ t ) for each step t ∈ { 0 , 1 , . . . , T -1 } . To look deeper, we draw a connection between the moment-matching objectives and step-wise TV distance using the following corollary.

Theorem 1 ( Relationship between moment-matching objective and TV distance). Under a constrain of uniform boundness on the class of Q -value functions for off-/on-policy learning: F Q = F Q ∗ = { f : ∥ f ∥ ∞ ≤ 1 } , the moment-matching objectives in Proposition 1 and Proposition 2 can be upper-bounded by the step-wise TV distance, Formally, we have

for the off-policy and on-policy perspectives, respectively.

Proof. See Appendix A.4 for the complete derivation.

We can observe from Theorem 1 that minimizing the step-wise TV distance can achieve suboptimal results compared to optimizing the moment-matching objectives L off ( π θ , f ) , L on ( π θ , f ) for off-policy and on-policy imitation learning, which are defined in Proposition 1 and Proposition 2, respectively. Thus, optimizing the moment-matching objectives can potentially achieve better optimization results for imitation learning.

Algorithm 1: Adversarial training procedure

Input: Dataset D xy with inputs and ground-truth outputs Teacher policy π ∗ ; Student policy π θ with initial parameters θ pretrained on D xy ; Off-policy Q -value function f ϕ 1 and on-policy Q -value function f ϕ 2 with initial parameters ϕ 1 and ϕ 2 , respectively; Step sizes K (outer), N (inner); Learning rate η ; Controlling factor α ; Off-/on-policy combination factor β Output: The optimized student policy π θ ∗ for k = 0 , 1 , 2 , . . . , K -1 do

3.3 Adversarial Training Algorithm

Optimization objective. As shown in previous work [14, 2, 21] incorporating both the off-policy and on-policy distillation benefits effectiveness and efficiency. We thus consider a training objective to jointly minimize the off-policy moment-matching objective in Proposition 1 and the on-policy moment-matching objective in Proposition 2. Both the off-/on-policy 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 initialize two small networks of a single-layer MLP to estimate the off-/on-policy Q -value functions, respectively. For example in a causal/seq-to-seq LM, the Q -value estimate module can be represented as f ϕ 1(2) ( y ≤ t , y ) = ( h π θ t + v off(on) y ) ⊤ w off(on) y for any action token y ∈ V . This estimates the Q -value function by taking the current t ∈ { 0 , 1 , . . . , T -1 } hidden step of a policy network h π θ t ∈ R H (for next token prediction) to combine with the feature vector of the token v off(on) y ∈ R H with a linear transformation by w off(on) y ∈ R H for off(on)-policy learning. Here, H represents the hidden size and the additional parameter cost is O ( H |V| ) for Q -value estimation. Finally, combining off- and on-policy objectives with a factor β ∈ (0 , 1) , the optimization problem can be represented as follows,

where L ( π θ , f ϕ 1 , f ϕ 2 ) represents the overall training objective. To minimize the objective w.r.t the policy parameters θ , we use a policy gradient approach and derive the policy gradient in Appendix A.5, formally represented as follows,

where ˆ Q f ϕ 2 : Y × V → R denotes the empirical Q -value defined in Eq. (21). Besides, we use stochastic gradient ascent (SGA) to maximize the objective of L ( π θ , f ϕ 1 , f ϕ 2 ) w.r.t. parameters of the on-policy Q -value function ϕ 1 and parameters of the off-policy Q -value function ϕ 2 .

Training procedure. The goal is to achieve an equilibrium between minimizing the objective w.r.t. the parameters of student policy θ ∈ Θ and maximizing the objective w.r.t. the parameters of on-policy and off-policy Q -value functions ϕ 1 , ϕ 2 ∈ Φ , formally defined as min θ max ϕ 1 ,ϕ 2 L ( π θ , f ϕ 1 , f ϕ 2 )

(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 D xy . In the training algorithm, we iteratively maximize the objective w.r.t. the parameters of Q -value functions f ϕ 1 , f ϕ 2 and simultaneously minimize the objective w.r.t. the parameters of student policy π θ . In each iteration of policy updating, we first perform N steps of stochastic gradient ascent (SGA) w.r.t. the parameters of Q -value functions ϕ 1 , ϕ 2 . Then, the parameters of student policy θ are updated by stochastic gradient descent (SGD) with the estimated policy gradient with sampling policy gradients.

3.4 Convergence Analysis

We further provide a convergence analysis for the algorithm proposed in §3.3. To deal with the challenges of non-convexity by certain reward structures, the algorithm is expected to obtain an ϵ -accurate stationary point of the policy parameters θ ∗ ∈ Θ , satisfying that E [ ∥∇L ( θ ∗ ) ∥ 2 ] ≤ ϵ . We focus on policy optimization and directly use the optimized off-/on-policy Q -value functions in each outer-loop iteration k ∈ { 0 , 1 , . . . , K -1 } . We denote ϕ 1 ( θ k ) = arg max ϕ 1 L off ( π θ k , f ϕ 1 ) , ϕ 2 ( θ k ) = argmax ϕ 2 L on ( π θ k , f ϕ 2 ) as the inner-loop optimized functions and use L ( θ k ) := L ( π θ k , f ϕ 1 ( θ k ) , f ϕ 2 ( θ k ) ) (def. in Eq. (9)) for simplicity in this section. We start with the following standard assumption [45].

Assumption 1. Suppose that the optimized Q -value functions and the parameterized policy π θ satisfy the following conditions:

(i) The uniformly boundness of off/on-policy Q -value functions optimized by Algorithm 1, i.e., ∥ f ϕ 1 ∥ ∞ , ∥ f ϕ 2 ∥ ∞ ≤ 1 .

(ii) The B -Lipschitzness and the L -smoothness of the parameterized policy, i.e., for any stateaction pair ( y ≤ t , y t +1 ) ∈ Y × V at any time step t ∈ { 0 , 1 , . . . , T -1 } ,

Theorem 2 ( Convergence rate of Algorithm 1 to stationary points). Let { θ k } 1 ≤ k ≤ K be the sequence of parameters of the policy π θ k given by Algorithm 1. Let the learning rate η = 2 B L √ 1 -γ T (1 -γ ) KL L . Under Assumption 1, we have

Proof. See Appendix A.6 for the complete derivation.

Theorem 2 illustrates that the output gradient norm square by Algorithm 1 can converge to a neighborhood around zero with the rate of 1 / √ K . Furthermore, leveraging a sufficient number of training iterations O ( ϵ -2 ) , Algorithm 1 can obtain an ϵ -accurate stationary point. This leads to the following corollary on the computational complexity of the training procedure.

Corollary 1 ( Computational complexity of Algorithm 1). We formalize the policy as a softmax function π θ with a linear transformation: softmax( θ y ≤ t ) for any y ≤ t ∈ R H , where θ ∈ R |V|× H and H denotes the hidden size. Then, to obtain an ϵ -accurate stationary point by Algorithm 1, the complexity of gradient computation is O ( ϵ -2 T |V| H ( N + T + |V| )) .

Proof. See Appendix A.7 for the complete derivation.

Corollary 1 shows that Algorithm 1 has a polynomial computational complexity w.r.t ϵ -2 , N , |V| , H and T , to obtain an ϵ -accurate stationary point for optimizing the training objective in Eq. (9).

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 D xy ; KD [16], which uses KL divergence on the supervised dataset D xy ; SeqKD [20], which applies SFT to the student model with teacher-generated outputs; ImitKD [24], which uses KL divergence on the student-generated outputs; MiniLLM [14], which uses RKL divergence with a policy gradient method; GKD [2], which uses JS divergence with an on-policy method; and DistiLLM [21], 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. [39]. All the reported results are the average across three random seeds.

4.1 Task-Agnostic Distillation

Experimental Setup. Wefollow the previous works [14, 21] for the implementation of the instructionfollowing 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 [36], VicunaEval [6], S-NI [37], and UnNI [17]. Following the previous works [14, 21], we also add the OpenWebText [13] 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 [14]. The evaluation metrics include ROUGE-L [25] and GPT-4 feedback with the same prompts as in [21]. More details on experimental setup refer to Appendix B.

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 [14], and JS divergence [2] shows effectiveness for KD in the instruction-following task. State-of-the-art methods [14, 2, 21] tend to combine the studentgenerated 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 GPT-4 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. [21] to conduct experiments on the SAMSum [12] dataset. For the machine translation tasks, we follow Ko et al. [21] to conduct experiments on the IWSLT'17 (en-de) [5] dataset. For the commonsense reasoning task, we conduct experiments on the StrategyQA dataset [11] with chain-of-thought augmentations

[38]. For all of the task-specific experiments, we use T5-XL [29] as the teacher model and T5-Large/Base/-Small as the student model. For the machine translation experiments, we employ a multilingual pretrained model, mT5 [43], to build the methods. For evaluation, we use ROUGE-L [25], BLEU [27], 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.

Main results. Table 2 displays the performances on three task-specific datasets. Since the original work of MiniLLM [14] 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.

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, offpolicy objectives as well as the joint form. Results on four tasks with a default combination factor β = 0 . 5 are shown in Figure 2. More instruct-following results are available in Appendix C.2 and results with different values of off-/on-policy combination factor are available in Appendix C.5. Compared with the KL divergence, RKL divergence, JS divergence and total variation

KL

RKL

JS

TV

Ours

20

22

24

26

28

30

32

ROUGE-L

on-policy

off-policy

mixed

(a) DollyEval.

KL

RKL

JS

TV

Ours

26

27

28

29

30

31

32

33

34

BLEU

on-policy

off-policy

mixed

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

KL

RKL

JS

TV

Ours

43

44

45

46

47

48

49

50

51

52

ROUGE-L

on-policy

off-policy

mixed

(b) SAMSum.

KL

RKL

JS

TV

Ours

54

56

58

60

62

64

on-policy

off-policy

mixed

(d) StrategyQA.

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.

0

2000

4000

6000

8000

10000

Step

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Loss (train)

training loss

on-policy distance

d

on

off-policy distance

d

off

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1.8

Distance (eval)

(a) Training loss and d on MM , d off MM against training step.

KL

RKL

JS

TV

2.05

1.84

1.78

1.53

Ours

0.75

2.79

3.28

2.62

1.94

1.45

off  on five test sets

2.21

2.07

2.37

2.15

2.42

1.11

2.18

1.68

1.13

0.72

2.02

1.67

1.43

1.23

0.62

2.23

2.27

1.93

1.65

0.93

3.0

2.5

2.0

1.5

1.0

DollyEval

VicunaEval

Average

SelfInst

UnNI

S-NI

4.5

4.0

3.5

3.0

2.5

2.0

1.5

1.0

DollyEval

VicunaEval

Average

SelfInst

UnNI

S-NI

(b) On-policy moment-matching distance d on MM on the test sets.

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 L ( π θ , f ϕ 1 , f ϕ 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 on MM and d off 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 on MM and the off-policy moment-matching distance d off 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 Theorem 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. Besides, the experiments are conducted with limited LLM architectures, such as OpenLLaMA and T5. Therefore, in future work, we aim to enhance the time and memory efficiency of our approach, and evaluate the proposed approach on a wider range of architectures.

KL

RKL

JS

TV

2.65

2.54

2.31

1.98

Ours

0.96

4.87

4.78

4.52

2.34

1.85

on  on five test sets

3.21

3.67

3.87

3.67

3.32

2.11

3.58

3.32

2.13

0.85

d

3.02

2.67

2.43

1.83

0.92

3.48

3.48

3.25

2.32

1.34

d

References

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.

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.

James Bagnell, Sham M Kakade, Jeff Schneider, and Andrew Ng. Policy search by dynamic programming. Advances in Neural Information Processing Systems , 16, 2003.

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.

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.

Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.

Kamil Ciosek. Imitation learning by reinforcement learning. In International Conference on Learning Representations , 2021.

Mike Conover, Matt Hayes, Ankit Mathur, Jianwei Xie, Jun Wan, Sam Shah, Ali Ghodsi, Patrick Wendell, Matei Zaharia, and Reynold Xin. Free dolly: Introducing the world's first truly open instruction-tuned llm, 2023.

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.

Xinyang Geng and Hao Liu. Openllama: An open reproduction of llama. URL: https://github. com/openlm-research/open_llama , 2023.

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.

Bogdan Gliwa, Iwona Mochol, Maciej Biesek, and Aleksander Wawer. Samsum corpus: A human-annotated dialogue dataset for abstractive summarization. In Proceedings of the 2nd Workshop on New Frontiers in Summarization , pages 70-79, 2019.

Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus. http://Skylion007.github.io/OpenWebTextCorpus , 2019.

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.

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.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2015.

Or Honovich, Thomas Scialom, Omer Levy, and Timo Schick. Unnatural instructions: Tuning language models with (almost) no human labor. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 14409-14428, 2023.

Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the Nineteenth International Conference on Machine Learning , pages 267274, 2002.

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.

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.

Jongwoo Ko, Sungnyun Kim, Tianyi Chen, and Se-Young Yun. Distillm: Towards streamlined distillation for large language models. In Forty-first International Conference on Machine Learning , 2024.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

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.

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 , pages 6121-6133, 2020.

Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Proceedings of the Workshop on Text Summarization Branches Out , pages 74-81, 2004.

Richard Yuanzhe Pang and He He. Text generation by learning from demonstrations. In 9th International Conference on Learning Representations , 2021.

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.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI Blog , 1(8):1-24, 2019.

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.

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.

Bharath K Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On integral probability metrics, \ phi-divergences and binary classification. arXiv preprint arXiv:0901.2698 , 2009.

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.

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.

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. https://github.com/tatsu-lab/stanford_alpaca , 2023.

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.

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. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 13484-13508, 2023.

Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganeh Kordi, Amirreza Mirzaei, Atharva Naik, Arjun Ashok, Arut Selvan Dhanasekaran, Anjana Arunkumar, David Stap, et al. Super-naturalinstructions: Generalization via declarative instructions on 1600+ nlp tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing , pages 5085-5109, 2022.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems , 35:24824-24837, 2022.

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.

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.

Taiqiang Wu, Chaofan Tao, Jiahao Wang, Zhe Zhao, and Ngai Wong. Rethinking kullbackleibler divergence in knowledge distillation for large language models. arXiv preprint arXiv:2404.02657 , 2024.

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.

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. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 483-498, 2021.

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.

Kaiqing Zhang, Alec Koppel, Hao Zhu, and Tamer Basar. Global convergence of policy gradient methods to (almost) locally optimal policies. SIAM Journal on Control and Optimization , 58(6):3586-3612, 2020.

A Proofs

A.1 Proof of Proposition 1

Proof. Similar to the proof of Performance Difference Lemma (PDL) [18, 3, 33], we have J ( π ∗ ) -J ( π θ )

where ( i ) follows from Bellman equation and noting that the transition probability T ( ·| y ≤ t , y t +1 ) is deterministic in an auto-regressive text generation problem. This completes the proof.

A.2 Proof of Proposition 2

Proof. Similar to the proof of Proposition 1, we have

A.3 Existing Step-Wise Distribution Distance for Distillation

Definition 4 ( Step-wise distribution distances for distillation [39]). Following Wen et al. [39], we define four groups of well-studied probability distribution distances as follows,

Total variation (TV) distance. The token-level TV distance between the probabilities of teacher policy π ∗ and student policy π θ given the current state y ≤ t can be defined by the ℓ 2 -norm as follows,

Kullback-Leibler (KL) divergence. The token-level KL divergence between the probabilities of teacher policy π ∗ and student policy π θ given the current state y ≤ t can be defined as follows,

Reverse Kullback-Leibler (RKL) divergence. The token-level RKL divergence between the probabilities of teacher policy π ∗ and student policy π θ given the current state boldsymboly ≤ t can be defined as follows,

Jenson-Shannon (JS) divergence. The token-level JS divergence between the probabilities of teacher policy π ∗ and student policy π θ given the current state boldsymboly ≤ t can be defined based on the KL divergence and RKL divergence as follows,

A.4 Proof of Theorem 1

Proof. We first derive an upper bound for the on-policy moment-matching objective of Eq. (3). Set F Q ∗ = { f : ∥ f ∥ ∞ ≤ 1 } , and by the definition of L ( π θ , f ) in Eq. (3), we have

Then, we have

where ( i ) follows from Jensen's inequality, ( ii ) follows from [31] and ( iii ) follows from the definition of TV distance.

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

which completes the proof of Theorem 1.

A.5 Derivation of Policy Gradient in Eq. (10)

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

Based on the definition of L off ( π θ , f ϕ 1 ) in Eq. (2), we have

Then, based on the definition of L on ( π θ , f ϕ 2 ) in Eq. (3), we have

where ( i ) follows from a standard derivation of gradient policy (c.f. [22]). For simplicity, set

as the empirical Q -value given any draw of trajectory τ ∼ π θ | y 0 = x , x ∼ p x in Eq. (20). Coming back to Eq. (18) and combining with Eq. (19) and Eq. (20), we have

Then, using the law of iterated expectations, we obtain the final formulation of policy gradient,

which completes the derivation of policy gradient in Eq. (10).

A.6 Proof of Theorem 2

Lemma 1. Let ˆ ∇L ( θ ) = -β ˆ G off ( τ, θ ) + (1 -β ) ˆ G on ( τ ′ , θ ) denote the empirical policy gradient given any trajectories x ∼ p x , τ ∼ π ∗ | x , τ ′ ∼ π θ | x , where L ( θ ) := L ( π θ , f ϕ 1 , f ϕ 2 ) (def. in Eq. (9)) denote the objective w.r.t. the policy parameters θ given any off-/on-policy Q -value functions f ϕ 1 and f ϕ 2 . Then, under Assumption 1, we have ∥ ˆ ∇L ( θ ) ∥ ≤ B L with

Proof. By triangle inequality, we have for any x ∼ p x , τ ∼ π ∗ | x , τ ′ ∼ π θ | x ,

By the formulation of off-policy gradient ∥ ˆ G off ( τ, θ ) ∥ in Eq. (10) under the condition of optimized off-policy Q -value functions f ϕ 1 by Algorithm 1, we have

By Jensen's inequality, we have

By Assumption 1, we have

Similarly, we can bound the on-policy gradient ∥ ˆ G on ( τ ′ , θ ) ∥ by Jensen's inequality as follows,

Based on the definition of ˆ Q f ϕ 2 ( y ≤ t , y t +1 ) in Eq. (21) and by Jensen's inequality, we have

Then, by Assumption 1 (i) that ∥ f ϕ 2 ∥ ∞ ≤ 1 , we have

Thus, we have with the constant

Proof. Under the definition of policy gradient in Eq. (10), for any θ 1 , θ 2 ∈ Θ , we have

Then, by Jensen's inequality and triangle inequality, we have

Based on the definition of off-policy gradient in Eq. (10) and using Jensen's inequality, we have for any x ∼ p x , τ ∼ π ∗ | x ,

Coming back to the bound of ∥ ˆ ∇L ( θ ) ∥ in Eq. (22), we combine it with Eq. (23) and Eq. (24). Then, we have

which completes the proof of Lemma 1.

Lemma2. Under Assumption 1, the objective function L ( θ ) is L L -smooth such that for any θ, θ ′ ∈ Θ ,

Then, by triangle inequality, we have for any t ∈ { 0 , 1 , . . . , T -1 } ,

By Taylor expansion of π θ ( y | y ≤ t ) , we have that for any t ∈ { 0 , 1 , . . . , T -1 } ,

where ˜ θ is a vector lying between θ 1 and θ 2 , i.e., there exists some λ ∈ [0 , 1] such that ˜ θ = λθ 1 +(1 -λ ) θ 2 . Then, combining with Eq. (27), yields

Then, combining with Eq. (26) yields

In addition, we can first bound I 2 using Jensen's inequality and triangle inequality,

By triangle inequality and the boundess of | ˆ Q f ϕ 2 ( y ≤ t , y t +1 ) | ≤ 2(1 -γ T ) 1 -γ , we further have,

By Taylor expansion of ∏ t -1 t ′ =0 π θ ( y t ′ +1 | y ≤ t ′ ) , we have

̸

where ˜ θ denotes a vector lying between θ 1 and θ 2 , i.e., there exists some λ such that ˜ θ = λθ 1 +(1 -λ ) θ 2 . Coming back to the boundness of I 2 in Eq. (29), we have

where the last inequality follows from the fact that

Then, combining Eq. (25) with the boundness of I 1 in Eq. (28) and the boundness of I 2 in Eq. (30), we obtain the final bound of

Next, we have for any θ, θ ∈ Θ ,

Then, by Eq. (31) and set θ 1 = θ ′ + t ( θ -θ ′ ) and θ 2 = θ ′ , we have

which completes the proof of Lemma 2.

We prove Theorem 2 as follows.

Proof of Theorem 2. Let θ t , θ t +1 , t ∈ { 0 , 1 , . . . , T -1 } be adjacent parameters of policy π θ t , π θ t +1 given by Algorithm 1. Then, using Lemma 2 by setting θ = θ k +1 , θ ′ = θ k for any k ∈ { 0 , 1 , . . . , K -1 } , we have

Following from the updating rule θ k +1 = θ k -η ˆ ∇L ( θ k ) and Lemma 1, we have

Then, we have

We introduce a probability measure space (Ω , F , P ) and then θ k : Ω → Θ , k ∈ { 0 , 1 , . . . , K -1 } can be viewed as a random variable on it. Let { σ ( θ k ) } 0 ≤ k ≤ K -1 denote a sequence of increasing sigma-algebras such that σ ( θ 0 ) ⊂ σ ( θ 1 ) ⊂ · · · σ ( θ K -1 ) ⊂ F , we define the conditional expectation E [ ˆ ∇L ( θ k ) | σ ( θ k )] as

where the second equality follows from the unbiased estimation property in Eq. (10). Then, taking the conditional expectation, we have

Taking total expectation, rearranging the terms and making average on k ∈ { 0 , 1 , . . . , K -1 } , we have

Let η = 2 B L √ 1 -γ T (1 -γ ) KL L and denote L ( θ k ) = L ( π θ k , f ϕ 1 ( θ k ) , f ϕ 2 ( θ k ) ) for simpilicity for any k ∈ { 0 , 1 , . . . , K -1 } , then we have

A.7 Proof of Corollary 1

Proof. Let the convergence rate in Eq. (32) satisfy that 2 B L √ (1 -γ T ) L L (1 -γ ) K ≤ ϵ , then we have

which indicates that when the iteration number of policy updating satisfies K := O ( ϵ -2 ) , it can reach an ϵ -accurate stationary point to optimize the objective in Eq. (9), such that

For simplicity, we define the policy as a softmax function with a linear transformation of y ≤ t ∈ R H with θ ∈ R |V|× H . Formally, for any trajectory τ and any timestep t ∈ { 0 , 1 , . . . , T -1 } , we have the probability of any y ∈ V ,

In the following, we will analyze the computational complexity in each policy updating iteration. First, we find that each inner-loop step of Q -value function updating has a gradient computation complexity of O ( T |V| H ) given the linear formulation of Q -value functions. Accordingly, N inner-loop steps in each policy updating iteration have a computational complexity of O ( NT |V| H ) . Second, for policy gradient computation, since computational complexity of ∇ log π θ ( y | y ≤ T ) is O ( |V| H ) , the computation complexity of policy gradient computation is O ( T |V| H ( T + |V| )) . Overall, the total gradient computational complexity is O ( ϵ -2 T |V| H ( N + T + |V| )) , which completes the proof of Corollary 1.

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 [28] and OpenLLaMA [10]. For the GPT-2 experiments, we use GPT-2 XL 2 with 1.5B parameters to construct the teacher policy and GPT-2 3 with 117M parameters to construct the student policy. For the OpenLLaMA experiments, we use OpenLLaMA-7B 4 with 6.7B parameters to construct the teacher policy and OpenLLaMA-3B 5 with 2.7B parameters to construct the student model.

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. [14]. For the implementation of compared baselines, we use the code by Ko et al. [21] and re-run the results. The optimization protocol for KD training largely follows the previous work [14, 21]. 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 α = 0 . 1 based on the development experiments. We use a default off-/on-policy combination factor β = 0 . 5 for main experiments while exploring other values for analysis. The hyperparameters for training are listed in Table 3.

B.2 Task-Specific Experiments

Base models. For the text summarization and commonsense reasoning experiments, we use T5XL 6 with 2.8B parameters to construct the teacher policy and construct the student policy with T5-Large 7 (770M parameters), T5-Base 8 (220M parameters) and T5-Small 9 (60M parameters). For the machine translation experiments, we use mT5-XL [43] to construct the teacher policy and mT5-Large/-Base/-Small to construct the student policy.

Training details. We initialize the corresponding teacher and student models using 10,000-step-finetuning 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. [21] 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 of the IWSLT'17 (en-de) dataset.

KL

RKL

JS

TV

Ours

20

22

24

26

28

30

32

ROUGE-L

on-policy

off-policy

mixed

(a) DollyEval.

KL

RKL

JS

TV

Ours

16

17

18

19

20

21

22

23

ROUGE-L

on-policy

off-policy

mixed

(b) SelfInst.

KL

RKL

JS

TV

Ours

15

16

17

18

19

20

21

22

23

ROUGE-L

on-policy

off-policy

mixed

(c) VicunaEval.

KL

RKL

JS

TV

Ours

30

32

34

36

38

40

ROUGE-L

on-policy

off-policy

mixed

(d) S-NI.

KL

RKL

JS

TV

Ours

26

28

30

32

34

36

38

40

ROUGE-L

on-policy

off-policy

mixed

(e) UnNI.

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.

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 or off-policy distillation. Besides, jointly using on-policy and off-policy moment-matching further improves the performances and achieves the best results on five instructionfollowing 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.

KL

RKL

JS

TV

Ours

43

44

45

46

47

48

49

50

51

52

ROUGE-L

on-policy

off-policy

mixed

(a) SAMSum.

KL

RKL

JS

TV

Ours

26

27

28

29

30

31

32

33

34

BLEU

on-policy

off-policy

mixed

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

KL

RKL

JS

TV

Ours

54

56

58

60

62

64

on-policy

off-policy

mixed

(c) StrategyQA.

0

2000

4000

6000

8000

10000

Step

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Loss (train)

training loss

on-policy distance

d

on

off-policy distance

d

off

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1.8

Distance (eval)

(a) Instruct-following.

0

2000

4000

6000

8000

10000

Step

0.4

0.5

0.6

0.7

0.8

0.9

1.0

1.1

Loss (train)

training loss

on-policy distance

d

on

off-policy distance

d

off

0.6

0.8

1.0

1.2

1.4

1.6

1.8

2.0

2.2

Distance (eval)

(b) SAMSum.

0

10000

20000

30000

40000

50000

Step

0.20

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

Loss (train)

training loss

on-policy distance

d

on

off-policy distance

d

off

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1.8

Distance (eval)

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

0

250

500

750

1000

1250

1500

1750

2000

Step

0.6

0.7

0.8

0.9

1.0

1.1

1.2

1.3

Loss (train)

training loss

on-policy distance

d

on

off-policy distance

d

off

1.0

1.2

1.4

1.6

1.8

2.0

Distance (eval)

(d) StrategyQA.

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 on MM and the off-policy moment-matching distance d off MM reduce as the number of training steps increases, which shows the effectiveness of our adversarial training approach for moment-matching.

DollyEval

SelfInst

VicunaEval

S-NI

UnNI

Average

KL

RKL

JS

TV

Ours

2.65

4.87

3.21

3.67

3.02

3.48

2.54

4.78

3.87

3.58

2.67

3.48

2.31

4.52

3.67

3.32

2.43

3.25

1.98

2.34

3.32

2.13

1.83

2.32

0.96

1.85

2.11

0.85

0.92

1.34

d

on  on five test sets

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

(a) On-policy moment-matching distance.

KL

RKL

JS

TV

2.05

1.84

1.78

1.53

Ours

0.75

DollyEval

2.79

3.28

2.62

1.94

1.45

off  on five test sets

2.21

2.07

2.37

2.15

2.42

1.11

SelfInst

2.18

1.68

1.13

0.72

2.02

1.67

1.43

1.23

0.62

2.23

2.27

1.93

1.65

0.93

3.0

2.5

2.0

1.5

1.0

VicunaEval

S-NI

UnNI

Average

(b) Off-policy moment-matching distance.

KL

RKL

JS

TV

Ours

Small

Base

Large

2.85

2.58

2.35

2.52

1.67

2.1

2.32

2.18

2.06

1.32

2.02

2.25

1.98

1.78

1.24

d

on  on five test sets

1.4

1.6

1.8

2.0

2.2

2.4

2.6

2.8

KL

RKL

JS

TV

Ours

Small

Base

Large

2.55

2.38

2.25

2.32

1.5

2.22

2.23

2.08

1.85

1.02

1.97

2.17

1.78

1.93

1.23

d

on  on IWSLT 17 (en-de)

1.2

1.4

1.6

1.8

2.0

2.2

2.4

KL

RKL

JS

TV

Ours

Small

Base

Large

2.61

2.31

2.52

2.18

1.87

2.28

2.37

2.33

2.12

1.61

2.42

2.16

2.25

1.87

1.39

d

on  on StrategyQA

1.4

1.6

1.8

2.0

2.2

2.4

2.6

KL

RKL

JS

TV

Ours

Small

Base

Large

2.85

2.78

3.17

2.77

1.82

2.4

2.28

2.37

2.18

1.21

2.21

2.32

2.27

2.02

1.34

d

off  on five test sets

1.25

1.50

1.75

2.00

2.25

2.50

2.75

3.00

KL

RKL

JS

TV

Ours

Small

Base

Large

2.08

2.51

2.14

1.95

1.12

1.65

2.32

1.87

1.23

0.72

1.87

2.1

1.72

1.32

0.68

d

off  on IWSLT 17 (en-de)

0.75

1.00

1.25

1.50

1.75

2.00

2.25

2.50

KL

RKL

JS

TV

Ours

Small

Base

Large

2.52

2.23

2.46

2.24

1.37

1.95

2.18

2.32

2.03

1.28

2.23

1.89

2.13

1.91

1.19

d

off  on StrategyQA

1.2

1.4

1.6

1.8

2.0

2.2

2.4

(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.

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 is more effective in minimizing the moment-matching distance than the distribution-matching methods.

C.5 Analysis on the Off-/On-Policy Combination Factor β

We study the impact of on-policy and off-policy objectives with the combination factor β ∈ { 0 . 00 , 0 . 25 , 0 . 50 , 0 . 75 , 1 . 00 } in Eq. (9), which denotes a linear combination coefficient of the on-policy and off-policy objectives. We observe that if β = 0 . 00 , only the on-policy objective contributes to policy learning. As it increases from 0 to 1, the influence of off-policy objective increases while that of the on-policy objective decreases. Finally, when β = 1 . 00 , only the off-policy objective contributes to policy learning. We conduct experiments across four datasets. Specifically, we evaluate ROUGE-L for OpenLLaMA2-3B on the DollyEval dataset, ROUGE-L for T5-base on the SAMSum dataset, accuracy for T5-base on the IWSLT'17 dataset and accuracy for T5-base on

d

the StrategyQA dataset. Results in Table 6 show that a combination of on-policy and off-policy objectives outperforms using either on-policy or off-policy objectives only across four datasets.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes] .

Justification: The claim of contributions in the abstract and introduction have been fully reflected in the sections of Methods and Experiments.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes] .

Justification: Limitations are discussed in section of the conclusion.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes] .

Justification: Complete proofs for the theoretical results are available in Appendix A. All proofs are based on Assumption 1.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes] .

Justification: Detailed experimental setups such as datasets, models and hyperparameters used in implementing proposed algorithms are all described in detail. See Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes] .

Justification: All the datasets used in this work are publicly available. The code and implementation details are released at this GitHub URL: https://github.com/ jiachenwestlake/MMKD .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes] .

Justification: We provide the details of experimental settings in Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes] .

Justification: All experimental results have error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes] .

Justification: Available in Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes] .

Justification: The research is conducted with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA] .

Justification: Our work mainly focuses on algorithm design and performance improvement, which has no relationship to societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA] .

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes] .

Justification: The paper has cited the original papers that produced the models, code packages or datasets.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA] .

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA] .

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA] .

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nTo improve knowledge distillation for large language models, the authors first motivate an RL-based formulation that aims to minimize the imitation gap while matching on and off-policy moment bounds, and then introducing an adversarial training algorithm that achieves this by posing it as a two-player minimax game. They showcase the efficacy of the approach with instruction-following and task-specific experiments.\n\nSTRENGTHS:\n**S1.** The problem is an important one and, as far as I can tell, the method is novel and well motivated.\n\n**S2.** The experiments are reasonable to establish the efficacy of the method.\n\n**S3.** The paper is well written and clear.\n\nWEAKNESSES:\n**W1.** *The paper is missing a more thorough comparison of the training costs of each of the compared methods.* The authors briefly mention in Section 5 that their method induces a larger computational and memory cost compared to some of the other baselines. This discussion should be fleshed out more and potentially backed by experimental evidence.\n\nMinor comments:\n- It does not have any practical influence, but the definition of $\\mathbf{y}$ in line 101 is not very elegant as $y_0$ will have a different dimensionality than $y_i$ for $i \\in \\\\{1,\\dots,T\\\\}$. It might be cleaner to define $\\mathbf{y} = \\mathbf{x} || (y_1,\\dots,y_T)$, where $||$ is a concatenation operator.\n\nQUESTIONS:\n- Is the comparison between SFT and the other methods fair in terms of compute budget? It might be interesting to see if increasing the number of epochs still leads to a performance improvement, or if it has plateaued.\n- Could there be any practical benefit to having a different $\\alpha$ for on and off-policy optimization in Algorithm 1? The intuition for this comes from the results of Section 4.3 and particularly Figure 2.\n\nLIMITATIONS:\nThe authors discuss the limitations appropriately in Section 5 of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper applies a reinforcement learning (RL) framework to the problem of auto-regressive text generation, framing knowledge distillation as a task of minimizing the imitation gap between teacher and student policies. The authors provide a theoretical analysis demonstrating that the proposed momentum-matching method offers a tighter bound on this imitation gap compared to traditional distribution-matching approaches, potentially leading to improved optimization. To efficiently optimize the momentum-matching target, the paper introduces an adversarial training procedure that alternates updates between the student policy parameters and the Q-value functions, which are used to assess the imitation gaps. Experimental results showcased within the paper indicate that the proposed momentum-matching method outperforms existing distribution-matching baselines in terms of effectiveness.\n\nSTRENGTHS:\n1. The paper is well-structured and presents a clear, enjoyable narrative, facilitating ease of understanding for the reader.\n\n2. Utilizing an RL framework, the paper theoretically demonstrates that the momentum-matching target provides a tighter bound for minimizing the imitation gap compared to conventional distribution-matching targets. Furthermore, an adversarial training procedure is proposed to effectively optimize the momentum-matching target, aiming to approximate a Nash equilibrium between the parameters of the student policy and the Q-value functions.\n\n3. Comprehensive experiments empirically demonstrate that the proposed method outperforms existing distribution-matching methods in performance. Additionally, the presented analysis of training loss curves illustrates the stability of the proposed adversarial training procedure.\n\nWEAKNESSES:\n1. Unlike distribution-matching methods, the RL-based momentum-matching adversarial framework requires significantly more computational resources and runtime due to the necessity of calculating policy gradients and updating the parameters of auxiliary networks involved in Q-value functions. While the authors acknowledge this limitation in the conclusion section, the paper lacks quantitative analysis concerning the computation of policy gradients. It does not detail the resource consumption of the overall procedure. This omission limits the reader’s ability to assess the practical applicability of the method.\n\n2. The presented experiments primarily examines knowledge distillation performance within similar or identical architectural frameworks. However, it does not demonstrate the method's generalizability across models with different architectures, thus leaving the robustness of the approach across diverse settings untested.\n\nQUESTIONS:\n1. Could you provide a detailed comparison of resource consumption and memory costs relative to the baseline methods? Additionally, can you discuss the impact of the number of samples used to estimate the policy gradient on the performance of your method?\n\n2. It seems that the performance of the auxiliary models used to compute Q-value functions significantly impacts the effectiveness of your method. However, the paper does not provide details on the training of these models. I am curious to know whether these models are trained from scratch during the adversarial procedure or if they are pre-trained on certain datasets before inclusion. Is it possible to directly fine-tune teacher models to compute Q-value functions?\n\nLIMITATIONS:\nThe paper briefly discusses the additional computational consumption and GPU memory costs associated with the proposed methods. However, it lacks a more thorough and detailed discussion of other potential limitations that may affect the applicability or scalability of the approach.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel approach to knowledge distillation for Large Language Models (LLMs) using an adversarial training method that incorporates both on and off-policy distillation. The method jointly learns a critic that estimates Q-values while updating both the Q-function and the student model to more closely match the teacher model. The authors employ a policy gradient method to update the student model.\n\nSTRENGTHS:\n+ As far as I am aware, a novel approach to knowledge distillation for LLMs -- although I am not an expert.\n+ A well-presented method, tying together some previous ideas on IRL into the distillation application.\n+ Demonstrates a boost to accuracy.\n\nWEAKNESSES:\n+ Some crucial details are unclear, particularly regarding the parameterization of the Q-value function (see questions section)\n+ Not very much discussion of the computational complexity or additional overhead in memory of having multiple models and requiring rollouts from the teacher and student model while training\n+ Lack of ablation studies -- the method is evaluated as a single monolithic method, when there are many variants that could be applied, such as a weighted combination of the two upper bounds. In particular, I'd like to see how the method using only the on-policy upper bound and the method using only the off-policy upper bound would compare against the method using the linear combination of the on and off policy upper bounds.\n\nQUESTIONS:\n+ How exactly is the Q-value function parameterized? Is it an extra head on the model, or a new model entirely?\n+ Regarding the use of policy gradients for training the student function:\n   a. Did you use a baseline to reduce variance?\n   b. What is the variance of these policy gradients? In applications such as RL, policy gradients typically have quite high variance compared to other methods.\n   c. Have you considered lightweight baseline methods, such as those presented in [1]?\n+ Can you provide an ablation of the different elements of the approach, such as investigating the relative importance of the two upper bounds?\n+ Can you provide an analysis -- even if it is brief -- on the computational cost and memory usage of using the additional Q-value critics while training?\n\n\n\n[1] Ahmadian, Arash, et al. \"Back to basics: Revisiting reinforce style optimization for learning from human feedback in llms.\" arXiv preprint arXiv:2402.14740 (2024).\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed an adversarial moment-matching approach for knowledge distillation of LLM. The idea is to reformulate the knowledge distillation from an imitation learning perspective and derive both on-policy and off-policy bounds for the imitation gap between the teacher and student models. The authors proposed an adversarial training algorithm to estimate and minimize the on-policy and off-policy moment-matching distances. The moment-matching distance is evaluated by the value function and the student is updated using policy gradients to minimize this distance. Experiments on instruction-following and task-specific datasets show that the proposed approach outperforms other knowledge distillation methods.\n\nSTRENGTHS:\nIt is novel to reformulate the knowledge distillation as a moment matching problem, where the matching distance is evaluated by the Q-value function.\n\nThe authors derive the imitation gap bound for both on-policy and off-policy setup, and optimize the gap to achieve the knowledge distillation.\n\nThe proposed method demonstrated good performance on both instruction-following and task-specific datasets. The seven baselines are either distribution matching based or supervised finetuned.\n\nI found the connection between the proposed moment-matching approach and distribution distances matching interesting. Specifically, the authors show that minimizing the total variation distance can achieve sub-optimal results for the moment-matching bounds\n\nWEAKNESSES:\n1. I would recommend having ablation studies and analysis of the impact of on-policy and off-policy objective. And analysis how each of them effect the overall performance.\n\n2. Solving Eq.(9) is involving optimizing the minmax problem. (a) First, the optimization requires additional computational steps for the inner-loop gradient update. How expensive is the computation? such as the time/memory cost (b) Is the optimization robust with respect of hyperparameter changes like $K$ and $\\alpha$.\n\n3. The method requires an auxiliary network for Q-value estimation. It make the training system even more delicated. What network is used for Q-value estimation? any analysis here?\n\nQUESTIONS:\nOverall, I think the idea is novel. However, as authors pointed out in the limitation part, the required time/memory cost/training efforts can assumed to be high.  I would recommend have a comparison here with other distribution matching methods or knowledge distillation methods.\n\nSee the weakness section for my questions.\n\nLIMITATIONS:\nSee above question and weakness sections.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}