Trasformatore a Stato a Blocchi

Abstract

I modelli a spazio di stato (SSM) hanno dimostrato risultati impressionanti in compiti che richiedono la modellazione di dipendenze a lungo raggio e si adattano efficientemente a sequenze lunghe grazie alla loro complessità di runtime subquadratica. Originariamente progettati per segnali continui, gli SSM hanno mostrato prestazioni superiori in una vasta gamma di compiti, sia nel campo della visione che dell'audio; tuttavia, gli SSM rimangono ancora indietro rispetto alle prestazioni dei Transformer nei compiti di modellazione del linguaggio. In questo lavoro, proponiamo uno strato ibrido denominato Block-State Transformer (BST), che combina internamente un sottostrato SSM per la contestualizzazione a lungo raggio e un sottostrato Block Transformer per la rappresentazione a breve termine delle sequenze. Studiamo tre varianti diverse e completamente parallelizzabili che integrano SSM e attenzione a blocchi. Dimostriamo che il nostro modello supera architetture basate su Transformer simili in termini di perplessità nella modellazione del linguaggio e si generalizza a sequenze più lunghe. Inoltre, il Block-State Transformer mostra un aumento di velocità di oltre dieci volte a livello di strato rispetto al Block-Recurrent Transformer quando viene impiegata la parallelizzazione del modello.

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.