Today we are posting the first part of our ongoing research project on Deep Equilibrium Models (DEQs). This first edition will focus on explaining the architecture and sharing potential angles to push the frontier of DEQ within a business context. The next part will show preliminary results. A critical note right off the bat: this is by no means a fully fledged academic proposal; rather, we choose to adopt our own Rapid Experimentation Framework (REF). Using REF, we aim to use a combination of experience, brute-force mechanisms and AI to build workable Proof-of-Concepts.

Introduction to Deep Equilibrium Models

Before we dive into the specifics, we will first briefly explain what Deep Equilibrium Models are and where they can be useful.

Most traditional Large Language Models (LLMs) are so-called explicit depth models. They work by stacking "layers", with each layer implementing a predefined (forward) computation. The major LLM breakthrough came by using residual transformers in these layers, allowing models to use self-attention and Multi-Layer Perceptrons (MLPs) to predict the next word. Due to this design, each layer "encodes" a different aspect of language comprehension. For example, the first layer builds broad, shallow relationships and low-level feature extraction, whilst layer 40 builds complex relations and factual retrieval and reasoning. The simplified formula looks as follows:

Implicit depth models (Deep Equilibrium / Recurrent LLMs) take a completely different approach. Instead of using multiple layers, a single transformer block is used. The model iterates repeatedly through this block until an equilibrium state is reached, denoted as z^*. In this framework, language comprehension does not happen across multiple layers; rather it "evolves" over time within the same layer. The first iteration forms simple syntax and word definitions, which are fed into the next iteration, further expanding the representation. Moving through the iterations, the changes made to the representation become smaller, eventually balancing out the residual between input and output. Once the state is stable, computation stops and the output is passed to predict the next word.

Each iteration follows approximately the same math as the explicit models:

Take note how the original input embedding x is injected. This is done to ensure the model does not converge to a standard state from which x is absent. Passing the hidden state repeatedly through the same non-linear transformations dampens the representation over time. This leads to information decay, essentially "losing" details of the original token if not injected each iteration.

As k approaches infinity, the representation stops changing (z_k+1 ≈ z_k = z^*). The system collapses into a single implicit fixed-point equation:

where f_θ​ is the combined transformer block operating on the equilibrium representation z^*.

Explicit versus implicit

Now that we understand the difference between explicit and implicit models, let us write down a few advantages and disadvantages of each framework.

• Compute per Token‍
  
  • Explicit models pass through each layer, executing the forward computation regardless of whether that token requires "more depth". This means compute per token is the same for a simple question or complex coding logic.
  • Implicit models only use as much compute as is needed to obtain the equilibrium state, meaning that simple queries, in theory, require fewer iterations than complex tokens.
• Memory Footprint
  
  • Explicit models build contextual depth across multiple layers, with each layer encoding a different representation change to the hidden state. Thus, in order to achieve depth, you would need more layers and hence more parameters. In turn, this scales the memory footprint, needing to load more parameters.
  • Implicit models use only one layer and build depth via iterations, meaning, in theory, fewer parameters are needed to achieve the same output. The trade-off here is speed, as iterating many times over significantly decreases token generation speed.
• Hardware Optimization
  
  • Explicit models feature a deterministic path across fixed layers, enabling compilers (PyTorch torch.compile, FlashAttention) to pre-allocate static VRAM layouts, fuse CUDA kernels, and achieve near-peak Tensor Core utilization through homogeneous, highly parallel matrix operations across entire batches.
  • Implicit models require non-linear numerical root-finding algorithms (Anderson, Broyden), with runtime convergence checks. This introduces dynamic loop bounds, breaking static CUDA execution graphs and causing a batch-level straggler problem: sequences in a batch that converge in 8 iterations sit idle waiting for harder sequences to reach tolerance or hit the maximum number of iterations.
• Stability
  
  • Explicit models prove to be stable when scaled up. During pre-training, stability is kept by utilising normalization techniques (RMSNorm, LayerNorm) and residual connections, even when building models with hundreds of layers and trillions of tokens.
  • Implicit models can, due to the use of numerical root-finders, drift into unstable states during training on large-parameter models. At this time, no proven method has been found to consistently maintain a stable, large DEQ model.
