Трансформер с блочной структурой состояний

Аннотация

Модели пространства состояний (SSM) демонстрируют впечатляющие результаты в задачах, требующих моделирования долгосрочных зависимостей, и эффективно масштабируются на длинные последовательности благодаря их субквадратичной сложности времени выполнения. Изначально разработанные для работы с непрерывными сигналами, SSM показали превосходную производительность в широком спектре задач, включая обработку изображений и звука; однако в задачах языкового моделирования SSM по-прежнему уступают трансформерам. В данной работе мы предлагаем гибридный слой под названием Block-State Transformer (BST), который внутренне объединяет подслой SSM для долгосрочной контекстуализации и подслой Block Transformer для краткосрочного представления последовательностей. Мы исследуем три различных, полностью параллелизуемых варианта, которые интегрируют SSM и блочное внимание. Мы показываем, что наша модель превосходит аналогичные архитектуры на основе трансформеров по перплексии в языковом моделировании и обобщается на более длинные последовательности. Кроме того, Block-State Transformer демонстрирует более чем десятикратное увеличение скорости на уровне слоя по сравнению с Block-Recurrent Transformer при использовании параллелизации модели.

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.