Adversarial Moment-Matching Distillation of Large Language Models

Authors: chen  jia
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 12:23:55

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.

Main paper: 
 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 [38] . 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 The comparison between the distribution-matching-based distillation and the action-value moment-matching distillation is outlined. π θ and π * denote the student policy and the teacher policy, respectively. For both on-policy (using student-generated outputs) and off-policy (using teachergenerated outputs) perspectives, our approach optimizes moment-matching of action-value functions (Q-functions) instead of minimizing the distribution distance measured by M = KL, RKL, TV, etc.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. 
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) [40, 41] . 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, 38] 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 [38, 14, 2, 21, 40] . Wen et al. [38] 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 [43, 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 [43] and TextGAIL [39] 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 MethodQuestion: 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. 
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}.We formulate text generation as a sequential decision-making 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 asV π θ (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 ,yt+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 performanceJ(π θ ) = E x∼px E τ ∼π θ |x T -1 t=0 γ t r(y ≤t , y t+1 ) .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:EQUATIONFrom 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:EQUATIONIn 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:EQUATIONIn 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 [38] . At the beginning, we draw a general formulation of the state-of-the-art methods for distillation of LLMs [38, 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,EQUATIONEQUATIONwhere M(•, •) denotes a distribution distance, consisting of total variation (TV) distance [38] 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. [38] 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. MM off(on) (π * (•|y ≤t ), π θ (•|y ≤t ))= E y∼π * (•|y ≤t ) f off(on) * (y ≤t , y) - E y∼π θ (•|y ≤t ) f off(on) * (y ≤t , y) ,Off-policy:EQUATIONwhere 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 haveEQUATIONEQUATIONfor 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.Input: Dataset Dxy with inputs and ground-truth outputs Teacher policy π * ; Student policy π θ with initial parameters θ pretrained on Dxy; 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 for n = 0, 1, 2, . . . , N -1 do Sample an input x ∼ Dx and generate an trajectoryτ off ∼ π * |x ϕ1 ← ϕ1 + αβη∇ ϕ 1 Loff (τ off , θ k , f ϕ 1 ) £ maximize L off (π θ k , f ϕ 1 )in Eq. ( 9) Sample an input x ∼ Dx and generate an trajectoryτ on ∼ π θ |x ϕ2 ← ϕ2 + α(1 -β)η∇ ϕ 2 Lon (τ on , θ k , f ϕ 2 ) £ maximize L on (π θ k , f ϕ 2 ) in Eq. (9) end Sample an input x k ∼ Dx and generate trajectories τ off k ∼ π * |x k and τ on k ∼ π θ |x k θ k+1 ← θ k -η -β Ĝoff (τ off k , θ k ) + (1 -β) Ĝon (τ on k , θ k ) £ minimize L(π θ , f ϕ 1 , f ϕ 2 )in Eq. ( 9) endOptimization 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 asf ϕ 1(2) (y ≤t , y) = (h π θ t + v off(on) y ) ⊤ w off(on) yfor 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,EQUATIONwhere 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,EQUATIONwhere Qf ϕ 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.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 9)) for simplicity in this section. We start with the following standard assumption [44] . Assumption 1. Suppose that the optimized Q-value functions and the parameterized policy π θ satisfy the following conditions:f off * = arg max f L off (θ * , f ) and f on * = arg max f L on (θ * , f ), respectively. We use L(θ) := L(θ, f off * , f on * ) (def. in Eq. ((i) The uniformly boundness of off/on-policy Q-value functions optimized by Algorithm 1, i.e.,∥f off * ∥ ∞ , ∥f on * ∥ ∞ ≤ 1. (ii)The B-Lipschitzness and the L-smoothness of the parameterized policy, i.e., for any stateaction pairEQUATION∥∇ log π θ1 (y t+1 |y ≤t ) -∇ log π θ2 (y t+1 |y ≤t )∥ ≤ L∥θ 1 -θ 2 ∥, for any θ 1 , θ 2 ∈ Θ (12)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 3B 1-γ 3K . Under Assumption 1 and the condition γ ≤ min{1 -ηL 2 , 2 -1 β }, we haveEQUATIONProof. 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 isO(ϵ -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).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. 
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 Table 1 : Comparison with state-of-the-art KD methods on the instruction-following dataset using fine-tuned OpenLLaMA-7B as the teacher and fine-tuned OpenLLaMA-3B as the student. We format the best, the second best and worse than SFT results. The results based on GPT-2 are available in Appendix C.1. 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. [38] . All the reported results are the average across three random seeds.Method DollyEval SelfInst VicunaEval S-NI UnNI GPT-4 R-L GPT-4 R-L GPT-4 R-L R-L R-L OpenLLaMA2-7B (Experimental Setup. We follow 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.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] . 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 [42] , 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. 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 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. 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.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. 
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. 
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). 
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).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. 
 
 Appendix: A.1 Proof of Proposition 1Proof. Similar to the proof of Performance Difference Lemma (PDL) [18, 3, 33] , we haveJ(π * ) -J(π θ ) = E x∼px τ ∼π * |x T -1 t=0 γ t r(y ≤t , y t+1 ) -E x∼px [V π θ (x)] = E x∼px τ ∼π * |x T -1 t=0 γ t r(y ≤t , y t+1 ) + V π θ (y ≤t ) -V π θ (y ≤t ) -E x∼px [V π θ (x)] = E x∼px τ ∼π * |x T -1 t=0 γ t r(y ≤t , y t+1 ) + γV π θ (y ≤t+1 ) -V π θ (y ≤t ) = E x∼px τ ∼π * |x T -1 t=0 γ t r(y ≤t , y t+1 ) + γE y ≤t+1 ∼T (•|y ≤t ,yt+1) V π θ (y ≤t+1 ) -V π θ (y ≤t ) (i) = E x∼px τ ∼π * |x T -1 t=0 γ t Q π θ (y ≤t , y t+1 ) -V π θ (y ≤t ) = E x∼px τ ∼π * |x T -1 t=0 γ t Q π θ (y ≤t , y t+1 ) - E y∼π θ (•|y ≤t ) Q π θ (y ≤t , y) ≤ sup f ∈F Q E x∼px τ ∼π * |x T -1 t=0 γ t f (y ≤t , y t+1 ) - E y∼π θ (•|y ≤t ) f (y ≤t , y) ,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.Proof. Similar to the proof of Proposition 1, we haveJ(π * ) -J(π θ ) = -E x∼px τ ∼π θ |x T -1 t=0 γ t r(y ≤t , y t+1 ) + E x∼px [V π * (x)] = E x∼px τ ∼π θ |x T -1 t=0 γ t V π * (y ≤t ) -r(y ≤t , y t+1 ) + V π * (y ≤t ) + E x∼px [V π * (x)] = E x∼px τ ∼π θ |x T -1 t=0 γ t V π * (y ≤t ) -r(y ≤t , y t+1 ) + γV π * (y ≤t+1 ) = E x∼px τ ∼π θ |x T -1 t=0 γ t V π * (y ≤t ) -r(y ≤t , y t+1 ) + γE y ≤t+1 ∼T (•|y ≤t ,yt+1) V π * (y ≤t+1 ) = E x∼px τ ∼π θ |x T -1 t=0 γ t V π * (y ≤t ) -Q π * (y ≤t , y t+1 ) = E x∼px τ ∼π θ |x T -1 t=0 γ t E y∼π * (•|y ≤t ) Q π * (y ≤t , y) -Q π * (y ≤t+1 , y t+1 ) ≤ sup f ∈F Q * E x∼px τ ∼π θ |x T -1 t=0 γ t E y∼π * (•|y ≤t )f (y ≤t , y) -f (y ≤t , y t+1 )Step-Wise Distribution Distance for KD Definition 4 (Step-wise distribution distances for distillation [38] ). Following Wen et al. [38] , we define four groups of well-studied probability distribution distances as follows,• Total variation (TV) distance. The token-level TV distance can be defined as follows,EQUATION• Kullback-Leibler (KL) divergence. The token-level KL divergence between the teacher policy π * and the student policy π θ can be defined as follows,EQUATION• Reverse Kullback-Leibler (RKL) divergence. The token-level RKL divergence between the teacher policy π * and the student policy π θ can be defined as follows,EQUATION• Jenson-Shannon (JS) divergence. The token-level JS divergence between the teacher policy π * and the student policy π θ can be defined based on the KL divergence and RKL divergence as follows,EQUATIONA.4 Proof of Theorem 1Proof. 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 havesup f :∥f ∥∞≤1 L on (π θ , f ) = sup f :∥f ∥∞≤1 E x∼px τ ∼π θ |x T -1 t=0 γ t E y∼π * (•|y ≤t ) f (y ≤t , y) -f (y ≤t , y) = sup f :∥f ∥∞≤1 E x∼px τ ∼π θ |x T -1 t=0 γ t E y∼π * (•|y ≤t ) f (y ≤t , y) - E y∼π θ (•|y ≤t ) f (y ≤t , y)Then, we havesup f :∥f ∥∞≤1 L on (π θ , f ) (i) ≤ E x∼px τ ∼π θ |x T -1 t=0 γ t sup f :∥f ∥∞≤1 E y∼π * (•|y ≤t ) f (y ≤t , y) - E y∼π θ (•|y ≤t ))f (y ≤t , y)(ii) = E x∼px τ ∼π θ |x   T -1 t=0 γ t y∈V π * (y|y ≤t ) -π θ (y|y ≤t )   (iii) = d on TV (π θ , π * ) (Def. 4),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,sup f :∥f ∥∞≤1 E x∼px τ ∼π * |x L off (π θ , f ) ≤ E x∼px τ ∼π * |x   T -1 t=0 γ t y∈V π * (y|y ≤t ) -π θ (y|y ≤t )   = d off TV (π θ , π * ) (Def. 4),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 haveEQUATIONBased on the definition of L off (π θ , f ϕ1 ) in Eq. ( 2), we have∇L off (π θ , f ϕ1 ) =∇ E x∼px τ ∼π * |x T -1 t=0 γ t f (y ≤t , y t+1 ) - E y∼π θ (•|y ≤t ) f (y ≤t , y) = -E x∼px τ ∼π * |x T -1 t=0 γ t ∇ E y∼π θ (•|y ≤t ) f ϕ1 (y ≤t , y) = -E x∼px τ ∼π * |x T -1 t=0 γ t y∈V π θ (y|y ≤t )∇ log π θ (y|y ≤t )f ϕ1 (y ≤t , y) = -E x∼px τ ∼π * |x T -1 t=0 γ t E y∼π θ (•|y ≤t )∇ log π θ (y|y ≤t )f ϕ1 (y ≤t , y)Then, based on the definition of L on (π θ , f ϕ2 ) in Eq. ( 3), we haveEQUATIONwhere (i) follows from a standard derivation of gradient policy (c.f. [22] ). SetEQUATIONfor any draw of trajectory τ ∼ π θ |y 0 = x, x ∼ p x in Eq. ( 20).Combining Eq. ( 18) with Eq. ( 19) and Eq. ( 20), we have∇L(π θ , f ϕ1 , f ϕ2 ) = -β E x∼px τ ∼π * |x T -1 t=0 γ t E y∼π θ (•|y ≤t )∇ log π θ (y|y ≤t )f ϕ1 (y ≤t , y)+ (1 -β) E x∼px τ ∼π θ |x T -1 t=0 γ t ∇ log π θ (y t+1 |y ≤t ) Qf ϕ 2 (y ≤t , y t+1 )Supposing that τ ∼ π * |x and τ ′ ∼ π θ |x are independent, then by rearranging the terms, we obtain the final formulation of policy gradient,∇L(π θ , f ϕ1 , f ϕ2 ) = E x∼px E (τ,τ ′ )∼ π * |x•π θ |x -β T -1 t=0 γ t E y∼π θ (•|y ≤t )∇ log π θ (y|y ≤t )f ϕ1 (y ≤t , y)+ (1 -β) T -1 t=0 γ t ∇ log π θ (y ′ t+1 |y ′ ≤t ) Qf ϕ 2 (y ′ ≤t , y ′ t+1 ) ,which completes the derivation of policy gradient in Eq. (10) .A.6 Proof of Theorem 2 Lemma 1. Let ∇L(θ) = -β Ĝoff (τ, θ) + (1 -β) Ĝon (τ ′ , θ) denote the empirical policy gradient given any trajectories x ∼ p x , τ ∼ π * |x, τ ′ ∼ π θ |x. Under Assumption 1, ∇L(θ) is 3B 1-γ -Lipschitz w.r.t. θ.Proof. By triangle inequality, we have for anyEQUATIONBy the formulation of off-policy gradient ∥ Ĝoff (τ, θ)∥ in Eq. ( 10) under the condition of optimized off-policy Q-value functions f off * by Algorithm 1, we haveĜoff (τ, θ) = T -1 t=0 γ t E y∼π θ (•|y ≤t ) ∇ log π θ (y|y ≤t )f off * (y ≤t , y) By Jensen's inequality, we have Ĝoff (τ, θ) ≤ T -1 t=0 γ t E y∼π θ (•|y ≤t )∇ log π θ (y|y ≤t ) f off * (y ≤t , y)EQUATIONSimilarly, we can bound the on-policy gradient ∥ Ĝon (τ ′ , θ)∥ by Jensen's inequality as follows,Ĝon (τ ′ , θ) ≤ T -1 t=0 γ t ∇ log π θ (y t+1 |y ≤t ) Qf on * (y ≤t , y t+1 )Based on the definition of Qf on * (y ≤t , y t+1 ) in Eq. ( 21) and by Jensen's inequality, we haveQf on * (y ≤t , y t+1 ) = T -1 t ′ =t γ t ′ -t E y∼π * (•|y ≤t ′ ) f on * (y ≤t ′ , y) -f on * (y ≤t ′ , y t ′ +1 ) ≤ T -1 t ′ =t γ t ′ -t E y∼π * (•|y ≤t ′ ) |f on * (y ≤t ′ , y)| + |f on * (y ≤t ′ , y t ′ +1 )|Then, by Assumption 1 (i) that ∥f on * ∥ ∞ ≤ 1, we haveEQUATIONComing back to the bound of ∥ ∇L(θ)∥ in Eq. ( 22) and combining it with Eq. ( 23) and Eq. ( 24) and under the assumption that γ ≤ 2 -1 β , we have∥ ∇L(θ)∥ ≤ βB 1 -γ + 2(1 -β)B (1 -γ) 2 ≤ 3βB 1 -γ ≤ 3B 1 -γ ,which completes the proof of Lemma 1.Lemma 2. Under Assumption 1, the objective function L(θ) satisfies that for any θ, θ ′ ∈ Θ,L(θ) ≤ L(θ ′ ) + ⟨∇L(θ ′ ), θ -θ ′ ⟩ + 6B∥θ -θ ′ ∥ + 1 2 L∥θ -θ ′ ∥ 2 1 -γ Proof.Under the definition of policy gradient in Eq. ( 10), for anyθ 1 , θ 2 ∈ Θ, supposing that τ ∼ π * |x, τ ′ 1 ∼ π θ1 |x and τ ′ 2 ∼ π θ2 |x are independent, we have ∥∇L(θ 1 ) -∇L(θ 2 )∥ = E x∼px E (τ,τ ′ 1 ,τ ′ 2 )∼ π * •π θ 1 •π θ 2 |x -β Ĝoff (τ, θ 1 ) -Ĝoff (τ, θ 2 ) + (1 -β) Ĝon (τ ′ 1 , θ 1 ) -Ĝon (τ ′ 2 , θ 2 )Then, by Jensen's inequality and triangle inequality, we haveEQUATIONBased on the definition of off-policy gradient in Eq. ( 10) and using Jensen's inequality, we have for anyEQUATIONThen, by triangle inequality, we have for any t ∈ {0, 1, . . . , T -1}, Then, following from Assumption 1 and the fact that |π θ1 (y|y ≤t ) -π θ2 (y|y ≤t )| ≤ π θ1 (y|y ≤t ) + π θ2 (y|y ≤t ), we haveE y∼π θ 1 (•|y ≤t ) ∇ log π θ1 (y|y ≤t )f off * (y ≤t , y) - E y∼π θ 2 (•|y ≤t ) ∇ log π θ2 (y|y ≤t )f off * (y ≤t , y)(I) ≤ B y∈V π θ1 (y|y ≤t ) + π θ2 (y|y ≤t ) = 2BNext, by Jensen's inequality and following from Assumption 1, we have(II) ≤ y∈V π θ2 (y|y ≤t ) ∇ log π θ1 (y|y ≤t ) -∇ log π θ2 (y|y ≤t ) f off * (y ≤t , y) ≤L∥θ 1 -θ 2 ∥Then, coming back to Eq. ( 26), we haveEQUATIONSimilarly, we have for anyEQUATIONComing back to Eq. ( 25) and combining Eq. ( 27) and Eq. ( 28) and assume that γ ≤ 2 -1 β , we haveEQUATIONNext, we have for any θ, θ ∈ Θ,L(θ) -L(θ ′ ) -⟨∇L(θ ′ ), θ -θ ′ ⟩ ≤ |L(θ) -L(θ ′ ) -⟨∇L(θ ′ ), θ -θ ′ ⟩| ≤ t∈(0,1) ⟨∇L(θ ′ + t(θ -θ ′ )), θ -θ ′ ⟩dt -⟨∇L(θ ′ ), θ -θ ′ ⟩ ≤ t∈(0,1) ∥∇L(θ ′ + t(θ -θ ′ )) -∇L(θ ′ )∥∥θ -θ ′ ∥dtThen, by Eq. ( 29) and setθ 1 = θ ′ + t(θ -θ ′ ) and θ 2 = θ ′ , we have L(θ) -L(θ ′ ) -⟨∇L(θ ′ ), θ -θ ′ ⟩ ≤ t∈(0,1) 6B∥θ -θ ′ ∥ + Lt∥θ -θ ′ ∥ 2 1 -γ dt = 6B∥θ -θ ′ ∥ + 1 2 L∥θ -θ ′ ∥ 2 1 -γ ,which completes the proof of Lemma 2.We prove Theorem 2 as follows.Proof. 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 haveL(θ k+1 ) ≤ L(θ k ) + ⟨∇L(θ k ), θ k+1 -θ k ⟩ + 6B∥θ k+1 -θ k ∥ + 1 2 L∥θ k+1 -θ k ∥ 2 1 -γFollowing from the updating role θ k+1 = θ k -η ∇L(θ k ) and Lemma 1, we have∥θ k+1 -θ k ∥ ≤ 3ηB 1 -γWe further assume that γ ≤ 1 -ηL 2 , and thus6B∥θ k+1 -θ k ∥ + 1 2 L∥θ k+1 -θ k ∥ 2 1 -γ ≤ 27ηB 2 (1 -γ) 2Then, we haveL(θ k+1 ) ≤ L(θ k ) -⟨∇L(θ k ), η ∇L(θ k )⟩ + 27ηB 2 (1 -γ) 2Taking expectations on both sides, we haveE[L(θ k+1 )] ≤ E[L(θ k )] -ηE[⟨∇L(θ k ), E x∼px E (τ,τ ′ )∼π * |x•π θ |x [ ∇L(θ k )]⟩] + 27ηB 2 (1 -γ) 2Rearranging the terms, and making average on k ∈ {0, 1, . . . , K -1}, we have1 K K-1 k=0 E ∥∇L(θ k )∥ 2 ≤ 1 Kη K-1 k=0 (E[L(θ k )] -E[L(θ k+1 )]) + 27ηB 2 (1 -γ) 2 = 1 Kη (E[L(θ 0 )] -E[L(θ K )]) + 27ηB 2 (1 -γ) 2 ≤ 4 Kη(1 -γ) + 27ηB 2 (1 -γ) 2 ,where the second inequality follows from the fact that L(θ) ≤ 2 1-γ for any θ ∈ Θ under Assumption 1 (i).EQUATIONwhich completes the proof of Theorem 2.Proof. Let the convergence rate in Eq. ( 30) satisfy that 6B 1-γ 3(1-γ)K ≤ ϵ, then we haveK ≥ 108B 2 (1 -γ) 3 ϵ 2 ,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 thatmin 0≤k≤K-1 E ∥∇L(θ k )∥ 2 ≤ ϵ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, π θ (y|y ≤t ) = exp(θ y y ≤t )EQUATIONIn 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(N T |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.We use NVIDIA A40 GPUs with 40GB RAM to conduct all the 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 3with 117M parameters to construct the student policy. For the OpenLLaMA experiments, we use OpenLLaMA-7B4 with 6.7B parameters to construct the teacher policy and OpenLLaMA-3B5 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 .Base models. For the text summarization and commonsense reasoning experiments, we use T5-XL6 with 2.8B parameters to construct the teacher policy and construct the student policy with T5-Large7 (770M parameters), T5-Base8 (220M parameters) and T5-Small9 (60M parameters). {0.90, 0.95, 0.99} {0.90, 0.95, 0.99} Combination Factor (β) {0, 0.5, 0.9, 0.99, 1.0} {0, 0.5, 0.9, 0.99, 1.0} Learning Rate (η) {5e -5 , 1e -4 , 5e -4 } {5e -6 , 1e -5 , 5e For the machine translation experiments, we use mT5-XL [42] 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. {0.90, 0.95, 0.99} {0.90, 0.95, 0.99} {0.90, 0.95, 0.99} Combination Factor (β) {0, 0.5, 0.9, 0.99, 1.0} {0, 0.5, 0.9, 0.99, 1.0} {0, 0.5, 0.9, 0.99, 1.0} Learning Rate (η) {5e -5 , 1e -4 , 5e -4 } {1e -4 , 5e -4 , 1e -3 } {1e -4 , 5e -4 , 1e Table 5 : Comparison with state-of-the-art KD methods on the instruction-following dataset using fine-tuned GPT-2 XL (1.5B) as the teacher model and fine-tuned GPT-2 (0.1B) as the student model. We format the best, the second best and worse than SFT results. Method DollyEval SelfInst VicunaEval S-NI UnNI GPT-4 R-L GPT-4 R-L GPT-4 R-L R-L R-L GPT-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.Step-Wise Distribution Distance 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 Results show that the proposed adversarial training algorithm is more effective in minimizing the moment-matching distance than the distribution-matching methods.We study the impact of on-policy and off-policy objectives with the combination factor β ∈ {0.00, 0.50, 0.90, 0.99, 1.00} in Eq. ( 9), which denotes a linear combination coefficient of the on-policy and off-policy objectives. We observe that if β = 0, 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, 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 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 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.• 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.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.In text generation, the state-transition is commonly assumed to be deterministic[43,26], i.e., T (y ≤t+1 |y ≤t , yt+1) = 1.https://huggingface.co/openai-community/gpt2-xlhttps://huggingface.co/openai-community/gpt2https://huggingface.co/openlm-research/open_llama_7bhttps://huggingface.co/openlm-research/open_llama_7bhttps://huggingface.co/google/t5-v1_1-xlhttps://huggingface.co/google/t5-v1_1-largehttps://huggingface.co/google/t5-v1_1-basehttps://huggingface.co/google/t5-v1_1-small 


Decision: This paper presents a RL framework to the problem of auto-regressive text generation, framing knowledge distillation of LLMs as a task of minimizing the imitation gap between teacher and student policies. After author rebuttal, it received scores of 5677. All the reviewers are generally happy about the paper, agreeing that (1) the proposed method is novel, (2) the paper is well written, and (3) the experiments are reasonable to establish the efficacy of the method.

Reviewers had also shown some common concerns, such as clarity regarding the parameterization of the Q-value function, and lack of quantitative analysis concerning the computation of policy gradients. The authors have done a good job in addressing these concerns during rebuttal. Overall, the AC would like to recommend accept of the paper.

Reviews: review: paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.

2. 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$.

3. 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?
questions: questions: Overall, 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.

See the weakness section for my questions.
limitations: limitations: See above question and weakness sections.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: + Some crucial details are unclear, particularly regarding the parameterization of the Q-value function (see questions section)
+ 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
+ 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.
questions: questions: + How exactly is the Q-value function parameterized? Is it an extra head on the model, or a new model entirely?
+ Regarding the use of policy gradients for training the student function:
   a. Did you use a baseline to reduce variance?
   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.
   c. Have you considered lightweight baseline methods, such as those presented in [1]?
+ Can you provide an ablation of the different elements of the approach, such as investigating the relative importance of the two upper bounds?
+ 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?



[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).
limitations: limitations: Yes
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.

2. 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.
questions: questions: 1. 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?

2. 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?
limitations: limitations: The 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.
paper_summary: paper_summary: To 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.
strength_weakness: strength_weakness: **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.

Minor comments:
- 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.
questions: questions: - 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.
- 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.
limitations: limitations: The authors discuss the limitations appropriately in Section 5 of the paper.

Comments: title: Response to Author Rebuttal
comment: I thank the authors for the additional results and appreciate their effort in the rebuttal. I believe these only serve to improve the strength of the paper (particularly the ones on the training costs), and will maintain my acceptance score.
Thank you for your thorough rebuttal.
Thanks for addressing all of my queries -- I know that probably took a lot of work, given the short timeframe.
Given that all the additional results raise no additional problems, and the remarks on the computational complexity illustrate that the overhead is relatively small, I will raise my score to 7.