SinKD: Sinkhorn Distance Minimization for Knowledge Distillation.

Area of Science:

• Artificial Intelligence
• Machine Learning
• Natural Language Processing

Background:

• Knowledge Distillation (KD) is crucial for compressing Large Language Models (LLMs).
• Existing KD methods using Kullback-Leibler (KL), reverse KL (RKL), and Jensen-Shannon (JS) divergences face limitations with overlapping distributions.
• These limitations lead to issues like mode-averaging, mode-collapsing, and mode-underestimation in logits-based KD.

Purpose of the Study:

• To propose a novel Knowledge Distillation method, Sinkhorn KD (SinKD), that overcomes the limitations of existing divergence measures.
• To enhance the precision and nuance in assessing distribution disparities between teacher and student models.
• To improve the effectiveness of KD for diverse Natural Language Processing (NLP) tasks.

Main Methods:

• Utilized Sinkhorn distance for a precise assessment of distribution disparities between teacher and student models.
• Introduced a batch-wise reformulation of KD, moving beyond sample-wise limitations.
• Captured geometric intricacies of distributions across samples in high-dimensional space.

Main Results:

• Demonstrated superiority over state-of-the-art (SOTA) methods on GLUE and SuperGLUE benchmarks.
• Achieved significant improvements across LLMs with encoder-only, encoder-decoder, and decoder-only architectures.
• Validated the comparability, validity, and generalizability of the proposed SinKD method.

Conclusions:

• SinKD provides a more effective approach to logits-based KD compared to traditional divergence measures.
• The batch-wise reformulation and Sinkhorn distance enable robust model compression.
• SinKD represents a significant advancement in efficiently distilling knowledge from large language models.