State space models (SSMs) have shown impressive results on tasks that require modeling long-range dependencies and efficiently scale to long sequences owing to their subquadratic runtime complexity. Originally designed for continuous signals, SSMs have shown superior performance on a plethora of tasks, in vision and audio; however, SSMs still lag Transformer performance in Language Modeling tasks. In this work, we propose a hybrid layer named Block-State Transformer (BST), that internally combines an SSM sublayer for long-range contextualization, and a Block Transformer sublayer for short-term representation of sequences. We study three different, and completely parallelizable, variants that integrate SSMs and block-wise attention. We show that our model outperforms similar Transformer-based architectures on language modeling perplexity and generalizes to longer sequences. In addition, the Block-State Transformer demonstrates more than tenfold increase in speed at the layer level compared to the Block-Recurrent Transformer when model parallelization is employed.

翻译：状态空间模型（SSM）在处理需要建模长程依赖关系的任务上表现优异，并因其次二次计算时间复杂度而能高效扩展到长序列。这类模型最初针对连续信号设计，已在视觉和音频领域的众多任务中展现出卓越性能；然而，在语言建模任务中，SSM的性能仍落后于Transformer。本文提出一种名为Block-State Transformer（BST）的混合层，其内部结合了用于长程情境化的SSM子层和用于序列短时表征的块Transformer子层。我们研究了三种完全可并行化的不同变体，将SSM与块级注意力机制相融合。实验表明，我们的模型在语言建模困惑度上优于类似的基于Transformer的架构，并能泛化至更长序列。此外，在采用模型并行化时，Block-State Transformer在层级的处理速度相比Block-Recurrent Transformer提升超过十倍。