Learning Hard Retrieval Decoder Attention for Transformers (2021)

Abstract: The Transformer translation model is based on the multi-head attention mechanism, which can be parallelized easily. The multi-head attention network performs the scaled dot-product attention function in parallel, empowering the model by jointly attending to information from different representation subspaces at different positions. In this paper, we present an approach to learning a hard retrieval attention where an attention head only attends to one token in the sentence rather than all tokens. The matrix multiplication between attention probabilities and the value sequence in the standard scaled dot-product attention can thus be replaced by a simple and efficient retrieval operation. We show that our hard retrieval attention mechanism is 1.43 times faster in decoding, while preserving translation quality on a wide range of machine translation tasks when used in the decoder self- and cross-attention networks.

## introduction the transformer translation model @xcite , which has outperformed previous rnn/cnn based sequence-to-sequence models @xcite @xcite , is based on multi-head attention networks. the multi-head attention mechanism, which computes several scaled dot-product attentions in parallel, can be efficiently parallelized at sequence level. in this paper, we investigate whether we can replace scaled dot-product attention by learning a simpler hard retrieval based attention that attends to a single token only. this simplifies computation and increases speed. we show that this can indeed be achieved by a simple and efficient retrieval operation while preserving translation quality. our contributions are as follows: ## background: the scaled dot-product attention the multi-head attention network heavily employed by the transformer translation model consists of h parallel scaled dot-product attentions, where h is the number of attention heads. the scaled dot-product attention mechanism takes three inputs: the query sequence q, the key sequence k and the value sequence v . it first compares each vector in q with all vectors in k by dot-product computation to generate the attention score matrix s: finally, the value sequence v is weighted by the attention probabilities p and accumulated as the attention result: ## experiment we implemented our approach based on the neutron implementation of the transformer @xcite . to investigate the impact on translation quality of our approach, we conducted our experiments on the wmt 14 english to german and english to french news translation tasks to compare with @xcite . we also examined the impact of our approach on the pre-processed data of the wmt 17 news translation tasks for 12 translation directions. the concatenation of newstest 2012 and newstest 2013 was used for validation and newstest 2014 as test sets for the wmt 14 english to german and english to french news translation tasks. we used the pre-processed data for wmt 17 news translation tasks. foot_0 ## conclusion we propose to learn a hard retrieval attention which only attends to one token rather than all tokens. with the one-to-one hard attention matrix, the matrix multiplication between attention probabilities and the value sequence in the standard scaled dotproduct attention can be replaced by a simple and efficient retrieval operation. in our experiments on a wide range of machine translation tasks, we show that using the hard retrieval attention for decoder attention networks can achieve competitive performance while being 1.43 times faster in decoding.