# Align-to-Distill: Trainable Attention Alignment for Knowledge Distillation in Neural Machine Translation

arXiv: https://arxiv.org/abs/2403.01479 · HTML: https://mlsanity.com/paper/2403.01479/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Knowledge distillation for LLMs, one of 12 directions in Knowledge distillation.
- Builds on 32 of its 38 references that are in the map.
- First posted on arXiv 3 Mar 2024.

## Placed in

- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 30 of 32)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Neural Machine Translation of Rare Words with Subword Units](https://mlsanity.com/paper/1508.07909/) (2015)
- [SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing](https://mlsanity.com/paper/1808.06226/) (2018)
- [A Call for Clarity in Reporting BLEU Scores](https://mlsanity.com/paper/1804.08771/) (2018)
- [FitNets: Hints for Thin Deep Nets](https://mlsanity.com/paper/1412.6550/) (2014)
- [Knowledge Distillation: A Survey](https://mlsanity.com/paper/2006.05525/) (2020)
- [MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers](https://mlsanity.com/paper/2002.10957/) (2020)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer](https://mlsanity.com/paper/1612.03928/) (2016)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Sequence-Level Knowledge Distillation](https://mlsanity.com/paper/1606.07947/) (2016)
- [Patient Knowledge Distillation for BERT Model Compression](https://mlsanity.com/paper/1908.09355/) (2019)
- [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://mlsanity.com/paper/2004.02984/) (2020)
- [Do Deep Nets Really Need to be Deep?](https://mlsanity.com/paper/1312.6184/) (2013)
- [Self-Distillation Amplifies Regularization in Hilbert Space](https://mlsanity.com/paper/2002.05715/) (2020)
- [Deep Encoder, Shallow Decoder: Reevaluating Non-autoregressive Machine Translation](https://mlsanity.com/paper/2006.10369/) (2020)
- [Learning to Translate in Real-time with Neural Machine Translation](https://mlsanity.com/paper/1610.00388/) (2016)
- [Pre-trained Summarization Distillation](https://mlsanity.com/paper/2010.13002/) (2020)
- [ALP-KD: Attention-Based Layer Projection for Knowledge Distillation](https://mlsanity.com/paper/2012.14022/) (2020)
- [Selective Knowledge Distillation for Neural Machine Translation](https://mlsanity.com/paper/2105.12967/) (2021)
- [Dynamic Knowledge Distillation for Pre-trained Language Models](https://mlsanity.com/paper/2109.11295/) (2021)
- [Distilling Linguistic Context for Language Model Compression](https://mlsanity.com/paper/2109.08359/) (2021)

## Built on by (most cited first, 5 of 5)

- [AI Safety in Generative AI Large Language Models: A Survey](https://mlsanity.com/paper/2407.18369/) (2024)
- [Scalable Model Merging with Progressive Layer-wise Distillation](https://mlsanity.com/paper/2502.12706/) (2025)
- [Attention Grounded Enhancement for Visual Document Retrieval](https://mlsanity.com/paper/2511.13415/) (2025)
- [IA2: Alignment with ICL Activations Improves Supervised Fine-Tuning](https://mlsanity.com/paper/2509.22621/) (2025)
- [Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning](https://mlsanity.com/paper/2609.33200/) (2026)

## Abstract (arXiv)

The advent of scalable deep models and large datasets has improved the performance of Neural Machine Translation. Knowledge Distillation (KD) enhances efficiency by transferring knowledge from a teacher model to a more compact student model. However, KD approaches to Transformer architecture often rely on heuristics, particularly when deciding which teacher layers to distill from. In this paper, we introduce the 'Align-to-Distill' (A2D) strategy, designed to address the feature mapping problem by adaptively aligning student attention heads with their teacher counterparts during training. The Attention Alignment Module in A2D performs a dense head-by-head comparison between student and teacher attention heads across layers, turning the combinatorial mapping heuristics into a learning problem. Our experiments show the efficacy of A2D, demonstrating gains of up to +3.61 and +0.63 BLEU points for WMT-2022 De->Dsb and WMT-2014 En->De, respectively, compared to Transformer baselines.