• Knowledge / Fact storage
  
  • Explicit models scale factual capacity alongside depth. Because every added layer introduces new, unique weight matrices, increasing physical depth directly expands the model's "memory". This effect allows larger LLMs to "memorize" vast amounts of knowledge, facts and specialized data.
  • Implicit models detach reasoning depth from factual storage. While running a single block through 20 or 50 iterations grants deep reasoning capabilities, it adds zero new parameters. Consequently, an implicit model's capacity to know and store world facts remains hard-capped by the parameter size of its single-layer block, regardless of how many iterations it executes.

In light of this comparison, we argue that DEQs do not function as a 1:1 drop-in replacement for standard LLMs. Rather, their real value lies in decoupling reasoning depth from parameter count: they allow models to "reason" within a smaller VRAM footprint, trading runtime compute and latency for hardware storage efficiency. This makes DEQs an intriguing option where parameter memory is strictly capped but reasoning depth is required, such as on-device edge computing, IoT microcontrollers, or embedded automotive systems. Furthermore, equilibrium models already showed prominence in the original paper by Bai, Kolter and Koltun, called "Deep Equilibrium Models". When referencing or building upon DEQs, please cite their original paper.

Our work

Our work focusses on building a PoC DEQ for a simple task in a limited VRAM environment. The original paper managed equivalent performance against a standard small language model; we opted to pursue the same goal whilst also adding JSON-schema parsing as a secondary focus.

The project itself is built via Cursor, using a custom stack in Python to build the solver and model. We will release the repo on GitHub later, once a stable version is attained.

The hardware specs are as follows:
GPU: RTX 3070 with 8GB of VRAM
CPU: Intel Ultra 7 265k
RAM: 64GB of DDR5

The model itself is called Mudeep, and for the first stage of building we've determined 3 prototypes.

• Mudeep Ripple: First iteration, trained on TinyStories, testing the hardware for training purposes
• Mudeep Undertoe: Focus on JSON schema parsing, trained with a combination of 3 datasets, customer-support functionality
• Mudeep Riptide: Not trained yet, further expansion of Undertoe, using more parameters

Mudeep Ripple and Undertoe share the same architecture, but have different training conditions.

Architecture

Mudeep uses a relatively simple architecture, with, as of yet, no major advancements in efficiency. We've decided to use a custom library for the solver compared to existing ones like TorchDEQ, so that we can experiment on a fundamental level in later iterations. The transformer block is built using PyTorch.

Text to Vector embedding

Mudeep uses the GPT-2 BPE tokenizer (via tiktoken) for embedding, with vocab size V = 50257 against a dimension of d = 256. A sequence length of T (Ripple T = 256, Undertoe T = 384) is embedded in a T x 256 matrix, with each token getting its own 256-dim vector.

1import tiktoken
2import torch
3import torch.nn as nn
4
5enc = tiktoken.get_encoding("gpt2")
6vocab_size = enc.n_vocab  # 50257
7d_model = 256
8
9self.tok_emb = nn.Embedding(vocab_size, d_model)
10self.drop = nn.Dropout(0.1)
11self.head = nn.Linear(d_model, vocab_size, bias=False)
12self.head.weight = self.tok_emb.weight

To give a clear picture of the output after embedding, we are looking at:

Transformer block

The input is first normalized, via RMSNorm (Zhang & Sennrich). This is done to limit large magnitude drifts that would occur as the network iterates. By scaling the token vectors down to an approximate root-mean-square of 1 and then multiplying by a learned dimension gain, you get a more stable representation.

1import torch
2import torch.nn as nn
3from torch import Tensor
4
5class RMSNorm(nn.Module):
6    def __init__(self, dim: int, eps: float = 1e-6):
7        super().__init__()
8        self.eps = eps
9        self.weight = nn.Parameter(torch.ones(dim))  # per-dim gain γ
10
11    def forward(self, x: Tensor) -> Tensor:
12        x_f = x.float()
13        rms = x_f.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
14        return (x_f * rms).type_as(x) * self.weight
15
16# in TransformerBlock.__init__ (d_model=256)
17self.ln_h = RMSNorm(d_model)    # after inject
18self.ln1 = RMSNorm(d_model)     # before attention
19self.ln2 = RMSNorm(d_model)     # before MLP
20self.ln_out = RMSNorm(d_model)  # before damp
21
22# in __init__
23self.ln_f = RMSNorm(d_model)    # before LM head

