Transformador de Estado de Bloques

Resumen

Los modelos de espacio de estados (SSMs, por sus siglas en inglés) han demostrado resultados impresionantes en tareas que requieren modelar dependencias de largo alcance y escalar eficientemente a secuencias largas, gracias a su complejidad de tiempo de ejecución subcuadrática. Originalmente diseñados para señales continuas, los SSMs han mostrado un rendimiento superior en una amplia variedad de tareas, tanto en visión como en audio; sin embargo, aún se quedan atrás en comparación con los Transformers en tareas de modelado de lenguaje. En este trabajo, proponemos una capa híbrida llamada Block-State Transformer (BST), que combina internamente una subcapa SSM para la contextualización de largo alcance y una subcapa Block Transformer para la representación a corto plazo de secuencias. Estudiamos tres variantes diferentes y completamente paralelizables que integran SSMs y atención por bloques. Demostramos que nuestro modelo supera a arquitecturas basadas en Transformers similares en términos de perplejidad en modelado de lenguaje y generaliza mejor a secuencias más largas. Además, el Block-State Transformer muestra un aumento de más de diez veces en velocidad a nivel de capa en comparación con el Block-Recurrent Transformer cuando se emplea paralelización de modelos.

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.