State Spaces Aren’t Enough: Machine Translation Needs Attention

State Spaces Aren’t Enough: Machine Translation Needs Attention Ali Vardasbi†∗ University of Amsterdam a.vardasbi@uva.nl Telmo Pessoa Pires† Robin M. Schmidt Stephan Peitz Apple {telmo, robin_schmidt, speitz}@apple.com Abstract Structured State Spaces for Sequences (S4) is a recently proposed sequence model with successful applications in various tasks, e.g. vision, language modeling, and au- dio. Thanks to its mathematical formula- tion, it compresses its input to a single hid- den 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 sen- tences. 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 in- troducing an attention mechanism. 1 Introduction The Transformer (Vaswani et al., 2017) is the most popular architecture for state-of-the-art Natural Language Processing (NLP) (Devlin et al., 2019; Brown et al., 2020; NLLB Team et al., 2022). How- ever, the attention mechanism on which it is built is not well suited for capturing long-range depen- dencies due to its quadratic complexity (Ma et al., 2023). Recently, Structured State Spaces for Se- quences (S4) was shown to be on par with the © 2023 The authors. This article is licensed under a Creative Commons 4.0 licence, no derivative works, attribution, CC- BY-ND. †Equal contribution. ∗Work done during an internship at Apple. Transformer on various sequence modelling tasks, including time series forecasting, language model- ing (Gu et al., 2022), and audio generation (Goel et al., 2022); and to surpass the Transformer on tasks requiring reasoning over long range dependencies, like the Long Range Arena (Tay et al., 2021). Internally, S4 keeps a state-space based represen- tation. Due to the way its weights are initialized, it is able to approximately “memorize” the input sequence, removing the need for an attention mech- anism. Indeed, the results from Gu et al. (2022) 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 recur- rence formula is used to generate the output step by step during inference. S4’s competitive performance in Language Mod- eling (LM) promises an alternative to the Trans- former for other sequence modeling tasks, such as Machine Translation (MT). In this work, we ex- plore 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 …

Content truncated to 500 words