Block-State-Transformer

Zusammenfassung

State-Space-Modelle (SSMs) haben beeindruckende Ergebnisse bei Aufgaben gezeigt, die die Modellierung langreichweitiger Abhängigkeiten erfordern und sich aufgrund ihrer subquadratischen Laufzeitkomplexität effizient auf lange Sequenzen skalieren lassen. Ursprünglich für kontinuierliche Signale entwickelt, haben SSMs in einer Vielzahl von Aufgaben in den Bereichen Vision und Audio überlegene Leistungen gezeigt; dennoch hinken SSMs in Sprachmodellierungsaufgaben der Leistung von Transformern hinterher. In dieser Arbeit schlagen wir eine hybride Schicht namens Block-State Transformer (BST) vor, die intern eine SSM-Unterebene für die langreichweitige Kontextualisierung und eine Block-Transformer-Unterebene für die kurzfristige Repräsentation von Sequenzen kombiniert. Wir untersuchen drei verschiedene, vollständig parallelisierbare Varianten, die SSMs und blockweise Aufmerksamkeit integrieren. Wir zeigen, dass unser Modell ähnliche Transformer-basierte Architekturen in Bezug auf die Sprachmodellierungsperplexität übertrifft und sich auf längere Sequenzen verallgemeinern lässt. Darüber hinaus demonstriert der Block-State Transformer im Vergleich zum Block-Recurrent Transformer bei der Modellparallelisierung eine mehr als zehnfache Geschwindigkeitssteigerung auf der Ebenenebene.

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.