Improving Language Model Distillation through Hidden State Matching
Authors: Sayantan Dasgupta
Trevor Cohn
Venue: ICLR 2025 Poster (2025)
Keywords: Knowledge Distillation
Centered Kernel Alignment
BART
mBART
T5
Primary area: applications to computer vision, audio, language, and other modalities

Abstract: Hidden State Matching is shown to improve knowledge distillation of language models by encouraging similarity between a student and its teacher's hidden states, as demonstrated by DistilBERT and its successors. This typically uses a cosine loss, which restricts the dimensionality of the student to the teacher's, severely limiting the compression ratio. We present an alternative technique using Centered Kernel Alignment (CKA) to match hidden states of different dimensionality, allowing for smaller students and higher compression ratios. We show the efficacy of our method using encoder--decoder (BART, mBART \& T5) and encoder-only (BERT) architectures across a range of tasks from classification to summarization and translation. Our technique is competitive with the current state-of-the-art distillation methods at comparable compression rates. It requires no pretrained student models, but rather can synthesize new student models from scratch through pretraining distillation. It can scale to students smaller than the current methods, is no slower in training and inference, and is considerably more flexible.

Published as a conference paper at ICLR 2025
IMPROVING LANGUAGE MODEL THROUGH HIDDEN STATE MATCHING
DISTILLATION
Sayantan Dasgupta Computing & Information Science University of Melbourne Melbourne, VIC, Australia fayandg@ure umi ch. edu
Trevor Cohn
Computing & Information Science
University of Melbourne
Melbourne, VIC, Australia trevor. crevor.cchn@uriimelle edu. edu
ABSTRACT
Hidden State Matching shown to improve knowledge distillation of language models by encouraging similarity between student and its teacher's hidden states, as demonstrated by DistilBERT and its successors. This typically uses cosine loss, which restricts the dimensionality of the student the teacher's, severely limiting the compression ratio. We present an alternative technique using Centered Kernel Alignment (CKA) to match hidden states of different dimensionality, allowing for smaller students and higher compression ratios. We show the efficacy method using mcoodr-ddcodee (BERT) architectures across a range of tasks from classification to summarization and translation. Our technique is competitive with the current state-of-the-an distillation methods at comparable compression rates and does not require already pretrained student models. It can scale students smaller than the current methods, is no slower in training and inference, and is considerably more flexible. The Code is available on githut"
(BART, mBART & T5) and encoder-only
INTRODUCTION
Modern LLM sizes have increased dramatically over the past few years, alongside their computational requirements. This gives rise to the need for knowledge distillation (KD) of language models with high compression ratio, in order to produce small, fast models for inference that capture the key capabilities of learning foundation models. An x D transformer with L. layers and D) hidden states usually has fully connected modules of dimension D x ODD), leading computational cost of O(D2) for every layer. With slight abuse of notation, the memory required for the inference of transformer is O(LD²), motivating the need for streamlined models with smaller D) for downstream inference on encource-constrinnee devices. Xue et al. (2023) demonstrated that deeper and narrower architectures typically yield the best performance for encoder-only models. Since the encoder plays pivotal role in Encoder-Decoode models, it provides the motivation to reduce the hidden state dimension of the teacher during compression rather than reducing only the number of layers.
Existing distillation methods use cosine loss between the hidden states, such as in DistilBERT (Sanh) 2019) 2019)) or SSleifer Rush (2020) on BART and mBART. This limits their application to students with the same hidden state dimensionality as the teacher, severely restricting the compression ratio. An exception is Jiao et al. (220)), which handles students with smaller dimensions using a linear projection to match the student and teacher's 's hidden states. This practice remains oftte-of-the-rr art and has recently been employed in Muralidharan et al. (2024). Our work aims to distill students with smaller dimensions than the teacher with compression ratio typically > 2x using a hidden loss based on Centered Kernel Alignment (CKA - Kornblith et al. (2019)). Existing methods for sequence- KD, such as Shleifer & Rush (2020). are limited to compression ratio of 2x. However, with the size of modern LLMs going into several billions of parameters, distillation with a low compression ratio has minimal impact. Other KD approaches include aligning the student and teacher attention matrices, such as in Wang et al. 2020). We instead start with benchmarks such as
Also at Google Research
https:: thub. com/ wwwm/aayan1//cci/5555555

---

