Improving Language Model Distillation through Hidden State Matching

1. "although the authors conduct experiments...."

We added the F1 scores on QQP and accuracy for QNLI in the revised version. If we have to add the results for superGLUE, we will exceed the page limit.

2. "the teacher models used in this paper are relatively small...."

We go up to the 3B version of Flan-T5. The 7B Llama or similar decoder models are pretrained with 1T or more tokens. To create a student through pretrained distillation from scratch for such models, the student must be distilled on at least billions of tokens. For example, the recent work here uses around 30B tokens https://arxiv.org/pdf/2410.17215

Compare that to our experiments, where we use around 100-125M tokens for the pretraining distillation BERT, mBART, or T5. It is easier to create baselines to show improvements on CKA for such a scale of experiments. Distilling models trained on 1T or more tokens poses a significant scaling challenge. It is impossible to train models on a new methodology and then show the improvement by creating baselines when each model has to be trained on billions of tokens. We must stick to models needing less pretraining to show methodological improvement.

3. "The authors appear to primarily focus on method......"

We mostly focus on hidden state matching here, which can improve KD regardless of the divergence used. We also derived some results on the BART 6x 640 student models using reverse KL divergence, which has recently been used in Agarwal et al ICLR 2024. KD with Reverse KL Divergence shows a similar improvement when using CKA loss across the hidden states.

For translation on the mBART 12x 384 model, the BLUE scores using Reverse KL

Thanks for the encouraging review. We performed experiments on Flan-T5 with 3B parameters (so-called “XL” size), a larger encoder-decoder. It can be seen as a larger version of mBART/ BART.

