UCLA

Transformer-backed Large Language Models and Large Vision-Language Models are among the most widely used machine learning models recently. However, due to their enormous size, their behaviors and dynamics become difficult to understand and explain, and their inference speed and cost become slow and forbidding. Therefore, understanding their latent neural patterns and distilling their latent neural representations into smaller, cheaper, and faster models are vital for creating more explainable, accessible, and sustainable machine intelligence. However, knowledge distillation algorithms face an inherent challenge in minimizing the divergence between the neural activation representations of a teacher Transformer and a student Transformer: The latent representation spaces of the smaller Student always have lower dimensionality than those of the larger Teacher. We will address this Transformer distillation dimensionality mismatch or reduction problem by analyzing the neural activation patterns within these latent representation spaces and by proposing corresponding novel algorithms for distilling each group of components of Transformers. Our ultimate goal is to advance the field of understanding and accelerating the inference process of large Transformers by studying their latent neural representations.

Specifically, we first develop novel algorithms for aligning the heads within the Multi-Head Attention components of these Transformers for latent neural representation distillation. We propose a new Attention Map Alignment Distillation (AMAD) method for Transformers with multi-head attention, which, unlike previous literature, can work for a Teacher and a Student with mismatched dimensionalities in terms of the numbers of attention heads. Specifically, we soft-align different heads in Teacher and Student attention maps using a cosine similarity weighting. The Teacher head contributes more to the Student heads for which it has a higher similarity weight. Each Teacher head contributes to all the Student heads by minimizing the divergence between the attention latent neural activation distributions for the soft-aligned heads.

We then advance the understanding and distillation of the latent neural representations of the Feed Forward Network and Layer Normalization components of these Transformers. Previous literature shows that a small subset of dimensions within the high-dimensional latent neural representation spaces of language Transformers emerge as “outliers”, with weights or latent neuron activations of unusually large magnitudes—several standard deviations from the mean—during pretraining, and seem to encode critical linguistic knowledge in a sparse way. We systematically extend these previous findings on emergent outliers to Encoder-Decoder Transformers and instruction-finetuned models at scale, for the first time. We show that emergent outlier dimensions of the Transformer latent neural representations contribute significantly more to zero-shot performance than non-outlier dimensions. Leveraging our findings, we further propose the Emergent Outlier Focused Distillation (EOFD) method, which prioritizes these critical outlier latent dimensions to address the trade-off among these representation dimensions during distillation. Specifically, we propose a weighted MSE loss to place greater emphasis on the emergent outlier dimensions of the latent neural representations, which are identified by the standard deviations of neuron activation magnitudes.

We demonstrate that our AMAD and EOFD distillation algorithms achieve state-of-the-art performance empirically and generalize well to different Transformer architectures. Specifically, for AMAD, we experiment on distilling VL-T5 and BLIP and apply the AMAD loss to their T5, BERT, and ViT sub-modules. We show that, under vision-language settings, AMAD outperforms conventional distillation methods on the VQA-2.0, COCO captioning, and Multi30K translation datasets. We further show that even without VL pre-training, the distilled VL-T5 models outperform corresponding VL pre-trained VL-T5 models that are further fine-tuned by ground-truth signals, and that fine-tuning distillation can also compensate to some degree for the absence of VL pre-training in BLIP models. As for EOFD, we empirically show that it outperforms state-of-the-art distillation methods on the General Language Understanding Evaluation (GLUE) benchmark. We further demonstrate that EOFD generalizes well across Encoder-only BERT, Decoder-only GPT-2, and Encoder-Decoder T5 architectures.