Published as a conference paper at ICLR 2025
Sanh et al. (2019) and Shleifer & Rush (2020), which the hidden layer loss in addition to KL Divergence and Masked or Causal modeling loss, respectively. Any gain from attention matching for methodology will also apply to the existing benchmarks.
The first attempt to solve a similar problem was DeepCCA (Andrew et al.) 2013), which can align the hidden states of the student and the teacher through projection. However, DeepCCA computationally expensive and difficult to scale when one of the dimensions is high. Instead, we use CKA to match the student and teacher's hidden states of different dimensionality and formulate a stochastic loss that can be scaled across mini-batches. This enables us to create streamlined student models with lower hidden state dimensions, which gives competitive results even from random initialization. In contrast. (aanh et al. (2019) and Shleifer & Rush (2020) achieve performance benefits by initializing the student layers with the teacher's 's weights, which is impossible when the student dimension is smaller.
LOGIT
LM HEAD
LOGIT
CKA
LM HEAD
CKA
CKA
dg
Embedding
Embedding
STUDENT
We show that CKA is also effective in pretraining distil- lation for encoder-decoder models like mBART and T5 for multilingual tasks. Encoder- unique ad advantage over decoder-only models in terms of KD: using the encoder as support, the decoder layers can be pruned to only a handful or even one layer to speed up inference (Shleifer & Rush 2020). However, distilling the encoders in such models requires pretraining on the unsupervised corpus. Existing work such as (Sheiferr Rush (2020) and et al.) (2022) performs end-to-end distillation for machine translation. They retain the teacher's 's entire encoder and distill only the decoder layers. This results in a low compression ratio, with the smallest student being only half the size of the teacher. We demonstrate how pretraining distillation multilingual corpora, utilizing a CKA-based hidden state loss, can eliminate the need to retain the teacher's s encoder.
TEACHER
models offer
Figure 1: CKA loss between the layers of the student and the teacher. The layers with solid color are trainable.
METHODOLOGY
We draw inspiration from Deep CCA in matching the hidden states of pair of neural networks (Andrew et al. 2013).. The algorithm attempts to match the representations of two networks, regardless of their dimensionality. In traditional Deep CCA. both networks are typically trained simultaneously to learn maximally correlated representations across modalities. In our adaptation, on the other hand, we keep the teacher network frozen while training the student network to match its hidden states.
Let us assume that the hidden states of the teacher and the students are hr E Rd and hş E Rds respectively, with dimensions dr and ds, with ds S dr. Let Hs, H E N x be the matrices with the hidden states of all the data points stacked together as rows. Canonical Correlation Analysis (CCA) takes into account the covariance and cross-correlation matrices between the hidden states, Ess H5 Hs. Ert HT H1 and Ers N H. Hs, where Hg Hs PHs and HT = HT pH, are the centered hidden states of the student and the teacher with u N Ei=1 (uz; and pag * ΣN liss as the mean of the teacher and student hidden states for N samples. The goal of CCA is to learn two vectors E Rr and b E Rd that maximize a Exab
ReCa LoTrro Esab'
CCA is usually computed through the Singular Value Decomposition of the matrix / ErS (Andrew et al., 2013). This makes the algorithm computationally very expensive. The scale of our experiments makes backpropagation with SVD unfeasible, requiring more efficient algorithm. We adopt Centered Kernel Alignment (Kornllihh et al. 2019) as an alternative to CCA for hidden state matching. Let define K and L as the kernels between the hidden states of the student and the teacher respectively, as Ki = ls (hs,, hs,) and Lij 1 (hT, hy,) for some kernel functions I: H x H R and similar for 1, with
H. being the hidden states of
2

---

Published as a conference paper at ICLR 2025
the token pairs with index (2,5). Then the Hilbert Schimdt Independent Criteria is defined as HSIC L) tr (TTLL))(NNNN)
with being the centering matrix defined as I I N 11
The authors do not note any improvement in accuracy for non-linear kernel over linear one (Kornblith et al. 2019). A linear kernel is also computationally less expensive, which is important for scaling the algorithm, especially for larger language models. We use linear kernel here, for which the HSIC between the teacher and the student states is
HSIC(Hss = (N (N-
1
Hss |
The Linear CKA between the hidden states of the teacher and the students is defined as,
| 1 CKA(Hs, V /HSIC(HT, HT) HSICCHS, Hs) |Er
HSIC(Hs, Hr)
(2)
be shown that 0 S CKA(HS, Hr) s (Proof in the Appendix). The authors of Kornblith et al. (2019) also show that CKA is invariant to orthogonal transforms and isotropic scaling. If the eigenvectors and eigenvalues of the covariance matrix Ess are ug, and Xs, respectively for
i E [ds] and similar for ETT, then CKA(Hs Hy) can relate to them as CKA(Hs. H7) = As,
(us;, Wry)² KKrnblith et al. 2019). Whereas if Roas is the dg A3, / EE A3;
ury)? Kornblith et al. 2019). It can be observed that CK A(Hs, Hr) turns into a quantity proportional to RecA when we simply replace each of Äs, and λT, with a constant value. In other words, CKA is the weighted sum of the same quantities, (us,, ) as the square of CCA with the weighting coefficient as the product of the normalized cigenvalues of the Gram matrices. From this rationale, we use the square root of CKA as proxy for CCA to match the hidden states of the student and the teacher. /CKA(Hs, Hy), i.e.
estimated value of the CCA, it can be shown that RcCA dg ΣE
The corresponding loss between the hidden states is defined
Pissle
LH
METT [Ess||
2. MINIBATCH CKA
CKA, as defined above, must be computed over the entire dataset. However, it is not feasible to compute it globally over all samples. We can estimate the covariance matrices for every single minibatch, but the sample size can be very low, leading to high variance. We try to include more samples in the estimation process and compute them over mini-batches. If the covariance matrices for minibatch b E [B] are Ers, ETT, & Ess, respectively,
then estimate the CKA from
them as the following, and then compute Lu = VC A(Hs, Hr), where
1 b= ETS,13 IE D= ETT, Fl| EE Ess, 1|
CkA(Hs, HT)
The hidden states of the transformers are accessible after the Layernorm module (Ba et al. 2016). so they have usually already been centered w.r.t. the mean of the batch. We incorporate the distillation loss between the teacher and the student probabilities (Lpist) typically defined in terms of KI. Divergence (Hinton et al. 2014). We finally add causal language modeling (CLM) loss for the student, making the final loss LCLM Dist LH. For pre-training distillation of such LMs, for document with tokens with x, being the one-hot vector for the token f, and Y being the target sequence for some supervised data, the CLM losses for unsupervised and supervised cases defined as,
T
T
(CLM(O)= = Σ
xe log P (#r(
(CLM(x.V) - Σ
ye log P ((Ndleee (5)
1=1
3

---

Published as a conference paper at ICLR 2025
Task
Teacher
#Params
Pre-training
Tpek-fp
Summarization
BART-large (24 1024)
440M
None
CNN, XSum
MT
mBART-large (24 x 1024)
610M
mC4
EN-RO, EN->FR
MT with Prompt
Flan-T5-3B (48 x 2048)
3B
mC4
EN -ES
Classification
BERT-base (12 X 768)
110M
C4
GLUE
Table 1: Details of the pre-training as well as supervised datasets used for different tasks
3 EXPERIMENTS
Here, we describe the experiments for KD with CKA for three different tasks: summarization (BART) in Section 3.1] machine translation (mBART in Section 3.2 and T5 Section B.3). and classification with an encoder-only model (BERT) in Section 3.4 We construct our baseline using linear projection (Lin) match the students' hidden state dimension to that of the teacher's, s, followed by an MSE loss, similar to Jiao al. (2200). Our distillation approach follows two stages,
1. Pretraining distillation of the teacher using an unsupervised corpus (except for BART)
2. Supervised distillation (BARTs & T5) or fine-tuning (BERT) using supervised dataset
The details of the datasets for pretraining and downstream tasks are mentioned in Table for different models following the order which they appear in this section. We keep the temperature at 1 unless mentioned otherwise and do not use hyperparameters to weigh the loss contributions. The experimental details, including learning rate, batch size, and the GPUs used, are discussed in the Appendix.
3.1 DISTILLATION FOR SUMMARIZATION
We begin by distilling BART-large (Lewis al. 2020 for the downstream task of single-document news summarization, using four student architectures as shown in Table We follow the experimental setup of 'Shleifer & Rush 2020),, who perform distillation for summarization on the CNN Daily Mail (Hermann er al. (01 and XSum (Narayan et al. 2018) datasets. For a document x and its summary y. the supervised loss is defined as Equation (5). The other two losses are the KL Divergence and the hidden state loss between the student and the teacher. We measure the performance using Rouge Lin 2004).
We distill students between 6 and 24 layers and hidden states dimensions between 640 768 (see Table We apply the CKA loss over each hidden layer the student, applied against uniformly spaced layers in the teacher to ac- commodate shallower student models. We also distill the same students using linear projection- based loss between the same pair of hidden states, then with no hidden loss Table 2 We do not distill any student with the same hidden state the teacher's (1024), as CKA would be pointless for this case. We also create distil- BART models with 4, and 12 layers, us- ing the same hidden dimensions as the teacher (1024), as Shleifer Rush (2020). These distilBART students are initialized by copying the alternate layers the teacher and distilled with cosine loss between the hidden layers using the same hyperparameters as the other students in Table 2 When we compare the Rouge scores against the student size for these two dissimilar architectures, when trained well, narrower also observed in encoder-only models in (Xue al. 2023)
tmge (CCN)
Boggl
km 2100i
20000
Eesturg p4p
34
150
306
Stadknt Size Ems
(a) ROUGE Scores CNN (b) ROUGE Scores XSUM
Figure 2: ROUGE scores size of the BART students
trained with CKA loss, and distilBART (Shleifer & Rush 2020) of around the same size trained with co- sine loss between the hidden layers and initialized with teacher's weights. CKA produces higher ROUGE scores.
narrower students with CKA loss perform better (Figure (2.. This shows that outperform their wider counterparts; a similar trend is
4

---

Published as a conference paper at ICLR 2025
Model
P(M)
C.R.
R2(CNN)
RL(CNN)
R2(XSum)
RL(XSum)
BART-lage (24 x 1024)
440
1.0x
21.0
30.6
21.8
36.5
KD wo H (6 x 640)
80
5,5x
15.1
25.8
13.5
27 .4
Lin-B (6 640)
80
5.5x
14.8
-03
25.6
0.2
12.7
-0.8
26.7
-07
CKA-B 640)
80
5.5x
16.8
+1.7
26.8
+..
15.0
+1.5
29.2
+1.8
KD we H (6 768)
100
4.4x
16.4
26.8
15.1
29.2
Lin-B (6 768)
100
4.4x
15.5
-0.9
26.2
-0.6
14.1
28.2
-10
CKA-B (6 768)
100
4,4x
17.7
+1.3
27.7
+0,9
16.5
+1.4
31.0
+1.7
KD wo H (12 768)
140
3.1x
17.7
27.7
17.6
32.0
Lin-B (12 x 768)
140
3,1x
17.7
+0.0
27.8
+0.1
17.7
+0.1
32.1
+0.1
CKA-B (12 X 768)
140
3.1x
18.5
+0.88
28.5
-0.8
18.7
+111
33.5
+1.5
KD wo H (24 x 768)
239
1.8%
19.0
29.1
20.3
34.7
Lin-B (24 x 768)
239
1.8x
19.2
+0.2
29.3
+0.2
20.7
+0.4
35.2
+0.5
CKA-B (24 x 768)
239
1.8x
19.5
+0..5
29.6
0.5
21.3
+1.0
35.8
+1.11
Table 2: ROUGE-2 (R2) and ROUGE-L (RL.) scores for different BART students on the CNN and XSUM datasets for KD with CKA. Every BART student has equal number of encoder and deceder layers. "KD wo H" stands for KD without a loss the hidden states, Lin-B for KD with the linear projection-based loss, and CKA-B for CKA loss. All the students are trained with the same hyperparameters. The numbers on the right of every column of Rouge score are the differences from the baseline ("KD wo H"), in green when positive and red when negative. C.R. is the compression ratio
We further experimented with distilling the BART-large into students with smaller hidden dimensions for all the baselines and CKA, as listed in Table [ The linear projection gives benefit to compression ratio of 3x, beyond which it degrades the results. The CKA method, however, improves the performance for every case when we study the ablation with respect to the hidden layer loss, and the margins of improvement increase with the compression ratio. For the highest compression ratios of 5,5x, CKA increases the Rouge score by at least 1.0, while the linear loss fails to improve the result.
3.2 DISTILLATION FOR MACHINE TRANSLATION
Next, we distill a multilingual mBART model (Liu et al. 2020) for machine translation. We choose deep and narrow student architectures with the settings 12 384, 12 x 512, 24 512, and 24 x 640, all having lower dimensions than the teacher (Table 3 As was the case with BART, we only consider students with smaller hidden dimensions than the teacher.
We used multilingual data from mC4 (Xue et al. 2020 for all the languages the teacher mBART model covers (details in Appendix). We used causal modeling loss on the input (Equation (5) and uniformly weighed the loss terms. We used a context size of 512 and trained the students for 25 epochs, each containing 40, 000 text samples of mC4, and computed the sum of CLM loss and KL. divergence on the validation set of mC4 at the end of every epoch (Figure 31). For the 24 x 640 and 24 512 models, we use the CKA loss between every pair of student and teacher hidden states. For 12 x 512 and 12 x 384, we use every alternate layer of the teacher.
The larger models converge faster, while the smaller students take much longer to converge. We plot the sum of the CLM loss and KL divergence in Figure Bh. and exclude the hidden loss since their values are incomparable. The loss converges faster than the KD with linear loss for the largest student (300M). KD with linear loss converges to a higher loss than C for the 173M student, whereas it does not even converge for the smallest student, 122M. We also pretrain a third set of models with no hidden loss to study ablation.
We distill the pre-trained mBART students for the downstream task translation from English to Romanian using the WMT16 dataset (Bojar et al. 2016). We use the supervised loss defined in Equation (3) for the sentence pair (I. y) where x is an English sentence and y is its Romanian translation. Table S presents the BLEU scores for EN RO translation using different student architectures, with the teacher benchmark result taken from Shleifer & Rush (2020). We also train two DiRt models with compression ratio of approximately 2x using cosine loss between the
5

---

Published as a conference paper at ICLR 2025
EN-+RO
EN EN-FR
Model
P(M)
C.R.
wolT
woFl
Lin
CKA
woPT
woll
Lin
CKA
mB-L(24x 1024)
610
1.0x
27.0
40.0
mB (12 384)
122
5.0%
8.9
8.8
8.0
-0.8
18.7|+9.9
26.3
34.5
30.9
-3.6
39.2
+4.7
mB (12 512)
173
3.5x
14.3
19.8
17.9
-19
22.3
+2.5
34.3
37.4
36,6
-0.8
40.2
+2.8
mB (24 512)
217
2.8x
19.5
21.6
21.7
+0.1
24.5
+2.9
37.2
38.6
40.0
+14
41.7
+3.1
mB (24 640)
300
2.0x
23.7
23.6
24.6
+1.0
26.3
+1.7
38.9
40.0
41.2
+1.2
42.3
+2.3
dmB (2x 1024)
287
2
15.5
31.5
dmB (4x 1024)
319
1.9x
21.5
39.3
Table 3: BLEU scores for different mBART student models for EN-RO and EN-FR translation. Every student mBART has equal number of encoder and decoder layers. woPT stands for KD with CKA but without pretraining. wolf for KD with no hidden states loss, Lin for KD with the linear hidden loss, and CKA for KD with CKA loss, all with pretraining distillation on mC4, The distil-mBAA with weights from the teacher layers and distilled using cosine Ioss between the hidden layers using the same hyperparameters the rest of the mB students. The numbers on the right of the Lin and CKA columns the differences from the baseline of KD with no hidden loss (woH), in green when positive and red when negative.
(dmB) students are initialized
hidden layers, as described in Shleifer & Rush (2020). The lowest number parameters a distilBART model can scale to is 287M, while we can easily scale down to smaller students. Smaller students make them accessible to practitioners with limited GPU resources. Furthermore, our CKA students achieve significantly better BLEU scores, even at compression ratio of 2 than the distilBARTs.
We further distill an mBART model fine-tuned for context-aware machine translation from English to French (Sarti al. 2024) on IWSLT2017 (Cettolo et al., 2017) with context comprising up to sentences. The authors also demonstrated that context- awarre tine-tuning enhances translation accuracy even without context, and we utilize their fine-tuned mBART translation without context. The training set used is a combination of million instances randomly sampled (without replacement) from the Frenich-French subset of the WMT14 dataset (Bojar et al.. 2014) and the training samples of the IWSLT dataset (232K), resulting in a total of 2.23 million training samples. The evaluation is performed student (24 x 640) outperforms the teacher at a compression ratio of 2 (Table 3). The performance benefit can be attributed the data augmentation from the WMT14 corpus. sumilar to the case of TinyBERT (Jiao et al., 2020), which also uses data augmentation during distillation and outperforms the teacher BERTT-base for MNLI (Williams et al., 2017) at a compression ratio 2x. Our smallest. student, with 122M parameters, produces BLEU score only 0.8 lower than teacher 5x larger.
teacher for distilling
the test of IWSLT2017 (8.6K). Our largest
When we compare the performance of CKA loss with that based on linear projection, CKA performs far better when we study the ablation w.r.t the hid- den loss. The maximum gain in performance comes at the highest compression ratio. Similar to the case of BART, linear loss degrades performance a high compression ratio of 5 5x. The variation of the BLEU score w.r.t the size of the students is shown in Figure 3 We further calculate the inference time of the distilled students on 40GB A100 GPU, All the CKA students achieve substantially lower infer- ence times (Table than the teacher for EN-RO or better BLEU scores at moderately lower inference
Model
EN-RO 312.4
EN-FR
mB- Large (24 1024)
108.0
mBART (12 x 384)
59.6
52.8
mBART (12 x 512)
61.9
56.5
mBART (24 x 512)
96.3
75.8
mBART (21 x 640)
102.0
81.2
Table 4: Inference time in ms for different mBART students, with the teacher the top
times, e.g., the 24-l students for EN-FR
Another area where our approach differs from Sequential KD used in SSleifer & Rush (020)) is the mancher-generated labels: very expensive to generate labels from the teacher through beam search. For example, it takes over 300 hours the EN-RO dataset of 620K with 80GB NVIDIA A100 GPU, with FLOP count of around 161 PFLOPs. Generating teacher labels for dataset with millions of training samples is extremely challenging, which precludes data augmentation during knowledge distillation (KD). Our pretraining-batt approach requires no teacher decoding but only one expensive pre-training stage on multilingual corpus (mC4), after which it can be fine-tuned for

---

Published as a conference paper at ICLR 2025
20
D4 del JINeay wHrk In XR (1THA) ceo In 3244 (1OS9)))
KL Divagggncc -cu Love
to
K0 [EKAI
11:00
30
80 9T)
(UM)
K00 (m0 F7)
28
Steps
Student Size (in
(a) Distillation Loss (mBART) for CKA and the linear loss (Lin)
(b) BLEU vs. mBART student size for EN-RO translation
(c) BLEU vs. mBART student size for EN-FR translation
Figure 3: Distillation loss for pretraining of mBART using CKA and the linear loss on the validation set of mC4 (left), with the arrows pointing to where grokking happens in the student models. BLEU score vs. parameter size of mBART for EN-RO and EN-FR translation in the next two, respectively. wo PT stands for KD with CKA but without pretraining, wo H for KD with no hidden states loss, Lin for KD with the hidden loss with linear projection, and CKA KD with CKA loss. KD without pretraining distillation (w/o PT) performs the worst
specific translation tasks. For example, we use the same pre-trained models for supervised distillation in EN-RO and EN-FR. The FLOP counts for the pretraining distillation on mC4 are 79, 84, 90, and 99 PFLOPs, respectively, for the students Table 4 in increasing order of size. In contrast, distil-mBART must repeat expensive sacher-bbaed decoding to generate pseudo labels for every task. Our technique is thus more economical and can augment as much data as necessary to improve performance for downstream tasks, as we have done for EN-FR translation.
3.3 DISTILLATION OF ZERO-SHOT MODEL
Instruction-tune tuned language models have become the workhorse of NLP. Here, we demonstrate our technique can be applied to distill Flan-T5-3B (Chung et al. 2024), an instruction tuned encoder- decoder model. The advantage of such models is that they can perform a wide range of tasks with reasonable accuracy without fine-tuning, which can be expensive for 3B model. Most of the KD performed on such models in the literature is based on teacher- geenrrtee labels, as in West et al. (2022). In contrast, perform generic KD on Flan-T5-3B, first by pretraining distillation followed by supervised KD, and skip the expensive step of generating teacher labels.
We first perform pre-training distillation on four student models: 12 768 (145M), x 768 (T5-Base 250M), 24 x 1024 (425M), and 48 x 1024 (TS-Large 780M). We use the same mC4 corpus for pre-training, with context length of 1024. However, since Flan-T5 is primarily trained on English tasks, we sample the English corpus of mC4 with a probability of 0.67 and add 33 other non-English language corpora, each with probability of 0.01 (details the Appendix). We used context size of 1024 and trained the students for 25 epochs, each containing 40, 000 multilingual text samples from mC4 using the unsupervised loss defined in Equation (5). The experiments with CKA loss are similar to those with mBART. However, the baseline with linear projection does not converge, regardless of whether pretraining used. Convergence difficult for CKA loss alone, and the x 768 and
FT
KD wo H
Model
P(M)
C.R.
Wo PT
CKA
Flan-T5-3B (48 x 2048)
2.85B
1.0x
28.0
T5 (12 x 768)
145M
19.7x
22,01-5,2
23.4
25.3
-1.
27.2
T5 (24 768)
250M
11.4x
24.3
-50
26.0
3.3
27.7
-16
29.3
TS (24 x 1024)
425M
6.7x
26.1
4.7
27.9
2.9
29.4
-14
30.8
T5 (48 x 1024)
780M
3.6x
28.0
3.8
29.3
-2.5
30.6
-1.2
31.8
Table 5: BLEU scores for different Flan-T5 15 student models. FT stands for the BLEU of the fine-tuned Flan-T5 models without KD (zero-shot for the teacher), PT stands for KD without pretraining. wo H stands for KD with pretraining but no hidden states loss, and CKA stands for KD with pretraining using CKA loss. The numbers on the right side in the KD columns are the difference from the score for CKA to study ablation
7

---

Published as a conference paper at ICLR 2025
24 X 1024 models converged only after initializing the weights from the converged 24 x 768 and 48 X 1024 models.
We further rum a supervised distillation on the pre-trained students for English-Spaniss translation using the WMT13 corpus (Allauzen et al. 2013) by adding the prompt "Translate from English to Spanish:" in front of every English sentence. We sample 3 million sentence pairs from the WMT13 corpus, which is 14.5 million sentences in size, without replacement for training, and then measure the BLEU score on the test set. In the absence of the linear baseline, we provide the result for KD with only KL Divergence and no hidden loss in Table [ We also fine-tune the student models on the same dataset that we use as a second baseline, and then add a third baseline for students with no pretraining distillation (Table 55. We use the Flan-T5 base (24 x 768) and large (48 x 1024) models from Wolf et al. (2019) for fine-tuning and create the other two models (12 x 768 and 24 x 1024) by removing the alternate layers from them. KD with CKA CKA-based hidden loss gives BLEU score gain of 1.1 (for 780M) to 1.9 (for 145M) over KD with no hidden layer loss.
3.4 DISTILLATION OF ENCODER-ONE ONLY MODEL
We finally apply CKA loss to the task-agnostic distillation of BERT. We discard the masked loss used in Sanh et al. (2019) and perform pure distillation using the combination of only KL. Divergence and the loss on the hidden layer, i.e., Dist + LH. We distill the BERT-base models into student models of several configurations: L x 512 D, chosen to have the same number of parameters as DistilBERT (67M); 8 x 512 slightly smaller than 4- layer DistilBERT (52M); and two smaller models (12 x 384 and 6 384) with reduced intermediate size 1536.
Similar to the case of mBART, we distill the stu- dent first using C4. We replace the cosine loss on the hidden layers of DistilBERT with the CKA loss. We add CKA loss between every pair of hidden states for the 12-layer student, skip every 3rd layer for the 8-layer student, and use every alternate layer for the 6-layer student. We train the model for 30 epochs, with each step involv- ing 320, 000 sample texts from the C4 training set, and compute the KI. Divergence for the C4 validation set at the end of every epoch. The KL. Divergence plots are shown in Figure 4 for the CKA loss compared to the baseline method with a linear projection for various student models. CKA performs better for students of all sizes.
0.1
CEA Li S12 BONE
832
12 an PIPE (DMN
tu 513 (310)
Divergence
8.00
8.08
8.00
20
8.66
10
20
Steps
Steps
(a) KL Divergence for stu-(b) (b) KL Divergence for stu- dents with D = 512 dents with D = 384
Figure 4: Difference in KL Divergence of KD with CKA Loss vs. the baseline of lincar projection (Lin) on the validation set C4 corpus. KD with CKA always results in a lower KL Divergence across all student sizes.
We further fine-tune the distilled student with CKA loss on downstream GLUE tasks, specifically: SST-2 (Socher et al.) (013)) for sentiment classification; MRPC (Dolan & Brockett 2005). QQP and STS-B for paraphrase similarity matching (Conneau & Kiela| 2018); MNLI (Williams et al. 2017).
QNLI (Rajpurkar et al. 2016) and RTE (Wang et al. 2018) for natural language inference; and COLA (Warstadt et al. 2019)) for linguistic acceptability. We report the Matthew correlation coefficient for COLA, the Fl score for MRPC and QQP, Spearman's 's rank correlation for STSB, and accuracy for the remaining datasets. CKA outperforms the linear baseline for all tasks, with the largest difference observed for COLA. We do not repeat the other baselines, as the benefits of pretraining or hidden state matching for BERT distillation are well established in works like Sanh et al. (2019) and Jiao et al. 2020).
The authors of DistilBERT initialize the students by copying the weights of the alternating layers from the teacher into the student model. However, we initialize the student with random weights due to the dimension difference. Our 12 x 512 model is competitive with MiniLM (6 x 768) of equal size (67M) and outperforms 6-layer DERt on almost every task except for SST2, where it is equivalent. outperforms attention
Since Jung et al. (2023) shows that KD with CKA is either competitive matching in MiniLM for the same student architectures, we do not repeat the same experiments. Our 8 x 512 model outperforms the layer DistilBERT and, for MNLI, even the 6 layer DistiIBERT.
8

---

Published as a conference paper at ICLR 2025
Task
P(M)
COLA
MRPC
RTE
MNLI-m
QNLI
QQP
Sumples
8.5K
67.3K
3.7K
2.5K
5.7K
390K
105K
364K
BERT base (12 x 768)
110
1.0%
52.1
93.5
88.9
66.4
87.1
84.6
90.5
71.2
LinBERT 384)
22.0
5.0x
27.0
89.2
80.4
52.7
78.2
80.1
84.9
68.1
CKABERT 384)
22.0
5.0x
29.6
53.8
81.0
68.3
LinBERT (12 384)
33.0
3.3%
41.1
90.2
83.0
58.4
81.7
81.1
85,8
69.2
CKABERT 384)
33.0
3.3x
44.8
91.0
83,9
61.2
82.9
82.0
87.1
69.7
DistiBERT (4 768)
52.2
2.1x
32.8
91.4
82.4
54.1
76.1
78.9
85.2
68.5
LinBERT (8 512)
49.8
2.2%
42.7
90.9
83.8
55.3
82.3
82.0
87.9
69.2
CKABERT x8x 512)
49.8
2.2x
45.3
91.8
86.1
58.5
83.4
83.0
88.5
69.7
DistilBERT x 768)
66.9
1.6x
49.0
92.3
86.9
58.4
81.3
82.6
88.8
69.6
MiniLM (6 x 768)
66.9
1.6x
49.2
92.0
88.4
65.1
85.0
83.0
90.1
69.9
LinBERT (12 512)
66.5
1.6x
46.5
91.4
87.0
61.0
83.3
83.0
89.6
69.6
CKABERT 512)
66.5
1.6x
50.2
92.3
87.8
63.0
84,9
88.5
90.0
70.0
Table 6: Results for different student encoder-only modols on the GLUE test set, with the teacher BERT-base the top. The students CKA and Linear loss (Lin) are distilled with the same hyperparameters. The DistilBERT results are taken from riao et al. 2020). The results of MiniLM are generated using the model from huggingface
4 RELATED WORK
4.1 KNOWLEDGE DISTILLATION OF SEQUENCE BASED LMS
There has been extensive work on KD for downstream classification tasks with BERT, Ture al. (2019) demonstrated that wo-sttge distillation typically yields better results for transformers such as BERT (Devlin al. 2018) or GPT-2 (Radford al.) (219) compared to single-stage distillation on downstream tasks. The first stage involves pretraining distillation a generic, unsupervised corpus, such as Wikipedia or the Open WebText dataset, and the students are then further distilled using supervised datasets for various downstream tasks.
The KD literature on language models can be categorized into two main areas. The first group aims to enhance the pretraining distillation of the initial stage. For example, Ture et al. (2019) uses no loss on the hidden layers, [ann et al. (2019) uses a cosine loss, and Wang et al. (2020) uses layerwise attention matching. Our work falls into this category. The second category uses the pre-trained models and focuses on downstream tasks. This includes [Sun et al. (011). which uses MSE loss on normalized hidden states, and Fu et al. (202T), which uses contrastive hidden state matching. However, both assume that the student's dimension is the same as the teacher's.
summarization, are usually more complicated than classification. Early work (Kim & Rush (016) suggested fine-tuning the students on labels generated by the teacher. Subsequently, Shleiter Divergence of the logits. Other works follow this approach, such as Li et al. (2022), which includes quantization with KD, or Wen al.) (2023). which replaces the KL Divergence with Shasnn-Shannon Divergence and Total Variation Distance. Recently, reinforcement learning has been used to improve divergence, such as on-policy distillation by AAarwal et al. (2024), which utilizes a reverse KI. Divergence. However, irrespective of the divergence between the student and the teacher. Other on-policy distillation work, such (Gu al.) 2024). also adopts a loss based on reinforcement learning, although they generate sequences from a mixture of teacher and student distributions. The generation step is expensive and limits their efficiency. and the largest dataset they contains 15,000 data points. Unlike Agarwal et al. (2024) and (Gu et al. 2024), which use the smaller pre-trained models of Flan-T5 or other LMs starting point, we derive our Flan-T5 students from scratch through pre-training distillation.
Generative downstream tasks, such as machine translation
Shllifee & Rush (2020) combined this with the KL
contribution focuses hidden state matching and will give equal benefits
4.2
CENTERED KERNEL ALIGNMENT
CKA proposed to measure the similarity between different layers of deep networks (Kornblith et al. 2019). However, it has been applied far beyond comparing layers between two similar networks, including measuring similarity between heterogeneous networks (e.g., Vision Transformers

---

Published as a conference paper at ICLR 2025
(Dosovitskiy et al. 2020) and Resnet (Raghu et al. 2021)) and in speech (Ollerenshaw et al. 2022 where its value has been shown to follow CCA closely. Raghu et al. (2019) study CCA and CKA scores of different layers during the inner loop iteration of meta-learning (Finn et al. 2017) and show that the two metrics follow a similar trend. [Saha et al. (2022) uses a similar CKA Loss for feature extraction for image classification
Iman--ma and CIFAR-100.
In NLP, CKA has been used to study the similarity between the intermediate layers of BERT (Sridhar Sarah 2020) and investigate the similarity between the layers of the original and fine-tuned models for BERT-style transformer models Phang et al.] (202T). Recently, Jung et al. (2023) used CKA extract structural features from BERT during distillation. However, unlike our work, they use the standard DistilIBERT (6 x 768) as the student with the same dimension as the teacher and do not reduce the dimension.
5 CONCLUSION
We proposed a novel hidden state matching using Centered Kernel Alignment for language model distillation. We perform experiments wide range of teachers from 110M BERT-base to 3B F5a Based on our experiments, we make the following key observations:
Hidden loss using CKA almost always improves the performance for both summarization and translation. The same does not hold for the linear baseline.
The lincar loss does not work beyond compression ratio of x for the mncodr-ddcoderrs
The generative tasks result in more complex hidden states during the decoding, and the linear projection cannot match hidden states that are too disparate.
Pretraining distillation on a multi- muuti-lingual corpus improves the performance of machine translation even without the hidden layer loss for both mBART and Flan-T5
The higher the complexity of a model, the more significant the performance gap between CKA and the linear baseline. The performance of the linear loss is much closer to that of CKA for simpler models, such as BERT, in classification tasks. However, unlike classifi- cation, the generative tasks for more complex a sequential structure. And there, the linear loss falls short. Flan-T5 is the most complex model we distill with the highest compression ratios, for which it does not converge.
use CLM loss based on
For the smallest BART student (80M), CKA produces at least +1.0 ROUGE score improvements. The linear baseline does not converge for the smallest 122 M student for mBART, whereas it does not converge for any model for Flan-T55. We get a BLEU score improvement of +9.9 for the EN-RO and +4.77 for EN-FR translation for the smallest mBART student (122M). For Flan-T5, the smallest student (145M) produces BLEU score only 0.8 lower than that of a 20x larger teacher.
5.1 WHY MULTI-LINGUAL PRETRAINING WORKS
The largest difference using pretraining with CKA occurs in the distillation of machine translation on mBART and Flan-T5 models. Why would similar method not also work for BART? The key difference between mBART and BART is that BART is trained exclusively on English data. The supervised datasets CNN and XSUM used for summarization are also exclusively in English and contain entire paragraphs as input prompts. As long as the input texts of CNN or XSUM are reasonably representative of all the word representations of BART, the student's encoders will learn the word features. We ran study on pretraining BART students using C4 but did not see a significant benefit in downstream performance.
The downstream translation tasks of mBART contain only a sentence or two of a specific pair of languages, which insufficient for extracting all the teacher's 's word representation features. The encoder plays the most significant role in synthesizing the word representation for mccodr-ddcooe models, while the decoder takes the features for the input sentence from the encoder through the cross-attention. Figure 3 shows the grokking during the pertaining distillation, and the grokked students perform much better than their counterparts without pertaining (Table with the smaller ones being worse. The smaller the student's 's encoder, the lower the capability to learn the teacher's complex word representation features. Other works like Agarwal et al. (2024) start with the smaller Flan-T5 models as the initial students, which are already trained multilingual datasets.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGEMENT
The authors thank Prof. Timothy Baldwin of the University of Melboure/MBBZAII for his valuable feedback on different experiments and the manuscript. This research was supported by The University of Melbourne's Research Computing Services (Spartan) and the Petascale Campus Initiative.
REFERENCES
Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos Garea, Matthieu Geist, and Olivier Bachem. On-policy distillation language models: Learning from self- generated mistakes. In The Twelfth International Conference Leurning Representations, 2024.
Alexander Allauzen, Nicolas Pécheux, Quoc Khanh Do, Marco Dinarelli, Thomas Lavergne, Aurélien Max, Hai-Son Le, and François Yvn.. Limsi @ wmt13. In Proceedings of the Eighth Workshop on Statistical Machine Translation, pp. 62-69 69, 2013.
Galen Andrew, Raman Arora, Jeff Bilmes, and Karen Livescu. Deep canonical correlation analysis. In International conference machine learning, pp. 1247-1255, PMLR, 2013.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.06455, 2016.
Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Kochn, Johannes Leveling. Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Aleš Tamchyna. Findings of the 2014 workshop statistical machine translation. In Proceedings of the Ninth Workshop on Statistical Machine Translation, pp. 12-58, Baltimore, Maryland, USA, June 2014. Association for Computational Linguistics. URL http: //ww/ /ww
ac aclweb. org http://nnhhllgyyy//om ww/p/////ll 4/114333333322
Ondrej Bojar, Rajen Chatterjee, Christian Federmann, Yvette Graham, Barry Haddow, Matthias Huck, Antonio Jimeno Yepes, Philipp Koehn, Varvara Logacheva, Christof Monz, et al. Findings of the 2016 conference on machine translation (wmt16). In First Conference on Machine Translation, pp. 131-198. Association for Computational Linguistics, 2016.
Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Nichues, 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, pp. 2-14, 2017.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddharthaa Brahma, et al. Scaling instruction-finttune language models. Jourmal of Machine Learning Research, 25(70):1-53, 2024.
Alexis Conneau and Douwe Kiela. Senteval: An evaluation toolkit for universal sentence representa- tions. arXiv preprint arXiv: 100..1033..5444 2018.
Sayantan Dasgupta, Trevor Cohn, and Timothy Baldwin. Cost-effective distillation of large language models. In Findings of the Association for Computational Linguistics; ACL 2023, pp. 7346-7354, 2023.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint
2018.
Bill Dolan and Chris Brockett. Automatically constructing corpus of sentential paraphrases. In Third International Workshop on Paraphrasing (IWP2005), 2005.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:20 /929, 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 12351 PMLR, 2017.

---

Published as a conference paper at ICLR 2025
Hao Fu, Shaojun Zhou, Qihong Yang, Junjie Tang, Guiquan Liu, Kaikui Liu, and Xiaolong Li. Lrc- bert: inteettrepresentatiinn contrastive knowledge distillation for natural language understanding. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 12830-12838, 2021.
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.
Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines to read and comprehend. Advances in neural information processing systems, 28, 2015.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. In NIPS 2014 Deep Learning Workshop. 2014. doi: (044850ARXXW.1500000555 httpts //axxiv.oo.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. TinyBERT: Distilling BERT for natural language understanding. Findings of the Association for Computational Linguisties: EMNLP 2020, November 2020.
Hee-Jun Jung, Doyeon Kim, Seung-Hoow Na, and Kangil Kim. Feature structure distillation with centered kernel alignment in bert transferring. Expert Systems with Applications, 234:1 10980, 2023.
Yoon Kim and Alexander M Rush. Sequence-level knowledge distillation. arXiv:1606.0793 07947, 2016.
arXiv preprint
Simon Komblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited, In International conference on machine learning, pp. 3519-3529. PMLR, 2019.
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising sequence-to-sequeee pre-training for natural language generation, translation, and comprehension. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 7871-7880, Online, July 2020. Association for Computational Linguistics. doi: 10. 1/25533/1/220.aa// a0tp:/0000...iii..0000 main. 703. URL https: setts::/aclanthhlooy: org/ 2020. acc-mmai 703.
Zheng Li, Zijian Wang, Ming Tan, Ramesh Nallapati, Parminder Bhatia, Andrew Arnold, Bing Xiang, and Dan Roth. Dq-bart: Efficient sequence-to- seuuence-to-sequenee model via joint distillation and quantization. arXiv preprint arXiv:2203. 11239, 2022.
Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 74-81, 2004.
Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, and Luke Zeutlemoye. Multilingual denoising pre-training for neural machine translation. Transactions of the Association for Computational Linguistics, 8:726-742, 2020. doi: 10.1 tacl 00343.
Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marein Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint
(tt::2407.146799 2024.
Shashi Narayan, Shay Cohen, and Mirella Lapata. Don't give me the details, just the sum- mary! topic-aware convolutional neural networks for extreme summarization. arXiv preprint arXiv: /808.00500 2018.
Thao Nguyen, Maithra Raghu, and Simon Kornblith. Do wide and deep networks learn the same things? uncovering how neural network representations vary with width and depth. arXiv preprint arXiv:2010. 15327, 2020.
12

---

Published as a conference paper at ICLR 2025
Anna Ollerenshaw, Md Asif Jalal, and Thomas Hain. Insights neural representations for end-to-end speech recognition. arXiv preprint SxXiv:2205.09456, 2022.
Jason Phang, Haokun Liu, and Samuel R Bowman. Fine-tuned transformers show clusters of similar representations across layers. arXiv preprint arXiv:2109.08400 009.00090080,,, 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl Blog, 1(8), 2019.
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning or feature reuse? http://9999999 2019.
towards understanding the effectiveness of maml. arXiv preprint
Maithra Raghu, Thomas Unterthiner, Simon Kornblith, Chiyuan Zhang, and Alexey Doso- vitskiy. Do vision transformers see like convolutional neural networks? In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 12116-12128. Curran Asso- ciates, Inc., 2021. URL http:tttee/pprreeciassseeeeeeeee neurips cc/pper files/paper) fi 11/111/55 652
pdf
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv: arXiv:/606.00 1606.05200 2016.
Aninda Saha, Alina Bialkowski, and Sara Khalifa. Distilling representational similarity using centered kernel alignment (cka). In Proceedings of the the 33rd British Machine Vision Conference (BMVC 2022). British Machine Vision Association, 2022.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, distilled version of BERT: smaller, faster, cheaper and lighter. CoRR, abs/1910.01108, 2019. BtLthttp//aaxx
org/abs/ /1900. 01108
Gabriele Sarti, Grzegorz Chrupata, Malvina Nissim, and Arianna Bisazza. Quantifying the plausibility of context reliance in neural machine translation. In The Twelfth International Conference on Learn- ing Representations, 2024.
net/ forum? htt.aeeeeetetoomoeeeeeeeeennann. wi==::ifiIIIidSI.
Sam Shleifer and Alexander M Rush. Pre-trained summarization distillation. arXiv preprint 1:Xiv:2010.13000 2020.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1631-1642, 2013.
Sharath Nittur Sridhar and Anthony Sarah. Undivided attention: Are intermediate layers necessary for bert? arXiv preprint 2rXiv:2012.11888 2020.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model com- 1908 09355
pression. arXiv arXiv: 1995509355 2019.
Iulia Ture, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arvi 1i..ii::00.0000 998.08962.. 2019.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: multi-task benchmark and analysis platform for natural language understanding. Proceedings of the 2018 EMNLP Workshop BlackboxNI.P: Analyzing and Interpreting Neural Networks for NLP, pp. 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 18.18653/v 11/533/1//I1/ 1/W 11-44446....... URL PRRlnttpss//acaathoooooo wrg.//looooor//iii 5446
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self- attention distillation for task-agnostic compression of pre-trained transformers. Advances in Neural Information Processing Systems, 33:5776-5788, 2020.
Alex Warstadt, Amanpreet Singh, and Samuel R Bowman. Neural network acceptability judgments.. Transactions of the Association for Computational Linguistics, 7:625-641, 2019.
13

---

Published as a conference paper at ICLR 2025
Yuqiao Wen, Zichao Li, Wenyu Du, and Lili Mou. f-divergence minimization for sequence-level knowledge distillation. arXiv preprint arXiv:2307.15190, 2023.
Peter West, Chandra Bhagavatula, Jack Hessel, Jena Hwang. Liwei Jiang, Ronan Le Bras, Ximing Lu, Sean Welleck, and Yejin Choi. Symbolic knowledge distillation: from general language models to commonsense models. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4602-4625, 2022.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprini arXiv: 000.17444405525 2017.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. HuggingFace' transformers: State-of-thh-an natural language processing. CoRR, abs/1910.03771. 2019. URL
http::/aaxxv..oree abs 1t///111111077777
Fuzhao Xue, Jianghai Chen, Aixin Sun, Xiaozhe Ren, Zangwei Zheng, Xiaoxin He, Yongming Chen, Xin Jiang, and Yang You. A study on transformer configuration and training objective. In International Conference on Machine Learning, pp. 38913-38925. PMLR, 2023.
Linting Xue, Noah Constant, Adam Roberts, Mihir Kale, Rami AI-Rfou, Aditya Siddhant, Aditya Barua, and Colin Raffel. mt5: A massively multilingual pre-trained text-to-text transformer. arXiv preprint wwXXvv22010.19955 2020.
A PROOF OF THE UPPER BOUND OF CKA
We derived that for the linear case,
EEs 11
CKA(Hs, HT) = |ET |F EssILE
(6)
1)2, where tr stands for the trace of a matrix (Equation in Kornblith et al. (2019)) Now, since the Gram matrices Hs and Hi HH, are both positive semi-delinite, using Cauchy-Schwal inequality for their trace, show that
can be observed that |E>s| = ur Hy
(N 1)2 li HT
1/2
(N 1)²
(7)
(N 113 (N 1r H 1/2
1/2
1
This proves |E>l|| < |Ess| | EEr || F. and shows that the value of CK CKA(Hs, Hy) is bounded above by 1. And being positive quantity, S CKA(Hs, Hr) < 1.
B ADDITIONAL EXPERIMENTAL DETAIL
B.1 SUMMARIZATION (BART)
We do not use hyperparameters to weigh the loss contributions for all the experiments. We use a batch size 16 and sum over batches for the computation of CKA and the other losses through gradient accumulation, making the effective batch size 256. We use the Adam optimizer with 17 le 4 and weight decay 5e 4. The context size used for the input document is 1024, while the context size for the summary is 128. All the experiments performed
A 1100 GPU with 80GB memory.
14

---

Published as a conference paper at ICLR 2025
B.2 TRANSLATION (MBART & T5)
We do not use hyperparameters to weigh the loss contributions. We use Adam Optimizer with " 3e 5 and weight decay 5e for all the pretraining distillation mc4 The context size for pretraining of mBART' is 512. We sample the languages with the following codes from mC4 with an equal probability: ar, es, de, en, es, et, fr, gu, hi, ja. kk, ko, Iv. my, ne, nl, ro, ru, si, tr, zh, af, az, bn. fa, he, id, ka, km, mk, ml, mn, mr, pl, pt, sv, sw, ta, te, th, uk, ur, xh, gl. sl.
While for Flan-T5, probability of 0.67 and 33 other languages with probability of 0.01 with the following codes: es, ja, fa, hi, fr, zh, bn, de, te, ar. pl, ta, pt, ur, gl, he, ko, th, nl, id, tr, vi, ru, sv, fi, sw, ro, cs, ms, so, el.
use context size of 1024. We sample the English corpus of mC4 with
In the downstream translation tasks for both models, we use context size of 256 for both source and target sentences. All the experiments are performed on an A100 GPU with 80GB memory.
B.3 CLASSIFICATION (BERT)
We use sequence length of 512 tokens during pretraining using C4 and and use the Adam optimizer with learning rate = 2e - 4 and weight decay 5e 4. We use a batch size of 32 for gradient computation and then accumulate the gradient for 40 batches, resulting in a large batch size of 1280. This is similar to using large batch sizes in Sanh et al. (2019). The covariance matrices are averaged over the 40 batches for CKA loss computation during the pretraining and added the final batch. We do not use hyperparameters to weigh the loss contributions. All the experiments are performed on an A40 GPU with 40GB memory.
The fine- tuningiing on GLUE tasks is done with the Adam optimizer with learning rate = 3e-5 5 to 1e-4 and weight decay 5e for batch size of 64. Since CKA loss gives better KL Divergence than the baseline, we fine-tuned only the students distilled with CKA for the downstream tasks. We did not use any hidden state loss during fine-tuning.
Phttss//huggingfacc.oo hhtnss//wuuggiiiiaaccoaattttrtt/ www.aasece/ogggggdddddseeeeeeee.
hhtpss//huggingffcc.co http:://huggnnghhccccoooattttt/s htwp://wwww......oooooo/onnanaaa
15


Reviews:
[{"review_id": "xM1eP3i516", "reviewer": "Reviewer_qCBt", "summary": "- This paper presents an approach to knowledge distillation that allows matching latent representations between different dimensionalities by using centered kernel alignment (CKA) instead of cosine similarity.\n- Comprehensive experiments on encoder-decoder models and masked language models confirmed that the proposed method consistently achieves higher performance compared to simple baselines, such as cases without added loss or with linear projection.", "questions": "- Linear CKA is equivalent to the RV coefficient and can be thought of as similar to Pearson correlation. Is it possible to provide any reasoning or justification for why this method is effective?\n- Additionally, is there a reason why linear CKA performs better than simple regularization techniques (e.g., weight decay on the student side), including methods that adjust divergence? In practice, does it empirically outperform basic or current regularization methods?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The experiments are comprehensive, covering encoder-decoder models (BART, mBART, T5) and masked language models (BERT). The task settings range from fine-tuning on specific tasks to instruction tuning. This experimental section is valuable for readers interested in knowledge distillation in similar settings.\n- The CKA used is simple, and the authors utilize a linear kernel (Line 108). This choice is expected to make the implementation easier to reproduce and facilitate scalability in training.", "weaknesses": "- There is limited mention of related work in knowledge distillation, making it difficult to assess the value of this study. For instance, there is no reference to research that incorporates modifications to divergence (e.g., Wen et al. ACL 2023, Gu et al. ICLR 2024) or to the distillation of causal language models, which is extensively discussed in Xu et al. (arXiv:2402.13116) and is currently of significant interest to readers.\n- The tasks addressed in this study do not align with the practical use cases of modern language models. For example, summarization is often performed using causal language models today, and instruction tuning is naturally suited for causal language models rather than encoder-decoder models. While papers with significant theoretical contributions may be accepted even if their experimental settings are somewhat “toy” or “outdated,” this paper does not focus on theory (there is no theoretical section). Instead, it proposes that replacing cosine similarity with linear CKA can improve knowledge distillation in various practical settings. The value of empirical studies is considerable, and this topic is indeed appealing. However, given that the experimental setup feels somewhat outdated, it may not align well with ICLR—a leading conference for representation learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fnw1AqxPHC", "reviewer": "Reviewer_J69D", "summary": "This paper looks at improving distillation between a teacher and student model using Centered Kernel Alignment (CKA). The main benefit of this is that it allows for the hidden dimension to be a different size than most other distillation methods used today and is much more flexible. Overall, this is a nice approach with interesting use of isotropy and math that makes it a valuable insight to the field.\n\nExperiments were done on encoder and encoder-decoder models (BART, mBART, BERT, T5) and the results on three different NLP tasks all justified the method. It would have been nice to see an experiment on a decoder-only architecture, as well as experiments on larger versions of the models discussed, but none of these downsides are enough for me to lower my review score. I’m also not super well-versed in the distillation literature so there is a chance that there are additional baselines that should have been considered that I’m not aware of, but the current experiments show the benefit of the method.", "questions": "I think of distillation being useful for very large models, but the current teacher models are relatively small. What happens if you try it on a larger model?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Really interesting application of CKA to fix a problem with distillation and make it much more general.", "weaknesses": "Larger versions of BART or mBART for at least one experiment would have been nice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7hsfzfFHaD", "reviewer": "Reviewer_n8kN", "summary": "The paper proposes a novel technique for knowledge distillation (KD) in LLMs using Centered Kernel Alignment (CKA) to address limitations in hidden state matching. Traditional methods rely on cosine loss, which restricts the student model's architecture to match the teacher's dimensions. This paper introduces CKA to match hidden states of different dimensionalities, enabling higher compression ratios and allowing more compact and efficient student models.\n\nThe authors tested their approach on various NLP tasks:\n- Summarization: Using BART on CNN and XSum datasets, evaluated by ROUGE scores. \n- Machine Translation: Tested mBART with multilingual data and evaluated on EN-RO (WMT16) and EN-FR (IWSLT2017) datasets with BLEU scores.\n- Classification: Used BERT on the GLUE benchmark for classification tasks, comparing with linear projection and cosine-based baselines.\nThe CKA-based student models consistently outperformed models with other hidden state matching.", "questions": "n/a", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The method is innovative, offering an effective alternative to traditional cosine loss by using CKA for hidden state matching, which enables distillation with flexible model architectures. This approach allows for higher compression ratios and more adaptable student models, overcoming limitations of previous distillation methods. \nAdditionally, while most other KD methods overlook computational complexity, the authors present techniques that reduce inference times while maintaining competitive accuracy. \nThe performance improvement is substantial.", "weaknesses": "- although the authors conduct experiments across three tasks, this is relatively limited compared to other KD studies. I would encourage the authors to expand their evaluation to cover more tasks within the GLUE or SuperGLUE benchmarks for a more comprehensive analysis.\n- the teacher models used in this paper are relatively small (e.g., BART, T5, and BERT), whereas the current research trend is increasingly focused on larger models (e.g., over 7 billion parameters) like LLaMA, which have stronger pre-training capabilities. It would be insightful to see whether the proposed method can sustain similar performance gains with larger models, as this would make the findings more applicable to the broader research community.\n- the authors appear to primarily focus on methods involving hidden state matching. I would suggest including comparisons with other methods based on KL divergence or alternative distillation losses to provide a more thorough context for the proposed approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)