Note how we use RMSNorm between every step within the transformer block, with each step getting its own learned gain parameter gamma. The theory behind this decision is to limit the chance of activation explosions during training. In this endeavour, we introduced a redundancy within Ripple and Undertoe, by normalising after inject and before attention. Whilst we haven't noticed degrading performance, further iterations will remove this step with a single before-attention normalization.

Our attention block is slightly modified to include custom RoPE. Mudeep uses 4 attention heads, meaning d_h = 64. The input, which looks like this after ln1:

is split into Q, K and V using fused linear projection, creating a 768-dim vector. The exact math here is unchanged, we simply use the trained weight vectors to compute Q, K and V. Further experimentation will focus on determining compatibility with obtaining an equilibrium state on QKV. We then reshape the total vector, and unbind each individual into Q, K and V.

Our code does this as follows:

1def forward(self, x: Tensor) -> Tensor:
2        b, t, c = x.shape
3        qkv = self.qkv(x).reshape(b, t, 3, self.n_heads, self.d_head)
4        q, k, v = qkv.unbind(dim=2)
5        q = q.transpose(1, 2)  # [B, H, T, Dh]
6        k = k.transpose(1, 2)
7        v = v.transpose(1, 2)

Q and K are rotated by position so relative distance is encoded. We use the standard split-half formulation (popularized by LLaMA and GPT-NeoX) rather than interleaved coordinate pairing. The 64-dimensional head vector is divided into two 32-dim halves, allowing the frequency tensor to rotate each coordinate pair across the 2D plane via standard rotation geometry.

