How Effective are State Space Models for Machine Translation?

TL;DR

This paper investigates whether modern linear recurrent models can match Transformers for machine translation, focusing on RetNet and Mamba, plus hybrids that integrate attention. Through extensive experiments in sentence- and paragraph-level MT, the authors show that Mamba is competitive with Transformers on sentence-level data, and that training distribution shifts toward longer sequences significantly reduce gaps on paragraph-level MT. Attention-enhanced hybrids (e.g., Mamba-MHA) deliver top translation quality and improve robustness to sequence length and named-entity recall, while pretrained priors further boost performance. Across models and regimes, Mamba demonstrates substantial efficiency advantages in memory and throughput, suggesting linear recurrent approaches as viable, scalable alternatives for MT with long contexts.

Abstract

Transformers are the current architecture of choice for NLP, but their attention layers do not scale well to long contexts. Recent works propose to replace attention with linear recurrent layers -- this is the case for state space models, which enjoy efficient training and inference. However, it remains unclear whether these models are competitive with transformers in machine translation (MT). In this paper, we provide a rigorous and comprehensive experimental comparison between transformers and linear recurrent models for MT. Concretely, we experiment with RetNet, Mamba, and hybrid versions of Mamba which incorporate attention mechanisms. Our findings demonstrate that Mamba is highly competitive with transformers on sentence and paragraph-level datasets, where in the latter both models benefit from shifting the training distribution towards longer sequences. Further analysis show that integrating attention into Mamba improves translation quality, robustness to sequence length extrapolation, and the ability to recall named entities.

Figures (6)

• Figure 1: Recall in recovering named entities on the WMT16 ro$\to$en dataset by their training set frequency: unseen entities do not appear in the training data, while regular and frequent entities appear $[1, 16)$ and $16+$ times.
• Figure 2: Sensitivity to input length, measured by the number of sources tokens, on the WMT23 de$\to$en datset, for models trained from scratch (top) and finetuned from a pretrained checkpoint (bottom).
• Figure 3: COMET scores per sequence length on WMT14 de$\to$en for trained-from-scratch models.
• Figure 4: Sensitivity to input length, measured by the number of sources tokens, on the WMT23 en$\to$de datset, for models trained from scratch (top) and finetuned from a pretrained checkpoint (bottom).
• Figure 5: Distribution of source length in 1) the training datasets: WMT23 de$\to$en (top left), WMT23 en$\to$de (top right), and 2) the test datasets: WMT23 de$\leftrightarrow$en (bottom left), our custom TED Talks de$\to$en (bottom right).