Transformador de Estado de Bloco

Resumo

Modelos de espaço de estados (SSMs, do inglês *State Space Models*) têm demonstrado resultados impressionantes em tarefas que exigem a modelagem de dependências de longo alcance e escalam eficientemente para sequências longas, graças à sua complexidade de tempo de execução subquadrática. Originalmente projetados para sinais contínuos, os SSMs têm mostrado desempenho superior em uma variedade de tarefas, tanto em visão quanto em áudio; no entanto, os SSMs ainda ficam aquém do desempenho dos Transformers em tarefas de modelagem de linguagem. Neste trabalho, propomos uma camada híbrida denominada *Block-State Transformer* (BST), que combina internamente uma subcamada SSM para contextualização de longo alcance e uma subcamada *Block Transformer* para a representação de curto prazo de sequências. Estudamos três variantes diferentes e completamente paralelizáveis que integram SSMs e atenção em blocos. Mostramos que nosso modelo supera arquiteturas baseadas em Transformers semelhantes em termos de perplexidade na modelagem de linguagem e generaliza para sequências mais longas. Além disso, o *Block-State Transformer* demonstra um aumento de mais de dez vezes na velocidade no nível da camada em comparação com o *Block-Recurrent Transformer* quando a paralelização do modelo é empregada.

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.