Implicit Language Models are RNNs: Balancing Parallelization and Expressivity

TL;DR

This paper introduces implicit state-space models that combine the parallelization benefits of transformers with the expressivity of RNNs, demonstrating superior performance on language modeling and reasoning tasks.

Contribution

The paper proposes implicit SSMs that implement RNN-like non-linear state transitions through fixed-point iteration, enabling scalable training with retained parallelization.

Findings

Implicit SSMs outperform transformers and explicit SSMs on language tasks.

Approximate fixed-point convergence suffices for effective training.

Largest implicit language model trained to date with 1.3B parameters.

Abstract

State-space models (SSMs) and transformers dominate the language modeling landscape. However, they are constrained to a lower computational complexity than classical recurrent neural networks (RNNs), limiting their expressivity. In contrast, RNNs lack parallelization during training, raising fundamental questions about the trade off between parallelization and expressivity. We propose implicit SSMs, which iterate a transformation until convergence to a fixed point. Theoretically, we show that implicit SSMs implement the non-linear state-transitions of RNNs. Empirically, we find that only approximate fixed-point convergence suffices, enabling the design of a scalable training curriculum that largely retains parallelization, with full convergence required only for a small subset of tokens. Our approach demonstrates superior state-tracking capabilities on regular languages, surpassing transformers and SSMs. We further scale implicit SSMs to natural language reasoning tasks and pretraining of large-scale language models up to 1.3B parameters on 207B tokens representing, to our knowledge, the largest implicit model trained to date. Notably, our implicit models outperform their explicit counterparts on standard benchmarks. Our code is publicly available at http://github.com/microsoft/implicit_languagemodels .