Transformeur à État de Bloc

Résumé

Les modèles à espace d'états (SSMs) ont démontré des résultats impressionnants sur des tâches nécessitant la modélisation de dépendances à long terme et s'adaptent efficacement à des séquences longues grâce à leur complexité temporelle sous-quadratique. Initialement conçus pour les signaux continus, les SSMs ont montré des performances supérieures sur une multitude de tâches, notamment en vision et en audio ; cependant, les SSMs restent en retard par rapport aux Transformers dans les tâches de modélisation du langage. Dans ce travail, nous proposons une couche hybride nommée Block-State Transformer (BST), qui combine en interne une sous-couche SSM pour la contextualisation à long terme et une sous-couche Block Transformer pour la représentation à court terme des séquences. Nous étudions trois variantes différentes et entièrement parallélisables qui intègrent les SSMs et l'attention par blocs. Nous montrons que notre modèle surpasse les architectures basées sur les Transformers en termes de perplexité en modélisation du langage et généralise à des séquences plus longues. De plus, le Block-State Transformer démontre une augmentation de plus de dix fois en vitesse au niveau de la couche par rapport au Block-Recurrent Transformer lorsque la parallélisation du modèle est employée.

English

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.