A Mixture of $h-1$ Heads is Better than $h$ Heads
This work addresses efficiency and performance in NLP models for tasks like machine translation and language modeling, but it is incremental as it builds on existing multi-head attention and mixture of experts concepts.
The paper tackles the overparameterization of multi-head attention in neural architectures by proposing a mixture of attentive experts model that reallocates heads to specialize on different inputs, achieving a 0.8 BLEU improvement over transformer-base on WMT14 English to German translation.
Multi-head attentive neural architectures have achieved state-of-the-art results on a variety of natural language processing tasks. Evidence has shown that they are overparameterized; attention heads can be pruned without significant performance loss. In this work, we instead "reallocate" them -- the model learns to activate different heads on different inputs. Drawing connections between multi-head attention and mixture of experts, we propose the mixture of attentive experts model (MAE). MAE is trained using a block coordinate descent algorithm that alternates between updating (1) the responsibilities of the experts and (2) their parameters. Experiments on machine translation and language modeling show that MAE outperforms strong baselines on both tasks. Particularly, on the WMT14 English to German translation dataset, MAE improves over "transformer-base" by 0.8 BLEU, with a comparable number of parameters. Our analysis shows that our model learns to specialize different experts to different inputs.