TLDR: This research paper introduces Procrustes distance and the Frobenius norm of the Feature Gram Matrix as superior distillation losses for transferring knowledge from large to small language models. It theoretically and empirically demonstrates that existing methods like Centered Kernel Alignment (CKA) and projection-based losses fail to fully preserve the geometric structure of features. Experiments with BERT and OPT models show that the new geometry-aware methods lead to statistically significant improvements in classification and instruction-following tasks, highlighting the potential of integrating feature geometry into knowledge distillation.
Large language models are incredibly powerful, achieving state-of-the-art results across many tasks in vision and language. However, their immense size and operational costs often make them inaccessible to a wider audience. This is where Knowledge Distillation (KD) comes into play, a technique designed to transfer the capabilities of a larger, complex ‘teacher’ model to a smaller, more efficient ‘student’ model.
Traditionally, knowledge distillation has focused on aligning the probabilistic outputs of teacher and student models. More recently, feature-based distillation methods have emerged, aiming to minimize differences between the hidden layer representations of these models. The core idea is for the student to mimic the intricate structure of the teacher’s internal feature space.
However, a recent study critically examines the effectiveness of current feature distillation techniques, such as projection-based mean squared loss and Centered Kernel Alignment (CKA). The researchers theoretically demonstrate that these widely used methods may not fully capture the essential feature structure, even when the loss appears to be zero.
To address this limitation, the paper introduces and advocates for the use of two alternative measures as distillation losses: Procrustes distance and the Frobenius norm of the Feature Gram Matrix. These measures are already recognized in the field of representational alignment for their ability to quantify how well different sets of features align geometrically. The authors provide a theoretical framework to justify their use, showing that they more faithfully capture the geometric alignment of feature representations compared to CKA and projection-based approaches.
The concept of ‘feature geometry’ in language models refers to the preserved relative representations, such as angles and inner products, between latent embeddings of models trained on similar data. This geometric structure is crucial for how language models organize and encode knowledge. The research questions whether task-specific feature distillation is directly linked to preserving this geometry between teacher and student models.
Empirical evaluations were conducted across different language model families, including BERT for classification tasks and OPT for instruction-following tasks. The results are compelling: feature distillation using the proposed Procrustes distance and the Frobenius norm of the Feature Gram Matrix showed statistically significant improvements in distillation performance. Specifically, improvements of up to 2 percentage points were observed in classification and instruction-following tasks.
Also Read:
- Concrete Score Distillation: A New Approach to Making Large Language Models More Efficient
- Evaluating Language Model Text Quality Through Internal Geometric Properties
This work highlights a significant advancement in knowledge distillation, emphasizing the importance of integrating feature geometry into the distillation process. By using more geometrically grounded loss functions, smaller models can more effectively learn and replicate the complex internal structures of their larger counterparts, paving the way for more efficient and accessible AI models. You can read the full research paper here.