1    def _rotate_half(x: Tensor) -> Tensor:
2        d = x.shape[-1]
3        x1, x2 = x[..., : d // 2], x[..., d // 2 :]
4        return torch.cat((-x2, x1), dim=-1)
5    
6    
7    def _rope_cos_sin(seq_len: int, dim: int, device: torch.device, dtype: torch.dtype, base: float = 10000.0):
8        half = dim // 2
9        inv_freq = 1.0 / (base ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
10        t = torch.arange(seq_len, device=device, dtype=torch.float32)
11        freqs = torch.outer(t, inv_freq)  # [T, half]
12        emb = torch.cat((freqs, freqs), dim=-1)  # [T, dim]
13        cos = emb.cos()[None, None, :, :].to(dtype=dtype)
14        sin = emb.sin()[None, None, :, :].to(dtype=dtype)
15        return cos, sin
16		
17      	#Within .forward
18        cos, sin = _rope_cos_sin(t, self.d_head, x.device, q.dtype)
19                      q = q * cos + _rotate_half(q) * sin
20                      k = k * cos + _rotate_half(k) * sin

The attention module follows the softmax formula:

QKT measures the relevance of a given token, in other words: whether the Q and K vectors point in similar directions. The root of Dh is used to scale the computed dot product down when the head dimension increases. This is needed to make sure that scores do not blow up. Furthermore, we apply causal masking to limit the model from "looking into the future". Lastly, we apply normalization (the softmax function) to get a probability distribution back, seen as attention weights A. Afterwards, the computed attention weights multiply the output by the value vector V, computing the Attn(Q, K, V) for a given token. Finally, the attention heads are merged back again, computing the output.

Our code uses PyTorch library to compute the self-attention:

1        cos, sin = _rope_cos_sin(t, self.d_head, x.device, q.dtype)
2        q = q * cos + _rotate_half(q) * sin
3        k = k * cos + _rotate_half(k) * sin
4
5        att = F.scaled_dot_product_attention(
6            q, k, v,
7            attn_mask=None,
8            dropout_p=self.attn_drop.p if self.training else 0.0,
9            is_causal=True,
10        )
11        y = att.transpose(1, 2).contiguous().reshape(b, t, c)
12        return self.resid_drop(self.proj(y))

Now, after self attention is computed, we add the result y back to the original stream and multiply by a coefficient α_attn.

α_attn is computed via the following code and trained. The same is done for α_mlp. We train these alphas to better work with the DEQ model. If we were to use standard static coefficients, loop-safety must be entirely learned within the weights, leading to larger residuals and unstable solvers.

1# Before training
2self._attn_logit = nn.Parameter(torch.tensor(-2.2))
3self._scale_max = 0.5
4
5# During training
6a_attn = self._scale_max * torch.sigmoid(self._attn_logit)
7# 0.5 * sigmoid(-2.2) ≈ 0.5 * 0.09975 ≈ 0.0499 ≈ 0.05

The multilayer perceptron block also remains unchanged. We used SwiGLU to obtain the following operation:

The SwiGLU formula itself looks as follows:

The code we use to compute this is based on Meta LLaMA, and uses torch. The SwiGLU up and down projections follow roughly 2/3d_ff, to keep parameter count near that of a classic MLP. This is a common choice when using LLaMA.

1import torch.nn as nn
2import torch.nn.functional as F
3from torch import Tensor
4
5class SwiGLU(nn.Module):
6    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
7        super().__init__()
8        h = int(2 * d_ff / 3)
9        h = max(64, ((h + 63) // 64) * 64)
10        self.w1 = nn.Linear(d_model, h, bias=False)
11        self.w3 = nn.Linear(d_model, h, bias=False)
12        self.w2 = nn.Linear(h, d_model, bias=False)
13        self.drop = nn.Dropout(dropout)
14
15    def forward(self, x: Tensor) -> Tensor:
16        return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x)))
17
18self.mlp = SwiGLU(d_model, d_ff, dropout)
19g = g + a_mlp * self.mlp(self.ln2(g))

Firstly, we expand W1 and W3 from 256 dimensions to 704. W1 passes through the SiLU activation function. SiLU acts as a filter, blocking low values and letting high values through. This builds a tensor, that can then be used in tandem with W3, acting as a gate. If a gate value is ≈ 0, it completely shuts off that feature. If a gate value is active, that feature passes through. This gated vector is then multiplied by W2, scaling the 704 dimension vector back down to 256. We then run the same residual calculation with α_mlp, and compute the output vector again.

Now in vanilla transformers, you would be done after this step, with an output g. Due to the fact that we run multiple iterations of this exact transformer block, we apply a damp function. This is a stabilization mechanism, that makes sure that the solver function stays contractive.

Our code, similar to the attention and mlp alpha, trains the damp logit.

1self._damp_logit = nn.Parameter(torch.tensor(0.0))
2
3def _damp(self):
4    return 0.15 + 0.70 * torch.sigmoid(self._damp_logit)
5
6return (1.0 - damp) * z + damp * g

This brings the full architecture to this:

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from torch import Tensor
5
6
7class RMSNorm(nn.Module):
8    def __init__(self, dim: int, eps: float = 1e-6):
9        super().__init__()
10        self.eps = eps
11        self.weight = nn.Parameter(torch.ones(dim))
12
13    def forward(self, x: Tensor) -> Tensor:
14        x_f = x.float()
15        rms = x_f.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
16        return (x_f * rms).type_as(x) * self.weight
17
18
19def _rotate_half(x: Tensor) -> Tensor:
20    d = x.shape[-1]
21    x1, x2 = x[..., : d // 2], x[..., d // 2 :]
22    return torch.cat((-x2, x1), dim=-1)
23
24
25def _rope_cos_sin(seq_len: int, dim: int, device, dtype, base: float = 10000.0):
26    half = dim // 2
27    inv_freq = 1.0 / (base ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
28    t = torch.arange(seq_len, device=device, dtype=torch.float32)
29    freqs = torch.outer(t, inv_freq)
30    emb = torch.cat((freqs, freqs), dim=-1)
31    cos = emb.cos()[None, None, :, :].to(dtype=dtype)
32    sin = emb.sin()[None, None, :, :].to(dtype=dtype)
33    return cos, sin
34
35
36class CausalSelfAttention(nn.Module):
37    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
38        super().__init__()
39        self.n_heads = n_heads
40        self.d_head = d_model // n_heads
41        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
42        self.proj = nn.Linear(d_model, d_model, bias=False)
43        self.attn_drop = nn.Dropout(dropout)
44        self.resid_drop = nn.Dropout(dropout)
45
46    def forward(self, x: Tensor) -> Tensor:
47        b, t, c = x.shape
48        qkv = self.qkv(x).reshape(b, t, 3, self.n_heads, self.d_head)
49        q, k, v = qkv.unbind(dim=2)
50        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
51
52        cos, sin = _rope_cos_sin(t, self.d_head, x.device, q.dtype)
53        q = q * cos + _rotate_half(q) * sin
54        k = k * cos + _rotate_half(k) * sin
55
56        att = F.scaled_dot_product_attention(
57            q, k, v,
58            attn_mask=None,
59            dropout_p=self.attn_drop.p if self.training else 0.0,
60            is_causal=True,
61        )
62        y = att.transpose(1, 2).contiguous().reshape(b, t, c)
63        return self.resid_drop(self.proj(y))
64
65
66class SwiGLU(nn.Module):
67    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
68        super().__init__()
69        h = int(2 * d_ff / 3)
70        h = max(64, ((h + 63) // 64) * 64)
71        self.w1 = nn.Linear(d_model, h, bias=False)
72        self.w3 = nn.Linear(d_model, h, bias=False)
73        self.w2 = nn.Linear(h, d_model, bias=False)
74        self.drop = nn.Dropout(dropout)
75
76    def forward(self, x: Tensor) -> Tensor:
77        return self.drop(self.w2(F.silu(self.w1(x)) * self.w3(x)))
78
79
80class TransformerBlock(nn.Module):
81    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1):
82        super().__init__()
83        self.ln_h = RMSNorm(d_model)
84        self.ln1 = RMSNorm(d_model)
85        self.attn = CausalSelfAttention(d_model, n_heads, dropout)
86        self.ln2 = RMSNorm(d_model)
87        self.mlp = SwiGLU(d_model, d_ff, dropout)
88        self.ln_out = RMSNorm(d_model)
89        self.inj = nn.Linear(d_model, d_model, bias=False)
90        self._attn_logit = nn.Parameter(torch.tensor(-2.2))
91        self._mlp_logit = nn.Parameter(torch.tensor(-2.2))
92        self._damp_logit = nn.Parameter(torch.tensor(0.0))
93        self._scale_max = 0.5
94
95    def _pos_scale(self, logit: Tensor) -> Tensor:
96        return self._scale_max * torch.sigmoid(logit)
97
98    def _damp(self) -> Tensor:
99        return 0.15 + 0.70 * torch.sigmoid(self._damp_logit)
100
101    def forward(self, z: Tensor, x: Tensor) -> Tensor:
102        a_attn = self._pos_scale(self._attn_logit)
103        a_mlp = self._pos_scale(self._mlp_logit)
104        damp = self._damp()
105
106        h = self.ln_h(z + self.inj(x))
107        g = h + a_attn * self.attn(self.ln1(h))
108        g = g + a_mlp * self.mlp(self.ln2(g))
109        g = self.ln_out(g)
110        return (1.0 - damp) * z + damp * g

Take note that we haven't included the solver yet, and that this is a fairly standard write up of a DEQ model. In later editions of this project, we will share more in depth advancements.

Solver

In order to reach the equilibrium state when the model gets an input, we use the Anderson solver. Now, there are many more solvers available, but for the purposes of getting a model working, the Anderson solver proved the most stable.

The solver first checks the residual of the input and output for a given iteration. If the residual is 0 (or an equivalent extremely small number), the system has converged to equilibrium and halts. This essentially means that the changes the transformer block makes to the representation, are very minimal and hence iterating further doesn't yield any gain.
Instead of keeping all residuals in memory, the algorithm uses a sliding window, in our case m = 5. The exact storage of these states is maintained in F and X history rings, holding the current z and the f(z) respectively. Note on the history buffer: We update columns via ring-buffer indexing (k % m) to avoid memory allocations and in-place tensor shifts during training, trading exact chronological ordering of the subspace for GPU throughput.
With this information, the solver attempts to answer the question: What linear combination of my last 5 states would make the residual as close to 0 as possible?

This is done by computing a vector of scalar weights, that solves a constrained linear least-squares problem. With m being relatively small and the hidden dimension d comparatively large, we can collapse the problem into a tiny matrix equation.

The previous function output is then multiplied with the computed scalar weight vector, creating the adjusted input.

1import torch
2from torch import Tensor
3from typing import Callable
4
5
6@torch.no_grad()
7def anderson_solve(
8    f: Callable[[Tensor], Tensor],
9    z0: Tensor,
10    *,
11    m: int = 5,
12    max_iter: int = 24,
13    tol: float = 1e-2,
14    lam: float = 1e-4,
15) -> Tensor:
16    """Solve z = f(z). Returns f(z*) ≈ z*."""
17    bsz = z0.shape[0]
18    dim = z0[0].numel()
19    device, dtype = z0.device, z0.dtype
20
21    X = torch.zeros(bsz, m, dim, device=device, dtype=dtype)
22    F = torch.zeros(bsz, m, dim, device=device, dtype=dtype)
23
24    def rel_res(z, fz):
25        num = (fz - z).flatten(1).norm(dim=1)
26        den = fz.flatten(1).norm(dim=1).clamp_min(1e-6)
27        return num / den
28
29    # probe 0
30    x = z0.reshape(bsz, -1)
31    fx = f(z0).reshape(bsz, -1)
32    X[:, 0], F[:, 0] = x, fx
33    if (rel_res(x.view_as(z0), fx.view_as(z0)) < tol).all():
34        return fx.view_as(z0)
35
36    # probe 1
37    x = fx
38    fx = f(x.view_as(z0)).reshape(bsz, -1)
39    X[:, 1], F[:, 1] = x, fx
40    if (rel_res(x.view_as(z0), fx.view_as(z0)) < tol).all():
41        return fx.view_as(z0)
42
43    H = torch.zeros(bsz, m + 1, m + 1, device=device, dtype=dtype)
44    H[:, 0, 1:] = 1.0
45    H[:, 1:, 0] = 1.0
46    y = torch.zeros(bsz, m + 1, device=device, dtype=dtype)
47    y[:, 0] = 1.0
48
49    for k in range(2, max_iter):
50        n = min(k, m)
51        G = F[:, :n] - X[:, :n]
52        H[:, 1 : n + 1, 1 : n + 1] = (
53            torch.bmm(G, G.transpose(1, 2))
54            + lam * torch.eye(n, device=device, dtype=dtype)
55        )
56
57        alpha = torch.linalg.solve(H[:, : n + 1, : n + 1], y[:, : n + 1])[:, 1 : n + 1]
58        x = torch.bmm(alpha.unsqueeze(1), F[:, :n]).squeeze(1)
59        fx = f(x.view_as(z0)).reshape(bsz, -1)
60
61        idx = k % m
62        X[:, idx], F[:, idx] = x, fx
63
64        if (rel_res(x.view_as(z0), fx.view_as(z0)) < tol).all():
65            return fx.view_as(z0)
66
67    return fx.view_as(z0)

Probe 0 and probe 1 serve as the warm-up steps (K₀ and K₁) before the history ring contains enough states to construct the least-squares system. We cap total iterations at 24 and set tol = 1e-2 (1% relative residual). While traditional numerical solvers aim for tighter tolerances (10^-4 to 10^-6), setting tol = 1e-2 prevents the model from spending excessive compute on diminishing returns, keeping token generation throughput viable.
Notice the check if (rel_res(...)).all(): return fx. Because PyTorch computes the batch synchronously, every sample in the batch must satisfy the tolerance threshold before execution halts. This is the exact code path responsible for the straggler latency noted earlier.

Forward

In order to help the Anderson solver stabilize, we implemented a warm start mechanism. This is simply a small network that guesses the z^* from the embeddings before the solver runs.

1class WarmStart(nn.Module):
2    """Shallow predictor g_phi(x) ≈ z* so Anderson starts near the attractor."""
3    def __init__(self, d_model: int):
4        super().__init__()
5        self.norm = RMSNorm(d_model)
6        self.fc1 = nn.Linear(d_model, d_model, bias=False)
7        self.fc2 = nn.Linear(d_model, d_model, bias=False)
8        nn.init.zeros_(self.fc2.weight)
9
10    def forward(self, x: Tensor) -> Tensor:
11        h = F.gelu(self.fc1(self.norm(x)))
12        return self.fc2(h)

Because z0 is detached prior to the solver, WarmStart does not receive gradients through the main language model loss. In our training loop (detailed in Part 2), WarmStart is supervised via an auxiliary loss, teaching it to track the attractor directly without unrolling the solver. Notice also that self.fc2.weight is initialized to exact zeros (nn.init.zeros_). This ensures that on step 0 of training, WarmStart outputs z0 = 0, reverting cleanly to the classical, provably stable starting point of DEQs, before gradually shifting toward warm attractor predictions as training progresses.

The full feed-forward looks as follows, labeled as MudeepLM:

1class MudeepLM(nn.Module):
2    def __init__(
3        self,
4        vocab_size: int = 50257,
5        d_model: int = 256,
6        n_heads: int = 4,
7        d_ff: int = 1024,
8        dropout: float = 0.1,
9    ):
10        super().__init__()
11        self.tok_emb = nn.Embedding(vocab_size, d_model)
12        self.drop = nn.Dropout(dropout)
13        self.block = TransformerBlock(d_model, n_heads, d_ff, dropout)
14        self.warm = WarmStart(d_model)
15        self.ln_f = RMSNorm(d_model)
16        self.head = nn.Linear(d_model, vocab_size, bias=False)
17        self.head.weight = self.tok_emb.weight
18
19    def forward(self, idx: Tensor, targets: Tensor | None = None):
20        x = self.drop(self.tok_emb(idx))
21        z0 = self.warm(x.float()).detach()
22
23        def f(z: Tensor) -> Tensor:
24            return self.block(z, x.float())
25
26        with torch.no_grad():
27            z_star = anderson_solve(f, z0)
28
29        z_star = z_star.detach()
30        z = f(z_star)
31        z = z + (z_star - z).detach()
32
33        logits = self.head(self.ln_f(z))
34
35        loss = None
36        if targets is not None:
37            loss = F.cross_entropy(
38                logits.view(-1, logits.size(-1)),
39                targets.view(-1),
40                ignore_index=-100,
41            )
42        return logits, loss

A critical technical detail is the forward pass gradient formulation:

1with torch.no_grad():
2    z_star = anderson_solve(f, z0)
3z_star = z_star.detach()
4z = f(z_star)
5z = z + (z_star - z).detach()

You might notice the line z = z + (z_star - z).detach(). Because the Anderson solver runs inside torch.no_grad(), we use this straight-through formulation to pass the exact equilibrium state z^* downstream in the forward pass, while routing gradients through a single unrolled application of f(z^*) during backpropagation.

Classical DEQs rely on the Implicit Function Theorem (IFT) to solve a secondary root-finding problem in the backward pass. Here, we deliberately bypass that backward solve in favor of a 1-step gradient approximation called Phantom Gradients, which stabilizes training on consumer hardware. We will dive deep into why this works, how it impacts loss convergence, and the exact stability trade-offs in Part 2.

This is the general setup of Mudeep, specifically ripple and undertoe. The premise here remains fairly simple, and in this stage of the project we are focused on getting the DEQ models into a stable state within the hardware constraints. In the next part, we will show the preliminary results and explain where we see improvements.

Early thoughts and angles to explore

Many of the LLM foundations and techniques are tailored to work with standard layer models. Within this first part, we reused much of the existing math within the equilibrium framework and as such, the results will not blow anyone away. But we are not deterred, as the recurrent depth mechanism allows for practical applications that a layer LLM simply can't provide (yet). However, we specifically run into 2 major challenges when building DEQs:

• Speed: Early results and testing show that token speed degrades heavily when scaling, to a point where the reduction in memory footprint does not outweigh the slowness. The preliminary solution is to achieve better solver rates or have dedicated hardware available. Both are not really practical as of yet, but the mathematical angle seems more promising.
• Knowledge depth: Scaling makes the model quite unstable, and hence adding more factual depth remains hard. Angles to consider here are scaling by means of Retrieval Augmented Generation (RAG) or fundamentally changing the MLP mechanism. However, most setups here fail due to Jacobian instability.

We hope to inspire more engineers / AI enthusiasts to try DEQs out and share results. Again, we want to be clear that this is a practical project and that it would not survive academic rigor. Rather, we want to share what we work on and hope to gather a community here. Until part 2!