Implicit Language Models are RNNs: Balancing Parallelization and Expressivity
Authors: Mark Schöne
Babak Rahmani
Heiner Kremer
Fabian Falck
Hitesh Ballani
Jannes Gladrow
Venue: ICML 2025 spotlightposter (2025)
Keywords: State-space models
deep equilibrium models
RNN
transformer
large-language models
sequence models
regular languages
Chomsky hierarchy
Primary area: deep_learning->algorithms
TL;DR: Implicit SSMs bridge RNN expressiveness and transformer parallelization by iterating transformations to approximate fixed points, enabling scalable training and improved performance on state-tracking tasks and large-scale language modeling.

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 github.com/microsoft/implicit_languagemodels

Implicit Language Models are RNNs: Balancing Parallelization and Expressivity
Mark Schöne 12
Babak Rahmani 22
Heiner Kremer?
Fabian Falck
Hitesh Ballani?
Jannes Gladrow
Abstract
Length Extrapolation S5
Layers Required for
Mamba2 layens)
Mamba2
State-space models (SSMs) and transformers dominate the language modeling landscape. How- ever, they are constrained to lower computa- tional complexity than classical recurrent neu- ral networks (RNNs), limiting their expressiv- ity. In contrast, RNNs lack parallelization dur- ing training, raising fundamental questions about the trade off between parallelization and expres- sivity. We propose implicit SSMs, which it- erate transformation until convergence a fixed point. Theoretically, show that implicit SSMs implement the non-linear state-transitions of RNNs. Empirically, we find that only ap- proximate fixed-point convergence suffices, en- abling the design of a scalable training curricu- lum that largely retains parallelization, with full convergence required only for small subset of tokens. Our approach demonstrates superior state-tracking capabilities regular languages, surpassing transformers and SSMs. We further scale implicit SSMs to natural language reason- ing tasks and pretraining of large-scale language models up to 1.3B parameters on 207B tokens- representing, our knowledge, the largest im- plicit model trained to date. Notably, our im- plicit models outperform their explicit counter- parts on standard benchmarks. Our code is pub- licly available at github. com/miie crosoft/ implicit. lecliiit__laane languagemodels.
Implicit Mamba2 (1 layet)
Implicit Mamba2
1.0
Min PLayers
16
Accuracy
0.5
0.0
16
64
128
32
128
Sequeniee length
Sequence length
Downstream Test-Time Scaling
LAMBADA
Hellaswug
Implicit Mamba2
Implicit Llama3
0.8
0.8
Accuracy
Accuracy
0.6
0.6
0.4
0.4
0.2
0.2
Self-iterations
Self-iterations
Language Model Scaling the Pile
-4- Mamba2 Llana3
Implicit Mamba2 + Implicit Llama3
14
Puplexity
12
8-
130M
370M
7 700M
1300M
Number parameters
Figure Top Left: Minimum layers required to solve the S; word problem, a theoretically hard formalization of state tracking, for different sequence lengths. Top Right: Length generalization for Mamba2 and our implicit Mamba2 trained on L = 32 and extrapolated up to L 128. Center: Down- stream task accuracy for a range of self-iterations at test- time for our implicit 1.3B parameter Mamba2 and Llama3 models. Dashed lines indicate explicit models' accuracy. Bottom: Scaling of language models pretrained on 207B tokens the deduplicated PILE.
1. Introduction
Transformers, despite their dominance on contemporary language benchmarks, exhibit fundamental limitations in computational expressiveness. Both thcoretically and empir-
'Equal contribution 'Chair of Highly-Parallel VLSI Sys- tems and Meuro-Microelectroniee TUD Dresden University of Technology, Dresden, Germany "Microsoft Research, Cam- bridge, United Kingdom. Correspondence to: Jannes Gladrow @microsoft.com>..
ically, they cannot fully recognize regular languages (Bhat- tamishra et al., 2020) or, equivalently, represent finite state machines (FSMs) (Merrill et al., 2022). This limitation is significant because FSMs form the backbone of many real- world state-tracking problems, including evaluating code,
<jannes, gladrow <jannes.gladrow
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
tracking object permutations (e.g., in games like chess or structured narratives), and modeling sequential dependen- cies in logic (Li et al., 2021), I), location tracking (Guan et al., 2023), games (Li et al., 2023) and scientific applications such as protein generation, genetics, and chemistry (Briand et al., 2023; Chowdhury et al., 2022; Boiko et al., 2023). This raises questions about the ability of transformers to maintain coherent world models based on transitions be- tween states (Vafa et al., 2024) and hence, their suitability for tasks requiring robust state-tracking. These shortcom- ings appear to stem from fundamental trade-off between parallclizability at training time and the ability to track state (Merrill & Sabharwal, 2023).
sparse set of tokens that cannot be resolved by transform- ers (and SSMs). Such non-solvable transitions critical for state tracking but remain intractable for the class of circuits representable by transformers and SSMs (Merrill et al., 2022; 2024). Exploiting these properties, we devise implicit models that combine the expressive power of RNNs with the parallelizability of transformers and SSMs (see Fig- ure 2). In contrast to conventional transformers and SSMs, implicit models can track state, even out-of-distribution (see Figure Right).. In contrast to RNNs, these models per- mit a much larger degree of parallelization as the depth of self-iteration is much smaller than the sequence length (see Figure 3Mid).
Surprisingly, recently emerging state-space models (SSM), a class of linear recurrent neural networks, are bound by the same trade-off: despite their seemingly sequential nature they cannot express some inherently sequential problems such as certain regular languages (Merrill et al., 2024). In contrast, non-linear recurrent neural networks (RNNs) are not bound by these restrictions on compute complexity and can track state (Siegelmann& & Sontag, 1992; Merrill, 2019) but lack parallelizability at scale. This raises the question: How much sequential processing does one have to accept to solve the state tracking problem?
Contributions. (a) We propose implicit SSMs and show theoretically that they represent non-linear and non-diagonal state-to-stun transitions similar to RNNs. (b) We confirm empirically that implicit SSMs can solve the S5 word prob- lem, which conventional SSMs and transformers fail to solve. (c) We show by constructing distributions with vary- ing difficulty level over the word problem that implicit SSMs as well as transformers require much fewer non- parallelizable transitions to learn word problems than RNNs (d) We demonstrate scalability of implicit models through a carefully chosen training curriculum that bounds the num- ber of iterations, training implicit SSM and transformers up to 1.3B parameters on 207B tokens of the deduplicated PILE (D-PILE) (Gao et al., 2020)- see Figure IBottom, the largest self-iterated model with dynamic halting condition to date, to the best of our knowledge. (e) We highlight a set of properties of our pretrained implicit language models such as favorable length generalization, and path-independent auto-regressive generation.
Previous attempts to address these limitations in trans- formers have leveraged non-linear transitions through self- iteration in the depth dimension (Dehghani et al., 2019; Banino et al., 2021). However, backpropagation through unrolled networks is computationally prohibitive at scale. Deep equilibrium (DEQ) models (Bai al., 2019), in con- trast, define function implicitly via the fixed-points of neural network; their output is the result of self-iteration until convergence. Training such networks requires back- propagation solely at the fixed point, eliminating the need to traverse the iterative path and thereby decoupling mem- ory usage from the depth of iterations. Emerging hard- ware promising rapid computation of fixed-points of neu- ral networks (Brunner et al., 2025) may tilt the hardware lottery (Hooker, 2020) in favor of such implicit models, making this an opportune moment to explore their potential.
2. Background
2.1. State-Space Models
SSMs are linear recurrent models which produce an output yi Rd given an input 2, E Rd and a sequentially updated hidden state vector he E RD via the recurrence
Our approach to balancing state tracking and parallelization relies on two key observations. First, we demonstrate that implicit models naturally adapt their compute load to the difficulty of the learning problem (see Figure 3eeft). At both training and test time, such models effectively inter- polate between their parallelizable form, when all tokens in the sequence are resolvable, and RNNs, when there are no resolvable tokens. Further, we show theoretically that implicit models have indeed non-linear token-to-token tran- sitions similar to RNNs. Second, based on the success of transformers on many practical language modeling prob- lems, we hypothesize that natural language contains only
ht A(xz)le-1+ + u((z))
yt = f(1n-1,21),
(2)
where u and f are possibly non-linear learned functions. The learned matrix A E Rax= is typically diagonal and can be constant (Gu et al., 2022; Smith et al., 2023) or an input-dependent matrix-valued function (Qin et al., 2023; Gu & Dao, 2023; Dao & Gu, 2024). A SSM combines a number of these blocks with non-linear feed-forward blocks. In contrast to non- non-linear RNNs, the linear state recurrence (1) allows for training parallelism along the sequence di- mension, and avoids the quadratic scaling of self-attention.
2

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
2.2. Limitations of Transformers and SSMs
3. Implicit Sequence Models
Efficient parallelization is one of the central features en- abling transformers and SSMs scale to large machine learning problems such as language modeling, Parallel cir- cuits, however, face fundamental trade-offs regarding the class of problems that they can address. In particular, trans- formers and SSMs theoretically fail to recognize certain reg- ular languages, or equivalently, to simulate FSMs (Merrill et al., 2022; 2024). Empirical studies have confirmed that neither of the models are capable of learning the algorithms constituting certain regular languages (Bhattamishra al., 2020; Sarrof et al., 2024). By contrast, the sequential nature of RNNs allows them to express all regular languages (Mer- rill, 2019). A detailed discussion is given in Appendix B.1.
3.1. Implicit State-space Models
The linear recurrence of SSMs shown in Equation (1) cannot resolve elaborate sequential problems (Merrill et al., 2024). Here, we propose to exploit self-iterations along the depth of neural networks to close the expressivity gap between SSMs and RNNs. Following the DEQ paradigm (see Sec- tion 2.3), we implicitly define a model via the fixed points of a SSM. Introducing the iteration variable 2!²) E Rd to Equations (1) and (2) yields the fixed point iteration
11 = A 22
(1(-11) ,001) M + u
( (s-1)
2t : 8t
(6)
=t = so (0-11) 11(11 xt
2.3. Deep Equilibrium Models
where =10) = 0 for t o,..., T and h/ =0 for
Most deep learning architectures explicitly parametrize a function x + y with a neural network. Deep Equilibrium Models (DEQ), in contrast, define a function implicitly via the fixed-pole pointo of an input-conditional neural network, i.e.,
O, O,...,S respectively. In practice, we rely on the Mamba2 (Dao & Gu, 2024) architecture and inject the ex- panded input 2₁ additively after the initial projection.
Computing the fixed-points and the gradient around the fixed points of Equations (6) and (7) requires to iterate the two loops
Fo(2",
(3)
1..... T along the input sequence and = 0....., S to converge to fixed points along depth. Fig- ure visualizes how these two loops give rise to two distinct modes of evaluation. The simultaneous mode simultane- ously finds the fixed points for all (see Figure 2A). and exploits parallelization strategies for SSMs (Dao & Gu, 2024). The sequential mode resolves the and loops in the transpose order, and processes sequences sequentially just like classical SSMs or RNNs (see Figure 2B). While the simultaneous mode allows for highly parallel training, the sequential mode enables efficient inference at constant memory, e.g. for language generation. The emergence of these two strategies for finding the fixed points sparks the question if they are interchangeable as visualized in Fig- ure 6. Can model be trained in the simultaneous mode and deployed for inference in the sequential mode? This ques- tion is answered in Section 5 and Figure for real language models.
where 2" is identified with the prediction y. Naively differ- entiating a loss function L(=*) with respect to the model pa- rameters 0 generally requires a costly differentiation through the employed fixed- fixed-point solver. Instead, to allow for gradi- ent computations with a constant memory footprint, DEQs utilize the Implicit Function Theorem:
Let Gz(z,x)= = Fo(z, If the Jacobian Ja, of G w.r.t. is non- singular in 2², then there exists an open set U around (ar 0) and a unique function t on U such that (xx = zd and G(O(ī, 0), x, 0) = 0 for all (ž, 0) E U. Furthermore, the derivative of &b w.r.t. θ is given by
0$ = - G. 00
oFa
(4)
00
A range of methods have been proposed to efficiently com- dc = 06 Oc using Equation (4) (Bai et al., 2019; ge 0. Geng et al., 2021). Here, we employ the Phantom Gradient approach of Geng et al. (2021) (see in the Appendix Fig- ure 5). The method is based on solving a smoothed version of the fixed point Equation (3) combined with a finite trun- cation of the von Neumann series of the Jacobian-vector- product in (4) given as
pute a0
For both modes, Equations (6) and (7) in the limit -> 20 read
h; A (2f 24) h (27,
(8)
2, = fo (2f 1;
(9)
0Φ
OF, 0e
Σ i=
OFo A 0z
+ +(1 +(1-1)7) > (5)
where 2 lime-roo (00) and h lim k((0) denote the fixed points. Equations (8) and (9) couple the outputs 2f with the states hj via the non-linear functions fo, A and u. Notably, what appears as minor technical modification of the original state-space model in Equations (1) and (2) leads to a fundamentally enhanced expressivity of the recurrent model.
00
where a small smoothing parameter x E (0, 1] helps main- taining a small condition number at the cost of increased fixed-point iterations and the truncation length & determines the accuracy of the approximation.
3

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Simultancous fixed-points
Sequential fixed-points
Simultancous vs Sequential Language Model
A
dog walked through
the
dog A
wallked
throughh
the
Absolute Dirt Perplexity
0.4
Implicit Mamba2
Implicit Liama
0.3
0.2
pea streling saps
0.1
Promoom Ginder Corperrien
Legend A/B
tokeea stregee
lokana
Flard poiet
0.0
256
1024
2048
Edkkt-aali care
Token Position
Figure 2: A: The simultaneous mode self-iterates the entire sequence such that trajectories interact during convergence. It exploits the parallelism of the backbone model. B: The sequential mode iterates each token individually. Only converged hidden states or kv-caches are passed on. This mode is used for generation. C: Difference in perplexity between the two modes for our 1.3B implicit models.
Theorem 1. An implicit SSM defined by Equations (8) and (9) with generic weights yields non-linear and non- diagonal state-to-state transition function hi-1 + hū- If further 2 fo (zī, "4t,hhi-1 1i) = 0 is non-singular w.r.t =i, then there exists a continuously differentiable function sp such that 27 = (hi. (2, 0). In this case, the state-to-stane Jacobian is given by
the attention operator as Attn( W WOR + Iz Win)), where WoRV Rdx3d produces the Q, K, V for the multi-head self-attention (Attn). and Winp E Rxd3d is the input projec- tion. Conventional transformers, with their finite number of layers, cannot learn certain formal languages outside of the TC circuit complexity class (Merrill et al., 2022; Strobl et al., 2024). However, chain of thought (CoT) mod- els ( (Wei et al., 2022) bypass this restriction by using an adaptive compute budget through recursive generation of intermediate tokens (Merrill Sabharwal, 2024). Implicit transformers (Bai et al., 2019) utilize adaptive compute budget differently, using fixed-point iterations that can be in- terpreted as sequences of latent thoughts (Hao et al., 2024), undergoing non non-linear updates similar to a non-linear RNN's hidden state.
dh;
0A 0Gp diag (hā-1)
A xz)
+
d
dzj ihi-1
+ Ou Oe
(10)
azt Ohi-1
Proof: We refer the reader to Appendix A.2 for the proof and an extended statement of the theorem. A numerical check of Equation (10) is provided in Figure 7.
DI
4. Implicit SSMs Adapt to Hard Languages
As discussed in Section 2.2, non-linear RNNs surpass trans- formers and linear SSMs in terms of circuit complexity. By the above construction, our implicit SSM exhibits the favourable computational properties of RNNs, lifting the illusion of state linear SSMs (Merrill et al., 2024). Fur- thermore, the gradients of fixed point iteration depend solely on the fixed point, and not on the path to the fixed point, by the implicit function theorem. This suggests that both modes resolving the two for loops yield functionally equivalent fixed points.
Implicit SSMs Lift the Illusion of State The illusion of state (Merrill et al., 2024) reveals that constant depth SSMs cannot simulate arbitrary finite state machines. A hard state tracking problem in the sense that all state track- ing problems can be reduced to it is given by the word problem for the symmetric group S3 (Barrington, 1989). The word problem for a monoid (M, o) is to resolve arbi- trary length products of the form m my m2 o my for mi, m m E M, k N. A comprehensive intro- duction to the word problem and our particular learning setting is provided in Appendix C.
These properties raise the following hypotheses, which we will investigate empirically in this work.
Hypothesis (Expressivity). Implicit SSMs can learn and express all regular languages.
We train a set of Mamba2 SSMs (Dao & Gu, 2024) to repro- duce the results of Merrill al. (2024). Figure Left high- lights that Mamba2 requires more layers as the sequences get longer. For example resolving sequences of 32 elements from S₁ requires a minimum of 16 layers. Extending the result of Merrill et al. (2024), Figure 1Right shows that the same Mamba2 model with 16 layers does not gener- alize beyond the training distribution when evaluated on sequences longer than 32 elements. Our implicit Mamba2, however, can utilize additional self-iterations at test-time to
Hypothesis 2 2aarassssaPPPPPlaaaaaaala (Parallelization). Implicit SSMs can be trained in simultancous mode and evaluated in sequential mode without loss in performance.
3.2. Implicit Transformers
Similar to implicit SSMs, one can define an implicit trans- former model (Bai et al., 2019) by injecting the input x, into

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
resolve longer sequences of up to 128 elements. This result establishes that implicit SSMs effectively learn to be RNNs. However, with naive unrolling in implicit SSMs, paralleliza- tion would still be challenging. In the following, we show a subtle yet important result: Implicit SSMs can adapt to word problems of varying difficulty even when trained with bounded depth.
focus our discussion on the implicit models in the following. The left panel of Figure 3 shows that implicit SSMs capture the underlying algorithm, as measured by out-of-distribttion evaluation with p = 0.5. even when trained on very few non-solvable tokens. While a fraction of 2 % hard tokens per sample (p- = 0.02) suffices for some configurations, reliable training can be observed from p = 0.1 on.
We are left with the question of how many self-iterations are required during training learn the algorithm intrinsic to the word problem. To answer this we trained a range of models with p = 0.1, setting a different upper bound on the number of self-iterations at training time. The number of self-iterations at test time is unbounded and solely defined by the fixed point iteration. The mid panel of Figure shows that small amount of down to 8 self-iterations at training time suffices to generalize from the distribution Do. 1 at training time to Do.s at test time. Interestingly, the number of test time self self-iterations is quite similar for the models trained with different upper bounds on the training time self- iteratioss hinting that the models learned similar algorithms. Note that the self-iterations required during training are significantly lower than the sequence length. For comparison, conventional RNN conducts L 256 non- non-parallelizable steps to solve the same problem, factor of 32 larger than the 8 self-iterations required by our implicit Mamba2. This comes at a cost: we need to self-iterate over every token. However, each self-iteration can be parallelized across the sequence dimension by the parallelization of the base model.
Languages with Sparse Non-Solvable Transitions SSMs excel in natural language processing tasks despite being theoretically constrained to the simple class of star- free formal languages (Sarrof et al., 2024). We conject that natural language is mostly composed of simple to com- prehend tokens, while harder tokens appear only sparsely. To study implicit models in a controlled learning environ- ment closer to natural language than the S₁ word problem, we construct a word problem that mixes simple and hard examples. Let M = M x G be a direct product of an ape- riodic monoid M and a non-solvable group G. A sequence mo..... ................................ My is sampled from M with replacement. To con- trol the number of hard examples and simple examples, we define a family of distributions Dp over M as follows. An element m E M is sampled uniformly at each step k; representing the presence of simple examples. On the other hand, we sample elements 9v E G \ {e} from G without the identify transformation, each with probability iap-1- The identity element gk = e E G is sampled with probability 1 --- p. The resulting transformations (m², 9s) are aperiodic at least when gk e, i.e. with probability 1- p.
In the right panel of Figure 3, we demonstrate that the phan- tom gradient is, in most cases, more effective method for gradient computation than backpropagation through the entire sequence of unrolling steps. To evaluate this, we train three variants of the Mamba2 model: (1) an implicit Mamba2, which self-iterates and employs phantom gra- dients; (2) an unrolled Mamba2, which backpropagates through all unrolling steps; and (3) an explicit Mamba2, conventional model. All models are trained on sequences of length I, 1= 256 sampled from Do. with depth constraint of 16 corresponding to 16 self-iterations for the implicit and unrolled models and 16 layers for the explicit model. Our result shows that a constant number of backpropagation steps using the phantom gradient method is enough to lear complex non-solvable transitions and generalize to difficult distributions at test time. Since phantom gradients require constant memory that is independent of the number of self-iteration steps, the training of larger language models appears feasible.
We will identify minimally sequential models that parallelize to a high degree and still capture all non-solvable transitions in language. Therefore, we apply our construction of word problem above to mix tokens from simple languages with tokens from non-solvable hard languages. This section studies a word problem over M = M x A5, where Ma is a simple aperiodic monoid with four elements and As C S₁ is the alternating group over 5 elements, the smallest non-solvable subgroup of S3- For details on the learning problem, we refer the reader to Appendix C.
Interpolating between SSMs and RNNs
We train Mamba2 and implicit Mamba2 models on range of mixtures of simple and hard tokens between p 0.0 and p- = 0.25, and in the case of the implicit models with varying self-iteration depths at training time between 2 and 128. All training sequences sample L- 256 tokens, and evaluation is conducted on the distribution Po.s, where half of the tokens is hard. The evaluation is hence an out-of- distribution (OOD) setting. We report averaged results over 10 random seeds with boostrapped 95 % confidence intervals as well as the best models per configuration. None of the conventional models got OOD accuracies beyond random chance as shown in the right panel of Figure 3, hence we will
CatbAbi: A benchmark requiring state tracking. To evaluate the state- state-tracking capabilities of SSMs on language tasks, we use the CATBABl dataset (Schlag et al., 2021), a modified version of the BABI dataset (Weston et al., 2015),
5

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Impact of Training Distribution
Impact of Self-iterations
Comparison of Models
OOD Accurncy 0p5)
OOD Accuracy (pm(1.5 bostze I mmm
OOD Accuracy
Training curves
bolt all
auplitt (mead)
implili (menn)
unnolld (mcan)
cxplict strolled inplict
(bee)
Test Accuncy 96
100
100
Test Accurse in %
100
Trai Avcuracy %
1000
75
Text Accuracy in
75
75
50
50
50
25
25-
0.02
0.25
128
Lont 1-134
pud. L=1124
pet. L-256
pit L-1014
4%
81
12k
16k
Fraction of hard tokens (Truining)
Max sull sell-iterations (Training)
Training steps
Figure 3: All models were trained and evaluated on sequences of length I, 256. The nut-o--distribution (OOD) evaluation conducted with p 50%. Left: Comparison of OOD accuracy for range of training distributions with hard token probabilities p. Mid: Comparison of OOD accuracy for range of self-iterations caps training time, trained with p 0.1. Right: Comparison of implicit Mamba2, unrolled Mamba2, and Mamba2 trained with 0.1. Unrolled Mamba2 unrolls a single layer with full backpropagation, while implicit Mamba2 receives only 4 Phantom Gradient steps. All models have a training depth of 16 (layers for Mamba2, self-iterations for implicit and unrolled).
consisting of 20 tasks within a 5M token corpus. These tasks, requiring various reasoning abilities like deduction, conference, or counting, involve short stories with embed- ded questions (Schlag al., 2021), and require state tracking in various degrees. We train our implicit SSM model, using Mamba2 as the core architecture, alongside the baseline Mamba2 model, both with up to three layers. Our findings show that the implicit Mamba2 model with a single layer outperforms its single-layer Mamba2 counterpart on most tasks. Additionally, more layers in the implicit model's backbone reduce the number of self-iteration steps needed to solve the tasks (see Appendix Figure 13a, Figure 13b). We furthermore evaluate the performance of the models for tasks sorted by increasing story length. We see how im- plicit models retain its performance as the lengths increases in Figure 13c at a slight increase in the number of iterations of the implicit models in Figure 13d.
of self-iterations to 24/32 followed by four steps of phantom gradient. We refer to these models as (24 + 4)/(32 + 4)- models for Mamba2/L model sizes: 125M, 350M 760M, and 1.3B. These models are pretrained in an autoregressive manner for next-token prediction across all sizes on the D-PILE (Gao et al., 2020) dataset, which consists of 207B tokens. For baselines, we use both Mamba2 (Dao & Gu, 2024) and Llama (Beck et al., 2024) models previously trained on a corpus of 300B tokens. Additionally, reproduce Mamba2" and Llamat as base- lines trained with the same code and data our implicit models. We evaluate the pretrained models the test set of the D D-PILE, examine their length extrapolation capabilities, and assess their common sense reasoning performance on downstream tasks. See Appendix D.3 for pretraining details.
respectively. We employ four
Pretraining Results and Downstream Performance. We report Table the next-token perplexity performance of all models trained on the entire 207B token corpus using a test split of the D-PILE'. We observe our implicit models consistently achieve a lower perplexity compared to their explicit counterparts see also Figure Bottom.. For details related to the dynamics of the implicit models on D-PILE, refer to Table 6. Additionally, we evaluate the models' per- formance on common sense reasoning tasks using the LM Evaluation Harness (Gao al., 2024). The results show that implicit Mamba2 outperform the explicit Mamba2". which are pretrained on the same number of tokens, on most tasks. This difference becomes more pronounced as the size of the models increases, specifically with the 760M
5. Implicit Large Language Models
We investigate whether implicit models can be effectively pretrained to function as language models. Motivated by the results of Section 4, we implement a pretraining strat- egy for implicit models with two stages of bounded and free self-iterations. Transformer (LLama) Touvron et al. (2023) and SSM (Mamba2) (Dao & Gu, 2024) architectures serve as the core backbones for our implicit models. In the bounded stage, we train with four self ceff-iterations and a single step of phantom gradient, which we refer to as the (4 + 1)- model. The (s + -)-nttition refers to s gradient- tape--rpee self-iteration steps and & phantom gradient steps. k refers to Equation (5), see also Figure 5. The free stage starts from checkpoint of the (4 + mo-el and increases the number
'The test split represents a random selection of 0.1 percent of the entire dataset. This size is in line with the proportion used for the PIL.E's validation set (Gao et al., 2020).
6

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
and 1.3B variants. Compared to the original Mamba2 base- line, trained on 1.5 times more data, the implicit models do better on HELLASWAC, PIQA, ARC-E, and are com- petitive in LAMBADA and ARC-C. C. Across all scales, the implicit Mamba2 models significantly outperform Mamba2 in the HELLASWAG task, yet they underperform in WINO- GRANDE and OPENBOOKQA.
see Figure 6), which we empirically demonstrate with our pretrained language models. Specifically, we utilize our 1.3B implicit Mamba2 (24+4) and Llama (32+4) models to compute next-token predictions on the D-PILE test split. The models are fed identical input tokens of length 2048 in batches of size 16 and predict outputs greedily in both si- multaneous and sequential modes. We observe token match rates of 97.6 % (on 3M tokens) between the outputs of the two modes for the implicit Mamba2, and 97. % (on 330K tokens) for the implicit Llama. Examples of these model predictions are provided in Appendix Table 8. The per-token perplexity differences in the predictions of the models are depicted in Figure 2. To our knowledge, this is the first demonstration of sequential evaluation with self-iterated models at constant memory in the number of self-iterations, enabling auto-regressive generation for this class of models.
It also noteworthy that our implicit Llama models sub- stantially outperform the baseline Llamas, including both the results reported in (Beck et al., 2024) and the Llamat- This improvement is consistent across all tasks and model sizes. Strikingly, we note that our implicit Llama (32+4) 760M is competitive to the explicit Llamat 1.3B.
10
Mamba2
Mambu2 (4+1)
Mamba2"
Mamba2 (24+4)
6. Related Work
Perplexity
Adaptive-Compute Time The idea of an adaptive com- pute budget goes back to (Schmidhuber, 2012) who employ a halting neuron to delimit the computation on partic- ular input. Graves (2017) generalized the idea and regu- larised the halting condition encourage the network to stop early. They implemented an adaptive-depth RNN and demonstrated the network adjusting the compute budget based on the difficulty of instances in a parity-check task. This idea was later applied to Transformers, resulting in "Universal Transformers" (UT) (Dehghani et al., 2019). UTs can either be unrolled to a fixed depth or augmented with a dynamic halting condition (DHC) per token. UT's were later shown to exhibit improved scaling laws compared to stan- dard transformers (Kaplan et al., 2020). PonderNet (Banino et al., 2021) introduced a principled probabilistic model for determining the halting condition. This approach im- proved on the UT on the BABI benchmark. Recently, a coxture-of-expertt (MoE) variant of the UT (MoEUT) was presented (Csordás et al., 2024) with IB parameters, seek- ing to improve the trrameterr-o-compeee
10'
0
4000
8000
12000
16000
Token Position
Figure 4: Length extrapolation performance of the per token perplexity on the the test split of the D-PILE of the original 1.3B Mamba2, our Mamba2", and our implicit Mamba2 with (4+1) and (24+4) self-iterations. Shaded gray area shows the in-distribution length.
Implicit-SSMs Demonstrating Length Extrapolation Ca- pabilities All implicit models in our study were trained on sequences of 2048 tokens. To assess their capability for length extrapolation, we evaluated the implicit models on the test split of the D-PILE, which was packed with longer sequences consisting of 4096, 8192, and 16384 to- kens. We compared these results with the baseline Mamba2 and Mamba2' in Figure 4, where the per-token perplexities are reported. Numerical values for all models at a range of token positions are reported in Table 7 in the Appendix. In comparison with the baseline Mamba2 models, our implicit Mamba2 models demonstrate stronger robustness longer sequences.
ratio of UTs. The MoFUT is an unrolled model with fixed iterations and does not employ DHC. While our models presented here are dense, they could, in principle, be turned into MoE. Gatmiry et al. (2024) show that looped linear transformers implement gradient-descent until convergence on the prediction loss defined by previous input-output examples in the context window. Lim et al. (2024) take the opposite approach to our work: Instead of augmenting SSMs or transformers. they propose an approach based on fixed- point iterations to enable parallel training of RNNs. However, their method incurs cubic cost in terms of state size, limiting the method to smaller models.
Effective Duality between Simultaneous Mode and Se- quential Mode Autoregressive generation, core func- tionality of contemporary language models, for implicit models requires that the sequential mode introduced in Sec- tion 3 and Figure is functionally equivalent to the simul- tancous mode used for pretraining. Effectively, the loops over and t in Equation (6) have to be interchangeable (also
Reasoning and out-of-distribution generalization. The ability of looped models to generalize better to input lengths
7

---

Implicit Language Models are RNNs: Balancing Parallelization and Expressivity
Table 1: Comparison of test set perplexity and downstream performance. We compare our implicit models, which have 4 self- self-iteration steps and phantom gradient step (denoted as 4+1), and those with 24/32 self-iteration steps and 4 phantom gradient steps (denoted as 24+4/32+4), with our baseline models Mamba2* and Llama These baseline models as well as the implicit models are trained on 207B tokens from the D-PILE dataset and range in size from 130M to 1.3B parameters. For further comparison, we include the original Mamba2 (Dao & Gu, 2024) (trained on 300B tokens of the PILE) and the Llama (trained on 300B tokens of the SLIMPAJAMAI from (Beck et al., 2024). The best performing model for each type is highlighted in bold, and the second- best is underlined. The training wall-clock time (WCT) is provided for implicit models relative to their explicit counterpart. For (24 + 4)- and (32 + 4)-iterated models, we provided the weighted average measured over the entire curriculum. We omit WCT' for the 370M-model series since the runs were carried out on various node sizes, making direct comparison difficult.
Model
Datasel/ Tokens (8)
Rel. WCT lngl/Kop
D-Pile ppl.
LAMBADA
LAMBADA
HellaSwag
PIQA
Arc-E scc]
Arc-C nec7
WinuGrande acc7
OpenbookQA acc1
Average a0e7
ppl:
ace1
secT
acct
Mumhu2
Pile/300
13.72
16.81
0.438
0.3525
0.649%
0.4739
0.3423
0.5233
0.300
0.4266
Mamba2'
D-Pile/007
13.05
18.51
0:1116
0.3527
0.6572
0.4815
0.2572
0.5130
0,300
0:1219
MamkkZZ(4+1)-oom
D-Pile/207
1.97
13.76
0,4118
0.3628
0.6485
0.4537
0.2287
0.5107
0.288
0,4149
Mambu2(2444 (24-b)222224))
D-Pile/207
3.27
12.86
18:03
0.4174
0.3673
0.6496
0.4604
0.2372
0.5178
0.200
0.4200
Lhmma
S0imaajamas/00
39,21
0.3151
0.3109
0.6515
0.4533
0.2363
0.5067
0:4178
Lhana!
D-Pile/207
12.77
17.08
0.4297
0.3513
0.6540
0,4794
0.2440
0.5122
0.280
0.4215
1.amma (-+1)-
0.90
12.73
15.54
0.4518
0.3206
0.6447
0.4823
0.2372
0
0,290
0.4308
Llama (32+1)-ours
D-Pile/207
2,08
11:73
0.4501
0.3958
0.6676
0.488
0.4423
MambuZ
Pile/300
10.55
8.00
0,5593
0,4692
0.7046
0.5476
0.2671
0,5564
0.324
0.8897
Momhaz"
D.Pile/2077
10.18
0,5333
0,4653
0,6942
0.5526
0.26%
0.5320
0.4790
Munks264+1)oouss
D-Pilo/207
10.02
8.79
0,5457
0,4681
0.6899
0.5358
0.269
0,5162
0.308
0.4722
370M
Marnk22(2444:cours
D-Plle/207
9.70
8:26
0.5575
114792
0,7044
0. 1454
0.2688
0,5351
0,316
0
Llama
SlimPajama/300
15:73
0,4419
0.4445
0.6915
0.5223
0.2628
0.5359
0.4832
Llama!
D- D-Pile/207
10.30
8.37
0.5621
0.1537
0.6844
0.5477
0.2577
0
0,318
0.1826
D-Pile/007
9.66
0.5898
0.5030
0.7024
0.5539
0.2611
0.5572
0.4973
I.lama (52 4)-ous
D-Pile/207
9.43
7.04
0.5956
0.5114
0.7078
0.5244
0.2705
0.5722
0.320
0.5003
Mumbo2
Pile300
9.23
5.86
0.6167
0.5192
0,7198
0.6103
0,2850
0.6030
0,362
0.535]
Mumba2"
D-Pile/207
8.98
6.24
0.125
0.5418
0,2231
0 066044
0.2855
0.5777
0.338
0.5262
Mamba2(44 13-ouss 1)-omns
2,05
8.60
6.15
0.6117
0.5869
0.7296
0.6077
0.3140
0,5509
0,336
0.5295
260M
Mambo22(2444)-our
D-Pile/007
3.43
5.90
0.6191
0.5098
0.7334
0.6090
0.3131
0,5730
0.338
0.5365
Llama
SlimPajama/300 300
9.501
0.514]
0.5216
0. 0.70M5
0.5644
0.2875
0,5667
0.5274
Llama!
D.Pile/207
8.88
066355
0.5448
0.7171
0.5908
0.2816
0.6054
0.338
0.5307
Liana (4+])-ours
D-Pile/207
1.64
8:27
5:15
0.6524
0.5853
0,7312
0.6052
0.3097
0.5967
0.356
0.5481
Iaama (3244)-ours
D-Plle/207
2.97
7.90
4.82
0.6783
115995
0.7416
1.6187
0.3012
0.5991
0.344
0,5535
Mumhu2
Pile/300
5:012
0.5995
0.7378
0.6418
0.3319
0.6117
0.378
0_5682
Mamba2"
D-Pise2007
1
8:28
5.12
0.6455
0,5439
0,7416
0.6145
0.3123
0.6117
0,352
0.5531
D.Pile/0077
1.91
7:97
06633
0.6136
0.7437
0.6343
0.3302
0.5746
0,5555
Mamba2(24 (2nb)aa2(244)))omm 43)-ors
D-Pile2007
3.19
7.70
4.99
0.6489
0.6267
0.7416
0.6423
0.3336
0.5988
0.352
0.5620
Liama
300mPajamm/300
7.23
0.5744
0.5781
0,7312
0.6279
0.3174
0.5904
0.5699
Hama!
D-Pile7007
7.99
4.95
0.6599
0.5936
0.7432
0.6385
0.3217
0.6062
0,5589
Llama (4x 1-ours
D. DiPile2007
1.96
7.66
4,40
0,6852
0.6392
0.7448
0.6338
0.3396
0,6575
0.360
0.5801
Llama (32+4)-ours
D-Pile2207
3.91
0.6901
0.7465
0.6654
0.3601
0.6401
0.5892
not seen during training is empirically well established: For example Yang et al. (2024a) show this for looped transform- ers, while Anil et al. (2022) demonstrate length generaliza- tion for DEQs, particularly when they are path independent. Du al. (2022) show that energy-based models trained to map energy- mergy-gradienttdeseenn steps algorithmic steps, can length generalize in summation, and complex algorithms such as shortest-path. On the theoretical side, The pio- neering work of Siegelmann & Sontag (1992) shows that iterated RNNs are Turing complete at infinite numerical precision. More recently, Deletang et al. (2023) studied a number of sequence models and report that grouping tasks by their rung in the Chomsky hierarchy is predictive of models ability to length-generale Merrill el al (Merrill, 2019; Merrill et al., 2020; Merrill & Sabharwal, 2023; Merrill et al., 2024), which we discuss inSection 2.2, showed that both transformers and SSMs are restricted to TCo; several studies sought to find more pre-
cise constraints. Weiss et al. (2021) observe that programs written in a specific language (RASP) can be mapped to transformer models of sufficient capacity. Zhou et al. (2023) then showed that transformers tend to -ength-generalise if the underlying Bata-generating process can be expressed in RASP. Sarrof et al. (2024) derived a similar refined con- straint for SSMs and showed that they can precisely express star-free regular languages. Grazzi et al. (2025) demon- strate that SSMs can track state in simple problems, such as parity, when their (diagonal) recurrence matrix A in Equa- tion (1) permits negative eigenvalues. Moreover, they illus- trate that variant of DeltaNet (Yang et al., 2024b) with (possibly) negative eigenvalues can solve the S5 problem when only swaps of two values are considered in the tran- sition. However, no variant of Mamba or DeltaNet was capable of learning S5 and achieving length generalization. To tackle the wtteeiiiaatiooeeeeesseeeeeeeeeee trade-off, Beck et al. (2024) propose two new LST -iT--irpd layer archi-
While the works of
8

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
tectures: the sLSTM and mLSTM layers. While the latter is parallelizable, the former is not and intended to enable the whole model to recognize regular languages. Finally, Soulos et al. (2024) survey strategies for chunking input sequences with transformers, maintaining parallelizability within each chunk and using RNN-like transitions between chunks. They find these architectures recognize regular languages for small chunk sizes with scaling remaining challenge.
numerous potential societal consequences consider, such as the energy consumption required for the training of large models and the additional iterations that our implicit models impose (see Table 4) to enhance computational expressiv- ity of language models. When developing training algo- rithms for our language modeling, computational efficiency was taken into account- see Appendix D.3. Nevertheless, further research is essential to maintain the power of the proposed models while continuing to reduce computational costs.
7. Discussion and Conclusion
A Acknowledgment
This work demonstrates that models implicitly defined by a fixed point iteration can solve hard state tracking problems that resist the capabilities of transformers and SSMs. We provide theoretical insight how implicit SSMs can deviate from pure diagonal and linear token-to-token transitions and effectively become an RNN in the limit. When trained with relatively small number of self seff-iterations our models seamlessly generalize from simpler to harder word problems (see Figure 3). This property is of special interest in lan- guage modeling where 'hard' sequences are rare but might occur clustered in applications requiring state tracking.
The authors of the paper would like to thank colleagues from the Analog Optical Computer (AOC) team at Microsoft Re- search Cambridge for their discussions and feedback during the project. Additionally, we acknowledge support from the Microsoft GCR team for providing the GPUs and prompt assistance in resolving issues faced during the training of large language models. MS was partially supported with funds from Bosch- sooschh-oorchungsssttunng im Stifterverband.
References
Our extensive study of synthetic state tracking problems informs pretraining schedule for large language models. The implicit Llama and Mamba2 models improve over the baselines in many cases, and prove particularly beneficial on downstream tasks such as HELLASWAG (see Table I). Performance on language modeling is typically primarily determined by parameter count which traditionally caused weight-shared models underperform (Tay et al., 2023). While implicit models lift the state-tracking limitations of explicit language models, state-of-the-an explicit models perform extraordinarily well on natural language. Self- iteration introduces additional cost that might only amortize on specific problems requiring state-tracking over long se- quences. However, emerging hardware that accelerates such self-iteration would alleviate this overhead (Brunner al., 2025). Furthermore, as LLMs make more progress on re- ducing perplexity, they may eventually face tokens requiring RNN-like transitions.
Anil, C., Pokle, A., Liang, K., Treutlein, J., Wu, Y., Bai, S., Kolter, J. Z.., and Grosse, R. B. Path independent equi- librium models can better exploit test-time computation. Advances in Neural Information Processing Systems, 35: 7796-7809, 2022.
Bai, S., Kolter, J. Z., and Koltun, V. Deep equilibrium models. In Advances in Neural Information Processing Systems (NeurIPS). 2019.
Banino, A., Balaguer, J., and Blundell, C. Pondernet: Learn- ing to ponder. (arXiv:2107.05407) September 2021. doi: 10.485500a2XX222177.054407
URL http: / //tt:ii/ /arxiv. httg/abs/1100 05407. arXiv:2107.05400 [cs].
Barrington, D. A. Bounded-width polynomial-size branching programs recognize exactly those lan- guages in ne1. Journal of Computer and System Sciences, 38(1):150-164, 1989. ISSN 0022-0000. doi: URL https::/www..ccieneedddaae witeance/rrticlle/pp////00000000
Finally, given the recent rise of test-time compute (Snell et al., 2025) and latent-space reasoning (Hao et al., 2024), models with adaptive depth per token deserve careful con- sideration potential bridgeheads for such techniques they natively offer adaptive depth and latent-space iteration.
Beck, M., Pöppel, K., Spanring, M., Auer, A., Prudnikova, O., Kopp, M. K., Klambauer, G., Brandstetter, J., and Hochreiter, S. xL.TTM: Extended long short-term mem- ory. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: / / //oppneeviee eww.. net/ forum? iettummid=ARAAR d=
Impact Statement
This paper presents work aimed at advancing the field of machine learning by developing sequential and language models that can track state and hence prove to be more pow- erful in reasoning. Regarding language modeling, there are
Bhattamishra, S., Ahuja, K., and Goyal, N. On the Ability and Limitations of Transformers to Recognize Formal

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Languages. In Webber, B., Cohn, T., He, Y., and Liu, Y. (eds.), Proceedings of the 2020 Conference on Empiri- cal Methods in Natural Language Processing (EMNLP), Pp. 7096-7116, Online, November 2020. Association for Computational Linguistics. doi: 10.18633/1120200 munlp-main.574 URL httpstttt//. https///aaaaatthhlo.ooo org/2020 emnlp main. m76n.576/.
Sunada, S., Tefas, A., Tossoun, B., Tsakyridis, A., Turit- syn, S. K., der Sande, G. V., Vaerenbergh, T. V., Veraldi, D., Verschaffelt, G., Vlieg, E. A., Wang, H., Wang, T., Wetzstein, G., Wright, L. G., Wu, C.. Wu, C., Wu, J.. Xia, F., Xu, X., Yang. H., Yao, W., Yildirim, M., Yoo, S. J. B., Youngblood. N., Zambrini, R., Zhang, H., and Zhang, W. Roadmap on neuromorphic photonics, 2025. URL https: www.iiii.oggg.o. aat/a/ass////..01 07917.
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Pro- ceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Chowdhury, R., Bouatta, N., Biswas, S., Floristean, A., Kharkar, A., Roy, R., Rochereau, C., Zhang, J., Church, G. M., Sorger,
K., and AlQuraishi, M. Single- sequence protein structure prediction using a language model and deep learning. Nature Biotechnology. 40:1617- 1623, 2022. URL https: www. nature com/ a articles 41188- 022-6---2..--
Boiko, D. A.. MacKnight, R., Kline, B., and
Gomes, G. N. Autonomous chemical research with large language models. Nature, 620:547- 552, 2023. URL https:// //ws://wwww http::/wwwwaaatu.c.coc www.naattrccommm com/ articles/s41585 023-06737770000000 06792-0.
Clark, P. Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:/803.05457 2018.
Briand, T., Rombach, C., Menden, M. P., Stegle, O., and Luecken, M. D. Dna language models are pow-
erful predictors of gedome-wide variant effects. Pro-
ceedings of the National Academy of Sciences, 120 (43):e2311219120, 2023. URL https: /www. www.pnas. org/ htgg/do/ do. /10. 1073 www.////////...a/ pras 2311219121
Csordás, R., Irie, K., Schmidhuber, J., Potts, C., and Man- ning, C. D. MoEUT: Mixture-of-eepprts universal trans- formers. In The Thirty- Annual Conference Neu- ral Information Processing Systems, 2024. URL https: / //openeeview. net/ www./eernon/tteee idd 2xVrk2VB
Brunner, D., Shastri, B. J., Qadasi, M. A. A., Ballani, H., Barbay, S., Biasi, S., Bienstman, P., Bilodeau, S., Bo- gaerts, W., Böhm, F. Brennan, G., Buckley, S., Cai, X., Strinati, M. C., Canakci, B., Charbonnier, B., Chem- nitz, M., Chen, Y., Cheung, S., Chiles, J., Choi, S., Christodoulides, D. N., Chrostowski, L., Chu, J., Clegg, H., Cletheroe, D., Conti, C., Dai, Q., Lauro, L. D., Diamantopoulos, N. P. Dine, N. U., Ewaniuk, J., Fan, S., Fang, I.., Franchi, R., Freire, P., Gentilini, S., Gigan, S., Giorgi. G. L., Gkantsidis, C., Gladrow, J., Goi, E., Goldmann, M., Grabulosa, A., Gu, M., Guo, X., Hejda, M., Horst, F., Hsieh, J. L., Hu, J., Hu, J., Huang, C., Hurtado, A., Jaurigue, L., Kalinin, K. P., Kopae, M. K., Kelly, D. J., Khajavikhan, M., Kremer, H., Laydevant, J., Lederman, J. C., Lee, J., Lenstra, D., Li. G. H. Y., Li, M., Li, Y., Lin, X., Lin, Z.,, Lis, M.,, Lüdge, K., Lugnan, A., Lupo. A., Lvovsky, A. I., Manuylvvic,, E., Marandi, A., Marchesin, F., Massar, S., McCaughan, A. N., McMa- hon, P. L.. Pegios, M. M., Morandotti, R., Moser, C., Moss, D. Mukherjee, A., Nikdast, M., Offrein, B. J., Oguz, I., Oripov, B., Shea, G., Ozcan, A., Parmigiani, F., Pasricha, S., Pavanello, F., Pavesi, L., Peserico, N., Pickup, L., Picrangeli, D., Pleros, N., Porte, X., Primav- era, B. A., Prucnal, P., Psaltis, D., Puts, L., Qiao, F., Rahmani, B., Raineri, F., Ocampo, C. A. R., Robertson, J. Romeira, B., Carmes, C. R., Rotenberg, N., Row- stron, A., Schoenhardt, S., Schwartz, R. L. ... T., Shainline, J. M., Shekhar, S., Skalli, A., Sohoni, M. M., Sorger, V. J., Soriano, M. C., Spall, J., Stabile, R., Stiller, B.,
Dao, T. and Gu, A. Transformers SSMs: General- ized models and efficient algorithms through structured state space duality. In Forty-first International Con- ference on Machine Learning, 2024. URL https: // //openrevi
weterrvieew... ew. net. forum? d=2tn t:d.
Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, L. Universal transformers. In In- ternational Conference on Learning Representations, 2019. URL wtt.tttps:://peeeeeeeeeeeee /ttp///eeeeeeeeeeeeee/eee.. net/ /frum? id=HyzdRiR9y7.
Deletang, G., Ruoss, A., Grau-Moya, J., Genewein, T., Wen- liang, L. K., Catt, E., Cundy, C., Hutter, M., Legg, S., Veness, J., and Ortega, P. A. Neural networks and the chomsky hierarchy. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: //openrevieww.om www.reee/ofootte formmfrrumioo http://eeeeoouuu....oo..ttaaaee. www.iddddbeieeeeeeeee
Du, Y., Li,, S., Tenenbaum, J., and Mordatch, I. Learn- ing iterative reasoning through energy minimization. In International Conference on Machine Learning, pp. 5570- 5582. PMLR, 2022.
Elman, J. L. Distributed representations, simple recurrent networks, and grammatical structure. Machine learning, 7:195-225, 1991.
10

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The Pile: An 800gb dataset of diverse text for language modeling, arXiv preprint wrXiv22101.0002: 2020.
Hägele, A., Bakouch, E., Kosson, A., Allal, L. B., Von Werra, L.,, and Jaggi, M. Scaling laws and compute- optimal training beyond fixed training durations. arXiv preprint arXiv:2405. 18392, 2024.
Hao, S., Sukhbaatar, S., Su, D., Li, X., Hu, Z., Weston, J., and Tian, Y. Training large language models reason in continuous latent space, 2024. URL https:// arxiv. org/abs/////ooooooooooo........r 2412.41121.00 06769.
Gao, L., Tow, J., Abbasi, B., Biderman, S., Black, S., DiPofi,
A., Foster, C., Golding, I.., Hsu, J., Le Noac'h, h,, A., Li, H., McDonell, K., Muennighoff, N., Ociepa, C., Phang, J., Reynolds, L., Schoelkopf, H., Skowron, A., Sutawika, L., Tang, E., Thite, A., Wang, B., Wang, K., and Zou, A. A framework for few-shot language model evaluation, 07 2024. URL https: /enodoo. org/ records/ 12608602.
Hooker, S. The Hardware Lottery. 2020. URL https:
//arxiv. wwwp:i.wwwr.....o.oe..o...obasaa / 911 05248.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. (arXiv:2001.06) -08311), January 2020. doi: 10.48550/arXiv. 2001.08361 URL ht /ttp:///aavvvv//. org/ / abs/ 2001. 08361. arXiv:2001.0836) [es].
Gatmiry, K., Saunshi, N., Reddi, S. J.. Jegelka, S., and Kumar, S. Can looped transformers learn to imple- ment multi-step gradient descent for in-context learn- ing? In Forty-first International Conference Machine Learning, 2024. URL https: / //te:://pppnnne /ppnneveeeeeeeeee net. forum? doommndd=ooaaoo dakkao
Kleene, S. Representation of events in nerve nets and finite automata. CE Shannon and J. McCarthy, 1951.
Li, B. Z., Nye, M., and Andreas, J. Implicit representations of meaning in neural language models. In Proceedings of the Association for Computational Linguistics and International Joint Conference on Natural Language Pro- cessing (Volume 1: Long Pupers), pp. 4744 4456. Associ- ation for Computational Linguistics, 2021. URL https: / /aclannthlooyy...oo. org/ 2021 aalalcllllag 309 143/.
Geng, Z., Zhang, X.-Y., Bai, S., Wang, Y., and Lin, Z. On training implicit models. Advances in Neural Information Processing Systems, 34:24247-24260, 2021.
Graves, Adaptive computation time for recurrent neural networks. 16f.::::603.00999 February 2017. doi: 10. 00530//aXXii.113.08999000 URL //tp://aaavv.///a///aa/eaaaaaaia org/ abs/ 08983. arXiv:1603.08963 [es].
Li, K., Hopkins, A. K., Bau, D., Viégas, F., Pfister, H., and Wattenberg, M. Emergent world representations: Exploring a sequence model trained on a synthetic task. In International Conference on Learning Representations, 2023. URL https://p//. / tetttps://ppereeeeeee http//ppnnveeeeeeeee//e/ee net.:wwwwwwweeeeeee.eee. net forum?
Grazzi, R., Siems, J., Franke, J. K., Zela, A., Hutter, F. and Pontil, M. Unlocking state-tracking in linear RNNs through negative eigenvalues. In The Thirteenth International Conference on Learning Representations, 2025, URL https: //tp:://pppeaaa http://ppppppeeeeeeeeeeeeeee net forum? id=UwTo3tVBk2.
id=DeG07_Te2n
Lim, Y. H.. Zhu, Q., Selfridge, J., and Kasim, M. F. Parallelizing non-linear sequential models over the se- quence length. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: / //openreview.o htttt//frummiidde /ttp:/rrrmmm---AAAAAAAA0A
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
Merrill, W. Sequential neural networks automata. In Eis- ner, J., Gallé, M., Heinz, J., Quattoni, A., and Rabusseau, G. (eds.), Proceedings of the Workshop on Deep Learn- ing and Formal Languages: Building Bridges, pp. 1- 13, Florence, August 2019. Association for Computa- tional Linguistics. doi: 1653 URL https: hrtps://acaatthooogge
Gu, A., Goel, K., and Re, C. Efficiently modeling long sequences with structured state spaces. In In- ternational Conference on Learning Representations, 2022. URL https::/ wwwprreiiiii.eeeeoooooeeeeeeeeee net/ id=uYLFozIvlAC.
Guan, L., Valmeekam, K., Sreedharan, S., and Kambham- pati, S. Leveraging pre-trained large language models to construct and utilize world models for basel task planning. In Advances in Neural Information Pro- cessing Systems, 2023. URL https::///// /tttpp:://axxx...//// waw.ixiiiiooo.. org/ 0oss/335..49903
Merrill, W. and Sabharwal, A. The parallelism trade- off: Limitations of log-precision transformers. Trans- actions of the Association for Computational Linguis- tics, 11:531-545, June 2023. ISSN 2307-387X. doi: 10.1162/tacl a 00562.

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Merrill, W. and Sabharwal, A. The expressive power of transformers with chain of thought. In The Twelfth International Conference on Learning Representations, 2024. URL bectpa://sperreee iN= Ph8Wh.
Sarrof, Y., Veitsman, Y., and Hahn, M. The expressive ca- pacity of state space models: A formal language perspec- tive. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: / //openrevi eww..wwww.....................e. net/ forum? d=
net nerum
Merrill, W., Weiss, G., Goldberg, Y., Schwartz, R., Smith, N. A., and Yahav, E. A formal hierarchy of rnn architec- tures. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 443-459, 2020.
Schlag, I., Munkhdalai, T., and Schmidhuber, J. Learn- ing associative inference using fast weight memory. In International Conference on Learning Representations, 2021. URL https: htttps://ppereviiiiee wpwp/ewpeeeeee.ee/oe/pe net net/forun? id=TuK6agbdt27.
Merrill, W., Sabharwal, A., and Smith, N. A. Saturated transformers are constant-depth threshold circuits. Trans- actions of the Association for Computational Linguistics, 10:843-856, 2022.
Schmidhuber, J.
Self-delimiting neural networks.
(arXiv:1210.011 22100018)) September 2012. doi: 10.48550/ urXiv.1210.00 URL http: /tt.//a........aa..//.. 1210. 0118. arXiv:1210.0118 [cs].
Schützenberger, M. trivial subgroups. (2):9901144, 1965.
On finite monoids having only Information and Control,
Merrill, W. Petty, J., and Sabharwal, A. The illusion of state in state-space models. In Salakhutdinov, R., Kolter, Z. Heller, K., Weller, A., Oliver, N., Scar- lett, J., and Berkenkamp, (eds.), Proceedings of the 41st International Conference on Machine Learn- ing. volume 235 of Proceedings of Machine Learn- ing Research, pp. 35502-35506. PMLR, 21-27 Jul 2024. https:s:: httpp//ooeeengss.e// mlr pprss// htmm
ISSN 0019-9958. http://ddaag////0000000009999999 doi:
URL httpp://ww..sccennedeeeeer
Siegelmann, H. T. and Sontag, E. D. On the computational power of neural nets. In Proceedings of the Fifth Annual Workshop on Computational Learning Theory. COLT '92, pp. 440 440-449, New York, NY, USA, 1992. Association for Computing Machinery. ISBN 089791497X. doi: 10. 10.1145/130385.100433 URL https://doi.cnt 10.
html.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can suit of armor conduct electricity? new dataset for open book question answering. arXiv preprint 100::080900000 2018.
130432.
Smith, J. T., Warrington, A., and Linderman, S. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https:/ wptp//eeeneeeeeeeee/ee net/ forum? id=Ai8HW3Axqks.
Omlin, C. W. and Giles, C. Extraction of rules from discrete-time recurrent neural networks. Neural Networks, 91):41-55,, 1996. ISSN 0893-6080. doi: http:://o..ogg///o//000099440000
URL https://www.ccceeeeeeeec...t iitenc//arccclee/11///3300038880
Snell, C. V., Lee, J., Xu. K., and Kumar, A. Scaling LLM test-time compute optimally can be more effective than scaling parameters for reasoning. In The Thirteenth International Conference on Learning Representations, 2025, URL https: /wpp/reveeeeeeeeeeeeeeeeeeeeeeee net forum? id=4FWAwZtd2r.
Paperno, D., Kruszewski, G., Lazaridou, A., Pham, Q. N., Bernardi, R., Pezzelle, S., Baroni, M., Boleda, G., and Fernández, R. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint 1rXiv:1606.06031, 2016.
Soulos, P., Terzic, A., Hersche, M., and Rahimi, A. Recur- rent transformers trade-off parallelism for length gen- eralization on regular languages. In The First Work- shop System-2 Reasoning Scale, NeurlPS 24, 2024. URL https:/ /wwpnrvveewweeee////////e net forum? i 1d=6P jZAAJvge.
Qin, Z. Yang, S., and Zhong, Y. Hierarchically gated re- current neural network for sequence modeling. In Thirty- seventh Conference on Neural Information Processing Systems, 2023. URL https: /ttps///ee./eeeee..... net/ Corum?id=P11cH11
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
Strobl, L., Merrill, W., Weiss, G., Chiang, D., and Angluin, D. What formal languages can transformers express? a survey. Transactions of the Association for Computa- tional Linguistics, 12:543-50 -61, 2024.
12

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Tay, Y., Dehghani, M., Abnar, S., Chung, H. W., Fedus, W., Rao, J., Narang, S., Tran, V. Q., Yogatama, D., and Met- zler, D. Scaling laws vs model architectures: How does inductive bias influence scaling? In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023. URL https: htttps://ope.rewe// wpw..eeiee.eeeeeeeeoeeeee. htt.//ee//tteeeeee forum?
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O., Susskind, J., Bengio, S., and Nakkiran, P. What algo- rithms can transformers learn? a study in length general- ization. arXiv preprint arXiv:2310. 2023.
id=E9dHOBPSVW.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302.n 13971, 2023.
Vafa, K., Chen, J. Y., Rambachan, A., Kleinberg. J., and Mullainathan, S. Evaluating the world model implicit in a generative model. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024. URL https: httttp:://ppe.reviiiee net/forum? id=avK4JFpegy.
Wei, J., Wang, X., Schuurmans, D.. Bosma, M., Xia, F., Chi, Le, Q. V., Zhou, D., al. Chain-of-though prompting clicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24888 2022.
Weiss, G., Goldberg, Y., and Yahav, Thinking like transformers. Meila, M. and Zhang, T. (eds.), Pro- ceedings of the 38th International Conference on Ma- chine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 11080-1 1090. PMLR, 18 24 Jul 2021. URL https: : httptttts/////oot// /ttpppeeeings.....///pp/p/// mlr ml..press/ v1 #139/weisse a html.
Weston, J., Bordes, A., Chopra, S., Rush, A. M., Van Merriënboer, B., Joulin, A., and Mikolov, T. Towards ai-complete question answering: A set of prerequisite toy tasks. arXiv preprint 1502
2015.
Yang. L., Lee, K., Nowak, R. D., and Papailiopoulos, D. Looped transformers are better at learning learn- ing algorithms. In The Twelfth International Confer- ence on Learning Representations, 2024a. URL https: //openreview wwt.:/iwwwwee.eeeeee..m.emeeeeee
Yang. S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. Parallelizing linear transformers with the delta rule over sequence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL httph:tpp//.......t. http////peeeeeeeeee/eeeeeee net/forum: id id-y8Rm4 VNRPH.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv: / 2019.
13

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Phantom gradients
pass
Fixed-point search
Shared Weights
Backward pass
Figure 5: Fixed-point iteration and phantom gradients: A neural network iterated until convergence in the forward pass. When employing the phantom gradient principle, only the final steps of the forward pass considered for the backward pass expressed in Equation (5). Therefore, the memory required for the backward pass is proportional to k and in particular not proportional to the number of steps in the forward pass.
Duality between Simultaneous and Sequential Modes
Simultancous mode
Sequential mode
for
h for 1,...,7 T 2; -
in 1,.... -- 0
Exchange for loops
for in 1, 1...., T
for in 1, D
h So (, 24)
fo (2t, h, he))
for
Figure 6: The simultaneous mode exploits the parallelism of state-space models or transformers, while the sequential mode well suited for language generation. State-space models can further utilize the sequential mode for processing with constant memory
any sequence length. The two modes emerge from exchanging the for loops the two variables t and in the DEQ iteration (6). We demonstrate in Section 5, and Figure 2, that 1.3B parameter language models trained with the simultancous mode show negligible difference in perplexity when evaluated with the sequential mode.
A. Implicit State-Space Models
A.1. Fixed Point Iteration
Figure 5 visualizes the forward and backward pass of implicit models. A fixed point search is conducted by unrolling the -(--)) falls below tolerance where we choose 0.05 for the language models model until the relative difference |---
and 0.01 for the state-tracking models.
The phantom gradient approach backpropagates only through a fixed number of steps, effectively decoupling the number of iterations the forward pass from the number iterations the backward pass.
Figure shows the two loops over the sequence and over depth as presented in Section 3.1. The equations involving the iteration variables are printed again below.
(0-1) 20a ((()) A (s-1) ,3se 11/2) ((((
(()) fe 2 (s-1) (1(() 11
14

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
A.2. Proof of Theorem
Theorem 1 (Extended). An implicit SSM defined by Equations (8) and (9) with generic weights yields a non-linear and non- diagonal state-to-state transition function & h2-1 hū. Let g (z, h,r,0)
fo (2, (2, h, x). If g (zt 1: 2t, 0) 0
and Jg.=l0t differentiable function yp on U s.t.
0g ä=
is non-singular, there exists an open set U with (hç-1:20. U and a continuously
12.0
g (1 (h, (,, =0 ().
In this case, the stale-to-stale Jucobian Jg given by
dhg
Ou Oip ON 0e diag (hi-1) + dzj dh; dzz ohū-1
A (27,
dh 25.0
Proof. Suppose that fixed points hi-1. zi-1 have been found for -- Our goal is back the above statements for 4. Let T (2) fo (z, 11),, With generic weights 8, c.g. in general linear position, a multi-layer neural network fo is non-linear and its input-output function T (z) has non-diagonal Jacobian. Finding the fixed point 2 can be expressed in terms of recursive application of the non-linear operator T as
lim lim 7
(0) 21
(1I)
80000
(0) =0. where we set 2; a non-linear function (hī 0) - z Once the fixed point z? is found, we can compute h by single pass through Equation (8).
0. By extension from lulti-layer networks to self-iterated multi-layer networks, generic weights yield
To proof Equation (10), we apply the implicit function theorem to the function
g (2, 2, x, = - fo (2,
(12)
is non-singular, the implicit function theorem guarantees the existence of open set U with (hā-1; x. 0) U and continuously differentiable function yp on U
If (=i, (i, 0) = O and 1
0g 82
g (yp (h, x, h, 2, (h,,I, E.
(13)
The derivative of yp at the fixed point given by
0
0fe
0f
I
(14)
0h
Oz at My
Oh
1..00
1 ,k 1.
By the above argument, yp is non-linear function if fo a non-linear function, and oي (hī Iz, =1:=iim,--os T (0). Now, consider Equation (8) at the fixed point
hī = A A a1) hi-1 xz)
(15)
where = = 1p (h Equipped with the derivative of sp, we can derive the state-to-state Jacobian the implicit SSM as
0). 0). If &o is a non-linear function of 1;- then Mi_1 + h is a non-linear function as well.
dhi
0h;
0A
du Th
A (2 v.t)
+
1111 +
dh:
0h;_1
.
0A 0y diag (hi_1)
Iu dzg
Öe
A (x(
Ozg ahe
Ōh
where we used the implicit dependency of 2 on hi_ via y to differentiate A and u.
15

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Autograd Jacobian
Theorem Jacobian
Relative Difference
Head index
Head index
Head index
15
15
10
10
10
Head index
Hend index
Head index
Autograd Jacobian
Theorem Jacobian
Relative Difference
Head index
Head index
Head index
10
10
10
Head index
Head index
Head index
10
Matzix clements
Mattix elements
Mattix elements
Figure 7: State-to-state Jacobian according to Equation (10) of a single layer implicit Mamba2. To simplify the visualization, only one state variable of each head is plotted. Left: Jacobian computed with automatic differenniation through the fixed point iteration. Center: Jacobian computed with Equation (10). Right: Relative difference Jangual Juceean
|Jungaa Jaccenn
Top:
Randomly initialized model. Bottom:
Trained model.
This equation (Equation (10) in the main text) highlights the non-diagonal corrections that the self-iteration of 2 introduces to the diagonal Jacobian Å of the explicit state-space Equation (1). We conduct numerical check of Equation (10) in
Figure on data from the synthetic task devised in Section 4. Therefore, the state h is created by processing a sequence of 8 tokens .u....,.2 Z7. Next, we loop over until 28
(a) ./g converge fixed point, where the stopping criterium is x 10-4, Figure compares the Jacobian of h; - hg when backpropagating through the iteration over s with PyTorch's automatic differentiation (left column) versus Equation (10) (central column). The right |Jasteenal - Jacsen
11
relative difference I=S
column shows the relative difference between the two methods as TTmsgaal + Jasmnl He with 10-16 The top row shows randomly initialized model, and the bottom row shows trained model. Both models are based on the Mamba2 architecture with single layer and dmuxtel 64, dhead 8, dsaae such that the model has 16 heads. A single state variable selected head, resulting in 16 x 16 grid per model and evaluation method.
B. Algebraic Structure of Finite State Machines
This section provides basic introduction to word problem and its relation simulating finite state machines (FSMs). We start with some results in the circuit complexity and then relate them the properties FSMs.
B.1. Circuit Complexity
Efficient parallelization is one of the central features enabling transformers and SSMs to scale to large machine learning problems such as language modeling. Parallel circuits, however, face fundamental trade-offs regarding the class of problems that they can address. Circuit complexity theory provides framework characterize the types of problems that parallel circuits can solve. TC is the class of circuits with constant depth and polynomial width composed of unbounded fan-in AND-gates, OR-gates, NOT-gates and MAJORITY-gmm depth circuits with a polynomial number of bounded fan-in gates. From the perspective of formal languages, NC is equivalent to the class of circuits recognizing the regular languages. Since the unbounded fan-in gates allowed in TC circuits can be constructed from log-depth circuits with bounded fan-in, follows that TC C. NC!. It is open if TC is a proper subset of NC!, and we will discuss a regular language for which no TC circuit construction known.
The second class of interest, NC1, is represented by logarithmic
16

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Both transformers and SSMs can be simulated by TC circuit families under mild assumptions (Merrill et al., 2022; 2024). If TC is a proper subset of NC², the leading sequence models today cannot recognize all regular languages. Consequentially, they cannot execute arbitrary finite state machines (FSMs), fundamental skill to execute tasks, or to form world models ( (Vafa et al.. 2024). Many empirical studies confirm these theoretical limitations of transformers and SSMs to learn regular languages (Deletang et al., 2023; Sarrof et al., 2024; Strobl et al., 2024). At the same time, recurrent neural networks are known to recognize regular languages (Kleene, 1951; Elman, 1991; Merrill et al., 2020), and to effectively implement internal FSMs to solve language problems (Omlin Giles, 1996).
B.2. Algebraic Concepts of Finite State Machines
Monoids and Groups There is a tight relationship between finite state machines and algebraic concepts such as monoid and groups. We define the relevant concepts for our state tracking problem described Section 4
Definition (Monoid). A set M and a binary operation o M M + M are called a monoid (M, if
1. there exists an identity element E M with e o m m o e = m for all m E M
2. the operation o is associative, i.e. (m² m2) o my = m. o (m2 ma) for all my, m2, my E M.
Straight forward examples for monoids are natural, rational or real numbers with multiplicatiin, or strings with string concatenation. Since monoid are associative, we can simplify notation and write m o m = m² and so on for all powers E N.
Definition (Aperiodic Monoid). A monoid (M, o) is called aperiodic if for all m E M there is a k E N s.t. m k-1
Monoid where each element has a a corresponding inverse element are called groups.
Definition 4 (Group). group (G. o) is a monoid with the additional property that for every g E there is g- 1 E G s.t. g gog pog-i=g-b
-- g
Examples for groups rational numbers with multiplication, or the orthogonal matrices with matrix multiplication. Notably, permutations set of k elements for k E N form a group, called the symmetric group Sp.
Our synthetic learning problem discussed in Section 4 will be constructed based on classical problem in computer science. Definition (Word Problem). Let M denote the set of all sequences over elements of M. The word problem on a monoid (M, is defined by the function
WP: M M
WP (mo, mi .... m&) mg mi mk
(16)
ie.. a word over M is resolved by composition to a single element in M.
The central question for our experiment will be which kinds of circuits can solve the word problem for arbitrary sequence lengths.
Theorem ((Brrington, 1989)). The word problem for any fixed non-solvable group G is complete for NC under AC" reductions.
Tracking the state of system can be formalized as executing a finite state machine (Merrill et al., 2024). To characterize the limits of certain FSMs, define few formal concepts. Definition 7 (FSM). A fiite state machine (FSM) consists of finite set of states Q, a finite set of input symbols Σ called the alphabet, and a transition function 8 Q x Σ Q.
Algebra of FSMs
Given an initial state Yo Q and sequence of symbols = 0102 E*, FSM transitions from the initial state into a final state.
Finite state machines naturally define a monoid.
Definition 8 (Syntactic Monoid), For each symbol a E Σ, define the function 8, : Q - generated by 5a, a E Σ and the composition of functions o, called the syntactic monoid (M, o) of the finite state machine.
+ Q. The transformation monoid M
17

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Yesterday Mary travelled to kitchen.
Yesterday Bill went to school.
2 Fred went to the park yesterday.
Mary went to the cinema yesterday.
3 This morning Fred journeyed to the bedroom.
This morning Julie went back to the kitchen.
4 Bill went back the bedroom yesterday.
Yesterday Julie went to the office.
5 Q: Where was Fred before the hedroom?
Q: Where Julie before the kitchen?
Table 2: Two examples of tracking entities in natural language from the CATBABI dataset (Schlag et al., 2021).
The algebraic structure of M is tightly coupled to the programs that the original FSM execute. Our investigation is based on the classical result stated in Theorem 6. The simplest example of a non-solvable group is the permutation group of five elements Sç. A corollary from theorem 6 is that the FSM whose syntactic monoid is Sş is complete in NC* and hence in the class of regular languages. We have thus identified a hard state tracking problem: Permutations of five elements.
Another classical result tightly related state-space models
Theorem 9 (SSchützenherger, 1965)). Let L be the regular language defined by a FSM, and let M be the syntactie monoid of the same FSM. Then L is a star-free language if an only if M is aperiodic.
is intuitive that the word problem for finite aperiodic monoids is in TC The maximal depth of the circuit is driven by the number of elements of the monoid and it's maximal & for the aperiodicity condition. Empirical studies have shown that transformers and SSMs can simulate a range of regular languages (Deletang et al., 2023; Strobl et al., 2024), but they struggle to learn the S₁ word problem in line with their characterization as TC" circuits (Merrill et al., 2024). SSMs can be further restricted to the star-free languages (Sarrof et al., 2024), i.e. those with aperiodic syntactic monoid.
C. Relaxed Non-Solvable Word Problems
The A5 or S5 word problems used in (Merrill et al., 2024) require to keep track of state for every single token. However, entities appear sparsely in natural language as shown in Table 2. Consequently, tracking state of every token is not an accurate representation of the state tracking problem in natural language. Here, we devise synthetic task to model tracking state of sparse entities. Let p the sparsity of entities. For example p = 0.9 refers 90 % sparsity, i.e. an entity density of 10 %. An example of a sparse sequence is presented in Figure 8. In particular, are interested in answering the following questions
1. Which density of entities is required learn the intrinsic algorithm of the state tracking task?
2. How many self-iterations do our implicit models need conduct to learn the intrinsic algorithm?
To construct a learning problem for sequence models. we represent each element of the monoid M as a to-
Star-free language
ken, and present a sequence m,....., m, of tokens to the
Legrnd
model. The ground truth at each position & = .,....,! L the token representing the element 170 ... o my- We then calculate the mean cross entropy loss over the en- tire sequence, providing learning signal each step A 1, L.
NC'-hard language
nan aon-sdvaable toker
Sparse NC -hard language
State-space models can learn the word problem for ape- riodic monoids (Sarrof et al., 2024), but fail so solve it for non-solvable groups such as S3 (Merrill et al., 2024). We confirm in Figure that implicit state-space models can in fact learn the word problem for Sy. We now want to interpolate between word problems for aperiodic and non-solvable monoids to test how much signal our im- plicit state-space model defined in Section 3 needs from the hard nobllolvable group word problem to learn it.
Figure 8: Examples of word problems with mixed solvable and volvablvable tokens. We refer to tokens from non-solvable group as non-solvable tokens, and to tokens from a star free language as solvable tokens.
Let M = M x G be a direct product of an aperiodic monoid M and a non-solvable group G. A sequence Mo..., my is sampled from M with replacement. To control the number of hard examples and simple examples, we define a family of
18

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Model Comparison
explicit
implicit
unrolled
100
Accuracy in %
75
50
25
Train p=90.5 L=256
Test p-0. L=256
Test p- L-259 L-256
Test p=055 L=1024
Test p .9 L=1024
Figure 9: Comparison of implicit Mamba2, unrolled Mamba2, and Mamba2 for p = 0.05. All models were trained and evaluated on sequences of length L = 256. Unrolled Mamba2 refers to single layer being unrolled multiple times with a full backpropagation trace, while the implicit Mamba2 receives only steps of Phantom Gradient. The training time depth of all models is limited to 16, i.c. 16 layers for Mamba2, and 16 self-iterations for the implicit and weight tied models. Implicit and unrolled models use unbounded test-time computation converge to fixed point. The comparison shows that the implicit model with 4 steps of Phantom Gradient succeeds over the unrolled model.
distributions D, over M as follows. element mf E M sampled uniformly at each step k. representing the presence of simple examples. On the other hand, we sample elements I G {e} from G without the identify transformation, each with probability -p-T- The identity element 9k
G is sampled with probability p. The resulting transformations (me,.00 9s) are aperiodic at least when gk = e, i.c. with probability p-p.
We'll call the tokens representing (m, e) simple tokens and the tokens representing (m.g) with g * e are called hard tokens. The names derive from the fact that SSMs resolve the word problem if it is only composed from simple tokens. Any non-zero probability p of sampling hard tokens renders the word problem unsolvable for fixed depth SSMs on arbitrarily long sequences.
Our construction of a distribution over a monoid allows us to test out-of-distribution generalization not only in terms of length generalization, the most common setting in the literature. By changing p between training time and test time, we construct training tasks and evaluation tasks with varying difficulty. This effectively offers OOD evaluation with the same number of tokens, but different mixtures of easy and hard tokens. While this property allows us to distil expressivity questions from length generalization properties, our construction is not limited to the same sequence length and could as well be used in the length generalization setup (see Figure 9).
D. Experimental Details
D.I. The Word Problem
Each data point in Figure and Figure 9 is based on 10 runs with different random seeds. We report the best run, mean accuracy and 95 % confidence interval for each data point. All word problem models were trained on sequences length = 256, and batch size of 512 on 32 GB V100s. The explicit models and self-iterated models with full backpropagation trace required gradient accumulation over two steps. The learning rate is set to 0.001. We disable dropout and weight decay, which appears to harm learning the word problem. The self leff-iterations are stopped upon convergence, which we define as a relative difference between two consecutive states of 0.01 for Figure or 0.05 for Figure 9. We trained a number of standard Mamba2 models with the same number of runs for multiple numbers of layers. These models struggle to capture the training distribution compared to self-iterated models, and none of them was able to generalize to a harder distribution or longer sequences.
19

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
D.2. CatbAbl
The models, both implicit and explicit, comprise up to three layers of the Mamba2 architecture with an embedding dimension of 256, a state dimension of 16 (expansion factor 2), and a head dimension of 32. We trained the models using batch sizes of 128 and 256, and learning rates of 0.0001, 0.0005, 0,001, and 0.005. The models were trained for 15,000 steps, with the implicit model specifically trained for 5,000 steps in unrolling mode, utilizing 32 steps with normal gradient checkpointing, followed by 10,000 steps of self-iteration fixed-point search. The self-iteration included a stop threshold of 0.03 and a training and testing maximum number of steps 50 and 200, respectively, and phantom gradient parameters of 6 steps with (A = 0.5). Data were packed in sequences of length 200 tokens as per (Schlag et al., 2021). Figures 10 and 11 illustrate the validation accuracy of the explicit and implicit Mamba2 models on the CATBABI dataset, respectively. Additionally, Figure 12 plots the number of iterations required for the implicit model to reach fixed point on the validation set of the CATBABI dataset.
85: 128. Layers:
BS: 128. Layers: 2
BS: 128, Layers: 3
1.0
1.0
0.9
0.9
0.5
Valldation Accuracy
Aecuracy
Velidation Accuracy
0.9
0.8
0.8
0.7
0.7
0.8
0.6
0.6
0.5
LR: 0,0155
1
0.5
0.5
L8. 5.033
Uk- 23-2
0.4
0.4
3.01
0.9
0.9
5000
10000
15000
0
5000
10000
15000
5000
10000
15000
Stegs
Slegrs
Stegs
85: 256, Layers:
BS: 256, Layers: 2
ES: 256, Layers:
1.0
1.0
0.3
0.0
0.0
Velizitton ACCUPHCY
0.3
0.8
Velicition Accuracy
0.8
0.7
Velidaaion
0.5
0.6
0.6
0.5
0.0055
3
0.5
0.5
0.033
0.4
18 3.03%
LR: 0.01
0.3
0.3
0.3
5030
10000
15000
5000
10000
15000
5000
10000
15000
Steps
Figure 10: Hyperparameter sweeps for explicit Mamba2 models over batch sizes 128, 256, layers 1,2,3 and various learning rates for training on the CATBABI dataset.
D.3. Language Modeling
Pretraining Details We have trained a suite of Implicit SSM models with the core architecture of Mamba2 and Implicit Transformer models with the core of Llama3. For each implicit model, we have corresponding weight-tied model that is also trained on the entire D-PILE dataset. We use the checkpoint from 80 percent of the way through training the weight-tied model to train the fully implicit model. We use four scales for the training of the models: 1.3B, 760M, 360M, and 130M. In all models, the LLM head weights are tied to the embedding weights. The implicit and weight-tied models have the same architecture as those of the Mamba2 and Llama models, except for an injection module, consisting of an MLP, which transforms the input into the domain of the mixer latent space. This module has a constant size equivalent to x demb +2 2 x d'state + Nheods in Mamba2, matching with diny.wg. and demb in Llama models, corresponding to the key, value, and queries, and is shared across all layers of the model. Details for each model provided in Table 3.
We followed the training recipe of Mamba2 and Llama models. In particular, we used weight decay of 0.1, no bias for the L.M head, AdamW hyperparameters 33 = (0.9, 0.95), RMSNorm instead of LayerNorm, and a linear warm-up step to the peak learning value, which is chosen as 5 times the value of the GPT-3 model. For the learning rate scheduler, we used a constant learning rate followed by a square root decay to a minimum value of 10-5 (Hägele et al., 2024). While this scheduling has also been shown to be compute-optimal (Hägele et al., 2024) alongside the cosine scheduling, it allows us to use intermediate checkpoints during training more conveniently without considering how the new learning rate affects
20

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
BS: 128, Layers;
BS: 128, Layers: 2
ES: 128, Layors:
0.3
0.0
0.9
Validation ACCUPICY
Velideloom Accuracy
0.3
0.8
0.8
0.5
0.6
0.6
3,911
0,0355
0.5
0.033
1K 3,012
LR: 0.02
0.3
0.3
0.3
5000
10000
15000
5000
10000
15000
5000
10000
15000
8S: 256. Layers:
BS: 256. Layers: 2
BS: 256. Layers:
1.0
1.0
1.0
0.9
0.9
0.9
valcelon Accirocy
0.3
0.8
Volidelon Acturesy
0.8
0.7
0.5
0.8
0.6
0.5
0.5
3.0355
0.4
0.031
Uh 0,015
0.3
10000
15000
0.3
5000
10000
0.3
10000
15000
Steps
Steps
Steps
Figure 11: Hyperparameter sweeps for implicit Mamba2 models over batch sizes 128, 256, layers 1,2,3 and various learning rates for training on the CATBABI dataset.
BS: 128. Layers: 1
BS: 128. Layers: 2
BS: 128. Layers: 3
175
175
16.0
150
150
140
velldation Steps
125
125
Velidation Steps
120
200
100
100
75
75
50
LR: 3.9911
50
LE: 2,0355
25
18. 0.003
25
0.005
0
5000
10000
15000
5000
10000
15000
0
500C
10000
15000
Steps
Steps
Steps
B5: 256. Layers: 1
8S: 256. Layers:
85: 256, Layers: 3
1.5
175
150
150
Valietion Steps
125
Stepss
125
Validdion Steps
100
100
75
75
50
50
50
L5 0.03E1
LR: 0.0025
25
25
3,003
0.005
0
5000
10000
15000
5000
10000
15000
0
5000
10000
15000
SAegrs
Sleep
Sleers
Figure 12: Hyperparameter sweeps for implicit Mamba2 models over batch sizes 128, 256, layers 1,2.3 and various learning rates for training on the CATBABI dataset.
the training implicit models, All models were trained with an effective batch size of IM tokens and training sequence length of 2048 tokens.
Downstream Task Details We evaluated our models as well as the baseline models on seven tasks: LAMBADA OPEENAL (Paperno et al., 2016), HELLASWAG (Zellers et al., 2019). PIQA (Bisk et al., 2020), ARC-EASY and ARC-CHALLENGE (Clark et al., 2018), WINOGRANDE (Sakaguchi et al., 2021), and OPENBOOKQA (Mihaylov et al., 2018). We used the
21

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
A
B
1.0
Imphot2 Imphol 1 lyst
ACCURIC
Thst Stagee
0.6
20
1rpec133 loyer
0.2
2
128772
THE TTETIOETTS
Noboea TLaKeRg
Compue COECOOENNN
boon
2m
C
Task
D
Tesk
1.0
pmpkitt Maye
50-
702 Acc//ary
0.7
40
Test
30-
ExpR0t 8.55
20
0.5
ExpRot 23505
Exppo11
10
0.4
2
Inpec11 Uyer
1/5 245 2/3
3/5 Mean Story Length
415 4/5
625
Figure 13: Small-scale reasoning advantages of implicit SSMs compared explicit SSMs on the CATBABI dataset. (a) Task-specific performance comparison, measured in accuracy of the models in predicting answers to questions within story, between the implicit Mamba2 and baseline Mamba2. The one-layer implicit Mamba2 model outperforms the one-layer explicit Mamba2 on most tasks. As the number of layers in the explicit Mamba2 increases, its performance approaches that the implicit Mamba2. Adding more layers the implicit Mamba2 benefits certain tasks, such as 'Basic induction' or 'Path finding,' where the implicit Mamba2 achieves the best performance. (b) The correlation between the number of self-iteration steps that the implicit Mamba2 takes solve task and the number of layers in the implicit model's 's backbone architecture, showing decrease in the required steps with additional layers. (c) The implicit Mamba2 retains its performance as the story length increases, whereas the explicit Mamba2's s performance declines. (d) The trend the number iterations needed by the implicit Mamba2 models as story length increases, indicating a modest rise in computational steps.
model checkpoints the 207B tokens of the D-PILE. For the evaluation of the downstream tasks, utilized the LM Evaluation Harness package (Gao et al., 2024) with the default specifications, i.e., a batch size of 256 and sequence length of 2048. All models were evaluated using one H100 80GB GPU.
Curriculum-Base Training of Language Models The training of implicit models achieved through curriculum training framework that is divided into two main phases: bounded phase and the free phase. In the bounded phase, the models are subjected to four steps of self-iteration, followed by a singular phantom gradient step store the activations necessary for backpropagation. This phase of training involves 80 percent of the dataset (the choice of this proportion and its influence on model performance are discussed below). Training then progresses to the free phase, wherein the model undergoes additional fixed point searches, capped at 24 self-iterations for SSMs and 32 for Transformers, and is followed by four phantom gradient iterations. A stopping criterion of c = 0.05 implemented during this second phase, allowing models to terminate the fixed point search once this threshold is met. For validation and testing on the D-PILE dataset, the limit on fixed-point iterations is set four times the number used in the free phase of training. The learning rate reduction
22

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Table 3: Architecture and pretraining details
Params.
SSM/Transfrmerr # layers
Dim, Emb.
Heads/ Dim.
Training Steps
Peak Learning Rate
Batch Size (Tokens)
1.3B
48/24
2048
32/64
197701
0.001
1M
760M
48/24
1536
16/96
197701
0,00125
1M
350M
48/24
1024
16/64
197701
0.0015
1M
130M
24/12
768
12/64
197701
0.003
1M
starts in phase two; for Transformers, this reduction begins immediately to prevent instability caused by their high spectral norm. Conversely, for SSMs, the learning rate cooldown starts after 90 percent of the overall training period has elapsed, due to their spectral norm being approximately one, which permits more substantial weight adjustments at a higher learning rate.
Duration of Bounded and Free Phases in Training Language Models In our exploration of the optimal duration for bounded and free phases of training, we aimed to find a balance between computational efficiency and the necessary nonlinear transitions each token must undergo through self-iteration. We tested models trained with 70, 80, and 90 percent of the bounded phase duration before starting the full fixed-point search in the free phase-refer Fig. 14a. For this evaluation, used a 130M Mamba2 model. Based the model's perplexity the validation split of D-PILE (2M examples of length 2048), observed that beyond certain extent of free phase training, the model's performance plateaus overfits. We determined that 20 to 10 percent of free phase training is optimal. Consequently, we applied 20 percent free phase training for the training of larger models. It is crucial to note that the 1 130M model, when trained with an effective batch size of 0.5M tokens, experienced overfitting. This overfitting was not present when we increased the batch size to 1M tokens, see Fig. 14b. Thus, we adopted a IM token batch size for the training of models across all scales.
Specification of Fixed Point Solver for Language Model Training Our experimentation with the fixed point solver in the free phase involved adjusting various parameters, including the maximum number of self-iterations for the fixed point search (16, 24, 32) and the number of gradient accumulation steps (2, 4). The 32 iterations have smaller stop threshold of 0.02, whereas the 16 and 24 iterations have stop threshold of 0.05 see Fig. 14b, We used 130M Mamba2 model for this evaluation with an effective batch size of 1M tokens. We measured the validation split perplexity during training. For the training of the implicit Transformer models, we used a maximum self-iteration cap of 32. This was due to the higher spectral norm of the Transformer models requiring more steps to reach a fixed point below the threshold.
90% (4+1)-10% (24+4)
65% (4+1)-5%% (18+2)
90% (4+1)-20% (26+4)
B
60% (4-1)-301% (32+4)
70% (4+1)-30% (24+4)
(4+1)-38%) (24+4)
14
(4-1)
12.6
100% (4-1)
Validation perplexity
Validation perplexity
12.5
13.6
12.4
13.4
13.2
12.3
12.2
12.8
12.1
12.6
200k
250k
300k
350k
120k
130k
140k
150k
160k
Training Steps
Training Steps
Figure 14: Left: Impact of Bounded Phase (4+1) 1) Training Duration on Model Performance: A comparison of perplexity on the validation split of the D-PILE obtained by 130M Mamba2 models trained with 70%, 80%, and 90% bounded phase
durations, with a batch size of 0.5M tokens. Right: Evaluation of Fixed Point Solver Specifications: The relationship between different maximum iterations (16, 24, 32) and gradient accumulation steps (2, 4) on the validation split perplexity of the D-PILE for the 130M Mamba2 model with a batch size of 1M tokens.
23

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Resource Allocation for Training and Evaluating Large Language Models We trained our suite of models cluster with AMD Instinct MI300X GPUs. Each node within the cluster comprises 8 GPUs, and employed distributed multi-node processing to train our models on up to 32 GPUs simultaneously, Table 4 details the number of GPUs allocated for the training of each model, as well as the total GPU hours consumed by each. The evalution of models on downstream tasks was achieved on one 80GB Nvidia H100 GPU.
Table 4: GPU resource allocation and utilization for training large language models. The GPU counts employed and the total GPU hours expended for the training of each Mamba2 and Llama model variant across different parameter scales-130M, 350M, 760M, and 1.3B is listed. The models were trained using a cluster of AMD Instinct MI300X GPUs, with GPUs per node, utilizing distributed multi-node processing with up 32 GPUs in parallel.
Model
GPU Counts
Total GPU Hours
Mamba2"
1620
Mamba2 (4+1)
8
3185.6
130M
Mamba2 (24+4)
8
2756.8
Llama
1 114..56
Llama (4 1)
32
1027.2
Llama (32+4)
32
1561.6
Mamba2"
32
1516
Mamba2 (4+1)
8
2427.2
Mamba2 (24+4)
2168.8
Llama
32
1324.8
Llama (4 1)
8
2515.2
Llama (32+4)
2335.2
Mamba2"
1920
Mamba2 (4+1)
16
3932.8
760M
Mamba2 (24+4)
16
3440
Llama
4636.8
Ilamaa (4 1)
16
7612.8
Llama (32+4)
32
7676.8
Mamha2"
32
3054.4
Mamba2 (4+1)
32
5820.8
1.3B
Mamba2 (24+4)
32
5084.8
Llama
32
3084.8
Llama (4 1)
32
6057.6 6
Llama (32+4)
32
7225.6
Sum
83132.16
E. Additional Results
E.1. Inference Dynamics and Convergence of Implicit Language Models
In Table Table 6, present the average number of steps that the implicit language models require to process sequence length of 2048 from the test split of the D-PIL. dataset during inference in simultaneous mode. Moreover, the table also shows the relative error difference of the solutions found by the language models. Notably, the (4+1) configurations achieve a fixed point, despite not being explicitly constrained to do so. Our observations further reveal that the implicit models trained with a full DEQ setup-(24+d) for Mamba2 and (32+4) for Llama consistently reach fixed points within their training stop thresholds, which are < 0.05 and below the inference step threshold of four times 24 and 32, respectively.
E.2. Length Extrapolation Capabilities in Pretrained Models
We evaluated the ability of our models to extrapolate to longer sequence lengths and compared their performance with baseline models. All our in-house trained models, including the Mamba2 and Lllaa baselines, were initially trained on
24

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
sequences of 2048 tokens and subsequently tested on sequences of 4096, 8192, and up to 16384 tokens. Table 7 presents the average perplexities across different model scales. We note that the original Mamha2 models, denoted Mamba2", were trained sequence lengths of 8192. Our observations indicate that in all instances, our implicit models, including the Mamba2 (4+1), and Mamba2 (24+4) as well as the Alama(32+4) configuration, maintain lower perplexities compared to the explicit Mamba2 and Llama respectively. We also found that the difference in perplexity between the longer and shorter sequences becomes more pronounced as the size of the models increases.
E.3. Effective Duality between Simultaneous Mode and Sequential Mode in SSMs
Empirically, we demonstrate that once implicit state-space models (SSMs) are trained in a simultaneous mode (parallelizable in token dimension), these trained models can be utilized in sequential mode. In sequential mode, self-iteration occurs for each token, thereby affirming the duality the two modes, We employ the 1.3B Mamba2 (24+4) model and 1.3B Llama (32+4) model, trained on 207B tokens from the D-PILE, evaluate the duality between the two modes using examples from the test split of the D- D-PILE. Table 8 presents some detokenized outputs of the model in both modes.
25

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Table 5: Sample types from the Catbabi (Schlag et al., 2021) dataset, a modified version of the bAbl dataset, categorized into 20 examples. These examples are derived from (Weston et al., 2015).
Task 1: Single supporting fact
Task 2: Two supporting fact
sandra went back to the hallway, john moved to the hallway. where is sandra? hallway johu travelled the bathroom. daniel travelled the office. where daniel? office john moved the office, sandia travelled the buthooen, where sanda?? bathroons daniel went back the bedrocen. duniel went back to the garden, where is daniel? garden sandra moved to the ballway. john went back the hathroom, where is daniel? garden.
many went back the bathroom. john went to the office. daniei grabhed the fonthall there. sundra travelled to the bathrooe. daniel left the foctball, sundra moved the bedroom. sandm journeyed to the kitches. sandra travelled to the garden. sandra wesr to the kitchen. daniel moved to the gardon. sandra moved to the garden. may went to the office. daniel went to the kitchen.
the bathroom. john tavelled
Task 3: Three supporting fact
Task 4: Two argument relations
john went back the bodroom. sands get the milk,. sandra discarded the milk sundra took the milk. john went back to the office. mary moved to the buthroom. john travelled to the kitchen damiel went back to the bethoom. john weni bock the office. murry look the apple. mary inavelled the office. danicl back drupped the milk there where was the apple before the buthroom? ? uffice
the badroom is west the buthmomm the kitchen is west of the hedroom. what west of the bathroum ? bedroom
kitchen. mary moved buthnoom.
Task 5: Three argument relations
Task 6: Yes/No questions
jeff journeyed the gorden. fred went to office. fred the ballway. fied got the upple there. maved to office. fred went the kncceen fred put down apple. took the apple there. many went back to the garden. fied travelled to the buthroo.. jetf grabbed the milk there. bill went to the office jeff journeyed to the bathrom. jeff put down the milk there. frod picked up the milk there. food paosed apple jeff. who gave the apple to fred
john went the office. sundra went the buthroum. is sandra the kitchen ! no sandra travelled to the gardes sandra went to hedroom. is sandra the hedrocmn yes sundra travelled the garden mary went to the kitchen sundra in the garden ? yes john grabbed the apple there dasiel joumeyee to the hedroom. is sandra in the bedroomm no jobn picked up the fonthall there. john took the milk there. may hodrocen ? no
Task 7: Counting
Task 8: List/Sets
mury gul the apple there. john went back to the kitchen. how manyy ubjects mary caryying" ? one john mowed to the garden mary left the apple there. how objects is mary carrging mone john moved the bathroom. mary grabbed the apple
john went the buthroom. daniel travelled to the garden. mary moved the hallway. sandra moved the bedroom. daniel journeyed to the hallway. mary picked the apple there what mary currying 7 apple sandra travelled to the bathroom. mayy drogped the upplo what is mary carrying dankel moved to the office. daniel gribhedd the foothall mere. what is mary carrying ? nothing danicl moved to bedrocms may got the apple there. what dasiel carrying Roothall
how many objects many carrying one john went to the gunlen, mary drogped the apple. how many ubjects is mary carrynng none sandi journeyed to the bedroom. may moved to the kitchen how many objects is mony coryiig 2 none
Task 9: Simple negation
Task 10: Indefinite knowledge
daniel bathroom. sandra is no longer the bedroom. dansel the kitchen no mary no longer in the bedroom. sandra is no longer the kitchen. is sandr in the kitchen ? mo sandra is in the hallway. dansel in the gandes. is sandra the hallway yes sandra in the bedroom. sandia the bathroom. is sandra in bedroom sandra travelled the hallway. . the kitchen. sandra the kitchen?
julse either the bedroom or cinema.. bill joumeyeed the cinema. bull the cinema yes bill in the office. julie went to the school, is julie in the park no julie travelled to the office mary is in the school is bill in the park no mayy is in the park. bill is either in the office or the kitchen. mary in kitchen fred travelled the office. mary went back the kitchen. ? yes
mury in
Task 11: Basic conference
Task 12: Conjugation
daniel journeyed the garden. afterwurds he journeyed sundra went back to the kitchen. then she moved to the office. where duniel kitchen john moved to the office. afterwards he joumeyed to the bedroom where john ? bedroom dansel journeyed the oftice. after that moved the hallway. where the hallway. . following that she journeyed the buhhrom.... where daniel ? hallway
kitchen. where daneel latchen
danicl and mary journeyed to the hedroom. mary and sandra moved to the battroom where is daniel hedroom sandra and joumeyed bedmomn where is sandra bedroom sandra and mary went to office sandm daniel went back to the hathroom where is sandra bathroom mary and danial went to the bedrooms. john and mary travelled the bathroom. where john 7 bathroom mary and daniel went to the hallony and clnnie back bodraom where daniel bedroom
hallway. mary and sandra
back the
bedroom sandra
Task 13: Compound conference
Task 14: Time reasoning
sundm and duniel joumeyed ed to the bedruom. following that they enmelled to the kiichen. where is duniel 7 kitchen john and mary invelled to the buthoom. then they went to the hallway where is john ? hallway daniel and sandra trovelled the office. following that where is sandra? buthroom damiel sandra moved to käichen. then they went to the buthroom. where is sandra?
bill eravelled to the office yesterday bill moved to the school this morning fred went bock to the kitchen yestenlay julii joumsyed to the school yesterda. where was bill before the school ? office this afternoon bill journeyed to the cinema. yestenday mary journeyed the park. where was before the cinema the kischan this morning. . where was bill before the cinems. ? school fred travelled to the park this afternoon.
the bulhroom
frod journeeed the bedruom this muming. julie went back
? bathroom mary and john went the kitchen. . then they went to the bedroom. where is john ? bedroum
Task 15: Basic deduction
Task 16: Basic induction
mice are afriid of cats, cats are afraid of wolves emily a cat. walves are afraid of cats jessica a mouse, gertrade a wolf, shoup are afrii of mice. winoea is wolf, what is winom afraid of 2 cut what is jessica afraid 2 cal what is jessica afraid ? what is jessica afraid 7
greg a swan. bernhard is a thino julius is frgg. bernhard is white, briun a thino, julius is green. grug is white. brion groen. lily is soan. what color is lily 9 white
Task 17: Positional reasoning
Task 18: Size reasoning
the rod squir is to the right of the ninngle the pink rectangle is helow the red square is the trinngle the lea of the pink rectingle ? yes the tringle the right of the pink rectangle? ? no pink mestangl?
The chocolote fins inside the container the box is bigger than the chest. the chest fits inside the container, the bot of chocolatee fiss inside the container, the chest bigger than the suitcase is the suitesse bigger than
the pink rectangle yes the pink reciangle to the right of the triungle ? yes the pink rectanple below the triungle ? yes the trinnll to the right of the pink restangle? ? mo is the pink recangle to the right of the triangle? ? yes
the triangle below
the trinngle
? the box fit sultcase does suiteuee no is the suitcase bigger than the box no is the suitcase bigger than the bux." no
the
Task 19: Path finding
Task 20: Agent's motivation
the buthwoom is south the garden. the kicche north of bedrocut. the hallway north of the garden. the hodroom is west of the garden. the office is west of the hallway. how do you go from the hallway to the bedroom ? s...
yaun tired. where will yana go bedroom jeson is bocod. where jasoss go garden antoine is bored. where will antoine go gurden antoine moved to the garden. why did antoine go to the garden ? bured antoine got the foetball there why did antoine get the football ? bored sumit is hungry. where will go lätchen sumit traeeledd to kitchen. why did sumit the kitchen hungry yann travelled to bedroem. why did yau go to bedroom tired
26

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Table 6: Inference dynamics and convergence characteristics of implicit language models across different scales. We show the performance of Mamba2 and Llama models at various parameter Scales-130M,, 350M, 760M, and 1.3B-whrd processing sequences of length 2048 from the test split of the Pile dataset. The average number of steps required for convergence during inference in simultaneous mode is detailed alongside the relative error difference of the solutions obtained by the models. Remarkably, the (4+1) configurations reach fixed point organically, without explicit constraints enforcing this behavior. The table further highlights that implicit models employing full DEQ setup-(24+4) for Mamba2 and (32+4) for Llama- demonstrate consistent convergence within their predefined stop thresholds, which are less than 0.05. These thresholds are also below the designated inference step threshold of four times 24 and 32 for the respective models.
Model
D-Pile Perplexity
Inference Steps
Rel. DifI.
Mamba2(4+1)
13.76
0.035
130M
Mumbu2(2444)
12.86
62
0.036
Llama(4+1)
12.73
13
0.014
Llama(32+44
11.73
0,033 53
Mamba2(441)
10.02
10
0.012
350M
Mumba2(2244)
9.70
0.024 49
Llama(4+1)
9.66
12
0.015
(lama(32+4)
9.43
0.037 $7
M4mba2(4+1)
8.60
10
0.013
760M
Mamha2(24+4)
8.35
0.025 45
Llama(4+1)
8.27
12
0.014
Llama(32+4)
7.90
0.044
M4mba2(6+1)
7.97
10
0.013
Mamba2(24-4)
7.70
0.029
1lama(4+1)
0.015
Llumu(32+4)
7.24
69
0.048
Table Length extrapolation performance of pretrained models on varying sequence lengths. We demonstrate the average Pile test perplexity results for Mamba2 (300B tokens) and our in-house trained Mamba2' (207B tokens) baseline models, as well as our Mamba2(4+1) and Mamba2(24+4) configurations, across a range sequence lengths-2004 4096, 8192, and 16384 tokens. Each model was originally trained on sequences of 2048 tokens, while the original Mamba2 was trained on 8192 tokens. The results highlight that implicit models consistently achieve lower perplexities than their explicit counterparts across all tested sequence lengths. Notably, model sizes increase, the discrepancy in perplexity between longer and shorter sequences grows more evident.
Model
2048
4096
8192
16384
ppl.
ppl
ppl.
ppl.
INDET
Mamba2
13.72
13.44
13.92
14.91
Mamba2"
13.05
12.72
12.76
12.94
Mamba2(4+1 I-ouus -ours
13.76
13.04
13.06
13.25
Mamba2((444)-ours
12.86
12.53
12.55
12.78
M37E
Mamba2
10.55
10.28
11.20
13.89
Mamba2"
10.18
9.94
10.07
10.34
Mamba2(4++)-ou
10.02
9.74
9.73
9.86
Mamba2(24++4)ouurs
9.70
9.45
9.44
9,59
760M
Mamba2
9.23
9.45
34.99
231.25
Mamba2-
8.98
8.94
10.38
12.56
Mamba2(4+1) -ours
8.60
8.38
8.40
8.55
Mambu2(24+4)-ouns
8.35
8.16
8.18
8.33
1.3B
Mamba2
8.40
8.47
12.69
25.37
Mamba2"
8.28
8.24
9.96
15.06
Mamba2(4+ -ours
7.97
8.36
9.53
Mamba2((4+4)-ouns
7.70
7.59
8.19
9.63
27

---

Implicit Language Models are RNNs: Balancing Parallclization and Expressivity
Table 8: Selected examples from the test split of the D-Pile dataset, showcasing the detokenized output of the next-token predictions by the 1.3B implicit Mamba2 model, which was trained in a parallel fashion and tested in both Simultaneous and Sequential modes. These outputs highlight the duality of the two approaches.
Ground Truth
Simultancous
Sequential
labour productivity and output would rise as a result it is essential to protect the professionalism of cer- tain categories of workers: the de- bates here centred on performance artists and female theatrical employ- ees engaged in highly physical and in- tensely emotional work * heavy phys- ical labour and strenuous exercise can lead to disruptions of the menstrual cycle women's 's physical and intel- lectual capacities are reduced during menstruation; women lose muscular strength and powers of concentration women's 's biological constitution and reproductive functions require spe- cific recognition in law Against the provision: employers are less likely to appoint women if they are guaran- teed paid time off work during men- struation
market. the. be. a result. The would a to ensure the interests and of the professions of workers profes- sions on arered on the management, the artists performers in the skilled work demanding competitive work the industry work is therenuous phys- ical are be to injuryions in the body cycle and+ the's 's bodies and emotional capacities are not by pregnancyation this are their strength and stam of en- durance menstru's 's menstrual clocks is menstrual capacity are special pro- tection the Thest this background of the should not likely to be women to they are not the matern off for for menstruation womenin un the com- mentators)
market. the be. a result of The would a to ensure the environment and of the professions of workers pro- fessions on arered on the manage- ment, the artists performers in the skilled work demanding demanding work the industry work is theren- uous physical are be to injuryions in the body cycle and- the's bodies and emotional capacities are not by pregnancyation this are their strength and stam of endurance menstru's 's menstrual clocks is menstrual capac- ity are special protection the Thest this background of the should not likely to be women to they are not the matern off for for menstruation womenin un the commentators)
form of "photon counting".. "This de-excitation is called 'fluorescence' and it is characterized by a lifetime of a few nanoseconds of the lowest vibrational level of the first excited state. De-excitation from the excited singlet state to the ground state also occurs by other mechanisms, such non-radiant thermal decay or 'phos- phorescence'. In the latter case, the chromophore undergoes a forbidden transition from the excited singlet state into the triplet state (intersystem crossing, ISC, Fig 2.4), which has non-zero probability, for example be- cause of spin orbit coupling of the electrons' magnetic moments" its a type of INTERSYSTEM CROSSING doing a search for Intersystem cross- ing, memristor brings up this link.. A composite optical microcavity, in which nitrogen vacancy (NV) centers in a diamond nanopillar
meaning of thethe"" is I'The iscept- factciting of thephotonorescence' and it is the by the photonphotone- time of the few hundredoseconds. the excited- level of the excited of -excitation is the first state state is the ground state is occurs, , as -radiative fluorescence' the case case, the ex- citedophore is a non transition to the ground singlet state to the ground state, 'ystem crossing), ISC), or.).I). which is lifetime-ne probability of but example, of the- coupling to the excited to spin moment. " a bit of fluorescenceC crossingOSSING, " google on "tersystem crossing Teor, up a mem material devicecavity is consisting which the- (NV) centers are diamond diamond crystalillar are coupled to aing gallery modes sili- con microsphere, is fabricated.
meaning of theto"" is IThhe iscept- factciting of thephotonorescence' and it is the by the photonphotone- time of the few hundredoseconds. the excited energy state of the excited of -excitation is the first state state the ground state is occurs, as -radiative fluorescence' the case case, the ex- citedophore is a non transition to the ground singlet state to the ground state, "ystem crossing), ISC), or.).1) I)). which is a lifetime-ne probability of but example, of the- coupling to the excited to spin moments. " bit of fluorescenceC crossingOSSING, google on "tersystem crossing, Teor, up a " mem material devicecavity is consisting which the- (NV) centers are diamond diamond crystalillar are coupled to aing gallery modes sili- con microsphere, is fabricated.
28


Reviews:
[{"review_id": "qsLqbEmex2", "reviewer": "Reviewer_gbY7", "summary": "This paper proposes implicit language models, which are RNNs defined implicitly via fixed-point iterations. Theoretically, the authors show that implicit models can represent non-linear and non-diagonal state transitions of RNNs, overcoming the limitations of transformers and state-space models (SSMs) which are restricted to simpler, linear transitions. Empirical results show that implicit models can solve a challenging state-tracking problem that transformers and SSMs fail to learn. In addition, the authors scale implicit models up to 1.3B parameters and show they outperform explicit counterparts on language modeling benchmarks, with favorable length generalization and auto-regressive generation capabilities.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback and positive assessment of our submission. Below we address raised concerns and assumptions:\n\n**Experimental details**: we will make sure to double check Appendix D to see if any detail is missing. Please see additional wall clock time and memory footprint in response to reviewer __AqCH__.  We are also releasing code with precise experiment configurations.\n\n**Essential references not discussed:** The Reviewer appears to make the key assumption that some or all of the models with linear attention or kernelized attention as discussed in (Tay et al. 2020), or models with gated linear attention and variations of gated state-space models such as Mamba (Gu & Dao 2024) already address the theoretical issues discussed in our manuscript. We would like to respectfully clarify this assumption. Yes, these models admit a recurrent inference mode. [1] discusses the structural similarities between these attention variants. However, they all face the same limitations discussed in (Merill et al. 2024, Sarrof et al. 2024) since their recurrence is not expressive enough. We discuss xLSTM as an exception, whose sLSTM module is a non-linear + non-diagonal RNN, which lacks parallelization, though. If we should be missing relevant related work, we are happy to take further suggestions.\n\n[1] https://arxiv.org/abs/2405.15731 \n\n**Practical usefulness:** Our primary goal is to highlight qualitative differences between implicit and explicit sequence models, without claiming universal superiority. Our experiments demonstrate that implicit models scale practically, as further confirmed by concurrent work (https://arxiv.org/pdf/2502.05171). Thus, implicit models are a viable choice when complex sequential processing (e.g., state tracking) is required. We offer a design point where you can get a qualitative jump in expressiveness at the expense of compute, a trade-off reminiscent of the test-time compute paradigm.\n\n> [...] however, the paper fails to motivate why that is useful.\n\nStatic depth models show exceptional performance. Yet, they even struggle with certain regular languages. This limits their capabilities to execute tasks that require state tracking (e.g. managing supply chain processes). We tested GPT 4o and o1 (API) on the S5 state-tracking task.\n|Sample Length|o1-mini Acc|GPT-4 Acc|  \n|-|-|-|  \n|5|0.967|0.200|  \n|15|0.967|0.100|  \n|32|0.100|0.133|  \n\n> 1.3B parameters do not provide concluding evidence that the state expressivity is actually needed for good performance\n\nWe agree that most downstream tasks do not require extensive state tracking, and we do not claim that it is required in all tasks. We do see increased performance in HellaSwag, a benchmark requiring limited state-tracking.\n\n# Questions\n\n1.Implicit models for language modeling and CatbAbI used next-token prediction loss with phantom gradients where backpropagating through fixed self-iterations (4 for language modeling, 6 for CatbAbI) after gradient-free searches (up to 24/32 iterations language modeling, 50 CatbAbI).  \n\n2.The LM (The Pile) curriculum balances cost and accuracy by applying more self-iterations to only n=20% of tokens (Table 1). To address Reviewer questions, we tested n=10% on the 1.3B Mamba2 (24+4) model, showing curriculum robustness.  \n\n|Metric|10%|20%|  \n|-|-|-|  \n|Tokens seen|207|207|   \n|LAMBADA |0.4186|0.4116|  \n|HellaSwag|0.3672|0.3527|  \n|PIQA|0.6502|0.6572|  \n|Arc-E|0.4596|0.4815|  \n|Arc-C|0.2372|0.2372|  \n|Wino|0.5170|0.5130|  \n|OpenQA|0.2900|0.3000|  \n|Average|0.4200|0.4219|  \n\n3.See initial statement\n\n4.Implicit models can implement non-diagonal + non-linear transitions. We show in Fig. 2 (and Sec. 5) that implicit models allow to conduct sequential inference, e.g. language generation, only carrying forward the converged hidden state (SSM) or KV-cache (Transformer). This allows memory allocation independent of the number of iterations. While RNN-based ACT (Graves, 2017) make the non-parallizability of RNNs even worse by sequentially iterating more steps per token, implicit models learn the adaptive budget for all tokens in parallel, which allows us to scale to 1.3B models.\n\n5.We will further discuss limitations and computational differences between implicit and explicit models. Larger models suffer less from state tracking (Fig. 1), but even the largest GPT models remain limited.\n\n# New experiments\n> \"Benchmarks\"\n\nWe were able to evaluate the models on an additional task, the MMLU. Please note that we selected our benchmarks to align with the Mamba2 and xLSTM papers.\n\n|Model|MMLU Accuracy|  \n|-|-|  \n|1.3B_implicit_mamba2 (24+4)|0.269|  \n|1.3B_implicit_llama (32+4)|0.258|  \n|1.3B_llama_baseline|0.2483|  \n|1.3B_mamba_baseline|0.2502|   \n\n**Computational benefits:**    \nNote that implicit models have qualitatively higher expressivity than explicit models as shown in our synthetic experiments. In addition, we show results depending on the compute budget in our response to __AqCH__.", "claims_and_evidence": "The study provides sufficient evidence to support its main theoretical and empirical claims: \n- The claim that implicit state-space models can represent non-linear and non-diagonal state transitions is proven theoretically in Sec 3.1 and Appendix B. \n- The authors reproduce the finding from prior work regarding the inability of transformers or SSMs to solve hard state tracking problems such as S5 and show that implicit SSMs can behave like RNNs, in Sec 4 and Fig 1. \n- The scaling of implicit models to large language modeling tasks up to 1.3B parameters is supported by the results presented in Section 5 and the detailed experimental setup in Appendix D.3 (9 tasks).\n- The claims regarding length extrapolation capabilities and duality between sequential and simultaneous modes are supported in Section 5, and Figs 4 and 2.", "methods_and_evaluation_criteria": "- The proposed method is well motivated and shown theoretically to address the previous models' inability to represent non-linear state interactions. The expectation is that overcoming this limitation it will lead to more expressive modeling and state tracking. \n- For evaluation, the authors focus on synthetic and real-world state tracking problems to systematically evaluate their theoretical findings. This is suitable experiment choice as it grounds the theoretical findings to empirical evidence, making the findings more trustworthy.\n- In addition, the evaluate performance on downstream performance carrying out experiments with language models of increasing size up to 1.3B. These results show whether the expressivity is actually useful on language modeling and downstream real-world tasks.", "theoretical_claims": "The main theoretical claim of the paper is captured in Theorem 1 and shows that the transition function defined by the implicit SSM is non-linear and non-diagonal. I checked the proof and found it to be logically correct; it applies the implicit function theorem to the function $g(z, h, x, \\theta) = z - f_\\theta(z, h, x)$ and then shows that the derivative of the implicit function $\\phi(h,x, \\theta)$ with respect to $h$ is  a non-linear function if $f$ is non-linear. Based on this then the state-to-state Jacobian is shown to be non-linear.", "experimental_designs_or_analyses": "Yes, I reviewed the experimental designs and analyses presented in the paper and found them to be sound and well-executed in general. Below I list a few non-major issues:\n\n1. State tracking experiments\n    - It would be useful to provide more details on how exactly the synthetic data distributions were created exactly and what is the intuition behind the chosen parameters. Providing some examples would also help. \n    - In addition, a few experimental details are missing on the hyper-parameters of Mamba2 model on both synthetic and CATBABI tasks (number of layer, learning rates, batch sizes, etc).  \n2. Language modeling experiments\n   - The scaling experiments would provide more confidence for practical impact if the mode size was up to at least 7B models; it is not guaranteed that the behavior observed below 1.3B will generalize and whether non-linear transitions are still useful.  \n    - It would help if experiments included more recent benchmarks for large language models such as MMLU, BBH, HELM. \n    - Provide more details on the training budget used for training the implicit models and state-space models and quantify the training + inference costs. It was not clear to me how the authors ensure equal budget for convergence and what is the exact computational benefit for implicit models.", "relation_to_broader_scientific_literature": "The contributions are generally well-situated within the broader scientific literature: \n- The paper builds upon previous theoretical work that has identified limitations of state-space models in capturing complex sequential states and recognizing certain formal languages (Merill et al. 2024, Sarrof et al. 2024). They proposed models aim to address these exact limitations.\n- The authors develop implicit models building on top of deep equilibrium models and implicit function theorem from prior work. The adaptive computation that is inherited from these models is a useful property that has shown to be useful in previous research (Graves 2017, Dehghani et al. 2019).\n- Provides additional evidence to the existing literature that looped models are able to generalize better to input lengths not seen during training (Yang et al. 2024a). \n\nIt would be useful to discuss the advantages of implicit models compared to recurrent-attention-based transformers or hybrid transformers that make use of full and recurrent attention mechanisms that are competitive in terms of quality and speed tradeoff.", "essential_references_not_discussed": "There are prior works that used simple and more advanced recurrent attention mechanisms based on the kernel-based view of attention for transformers (Tay et al. 2020). With such formulations of attention the transformers are converted into RNNs which do not have the problems as pointed out in this paper. It would be essential to discuss the unique advantages of implicit models beyond implementing non-linear state-transitions of RNNs which has been addressed in the past.", "other_strengths_and_weaknesses": "Strengths:\n- The idea of using fixed-point iterations to combine the expressive power of RNNs and the parallelization benefits of transformers is quite interesting and useful, since the training becomes more efficient due to the computation of gradients with a constant memory footprint.\n- Provides a convincing theoretical analysis of implicit models and makes a solid connection between the benefits of RNNs and Transformers.\n- Experimentation is thorough and shows promising results up to 1.3B parameter models. The evaluation covers both tasks that require state tracking and tasks used in large language modeling studies.\n\nWeaknesses:\n- There is lack of discussion and comparison to hybrid models that leverage the benefits of RNNs and Transformer models through a combination of recurrent and softmax attention mechanisms. \n- The family of implicit models provides an appealing solution for the lack of non-linear state transitions, however, the paper fails to motivate why that is useful in practical real-world tasks where instruction-following models based on Transformers perform exceptionally well. \n- The scaling results up to 1.3B parameters do not provide concluding evidence that the state expressivity is actually needed for good performance on downstream tasks. In addition, the performance improvements also are not very consistent across different model sizes which further casts doubt on the generalizability to larger model sizes.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Could you elaborate on how the implicit model are trained exactly? A more detailed explanation of the training details (e.g. masking, next-token prediction, etc) would help the reader better understand how this approach is implemented and what is the training cost involved. \n2. Regarding the curriculum-based approach and the duration of different phases, what are the key considerations for the different design choices for different datasets? It would be good to specify how sensitive is the final performance on these choices. \n3. Previous studies have devised hybrid architectures that leverage the benefits of both RNNs and Transformers. Are the expressivity issues applicable to them? How do implicit models compare to them? \n4. What are the unique contributions of this work in the area of adaptive computation? Reading through the related work, it is not very clear what is the contribution of this work.\n5. The implicit models are shown to have good properties but there is little emphasis on their limitations. I wonder if the authors could provide some additional discussion and analysis on which tasks or settings the implicit models have difficulties or underperform other models. Also, do the authors expect their results to generalize to even larger model sizes?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "KCmObfA3ud", "reviewer": "Reviewer_AqCH", "summary": "1. The authors propose a DEQ-ified version (referred to as implicit models) of state space models like Mamba2. \n2. This is motivated by the fact that the diagonal (and real) state transition matrix of these models is not expressive enough for state tracking. They show an implicit model has as a non-linear and a non-diagonal state transition matrix which is required for state tracking. \n3. They test implicit Mamba2 model on the CatbAbi synthetic benchmark and on language modeling on D-PILE dataset.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the Reviewer for the positive perception. Below we try to answer to the reviewer’s remaining questions. \n\n__Q1: Can authors provide a wall-clock time analysis for their method against vanilla Mamba-2__\nThanks. We will add the following to the paper comparing throughput and wall clock time (WCT) relative to the explicit model for the largest models (760M and 1.3B).\n### Training Token Throughput and Wall Clock Time\n|Model|760M Tok/s|760M Rel T-put|760M Rel Time|1.3B Tok/s|1.3B Rel T-put|1.3B Rel Time|  \n|-|-|-|-|-|-|-|  \n|Mamba2*|1872|-|-|588|-|-|  \n|Mamba2(4+1)|914|49%|205%|309|53%|191%|  \n|Mamba2(24+4)|209|11%|180%|71|12%|166%|  \n|ImpMambaAvg|546|29%|343%|184|31%|319%|  \n|Llama†|775|-|-|582|-|-|  \n|Llama(4+1)|472|61%|164%|237|41%|196%|  \n|Llama(32+4)|47|6%|166%|50|9%|234%|  \n|ImpLlamaAvg|131|16.8%|297%|149|26%|391%|  \nThe averaged numbers for the implicit models take the curriculum into account.\n\n## Inference WCT Measurements (Time per token in milliseconds, averaged over 2048 tokens generated)\n|Steps|Llama 130M|Llama 1.3B|Mamba 130M|Mamba 1.3B|  \n|-|-|-|-|-|  \n|expl|46.7|97.9|23.5|45.7|  \n|1|62.6|120.7|29.2|55.6|  \n|2|130.8|236.7|53.5|102.9|  \n|4|200.3|421.6|93.6|182.0|  \n|8|440.3|745.2|180.2|3612.0|  \n|16|748.8|1294.0|356.1|705.0|  \n|32|1356.3|3204.8|710.5|1414.7|  \n\n## Memory Usage in Inference [MB] (Implicit / Explicit)\n|Model|Llama 130M|Llama 1.3B|Mamba 130M|Mamba 1.3B|  \n|-|-|-|-|-|  \n|Implicit / Explicit|871 / 511|10216 / 5281|935 / 547|10592 / 5488|  \n\n## Word Problem Inference\nWe evaluate 4k samples at batch size of 512. For the single layer implicit Mamba2\n- setting an upper limit of 4 self-iterations takes 26s and gets 97.8% accuracy on the p=0.5 distribution\n- setting an upper limit of 16 self-iterations takes 35s and gets 99.8%. Note that iterations terminate after 6 steps on average due to convergence.\n\nThe explicit model with 16 layers takes 37s and gets 1.5%. This explicit model has the same dimensions per layer, and hence 16x the number of parameters compared to the implicit model. \n\n## Catbabi WCT Training\n|Model|GPU-Hrs|Rel T-put|Rel Time|  \n|-|-|-|-|  \n|expl-1-lyr|1.02|-|-|  \n|impl-1-lyr|1.83|55.74%|179.41%|  \n|expl-2-lyr|1.81|-|-|  \n|impl-2-lyr|2.96|61.15%|163.54%|  \n|expl-3-lyr|2.58|-|-|  \n|impl-3-lyr|4.51|57.21%|174.81%|  \n\n## Catbabi WCT Inference (Time per token in milliseconds, averaged over 50 tokens generated)\n\n|Model|Step-1|Step-2|Step-4|Step-8|Step-16|Step-32|  \n|-|-|-|-|-|-|-|  \n|expl-1-lyr|1.879||||||  \n|impl-1-lyr|2.966|4.8321|8.6203|16.0183|30.9154|60.5825|  \n|expl-2-lyr|3.1831||||||  \n|impl-2-lyr|4.7185|8.2655|14.8508|27.6253|51.3367|105.8571|  \n|expl-3-lyr|4.4807||||||  \n|impl-3-lyr|6.3261|11.0607|20.2967|38.7703|75.5918|149.1918|  \n\n__Q2:  Could authors do a language modelling experiment with a control on the FLOPs.__\n\nOur primary goal is not FLOPs efficiency, but rather exploring a fundamental trade-off: how much true recursion is necessary for language modeling and reasoning, balancing parallelizability and expressiveness. While a single iteration of our implicit model approximately matches the explicit model in FLOPs, our intention isn't to claim FLOPs optimality. Instead, we aim to investigate qualitative differences, accepting increased representational power at the expense of greater and dynamic depth. Below, we provide a test-time compute table which allows for a FLOPs-matched comparison and highlights the flexibility of implicit models to trade off compute vs performance at test time.\n\n| Task    | Model                | f:1  | f:2  | f:4  | f:8  | f:16 | fixpt |  \n|---------|----------------------|------|------|------|------|------|-------|  \n| Avg Acc over tasks in Table1 | Implicit Mamba 1.3B|0.31|0.38|0.52|0.56|0.56|0.56|  \n| Avg Acc over tasks in Table1 | Implicit Llama 1.3B|0.30|0.30|0.42|0.57|0.59|0.59|  \n\n\n__Q3: Could the authors comment and contrast their method against SSM with complex diagonal values/negative eigen values [1]__?\n\nWe discuss the paper that the Reviewer is referencing in the Related Work section. It is important to see that even with negative eigenvalues/complex diagonal values, the models in that reference are not able to solve the full S5 problem but only S5 restricted to transitions of (two-element) swaps (see Fig. 4 in Grazzi et al.). \n\n__Q4: Is it possible to get a characterization of the class of transition matrices that implicit model admits__?\n\nWhile a full characterization of the transition matrices goes beyond our study, Eq. (14) provides first insights. $\\frac{\\partial \\varphi}{\\partial h_{t-1}^*}$, the derivative of the fixed point w.r.t. hidden state in Eq. (12), is a source of general non-diagonal entries that depends on the implicit function $\\varphi$. The hidden state is propagated through fully connected layers in the forward pass during the self-iteration. This leads to non-linear and non-diagonal contributions comparable with RNNs or multi-layer feed-forward networks.", "claims_and_evidence": "For theoretical claims and issues, please see \"Theoretical Claims*\nFor methodological claims and issues, please see \"Methods And Evaluation Criteria\"", "methods_and_evaluation_criteria": "Overall, **I think this is a good paper**, however, the only methodological/evaluation concern I have is on the **wall-clock time of this method**. Please correct me if I am wrong, but I expect that this method's (4+1) version takes atleast 4x the number of FLOPs in the forward pass than vanilla Mamba2 and that it's (32+4) version would take 32x the number of FLOPs than vanilla Mamba2. To probe this a little more, I suggest the following experiments:\n\n1. Can the authors provide a wall-clock time analysis for their method against vanilla Mamba-2\n2. I believe the the authors have currently controlled for number of parameters; could they also do a language modeling experiment with a control on the FLOPs. I am curious if the model's superior performance (on sizes > 350M) is actually due to the increased computation rather than the ability of the model to do state tracking.\n\n\nFurthermore, I think the wall-clock time becomes even more relevant as recent works like [1] have shown that Mamba2 can do state tracking if the transition matrix has complex eigenvalues. I think associative scan implemented for Mamba-1 already supports diagonal matrices with complex eigenvalues and it might be a more practically efficient solution to this problem. **Could the authors comment and contrast their method with this solution?** *NOTE: I would not hold this fact against this paper since [1] is a recent work.*\n\n\n[1]: Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. Riccardo Grazzi, Julien Siems, Jörg K.H. Franke, Arber Zela, Frank Hutter, Massimiliano Pontil", "theoretical_claims": "In light of my comment on [1] in \"Methods And Evaluation Criteria\", I am curious if the authors can compare/comment on the difference in expressivity of their method and an SSM with a complex valued diagonal transition matrix. Is it possible to get a characterization of the class of transition matrices that the model admits?\n\n[1]: Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. Riccardo Grazzi, Julien Siems, Jörg K.H. Franke, Arber Zela, Frank Hutter, Massimiliano Pontil", "experimental_designs_or_analyses": "Please see \"Methods And Evaluation Criteria\"", "relation_to_broader_scientific_literature": "Tries to fix the problem that SSMs cannot do state tracking which is known in the community.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A\n\nRunning the model on more state tracking tasks like Parity and Arithmetic Mod (w or w/o brackets) might be help strengthen the paper. \n\n[1]: Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues. Riccardo Grazzi, Julien Siems, Jörg K.H. Franke, Arber Zela, Frank Hutter, Massimiliano Pontil", "questions_for_authors": "Please see \"Theoretical Claims\" and \"Methods And Evaluation Criteria\"", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XvKiXFzATZ", "reviewer": "Reviewer_oVp8", "summary": "This paper describes an implicit approach to training state-space models with arbitrary depth by having the models evaluated in a fixed point and implicitly differentiating using the implicit function theorem, like DQEs. They find that on certain tasks, implicit SSMs outperform SSMs, which are unable to learn these tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# General comment\nWe would like to thank the reviewer for their comments. We are pleased that the reviewer finds our methods, evaluation criteria, experimental design, and analysis satisfactory. We appreciate the suggested clarifications and would like to take this opportunity to elaborate further on our theoretical results.\n\n\nWe rigorously derive the Jacobian of the hidden state-to-state in Eq. (14) by taking the derivative of the fixed point condition Eq. (8). Incorporating the reviewers feedback, we will add the corresponding fixed point condition for $z^*$ as well along Eq. (8) as\n$$z_t^* = f_\\theta\\left(z_t^*, h_{t-1}^*, x_t\\right)$$\n\nIf\n$M = \\frac{\\partial g_\\theta}{\\partial z}\\vert_{z_t^*, h_{t-1}^*}$\nis non-singular, we can further apply the implicit function theorem to replace $\\frac{\\partial \\varphi}{\\partial h}$ in Eq. (14) with Eq. (12). While there might exist $\\theta$ such that $M$ is singular, we numerically verified that $M$ is non-singular for randomly initialized networks, and we did not observe singular $M$ during our training experiments.  \n  \nEq. (14) contains products of matrices. As the reviewer points out, there is no guarantee that these products will not cancel out all non-diagonal terms, which could effectively lead to a diagonal Jacobian. As an example, the probability of two random Gaussian matrices multiplying to a diagonal matrix is zero. For the case of Eq. (14), we did not provide a rigorous proof that the probability is zero. Yet, we numerically checked that the Jacobian is non-diagonal both using autograd as well as Eq. (14).  \n  \nTo rule out any concerns about the rigor of our theoretical contribution, we suggest revising Theorem 1 to mention only what we can rigorously prove:  \n  \nTheorem 1:    \nThe Jacobian of the implicit SSM is given by Eq. (14). If   \n\n $$ \\frac{\\partial g_\\theta}{\\partial z}\\vert_{z_t^*, h_{t-1}^*}$$\n\nis non-singular, then $\\frac{\\partial\\varphi}{\\partial h}$ is given by (Eq. (12)).\n\n\nRemark 2: In contrast to the explicit state-space model in Eq. (1) and (2) the implicit state-space model allows for non-linear and non-diagonal state-to-state transitions. We empirically observe that $\\varphi$ is non-linear and that the Jacobian is non-diagonal.", "claims_and_evidence": "This paper claims “Notably, our implicit models outperform their explicit counterparts on standard benchmarks”, which is supported by Table 1, and can in fact model stateful systems, which is supported by Figure 1.", "methods_and_evaluation_criteria": "Several benchmarks are used for evaluation, which all seem reasonably well-suited to the task.", "theoretical_claims": "Theorem 1 claims that the transition function in equation (8) is non-linear and non-diagonal. Appendix B contains the proof.\n\nI am somewhat unfamiliar with this style of proof, so perhaps this is justified by an “almost always” that is implicit, or I am completely missing something. However, several times the compositionality of non-linear functions, non-diagonal matrices, etc is used, which does not hold in general.\n\nFor example, the line following equation 12 suggests that because (partial phi)/(partial h) = - (I - (partial f)/(partial z))^{-1} (partial f) / (partial h), it must be the case that f being nonlinear implies phi is nonlinear. However (again, I could be completely wrong), this is not guaranteed. For example, the equation f(h, z) = e^h e^z + z would have (partial phi)/(partial h) = - (1 - (e^h e^z + 1))^{-1} e^h e^z = - (e^h e^z)^{-1} e^h e^z = -1, so phi would be linear despite f being nonlinear. Obviously, this example is contrived specifically to be a counterexample, and the property holds in almost all cases, but unless there is something I am missing, it seems like this should be noted in the theorem statement.\n\nSimilar “two nonlinear functions exactly cancelling each other out”-style counterexamples could surely be found for the claim on line 785 and a similar style of argument could be made to find a counterexample to the claim on line 796 that this Jacobian is necessarily non-diagonal.", "experimental_designs_or_analyses": "The experimental designs and analyses seem appropriate.", "relation_to_broader_scientific_literature": "I am not familiar enough with the existing literature to fill this section.", "essential_references_not_discussed": "I am not familiar enough with the existing literature to fill this section.", "other_strengths_and_weaknesses": "All points are addressed elsewhere.", "other_comments_or_suggestions": "Minor nits:\nEquation (8) should probably be paired with an additional equation z^* = f_theta(z^*, h_t^*, x_t) in order to make the fixed-point nature of the definition clear.  Additionally, in the text immediately following this equation, “The fixed point z^*_t depends on h^*_t, and hence by equation (7) on h^*_{t-1}” I believe you intended to write “by equation (6)” because that is where the dependency with h_{t-1} is established.\n\n216 right column: A5 should be A_5.", "questions_for_authors": "All points are addressed elsewhere.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "htaiWkPUOR", "reviewer": "Reviewer_e8CX", "summary": "This paper introduces implicit SSMs, which are a parameter tied form of SSM that can be run for arbitrarily many self-iterations until convergence. They propose training implicit SSMs in a scalable manner using phantom gradients from the implicit layers literature. They demonstrate the ability of implicit SSMs to state track on hard OOD settings and in language modeling.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# General comment\nWe thank the reviewer for lots of insightful feedback, which significantly helps us to revise our manuscript. We appreciate the positive evaluation (\"On balance, I think this is a great paper\") as well as the engagement expressed by many detailed questions. In face of the space limit (5k characters), we have to focus on selected questions. \n# Theoretical contributions\nWhile we are not aware of theoretical convergence guarantees, we checked that our models converge to fixed points on all datasets that we present in the paper. Models that reach fixed points have the properties claimed in Section 3.\n\n> $h^*_{t-1}$ vs $h^*_t$ in $\\varphi$. \n\nThere is a typo in Eq. (7), which should be analogous to Eq. (2), where the dependency on $h_{t-1}$ is correctly stated in line with standard RNN formulations. Consequentially $\\varphi$ takes $h_{t-1}^*$ as an argument.\n\n> Is the minus sign in equation (4) correct? \n\nNo, thanks for spotting this!\n\n> Please include such a numerical check in your rebuttal.\n\nIncorporating the correct sign, we compared the RHS of Eq. (14) at the fixed point with the unrolled autograd (AD) Jacobian. The absolute difference between the AD Jacobian and Eq. (14) is three orders of magnitude smaller than the values in the Jacobian.\n\n# Addressing wall-clock time \nWe have included the WCT for all datasets and the memory of language models in our response to __AqCH__.\n## Synthetic state tracking\nThe halting condition is convergence of $z$ (relative diff. of 5%). \n- setting an upper limit of 4 self-iterations takes 26s on the test set (4k examples) and gets 97.8% accuracy on the p=0.5 distribution. \n- an upper limit of 16 self-iterations takes 35s and gets 99.8%. Note that iterations terminate after 6 steps on average due to the above halting condition. \n- the explicit model with 16 layers takes 37s and gets 1.5%. \nWe believe that this is a fair baseline.\n\n> this OOD task with increasing the number of $S_5$ tokens was extremely clever and a great addition to the literature.\n\nthanks!\n## catbAbI\nThe 1-layer implicit model trains faster than the 3-layer explicit model, but takes slightly more time during inference (see  __AqCH__).\n\nThe discrepancy between Figures 10a and 10c arises because Figure 10a averages accuracy across all story lengths per task, obscuring differences related to the distribution of story lengths that varies a lot across tasks and is also non-uniform per task.\n\n> how would a parameter matched explicit 6 layer do in comparison\n\nAccording to Fig. 1, a deep enough model will sufficiently track state, which should also hold for catbAbI. We view catbAbI as an intermediate sanity check between our synthetic task and the larger language models, and hence did not investigate this further. \n\n## Language Modeling\nWe will add the values reported to __AqCH__ to table 1 and table 6.\n# Questions for authors\nQ1: It is an empirical contribution of this work that passing on $h_t^*$ suffices for inference, which enables const memory language generation. Intuitively, DEQs are all about fixed points, and “path independence” has been observed in prior works.\n\nQ2: Self-iteration introduces additional FLOPs. Many practical problems might already be sufficiently addressed by the baseline models. Similar to test-time computation, self-iteration pays off only for certain problems.\n\nQ3: There are multiple runs that  overlap with the best run. Perhaps a box-plot would be a more appropriate choice. Left was trained with 32 iterations.\n\nQ4:\nabout 6 iterations (5-7 for different models) per token.\n\nQ5:\nIt seems that differentiating locally around the fixed point and not along the full trajectory provides a stronger bias towards learning sequential problems.\n\nQ6:\nPhantom gradients are similar to the truncated backpropagation method used by Geiping et al., but uses the update rule $z_{t+1} = (1 - \\lambda) z_t + \\lambda f(z_t, x)$. For $\\lambda=1$, the two algorithms match.\n\nQ7:\nWe found that continuing the 4x unrolling for the complete training is not enough to achieve competitive accuracy.\n\nQ8:\n'32' is indeed a typo. Should be 4x unrolling till step 5000.\n\n# Further comments\nAs requested by multiple reviewers, we will add a new section on limitations to the manuscript to discuss the moderate downstream task improvements in face of the larger wall-clock time.\n\n> Can you provide any theoretical perspective on why implicit models are better at extrapolation? \n\nPrior works indicate that implicit models have higher robustness to noise due to attractor dynamics. Viewing uninformative tokens as a source of noise (e.g. at longer sequence length) might explain the robustness.\n\n> provide evidence that the explicit models are being trained to the utmost on the synthetic state tracking task\n\nFig. 1 (top left) shows that the number of layers required by the explicit model to solve the S5 problem grows linearly with the sequence length showing that explicit models are limited by their depth.", "claims_and_evidence": "* The discussion around equations (6) and (7) implies that these equations converge to a fixed-point. However, is this correct? Not all iterative equations converge, and no proof of this claim is given.\n* > Theoretically, we show that implicit SSMs implement the non-linear state transitions of RNNs\n    * This claim is formalized in Theorem 1\n    * I am concerned about the shift from $h_t^*$ to $h_{t-1}^*$ in the proof, see my comments in the \"Theoretical Claims\" section\n* > Empirically, we find that only approximate fixed-point convergence suffices\n    * Figure3, Mid supports this claim\n* Lines 1080-1 (Figure 10 caption) \n> The implicit Mamba2 retains its performance as the story length increases, whereas the explicit Mamba2's performance declines\n    * This claim is contradicted by the performance of the explicit 3 layer in Figure 10c, which on the balance appears to have better performance as story length increases.\n* Lines 331-4 \n> The implicit Mamba2 models maintain their perplexity as sequence length increases, whereas the baseline Mamba2 models exhibit an increase in perplexity with longer sequences\n    * Yes, demonstrated in Figure 4\n* Effective Duality between Simultaneous Mode and Sequential Mode\n    * Yes, this claim is supported", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation are excellent.\n* The OOD task in Figure3 was a great addition to the literature\n* The use of the CatbAbI dataset was a good idea", "theoretical_claims": "_Theorem 1 and Proof in Appendix B_\n\n* This is a bit of a nit / more of a notational question, but in equation (14), wouldn't it be more correct to write $\\dfrac{d h_t^*}{d h_{t-1}^*}$ (as opposed to $\\dfrac{\\partial h_t^*}{\\partial h_{t-1}^*}$, as currently written)? Looking at equation (13), it seems to me like $\\dfrac{\\partial h_t^*}{\\partial h_{t-1}^*} = \\Lambda(z_t^*, x_t)$, while the full derivative contains the off diagonal correction terms.\n* **A more important point that needs to be discussed more precisely and clearly** is the choice of $h_{t-1}^*$ as opposed to $h_t^*$ as the argument of $\\varphi$. In particular, on line 785, it is stated without justification that $z_t^* = \\varphi(h_{t-1}^*, x_t, \\theta)$. However, in equation (7), $z_t^{(s)}$ is a function of $h_t^{(s)}$ and not of $h_{t-1}^{(s)}$. This distinction between $h_{t-1}^*$ as opposed to $h_t^*$ is important because if $\\varphi$ is actually a function of $h_t^*$, then the off-diagonal terms in equation (14) disappear and the Theorem is incorrect as stated. Therefore, I think it is very important for the authors to **add a lot more rigorous detail** about why $\\varphi$ takes $h_{t-1}^*$ as an argument instead of $h_t^*$, especially because this shift is different from the set up for equations (6) and (7)\n    * I also think it is very important that the authors provide a numerical check of equation (14) in their trained models. I.e, in their trained models, what actually are the Jacobians $\\dfrac{\\partial h_t^*}{\\partial h_{t-1}^*}$. Do they correspond within numerical tolerance to the RHS of equation (14)? Or not? Please include such a numerical check in your rebuttal.\n\n_Phantom Gradient (Section 2.3)_\nIs the minus sign in equation (4) correct?\nAs I understand things, we know that\n$$ G(\\Phi, x, \\theta) = 0.$$\n\nThus, taking derivatives wrt $\\theta$, it follows that\n$$ \\dfrac{\\partial G}{\\partial \\theta} + \\dfrac{\\partial G}{\\partial z} \\dfrac{\\partial \\Phi}{\\partial \\theta} = 0.$$\n\nNow, we know that $\\dfrac{\\partial G}{\\partial \\theta} = - \\dfrac{\\partial F}{\\partial \\theta}$, and so plugging in it follows that\n$$   \\dfrac{\\partial G}{\\partial z} \\dfrac{\\partial \\Phi}{\\partial \\theta} =  \\dfrac{\\partial F}{\\partial \\theta}.$$\n\nTherefore, it would seem that there is a sign error in equation 4.", "experimental_designs_or_analyses": "* Will the code be published? It is difficult to check the experiments without code.\n* How was $\\lambda$ set for the Phantom gradients (see equation 5)? This choice of hyperparameter does not seem to be discussed anywhere in the paper, even though there is a very good treatment of other experimental design choices. What happens if this hyperparameter $\\lambda$ is varied?\n* In Figure3 Right, I am concerned that the explicit models did the worst on train accuracy. Is there an explanation for this behavior? I would have thought that 16 layers of mamba would be enough to memorize to memorize a sequence length of length 256, i.e. get to at least 90% train accuracy. Is there any way to explain this phenomenon or provide evidence that the explicit models are being trained to the utmost on the synthetic state tracking task (Figure 3). \n* I would really like to see more reporting of wall clock time and memory usage in the experiments. I discuss wall clock time at length in \"Other Strengths and Weaknesses\" Section. As for memory, I was extremely impressed by the batch size of 1M tokens for the language modeling experiments. I want to know how much memory was required for training with a batch size of 1M tokens. Such an addition of max memory needed for training should be added to Table 6.\n* Broadly speaking on the Language Modeling experiments (Table 1, Table 5, etc), I wasn't sure if a proper ablation was done for depth. For example, what would happen if the explicit models were made deeper but less wide (to preserve parameter matching). Their depth could be scaled by the number of inference steps reported in Table 2 (are those inference  steps means over tokens in Table 2?). \n* The large scale language modeling tasks in Table 1 are good and show modest improvement of implicit over explicit models, but they do not blow away the explicit models. Are there any tradeoffs, i.e. downsides of using an implicit model, in terms of memory, compute, wallclock time, or some other metric? An explicit bolded paragraph on \"Limitations\" would be a nice contribution to contextualize the method and help practitioners.", "relation_to_broader_scientific_literature": "This paper builds on the implicit layers line of work to create an implicit model that can actually perform well on language tasks!", "essential_references_not_discussed": "The authors cite Lim et al '24, but note that their method incurs cubic cost in terms of state size. The authors may also wish to cite\n* Gonzalez et al '24, \"Towards Scalable and Stable Parallelization of Nonlinear RNNs,\" https://arxiv.org/abs/2407.19115\nwhich is an extension of Lim et al but uses quasi-Newton methods to avoid the cubic cost in state size.\n\nThe authors may also consider citing the following paper which started the deep SSM line of work\n\n* Gu et al '21, \"Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers,\" https://arxiv.org/abs/2110.13985\n\nIn particular, see Appendix C.\n\nIn this paper, Gu et al prove that many layers of an SSM can approximate a Picard iteration (a fixed point iteration). The methods in this proposed paper are effectively doing a Picard iteration as I understand it, so some comparison with the theory developed by Gu et al '21 may be useful to the academic community.\n\nWhile not required because of the ICML policy on concurrent work, the authors might considering citing and discussing the concurrent work\n* Geiping et al '25, \"Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach.\" https://arxiv.org/abs/2502.05171\nThey also uses self-iterations, but with attention layers instead of SSM layers. They also apply to language modeling. A robust discussion of similarities and differences in approach would be a great resource for the community.", "other_strengths_and_weaknesses": "On the balance, I think this is a great paper. However, I have various minor concerns scattered throughout this review that I think should be addressed. More importantly, I have a major concern about **wall-clock time** that I elaborate on in this section. I really need to see this question of wallclock time directly addressed before I can advocate for publication.\n\n_Synthetic State tracking task (Figure 3)_\n\nFor example, in all of Figure 3, the implicit models are granted unbounded test iterations (what actually is the halting condition? I do not think the halting condition is stated explicitly in the paper).\n\nWhat is the wall-clock time for inference of the implicit models at test time in Figure3, compared to the wall clock time of Mamba2 (explicit) for inference?\n\nAlso, I'm curious about the fairness of the Mamba2 baseline. In Figure3 right, the explicit and implicit models have matched train time depth, which is very good. But what happens if we match test time depth? I.e., report the average number of test time iterations used by implicit mamba in Figure3 right. And then give that many layers in depth for both training and test to explicit Mamba2 (go ahead and parameter match still). **I would really like to see this fair baseline for Mamba2 before concluding that explicit layers are limited on this task.**\n\nHowever, this OOD task with increasing the number of $S_5$ tokens was extremely clever and a great addition to the literature.\n\n_catbAbI task (Appendix D.2)_\n\nWhat is the wall clock time, both for training and for test, of the implicit and explicit mamba2 models (1,2, and 3 layers) on the catbAbI dataset. As demonstrated in Figure10a, a 3 layer explicit mamba2 has almost identical performance to a 1 layer implicit mamba2. However, if implicit mamba2 takes dramatically longer on wallclock time, it is difficult to recommend the use of 1 layer implicit (with a large number of fixed-point iterations) over a 3 layer explicit mama2.\n\nHow should I reconcile Figure 10a and Figure 10c however? On 10a it looks like explicit and implicit 3 layers almost always have similar performance; while on Figure10c it looks like implicit 3 layer is always better than explicit 3 layer. But how would a parameter matched explicit 6 layer do in comparison, both in test accuracy and on **training and test time wall clock time.**\n\n_Language Modeling task_\n\nIt would be helpful to add wallclock time (for both training and inference) to Table 6 (in addition to peak memory requirement for training and inference, see discussion in \"Experimental Designs or Analyses.\")\n\n_Extrapolation Advantage_\n\nA main selling point of this paper (especially on the synthetic state tracking experiment and on the length extrapolation aspect of language modeling) is that implicit models seem to be better on extrapolation tasks (proportion of hard tokens, length of sequence) at test time. **Can you provide any theoretical perspective on why implicit models are better at extrapolation?** Doing so would really strengthen the paper!", "other_comments_or_suggestions": "Style\n* The capitalization in this paper is a bit nonstandard at times, i.e. sometimes too generous with capitalization (eg \"Illusion of State\", \"Word Problem\", \"Implicit Function Theorem\"; also capitalization of a sentence fragment after a colon), but then other times doesn't capitalize \"theorem\" when it should. The authors may wish to review standard English capitalization style guides.\n\nTypos\n* In equation 2, I think it should be $h_t$ and not $h_{t-1}$\n* Line 434: the quotes around 'hard' are a bit ugly.\n* Line 662: \"it's\" should be \"its\"\n* Line 699\n     >  Monoid whose elements can be inverted have a particularly right structure\n    * Perhaps this sentence should read: \"Monoids where every element has an inverse are called \\emph{groups}.\"\n\nSmall Suggestions\n* Another nit, but in lines 199-202, the authors write\n> The Illusion of State reveals that SSMs cannot simulate arbitrary finite state machines\n    * wouldn't it be better style to not capitalize illusion of state?\n    * I think SSMs can in fact simulate arbitrary finite state machines, just not without depth growing in the sequence length. I think the sentence should instead read \"\\citet{Merrill24} shows that SSMs cannot simulate arbitrary finite state machines with constant depth.\"\n* Another nit: line 202-4\n> A hard state tracking problem in the sense that all state tracking problems can be reduced to it\n    * not the most elegantly phrased sentence\n    * strictly speaking, not _all_ state tracking problems can be reduced to $S_5$ (consider for example $\\mathbb{Z}_7$).", "questions_for_authors": "1. How is Figure 2B related to equation (8)? It would seem that every token in depth still needs to propagate across the sequence length in order to simulate equations (6) and (7), which converge to (8). Why is only $h_t^*$ needed to be passed on, instead of all $h_t^{(s)}$?\n2. Would you please elaborate on the following sentence in your conclusion?\n> While implicit models lift the limitations of state-of-the-art language models, self-iteration comes at a cost that only amortizes over the long tail of natural language\n3. In Figure 3 (Left and Mid), there is a large discrepancy between the best run and all runs. Would you please elaborate on what is causing this difference? Is there any way to see from train performance alone which run will yield the highest test accuracy? Also, how many self-iterations are used on average at test time in Figure 3 (all 3 panels)? And how many self-iterations are used during training for Figure 3 Left?\n4. Could more detail be provided about the following sentence on lines 251-5?\n> Interestingly, the number of test time self-iterations is quite similar for the models trained with different upper bounds on the training time self-iterations, hinting that the models learn similar algorithms\nHow many test time self-iterations is this?\n5. Why in Figure3 Right does unrolled have better training performance than phantom gradients, but the opposite relation is true for test?\n6. For Figure3 Right, an important missing baseline would seem to be a truncated backpropogation, i.e. only backpropogate through 8 steps (as is done in Geiping et al 2025, \"\"Scaling up Test-Time Compute with Latent Reasoning\", https://arxiv.org/abs/2502.05171). This approach may be somewhat intermediate in the train/test tradeoff between phantom gradients and full unrolled backpropogation, and would also be constant memory. I think you are already doing this on the CatbAbi task, see lines 927-8.\n7. In the catbAbi experiments (D.2), did you ever try an ablation where you unrolled the entire way through, and never switch to self-iteration fixed-point search? In Figure 8, there is a sharp dip in validation accuracy at 5000 steps when you switch schedules;  but otherwise, it doesn't look like the validation trajectory was changed very much as a result of the change in training procedure. Moreover, I thought that the point of Figure 3 Mid was that fixed number of self-iterations at train time was an acceptable way to train the model.\n8. I must not be understanding something about Figure 9 (maybe the y-axis should be renamed from \"Validation steps\" to \"number of self-iterations\"). Still, I thought that based on lines 927-8 that 32 self-iterations would be used on the first 5000 gradients steps, but it looks more like 4. How was this number chosen? And was does \"trained for 5000 steps in unrolling mode, utilizing 32 steps with normal gradient checkpointing\" mean?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors propose implicit language models, RNNs defined implicitly via fixed point iteration. They show theoretically that this architectures can simulate non-linear RNNs, therefore from a representational point of view they can represent functions SSM cannot. Particularly in the rebuttal the theoretical statement is made a bit more sharp. The authors also show empirical evidence for the performance of their proposed model, up to size 1.3B.  

I think this work is a step forward towards better understanding recurrent models and also making them a valid contender to transformers. I believe the authors have provided a very detailed rebuttal, answering most of the question raised, and have improved their paper in the process (including improving the theoretical result by making it more sharp). The empirical evaluation is also thorough and provides a clear signal that this direction is worth pursuing. Overall I think this is a strong contribution to the field. 
I do urge the authors to include all provided details in the rebuttal into the final paper, since I think the rebuttal made the paper stronger.
Decision: Accept (spotlight poster)