Implicit Language Models are RNNs: Balancing Parallelization and Expressivity (2502.07827v3)

Abstract: State-space models (SSMs) and transformers dominate the LLMing 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 LLMs 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 .

• The paper introduces implicit state-space models (SSMs), leveraging deep equilibrium models, to bridge the gap between RNN expressivity and transformer/explicit SSM parallelization.
• Implicit SSMs achieve non-linear state transitions, enabling performance gains on state-tracking tasks and scaling to large language models with competitive perplexity.
• This work suggests implicit models could unify sequence modeling architectures, offering practical benefits for training large models and opening avenues for hardware acceleration and broader applications.

Implicit LLMs are RNNs: Balancing Parallelization and Expressivity

The paper "Implicit LLMs are RNNs: Balancing Parallelization and Expressivity" presents a novel approach to LLMing that utilizes the structure of recurrent neural networks (RNNs) to enhance expressivity while maintaining the parallel computational advantages of transformers and state-space models (SSMs). The authors introduce implicit SSMs, leveraging the idea of deep equilibrium models (DEQs), to theoretically and empirically bridge the gap between the expressivity of RNNs and the training efficiency of more recent LLMs like transformers.

Key to this advancement is the balance between expressivity and parallelization. Classical RNNs, while expressive, struggle with parallel processing, limiting their scalability for large-scale machine learning problems. Transformers and SSMs, on the other hand, although parallelizable, cannot effectively capture the sequential dependencies required for tasks that necessitate robust state tracking. By allowing self-iterations in the depth dimension, implicit SSMs achieve convergence to fixed points akin to the state transitions in RNNs, thus enabling these models to solve complex sequential problems that transformers and explicit SSMs inherently struggle with.

Theoretical Insights and Empirical Validation

The foundational theoretical contribution of this paper is demonstrating that implicit SSMs can achieve non-linear state-to-state transitions that are not constrained by the diagonal Jacobian seen in linear SSMs. This non-linearity is essential for representing finite state machines (FSMs) and, consequently, recognizing certain classes of regular languages beyond what transformers or linear SSMs can manage. The authors leverage the fixed-point iteration approach from DEQ models, showing that only an approximation to the fixed-point convergence is necessary, thereby allowing a scalable training procedure that retains parallelizability.

Empirical results highlight the capabilities of implicit SSMs in state tracking, outperforming transformers and explicit SSMs on regular language tasks. Furthermore, the authors pretrain large implicit LLMs, scaling up to 1.3 billion parameters on 207 billion tokens, achieving lower perplexities than their explicit counterparts. This not only demonstrates scalability but also underscores the competitive performance of implicit models on standard NLP benchmarks.

Implications and Future Directions

The implications of this work are manifold. Practically, the ability to balance parallelization and expressivity without sacrificing performance offers a significant advantage for training LLMs. Theoretically, this research suggests that implicit models could serve as a unifying framework, potentially inspiring further advances that refine the capabilities of neural networks for sequence modeling.

The introduction of implicit SSMs invites exploration into hardware accelerations specific to self-iteration computations, particularly as new neuromorphic and photonic hardware emerges. Additionally, the results on downstream NLP tasks indicate that implicit models could be a pivotal step in addressing the inherent limitations of current LLMing architectures, such as their abilities in reasoning and in-context learning. Further research could explore optimization strategies for implicit models, investigate their performance on fine-tuning tasks, and expand their application scope beyond natural language processing to other domains like protein folding and autonomous robotics.

Overall, this work represents a significant addition to the discourse on neural network architectures, empirically validating the benefits of revisiting RNN-like structures within a modern framework, bridging the longstanding gap between parallel computation and expressive capacity.