Dispatcher: A Message-Passing Approach to Language Modelling (2022)

Abstract: This paper proposes a message-passing mechanism to address language modelling. A new layer type is introduced that aims to substitute self-attention for unidirectional sequence generation tasks. The system is shown to be competitive with existing methods: Given N tokens, the computational complexity is O(N logN) and the memory complexity is O(N) under reasonable assumptions. In the end, the Dispatcher layer is seen to achieve comparable perplexity to self-attention while being more efficient.

## introduction the introduction of self-attention @xcite has produced a considerable surge of language models @xcite @xcite . originally, selfattention had been envisioned as a three elements algorithm (key, query, and value) to be applied onto an encoder-decoder framework for machine translation. it soon became evident that the transformer architecture can successfully master the most relevant nlp tasks @xcite with some marginal modifications. one key application of self-attention has been language generation @xcite , where typically the model attempts to predict the next token given a limited window of prior elements. a full text can thus be generated word by word. self-attention -bidirectional in nature -needs to be masked in order to avoid backward propagation of information. until a few years ago recurrent models @xcite @xcite outperformed every other method for language modelling and generation. this has changed with the introduction of masked self attention (msa) models, which have achieved the state-of-the-art in language generation, culminating in some unexpected results for multi-task zero-shot learning @xcite as well as intriguing few-shot abilities @xcite . the main argument of this work is to show that language modelling can efficiently rely on a message passing approach to perform, proposing a method that does not leverage upon self-attention. instead, the system builds a tree-like structure of forward message passing weighed by dispatching coefficients. in the end, the dispatcher architecture can generate texts as well as the original transformer model, more efficiently. the main contributions of the paper are to introduce the novel algorithm as well as compare perplexity to the "standard" self-attention on the task of language modelling. ## model the original transformer architecture is composed of a number of self-attention, skip connection, and feed-forward layers. given n tokens, the selfattention block has a computational and memory complexity of o(n 2 ) and is therefore problematic for long sequences. here we propose to substitute each self-attention layer with a different algorithm. within the dispatcher layer, information is pushed forward onto the next tokens in a recursive fashion. the algorithm is given a list of embeddings as input, with the aim to create output embeddings that contain a mixture of the tokens that precede them, without any leakage from the tokens that follow. the system achieves this goal by summing the tokens with themselves shifted by a power of two, iteratively. each of these steps is labelled shift and sum in fig. 1 . in the pseudo-code shown in alg. 1 the dis- the message coming from the prior tokens follows a binary tree structure, as depicted in fig. 1 . the sum is weighed by the dispatching coefficients, which effectively decide whether information coming from the left of the tree should propagate further, and by what amount. these weights are computed through a dense layer applied to the original tokens. a constant mask is applied to the tensor c after it has been computed to avoid leakage after the rollright operation. the algorithm presented above describes a single-head unit. as with self-attention, this layer can be split into a set of dispatcher heads to improve performance. the number of heads then becomes another hyper-parameter to tune during training. finally, if the number of input embeddings is not a power of two, the loop stops when the shift value is greater than the input length. ## conclusions a novel architecture dedicated to language modelling is introduced and shown to achieve comparable perplexity with self-attention based models, requiring less computational and memory resources. a larger number of tokens can allow for a wider context window and more detailed grounding of the model's predictions. finally, low perplexity in the task of language modelling is often predictive of high-quality text generation. this intriguing possibility will be pursued in a future work.