A Mixture of h - 1 Heads is Better than h Heads (2020)

Abstract: 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.

## introduction the transformer architecture and its variants achieve state-of-the-art performance across a variety of nlp tasks, including machine translation (vaswani et al., 2017; ott et al., 2018), language modeling (radford et al., 2018; baevski and auli, 2019), semantic role labeling (strubell et al., 2018), and more (devlin et al., 2019; liu et al., 2019b; yang et al., 2019b). under the hood, multihead attention provides the driving force: multiple separately parameterized attention functions act in parallel to contextualize the input representations; their outputs are then gathered by an affine transformation, and fed to onward computation. experts: input attention heads: figure 1: illustration of mae: a mixture of attentive experts. each h i box is an attention head in a given layer; there are h of them in total. experts are groups of h -1 attention heads. mae learns an input-dependent distribution of the experts (g). at each training step, a single expert is selected and updated (solid line); during the evaluation, experts' outputs are linearly combined with weights produced by g. recent efforts by voita et al. (2019) and michel et al. (2019) suggest that typical transformer networks are overparameterized, in the sense that at test time, many of the heads, or even a full layer (fan et al., 2020), can be removed without significant loss in performance. @xcite in response to this observation, they propose to prune the unimportant attention heads in the model after it is trained, aiming for faster inference. in this paper, we ask whether, instead of reducing the model capacity, we can use it more effectively. we propose mixture of attentive experts (mae). mae retains all attention heads, and learns to activate different heads on different inputs (see illustration in figure 1 ). we start by showing that multi-head attention can be seen as an uniform, input-agnostic mixture of experts (jacobs et al., 1991), by grouping a subset of atten-tion heads as an expert ( §2.2). we then introduce mae, which instead of uniformly weighting the experts, complements the experts with a learned, input-dependent function that assigns their responsibilities ( §2.3). to train mae, we propose a two-step algorithm based on block coordinate descent ( §3), which alternates between updating the experts' responsibilities and their parameters. we evaluate mae on machine translation and language modeling ( §4). our approach outperforms strong baselines on both; on the wmt14 english to german mt dataset, mae outperforms transformer-base (vaswani et al., 2017) by 0.8 bleu with a negligible increase in the number parameters. our analysis shows that mae learns to encourage different experts to specialize on different inputs ( §5). ## mae: mixture of attentive experts this section describes mae in detail. it is inspired by a mixture-of-experts view of multi-head attention, which we present in §2.2. specifically, we show that multi-head attention can be viewed as a mixture of uniformly weighted experts, each consisting of a subset of attention heads. based on this observation, we propose mae, which learns to weight the experts ( §2.3) depending on the input. we begin by laying out notation and necessary background in §2.1. ## training mae with block coordinate descent it is straightforward to jointly train the experts and the gating functions in an mae model using backpropagation. however, in line with previous observations (shen et al., 2019), we empirically observe that this is prone to degenerate solutions where the gating functions tend to learn to similarly weight the experts (see §5.1). @xcite as a remedy, we propose a block coordinate descent (bcd) training. at a high level, training is decomposed into two interleaving steps: a g step updates the gating function g(•; φ), fixing the experts; an f step fixes the gating function and updates one randomly selected expert f i (•; θ i ). @xcite the computations for g and f steps differ: forwardprop with z and calculate l. ## experiments we empirically evaluate mae on machine translation ( §4.2) and language modeling ( §4.3) benchmarks. we first introduce the compared models ( §4.1). ## analysis this section first empirically confirms that mae learns to activate different experts on different inputs in §5.1. we then run a synthetic experiment to explore mae's potential in transfer learning ( §5.2). ## conclusion we presented mae. it is inspired by a mixture-ofexperts perspective of multi-head attention. with a learned gating function, mae activates different experts on different inputs. mae is trained using a block coordinate descent algorithm, which alternates between updating the responsibilities of the experts and their parameters. our experiments show that mae outperforms the transformer baselines on machine translation and language modeling benchmarks. the analysis shows that mae learns to activate different experts. the code is publicly available at https://github.com/ noahs-ark/mae .