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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Jin, Heegon, Son, Seonil, Park, Jemin, Kim, Youngseok, Noh, Hyungjong, Lee, Yeonsoo
Format: Preprint
Veröffentlicht: 2024
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866917621408989184
author Jin, Heegon
Son, Seonil
Park, Jemin
Kim, Youngseok
Noh, Hyungjong
Lee, Yeonsoo
author_facet Jin, Heegon
Son, Seonil
Park, Jemin
Kim, Youngseok
Noh, Hyungjong
Lee, Yeonsoo
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2403_01479
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Align-to-Distill: Trainable Attention Alignment for Knowledge Distillation in Neural Machine Translation
Jin, Heegon
Son, Seonil
Park, Jemin
Kim, Youngseok
Noh, Hyungjong
Lee, Yeonsoo
Computation and Language
Artificial Intelligence
68T50
I.2.7
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.
title Align-to-Distill: Trainable Attention Alignment for Knowledge Distillation in Neural Machine Translation
topic Computation and Language
Artificial Intelligence
68T50
I.2.7
url https://arxiv.org/abs/2403.01479