The decoder models are pre-trained with 1T or more tokens, and to distill them, we need to train the students on at least a billion tokens. For example, the recent work uses around 30B tokens (https://arxiv.org/pdf/2410.17215). Compare that to our experiments, where we use around 100-125M tokens for the pretraining distillation BERT, mBART, or T5. Creating multiple baselines to show improvements in CKA is easier for such a scale of experiments.

The linear baseline is still the SOTA, and has recently been used in https://arxiv.org/html/2407.14679v1 (Neurips '24)

1. "There is limited mention of related work .... currently of significant interest to readers."

We did cite Wen et al n line 467, “or Wen et al. (2023), which replaces the KL Divergence with Jensen-Shannon Divergence.” We are not distilling causal language models like (Gu et al.). Nonetheless, we have added a cite to this paper in the revision (line 470). However, the works on causal modeling (arXiv:2402.13116) would be of limited utility, given the different settings and the page limit.

2. ``The tasks addressed in this study .... encoder-decoder models."

We respectfully disagree. In summarization, works on CNN and XSUM are still benchmarked against BART, e.g. Ryu et al 2024.acl-long.319. BARTScore is widely used to evaluate the quality of the summarization by comparing the summary produced by a model with that of BART. Further, modern decoder models, such as Llama-2 or Mistral are much worse in summarization than BART. See Table 5 (page 2778) of ravaut-etal-2024-context (https://aclanthology.org/2024.acl-long.153). Llama-2-13B produces Rouge-2 scores of 14.10 and 8.61, respectively, for CNN and XSum, compared to 21.0 and 21.8 for BART. The numbers for Mistral-7B or Llama-2-7B are similar to Llama-2-7B, and are much worse than BART.  A weaker teacher will lead to worse performance on KD, and we will struggle to identify whether it is because of the lack of teacher supervision or a drawback of our algorithm. Our distilled BART student of 80M parameters produces a Rouge-2 > 15 and a Rouge-L> 26 for both CNN and XSum. Please let us know if you find a causal model with less than 100M parameters producing a similar Rouge score.

Further, during the inference in a causal model, it runs a forward call on the entire prompt to generate every token. In an encoder-decoder model, the encoder runs a forward call on the prompt only once. The decoder runs a forward call only on the output (e.g., the summary) for every token. Since the prompt for summarization is usually very long because it contains the entire document, it is prohibitively slow to generate a summary in a causal model. Compared to that, an encoder-decoder model calls the repeated forward passes only on the summary during inference. The most significant purpose of KD is to improve the inference time. While we agree that LLMs are a hot topic, we do not see why distilling LLMs, especially for summarization. It would have a much higher inference time than Enc-Dec students of the same size.

Similarly, instruction tuning is possible for any model with generative capabilities, and Flan-T5 performs rather well compared to its Decoder counterparts of a similar size. The MMLU score for Flan-T5-3B is 49.3, and the score on BBH is 41.0. They both are higher than those of Llama2-7B (MMLU 45.3, BBH 32.7). Overall, while LLMs are making significant advances in many tasks, this is only sometimes the case; there is still a role for encoder-decoder (and encoder-only) architectures. Further, LLMs pose significant risks, passed on during KD to the students, and rejecting a paper demanding that the LLMs must be distilled, ignoring the risks, seems a little overstretched. 

3. "While papers with....for representation learning."

Our paper's main contribution is not theoretical but methodological: CKA loss improves KD for sequence-level classification and generative tasks such as summarization and translation. We submitted our paper in the "Application" track. We focused on particular NLP tasks here and used standard SOTA teachers (as argued above). KD is significantly more complex than low-rank fine-tuning or even quantizing LLMs, and relatively little work standardizes even the Vanilla KD (w/o hidden states) for decoder models/LLMs. We cannot apply our methodology to show the improvement in such a non-standard setting. We create our model from scratch, whereas works like MiniLLM use already pre-trained models.

The summarization datasets are standard and have recently been used in ravaut-etal-2024-context, ACL 2024. We have already explained how significant BART is for summarization. The mBART models are still in use, e.g. Liang et al in 2024.acl-long.588.pdf. The same goes for any work on BERT, which is still benchmarked using the GLUE dataset, such as Kim et al. 2024.emnlp-main.907.  Overall, we do not agree that these are all outdated or toy experiments

4. Q1:

Our approach maximizes the correlation between the hidden states of the respective layers of the student and the teacher. The hidden states carry a lot of information, as explained in Azaria et al., 2023, findings-emnlp.68. So, it is natural that incorporating any correlation between them will improve performance.

5. Q2:

CKA and CCA are correlation measures between a pair of multivariate random variables. This is vastly different from regularization techniques such as weight decay.

Thanks for the response and the clarification; highly appreciated !!! It would also be appreciated if you could adjust the ratings (3 before) according to your withdrawal of the criticism of weakness 2. The scores of reviewer n8kN and J69D corroborate their criticism of our paper. Reviewer n8kN also asked for specific results, which we added, e.g., the two classification tasks (QNLI and QQP) on BERT in the main paper and some results on reverse KL divergence for BART/mBART. Reviewer n8kN did not hold the lack of results on LLM as a reason for rejection. This review is the odd one out, suggesting a rejection. Also, you mentioned that you have a clearer understanding of Encoder/Enc-Dec models from the discussion on the paper, something you probably lacked before. It is unfair to the other reviewers if you assign a high confidence score of 4, since both are more familiar with BERT/BART/mBART, and their distillation.

We used standard baselines from the existing literature. The first baseline is Vanilla KD without any hidden state. The second one, i.e., the linear baseline, is the SOTA for LM distillation, which has recently been used in https://arxiv.org/html/2407.14679v1 (Neurips '24). We are unaware of further advancements in the literature that match dissimilar hidden states. We explained to reviewer n8kN that distilling the modern causal models requires around a billion tokens, and it is impossible to benchmark on various architectures when every experiment takes a billion tokens. We typically train for 100-125M for BERT, mBART, or T5, which is adequate to show the efficacy of our method. We are from academia and lack the infrastructure to distill models on billions of tokens.

Regarding Q2, we aim to establish a correlation between the teacher and student representation. But there might be more than one interpretation of it. The paper you mentioned uses HSIC regularization to increase adversarial robustness, and our scope is limited to knowledge distillation. Going by the argument, Every KD work constrains the student model's function by constraining it to the teacher's function space. However, that does not mean we should expand the KD experiments into regularization space; neither does any KD work on LM in the literature.