# Autocorrelation Matrix Knowledge Distillation: A Task-Specific Distillation Method for BERT Models
## 3 Method
### Hidden States Distillation  
In the Transformer model, the hidden states are an important part of the intermediate representation of each layer, capturing the semantic information and features of the input sequence. In traditional knowledge distillation methods, the student model usually needs to use a projection layer to solve the problem of mismatching the hidden states dimension with the teacher model. The hidden states matrix \(\mathbf{H}_{i}^{S}\in\mathbb{R}^{l\times D^{\prime}}\) of the student model and the hidden states matrix \(\mathbf{H}_{j}^{T}\in\mathbb{R}^{l\times D}\) of the teacher model cannot be directly compared due to their different dimensions. Here, \(l\) represents the length of the input sequence, \(D\) and \(D^{\prime}\) represent the hidden states dimensions of the teacher model and the student model, respectively, and \(D^{\prime}\) is usually smaller than \(D\). To solve this problem, traditional methods introduce a projection layer \(\mathbf{P}\in\mathbb{R}^{D^{\prime}\times D}\) to project the hidden states of the student model to a dimension that matches the teacher model:  
\[\mathbf{H}_{i}^{S}\mathbf{P}\in\mathbb{R}^{l\times D} \tag{10}\]  
However, this method has two major limitations: first, the projection operation adds additional computational complexity, especially in large-scale model scenarios, where the computational overhead increases significantly; second, the dimension conversion process may lead to information loss, meaning that the student model is unable to fully capture the characteristic performance of the teacher model.  
To overcome these problems, AMKD takes a different approach: by calculating the autocorrelation matrix of the hidden states, explicit dimensional projection is avoided. Specifically, AMKD calculates the autocorrelation matrix of the hidden states matrix of each layer, converting it from the \((length,hidden\_size)\) shape to a unified \((length,length)\) shape; that is,  
\[\mathbf{C}_{i}^{S}=\mathbf{H}_{i}^{S}\cdot(\mathbf{H}_{i}^{S})^{\mathsf{T}}\in\mathbb{R} ^{l\times l},\quad\mathbf{C}_{j}^{T}=\mathbf{H}_{j}^{T}\cdot(\mathbf{H}_{j}^{T})^{\mathsf{ T}}\in\mathbb{R}^{l\times l}, \tag{11}\]  
where \(\mathbf{C}\) represents the autocorrelation matrix. By calculating the autocorrelation matrix, AMKD effectively solves the dimension mismatch problem. No matter how different the hidden states dimensions of the student model and the teacher model are, after autocorrelation matrix conversion, they can be compared on a unified \(l\times l\) dimension.  
The autocorrelation matrix can not only capture the performance of a single feature, but also capture the high-order relationship between features. The formula is as follows:  
\[C_{uv}^{T}=\sum_{k=1}^{D}H_{uk}^{T}H_{vk}^{T},\quad C_{uv}^{S}=\sum_{k=1}^{D^ {\prime}}H_{uk}^{S}H_{vk}^{S}, \tag{12}\]  
where \(C_{uv}\) represents the inner product between the \(u\)-th row and the \(v\)-th row of the hidden states matrix \(\mathbf{H}\). As shown in Figure 2, the left figure shows the autocorrelation matrix of the teacher model, and the right figure shows the autocorrelation matrix of the student model. By comparing these two matrices, we can see how the student model learns the complex feature relationships in the teacher model through AMKD. The color depth represents the correlation between features, and the darker the color, the stronger the correlation.  
In AMKD, we use mean squared error (MSE) to measure the difference between the autocorrelation matrices of the student model and the teacher model, and the hidden states distillation loss is defined as  
\[\mathcal{L}_{\text{hidden}}=\sum_{i=1}^{N}\text{MSE}\Big{(}C_{i}^{S},\mathbf{C}_{ j}^{T}\Big{)}, \tag{13}\]  
By minimizing the MSE of the autocorrelation matrix, AMKD ensures that the student model can not only learn the individual features of the teacher model, but also capture the complex relationships between the features. This significantly improves the performance of the student model in complex tasks.