MiniLM (Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers)

A task-agnostic knowledge distillation framework for compressing pre-trained Transformer language models (Wang, Wei, Dong, Bao, Yang, & Zhou, NeurIPS 2020) that trains an M-layer student model with arbitrary hidden dimension d'_h to deeply mimic the self-attention module of only the last (L-th) Transformer layer of the teacher by minimizing the sum of two Kullback-Leibler divergences: L = L_AT + L_VR, where L_AT matches multi-head query-key attention distributions A_{L,a}^T = softmax(Q_{L,a}^T (K_{L,a}^T)^T / sqrt(d_k)) and L_VR matches multi-head scaled dot-product value relations VR_{L,a}^T = softmax(V_{L,a}^T (V_{L,a}^T)^T / sqrt(d_k)), optionally bridged by an intermediate-width Teacher Assistant when both depth and width are halved (M <= L/2, d'_h <= d_h/2).