TALAS: Teacher-Anchored Layer Alignment with Adaptive Sharpness-Aware Minimization for Embedding Distillation

2026-06-19

AI Summary

Existing knowledge distillation methods often incur substantial computational overhead and suffer from performance bottlenecks due to capacity gaps, as they compel student models to fully mimic the teacher’s embeddings or intermediate features. To address these limitations, this work proposes the TALAS framework, which innovatively integrates three key components: a teacher-anchored mechanism that distills only high-level sentence embeddings, bottom-layer alignment via self-distillation to preserve semantic relationships, and adaptive sharpness-aware minimization (ASAM) enhanced with geometric constraints in the embedding space. This approach significantly reduces both training computation and memory requirements while achieving state-of-the-art performance on standard sentence embedding benchmarks. Moreover, it consistently improves model generalization and robustness without relying on exhaustive feature imitation.

Abstract

Knowledge Distillation (KD) has established itself as a pivotal technique for compressing large pre-trained language models. However, existing methods that force a student to strictly mimic the teacher's sentence embeddings or internal features often incur prohibitive computational costs and yield suboptimal performance due to the inherent capacity gap. To address these challenges, we propose TALAS (Teacher-Anchored Layer Alignment with Sharpness-aware minimization), a unified framework that synergizes hierarchical (multi-layer) alignment with robust optimization. First, we introduce a Teacher-Anchored mechanism that selectively distills final sentence embeddings only into the student's upper layers, thereby reducing overhead while respecting capacity constraints. Second, we bridge the semantic gap in lower layers via Layer-Aligned Self-Distillation, which propagates knowledge top-down using internal geometric relational constraints in the embedding space. Finally, to prevent the student from memorizing point-wise teacher noise, we integrate Adaptive Sharpness-Aware Minimization (ASAM) into the training objective, guiding the model towards flat minima for enhanced generalization. Empirical results on standard sentence embedding benchmarks demonstrate that TALAS consistently outperforms strong distillation baselines while achieving superior training efficiency in terms of computational cost and memory footprint.

Problem

• Knowledge Distillation
• Model Compression
• Embedding Distillation
• Capacity Gap
• Computational Cost

Innovation

• Teacher-Anchored Layer Alignment
• Adaptive Sharpness-Aware Minimization
• Embedding Distillation
• Knowledge Distillation
• Flat Minima