Blok-Toestand Transformator

Samenvatting

State space-modellen (SSM's) hebben indrukwekkende resultaten getoond bij taken die het modelleren van langetermijnafhankelijkheden vereisen en efficiënt schalen naar lange sequenties vanwege hun subkwadratische rekentijdcomplexiteit. Oorspronkelijk ontworpen voor continue signalen, hebben SSM's superieure prestaties getoond bij een breed scala aan taken, zowel in visuele als auditieve domeinen; echter blijven SSM's achter bij de prestaties van Transformers bij taalmodelleertaken. In dit werk stellen we een hybride laag voor, genaamd Block-State Transformer (BST), die intern een SSM-sublaag combineert voor langetermijncontextualisatie en een Block Transformer-sublaag voor kortetermijnrepresentatie van sequenties. We bestuderen drie verschillende, en volledig paralleliseerbare, varianten die SSM's en bloksgewijze aandacht integreren. We tonen aan dat ons model vergelijkbare Transformer-gebaseerde architecturen overtreft wat betreft perplexiteit bij taalmodelering en generaliseert naar langere sequenties. Daarnaast toont de Block-State Transformer een meer dan tienvoudige snelheidstoename op laagniveau vergeleken met de Block-Recurrent Transformer wanneer modelparallelisatie wordt toegepast.

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.