A Scalable Recurrent Architecture - Rewritten for Easier Understanding

After

3.1. Macroscopic Design

The model suggests 3 new functional groups: prelude \(P\), the core recurrent block \(R\), and coda \(C\). While keeping the stucture primarily around decoder-only transformer blocks (Vaswani et al., 2017; Radford et al., 2019). The core block is set between the prelude and coda blocks. And by looping the core we can put an indefinite amount of "verses".

The idea behind the 3 groups.

• \(P\), the prelude, embeds the input data into a latent space using multiple transformer layers.
• \(R\), the core recurrent block, is the central unit. It's responsible for recurrent computation modifying states \(\mathbf{s} \in \mathbb{R}^{n \times h}\).
• \(C\), the coda, has 2 main responsibilities. It un-embeds from latent space using several layers and contains the prediction head of the model.

Motivation for this Design.

The minimal setup required for learn stable iterative operators includes the following.

• An initial starting state (a random starting point, \(s_0\)).
• A previous state to build upon (\(s_{i-1}\)).
• The input data (\(e\)) injected at every single step.

A good example is gradient descent of a function \(E(x, y)\), where \(x\) may be the variable of interest and \(y\) the data. Gradient descent on this function starts from an initial random state, here \(x_0\), and repeatedly applies a simple operation (the gradient of the function it optimizes), that depends on the previous state \(x_k\) and data \(y\). Note that we need to use \(y\) in every step to actually optimize our function.

The formal architecture.

Given a number of recurrent iterations \(r\), and a sequence of input tokens \(\mathbf{x} \in V^{n}\) these groups are used in the following way to produce output probabilities \(\mathbf{p} \in \mathbb{R}^{n \times |V|}\). As shown in Figure 2.

• The prelude produces the embedded input data \(\mathbf{e} = P(\mathbf{x})\).
• The initial random state \(s_0\) is drawn from some standard deviation \(\sigma\) distributed normally. \(\mathbf{s}_0 \sim \mathcal{N}(0, \sigma^2 \mathbf{I}_{n \cdot h})\)
• The core recurrent block repeatedly produces a new state \(s_i\) while still using the original embedded input data \(\mathbf{e}\). \(\mathbf{s}_i = R(\mathbf{e}, \mathbf{s}_{i-1}) \quad \text{for } i \in \{1, \dots, r\}\).
• The coda produces the final output probabilities for predicting the next token from the last state. \(\mathbf{p} = C(\mathbf{s}_r)\).

Notice we repeatedly inject the data \(\mathbf{e}\) in our set-up in every step of the recurrence. Similarly to what's mentioned in the motivation. If \(\mathbf{e}\) was provided only at the start, e.g. via \(\mathbf{s}_0 = \mathbf{e}\), then the iterative process would not be stable¹, as its solution would depend only on its boundary conditions.

This architecture is based on deep thinking literature, where it is shown that injecting the latent inputs \(\mathbf{e}\) in every step (Bansal et al., 2022) and initializing the latent vector with a random state stabilizes the recurrence and promotes convergence to a steady state independent of initialization, i.e. path independence (Anil et al., 2022).

The inspiration for the architecture.

The structure of using several layers to embed input tokens into a hidden latent space is based on empirical results analyzing standard fixed-depth transformers (Skean et al., 2024; Sun et al., 2024; Kaplan et al., 2024). This body of research shows that the initial and the end layers of LLMs are noticeably different, whereas middle layers are interchangeable and permutable. For example, Kaplan et al. (2024) show that within a few layers standard models already embed sub-word tokens into single concepts in latent space, on which the model then operates.