State Spaces Aren’t Enough: Machine Translation Needs Attention (2023)

Abstract: Structured State Spaces for Sequences (S4) is a recently proposed sequence model with successful applications in various tasks, e.g. vision, language modelling, and audio. Thanks to its mathematical formulation, it compresses its input to a single hidden state, and is able to capture long range dependencies while avoiding the need for an attention mechanism. In this work, we apply S4 to Machine Translation (MT), and evaluate several encoder-decoder variants on WMT’14 and WMT’16. In contrast with the success in language modeling, we find that S4 lags behind the Transformer by approximately 4 BLEU points, and that it counter-intuitively struggles with long sentences. Finally, we show that this gap is caused by S4’s inability to summarize the full source sentence in a single hidden state, and show that we can close the gap by introducing an attention mechanism.

## introduction the transformer @xcite is the most popular architecture for state-of-the-art natural language processing (nlp) @xcite @xcite . however, the attention mechanism on which it is built is not well suited for capturing long-range dependencies due to its quadratic complexity @xcite . recently, structured state spaces for sequences (s4) was shown to be on par with the transformer on various sequence modelling tasks, including time series forecasting, language modeling @xcite , and audio generation @xcite ; and to surpass the transformer on tasks requiring reasoning over long range dependencies, like the long range arena @xcite . internally, s4 keeps a state-space based representation. due to the way its weights are initialized, it is able to approximately "memorize" the input sequence, removing the need for an attention mechanism. indeed, the results from @xcite show that the self-attention layers can be replaced by s4 layers without losing accuracy, and that it is able to effectively model long-range dependencies in data. moreover, one of the key advantages of the s4 kernel is that its forward step can be formulated both as a convolution and as a recurrence formula, allowing fast implementation during training, when the convolution method is used, while the recurrence formula is used to generate the output step by step during inference. s4's competitive performance in language modeling (lm) promises an alternative to the transformer for other sequence modeling tasks, such as machine translation (mt). in this work, we explore s4-based architectures for mt. our goal is to find the best performing s4 architecture, and we study the impact of several architectural choices on translation accuracy, namely the effect of model depth, the number of s4 blocks, and the importance of the encoder. despite our best efforts, our top performing attention-free s4 model lags significantly (∼ 4 bleu points) behind the transformer, with the gap increasing with input length. we hypothesize this is due to the fact that s4 compresses the source sentence to a fixed-size representation, and thus lacks a way to access the token-level states of the source, which is important for mt. as the input length increases, it becomes increasingly hard for the model to accurately store the full source sentence in a single hidden state. in contrast, the decoder cross-attention in the transformer acts as a retrieval mechanism, allowing to accurate retrieval of the source sentence during decoding. armed with this observation, we enhance s4 with crossattention, and show this is enough to close the gap to the transformer. finally, we combine the transformer and s4 into an hybrid architecture that outperforms both of them. to summarize, the main contributions of the present work are: ## background in this section, we provide a brief overview of s4 and machine translation. ## results in this section, we describe the experimental setup, and discuss our results. ## conclusion and future work in this work, we explored the application of s4 to machine translation and conducted an investigation into the best architecture and hyperparameters. despite our efforts, we found that s4's translation accuracy lagged behind the transformer, and the performance gap widened for longer sentences. we then showed that this was due to the limitations of the fixed-size representation used by s4, which had to compress the entire prior context, including the source sentence and previous output tokens. finally, we showed that the performance gap can be closed by incorporating attention. since we did our investigation into s4, numerous new ssm models have been proposed. of particular note are s5 @xcite , which utilizes a multi-input multi-output ssm, instead of one single-input single-output ssm per feature as s4 does, and h3 @xcite , which is faster and better at lm than s4. we hope future research explores how well these models perform on mt. additionally, it is worth noting mega @xcite , which incorporates ssm's into the transformer attention, and is effective in mt, albeit at the expense of quadratic complexity.