State-space models can learn in-context by gradient descent

Abstract: Deep state-space models (Deep SSMs) are becoming popular as effective approaches to model sequence data. They have also been shown to be capable of in-context learning, much like transformers. However, a complete picture of how SSMs might be able to do in-context learning has been missing. In this study, we provide a direct and explicit construction to show that state-space models can perform gradient-based learning and use it for in-context learning in much the same way as transformers. Specifically, we prove that a single structured state-space model layer, augmented with multiplicative input and output gating, can reproduce the outputs of an implicit linear model with least squares loss after one step of gradient descent. We then show a straightforward extension to multi-step linear and non-linear regression tasks. We validate our construction by training randomly initialized augmented SSMs on linear and non-linear regression tasks. The empirically obtained parameters through optimization match the ones predicted analytically by the theoretical construction. Overall, we elucidate the role of input- and output-gating in recurrent architectures as the key inductive biases for enabling the expressive power typical of foundation models. We also provide novel insights into the relationship between state-space models and linear self-attention, and their ability to learn in-context.

• The paper introduces GD-SSM, a state-space model variant that mimics one-step gradient descent using a diagonal recurrent layer as a gradient accumulator.
• The paper shows that stacking GD-SSM layers and integrating MLPs enables effective multi-step and non-linear regression, highlighting scalability and robustness.
• The paper’s empirical results reveal that a single-layer GD-SSM can match transformer performance on synthetic linear tasks, underscoring its efficiency.

In-Context Learning in State-Space Models via Gradient Descent

The paper "State-space models can learn in-context by gradient descent" explores the capabilities of deep state-space models (SSMs) in in-context learning tasks, investigating their potential as efficient alternatives to transformers. The authors argue that SSMs, when constructed with specific architectural features, can mimic the gradient descent mechanism commonly associated with in-context learning in transformer-based models.

Architectural and Theoretical Insights

The study introduces a variant of SSMs, termed GD-SSM, which integrates a single structured state-space model layer augmented with local self-attention. This layer is demonstrated to replicate the outputs of an implicit linear model after one gradient descent step. The core insight is that the diagonal linear recurrent layer within the SSM acts as a gradient accumulator, effectively aligning with the parameters of the implicit regression model.

The authors extend their theoretical exploration to multi-step and non-linear regression tasks. They establish that stacking layers in the GD-SSM enables multi-step gradient descent, while the introduction of Multi-Layer Perceptrons (MLPs) facilitates handling non-linear tasks. The SSMs with these enhancements remain competitive across various regression problems, showcasing their expressiveness and scalability.

Empirical Validation

Empirical results support the theoretical framework, with trained GD-SSMs on synthetic linear regression tasks exhibiting losses that match those calculated from analytical constructions. These results persist even when tasks deviate from those encountered during training, underscoring the robustness of the model in generalizing learning rules.

Comparative performance evaluations further highlight the competitiveness of GD-SSMs against traditional transformers and other recurrent networks. Notably, while transformers may require multiple layers to replicate similar tasks, a single-layer GD-SSM suffices, emphasizing the model's efficiency.

Implications and Future Directions

The findings illuminate the potential for SSMs to serve as efficient and scalable alternatives to transformer architectures in tasks requiring in-context learning. This points to broader implications in designing models with intrinsic support for gradient-based updates, extending beyond simple autoregressive tasks.

The study suggests several directions for future exploration. These include scaling GD-SSMs in more complex and higher-dimensional tasks, integrating additional model components for enhanced capabilities, and examining the architectural features that contribute to efficient in-context learning further.

In conclusion, this paper contributes significantly to understanding the architectural and functional capacities of state-space models in in-context learning, providing a foundation for future research and practical implementations in AI systems.