Adversarial Moment-Matching Distillation of Large Language Models

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.

Strengths

**S1.** The problem is an important one and, as far as I can tell, the method is novel and well motivated. **S2.** The experiments are reasonable to establish the efficacy of the method. **S3.** The paper is well written and clear.

Weaknesses

**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

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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors discuss the limitations appropriately in Section 5 of the 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.

Strengths

1. The paper is well-structured and presents a clear, enjoyable narrative, facilitating ease of understanding for the reader. 2. 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. 3. 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.

Weaknesses

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

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?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

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.

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.

Strengths

+ As far as I am aware, a novel approach to knowledge distillation for LLMs -- although I am not an expert. + A well-presented method, tying together some previous ideas on IRL into the distillation application. + Demonstrates a boost to accuracy.

Weaknesses

+ 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

+ 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).

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

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.

Strengths

It is novel to reformulate the knowledge distillation as a moment matching problem, where the matching distance is evaluated by the Q-value function. The authors derive the imitation gap bound for both on-policy and off-policy setup, and optimize the gap to achieve the knowledge distillation. The proposed method demonstrated good performance on both instruction-following and task-specific datasets. The seven baselines are either distribution matching based or supervised finetuned. I 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

Weaknesses

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

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.

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

See above question and weakness sections.

Response to Author Rebuttal

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.