Comparing transformers and state-space models for sequential learning

A Comparative Analysis of Attention versus State-Space Models for In-Context Learning

Machine Learning

Summary

This paper looks at two ways computers learn from sequences of information: transformers, which use attention, and state-space models (SSMs), which handle memory differently. The authors create a new way to compare these models based on how well they can gather evidence, keep track of what they know, and use past clues. They find that SSMs are better at some memory tasks, but attention models have advantages in how they select important information. Their findings are tested both theoretically and practically, showing insights that could guide better design of learning systems.

What this means in practice

Authors

Enes Arda, Semih Cayci, Atilla Eryilmaz

Abstract

Transformers and state-space models (SSMs) are two prominent sequential learning architectures, yet their comparison remains largely empirical and existing theoretical analyses are typically task-specific or architecturally restricted. In this paper, we develop belief geometry, a unified analytical framework for comparing the representational capabilities of broad classes of attention and SSMs. Starting from a generalized formulation of in-context linear regression and using cumulative Bayes regret as our measure, we abstract three capabilities required by many sequential learning problems in our belief geometry: evidence assembly, belief maintenance, and addressing. We then study three cases of our formulation that isolate these capabilities and yield sharp architectural lessons: For belief maintenance, SSMs attain the optimal regret over stationary aggregation kernels; for positional assembly, SSMs have a memory advantage; and for content addressing, softmax attention has an exponential width advantage over sigmoid-selective SSMs. Experiments with LLaMA-type Transformers and Mamba-2 show that these architectural insights extend beyond our analytically tractable classes and linear-regression testbed.