The Context-Ready Transformer

We first describe the architecture during sequential left-to-right generation—the setting where context-ready inference is exact. Let $T$ denote the sequence length, $C$ the embedding dimension, and $V$ the vocabulary size. The core component is a $D$-block unit: $D$ transformer layers (each consisting of causal self-attention and a feed-forward network with residual connections), described in detail below. Processing token $t$ requires the outputs $z_{0},\ldots,z_{t-1}\in\mathbb{R}^{C}$ of this block unit from all preceding tokens. Let $e_{t}\in\mathbb{R}^{C}$ denote the token embedding at position $t$, and define $z_{-1}=\mathbf{0}$.

Correction. A dedicated correction FFN generates the correction for position $t$ from the cached block output $z_{t-1}$ and the current token embedding $e_{t}$:

The correction FFN is a feed-forward network (Linear($C\to 4C$) $\to$ GELU $\to$ Linear($4C\to C$)) with its own weights, separate from the block’s FFN. The correction is token-aware: it depends on both $z_{t-1}$ (context of tokens $0,\ldots,t{-}1$) and $e_{t}$ (the current token embedding). Since both inputs are available when token $t$ arrives, the correction is causal—it depends on no future tokens.

Contextualization. The new token enters the block with the correction added to its raw embedding:

Block processing. A $D$-block unit applies $D$ transformer blocks with separate weights and standard residual connections:

Each $\texttt{Attn}_{i}$ is causal self-attention with Rotary Position Embeddings (RoPE) (Su et al., 2024). Each $\texttt{FFN}_{i}$: Linear($C\to 4C$) $\to$ GELU $\to$ Linear($4C\to C$). The parameter $D$ controls inference depth.

Prediction.

After prediction, $z_{t}$ is cached and the KV caches $\mathcal{K}_{i},\mathcal{V}_{i}$ are updated for future tokens.

Sequential inference is exact but inherently serial. For training, we unroll the correction process $K$ times, processing all $T$ positions in parallel at each step. Gradients flow through $K$ unrolling steps rather than $T$ time steps, making the architecture trainable like a transformer despite being recurrent at inference.

Given token embeddings $e=(e_{1},\ldots,e_{T})$, with $z_{0}^{(k)}=\mathbf{0}$ for all $k$ (the initial cache from Section 3.1), initialize $\texttt{correction}^{(0)}=\mathbf{0}$. For $k=1,\ldots,K$:

The $D$ transformer blocks share weights across iterations $k$ but have separate weights across layers $i$.

Non-cumulative correction. Each iteration replaces the previous correction entirely: $\tilde{x}^{(k)}=e+\texttt{correction}^{(k)}$, not $\tilde{x}^{(k)}=\tilde{x}^{(k-1)}+f(\tilde{x}^{(k-1)})$. Only the last correction matters.

Past-only correction. The correction at position $t$ uses $z_{t-1}^{(k)}$. Corrections propagate left to right: position $0$ converges after one iteration, position $1$ after two, and so on.

Random-depth training ($k_{\min}$). We sample $K\sim\text{Uniform}(k_{\min},K_{\max})$ each batch with $k_{\min}=2$, forcing the model to produce good predictions at any depth, which empirically encourages contraction.

Loss and dropout. The training loss (cross-entropy) is computed on the logits from the final iteration $K$ only. Dropout masks are resampled independently at each unrolling step $k$.

When a new token arrives with embedding $e_{t}$, the model computes the correction from the cached $z_{t-1}$ and $e_{t}$, passes the corrected embedding through the $D$-block unit, caches the output $z_{t}$, and predicts. This is one forward pass—no iteration over $K$ steps—regardless of the training depth $K$.

Why inference needs no iteration. During training, $K$ unrolling steps refine the corrections. At inference, this iteration is unnecessary: since earlier positions are already computed and cached, the correction for token $t$ is fully determined by $z_{t-1}$ and $e_{t}$ in a single pass (Theorem 2). The training approximation matters only during training: the first $K$ positions converge exactly after $K$ steps (Lemma 1 in Appendix A.2), and for later positions, the approximation error shrinks geometrically with $K$ when the correction operator is contractive (Theorem 3).

Data. OpenWebText (Gokaslan and Cohen, 2019) with byte-pair encoding (BPE, vocabulary 32,000). Context lengths: 64, 256, 512, and 1024 depending on the experiment. Ablations use English Wikipedia (BPE 16k); additional Wikipedia results in Appendix C.8. All results are validation perplexity (PPL) on held-out splits.

Training. AdamW optimizer (Loshchilov and Hutter, 2019), gradient clipping at 1.0. Learning rate $2\times 10^{-4}$ unless noted. Training depth $K=5$ with $k_{\min}=2$ by default; $K=10$ where noted. Dropout 0.2, FFN expansion factor 4. Full hyperparameters in Appendix C.

FLOP accounting.¹¹1Following standard convention in the literature, we count multiply-accumulate operations and label them FLOPs; actual floating-point operations are roughly $2\times$. We report total FLOPs per token, including the transformer blocks, correction FFN, and prediction head ($W_{\text{head}}\in\mathbb{R}^{V\times C}$, costing $VC$ FLOPs/token). Each transformer block costs $12C^{2}$ FLOPs ($4C^{2}$ for attention projections, $8C^{2}$ for the FFN). A context-ready model with $D$ blocks costs $D\times 12C^{2}+8C^{2}+VC$; a standard $N$-layer transformer costs $N\times 12C^{2}+VC$. Same-width comparisons (Tables 2, 3) share the same prediction head and are FLOP-matched. Cross-width comparisons (Table 1) explore depth-width tradeoffs: wider models pay a larger prediction head, so total FLOPs differ. Despite higher total FLOPs, the wider, shallower context-ready models deliver lower wall-clock inference time because fewer sequential layers dominate latency on modern GPUs (Section 5.4).

Baselines. Standard transformers with separate weights per layer and RoPE attention (Su et al., 2024). All results are single runs. The breadth of the evaluation—across widths ($C{=}50$–$2048$), depths ($D{=}1$–$23$), block sizes (64–1024), two datasets, a synthetic task, and multiple training strategies—provides stronger evidence than multi-seed runs on a single configuration: the context-ready architecture wins consistently across all these axes.

Table 1 compares context-ready models against standard transformers across depth-width tradeoffs at two compute scales. At the smaller scale (block size 256, 100K iterations), context-ready $D{=}5$ at $C{=}1120$ achieves 36.38 PPL, beating both roformer $N{=}6$ at $C{=}1088$ (37.76, $\Delta=-1.38$) and roformer $N{=}12$ at $C{=}768$ (37.83, $\Delta=-1.45$). At the larger scale (block size 256, 200K iterations), context-ready $D{=}6$ at $C{=}2048$ (29.04) matches roformer $N{=}12$ at $C{=}1536$ (29.01), despite using a much higher width-to-depth ratio. Depth has diminishing returns: going from $N{=}12$ to $N{=}24$ gains only 0.33 PPL.

Table 2 isolates the value of the correction mechanism at fixed width ($C=1024$), so the prediction head is identical across all rows and the comparison is FLOP-matched. $D=12$ context-ready ($152C^{2}$ FLOPs) beats roformer $N{=}13$ ($156C^{2}$) by 0.54 PPL and matches roformer $N{=}14$ ($168C^{2}$). The correction FFN adds only $8C^{2}$ FLOPs yet provides a genuine PPL improvement at the same parameter budget. At deeper scale, $D=23$ ($284C^{2}$) edges out $N=24$ ($288C^{2}$) by 0.53 PPL—directionally consistent with the $D{=}12$ result, but a single-run margin that should be read as evidence of parity rather than robust superiority.

Table 3 tests whether the correction advantage is a small-scale artifact. Comparing $D{=}2$ vs. $N{=}2$ at the same width with Chinchilla-matched token budgets, the relative improvement grows from 9.6% at $C=256$ to 16.5% at $C=1024$.

Token-matched results. At $C=1024$ with block size 64, $D{=}x$ beats $N{=}x$ at every depth tested once training progresses past a crossover point: $D{=}1$ beats $N{=}1$ by 34.4 PPL (crossover at ${\sim}424$M tokens), $D{=}2$ by 7.6 PPL (${\sim}565$M), $D{=}3$ by 3.0 PPL (${\sim}835$M), and $D{=}6$ by 0.7 PPL (${\sim}1{,}032$M). All gaps are still growing at the end of training. The correction mechanism provides a consistent advantage at every depth; the advantage is largest when depth is small, consistent with the correction doing the most work when the block has the fewest layers.

Inference latency. We measure autoregressive generation speed on an A100 over 10,000 tokens with KV caching. $D{=}1$ $C{=}2048$ (149M FLOPs/tok) generates at 919 tokens/s vs. 351 tokens/s for roformer $N{=}6$ $C{=}1088$ (120M FLOPs/tok)—a $2.6\times$ speedup despite higher total FLOPs. $D{=}5$ $C{=}1120$ (121M FLOPs/tok) generates at 349 tokens/s vs. 201 tokens/s for roformer $N{=}12$ $C{=}768$ (110M FLOPs/tok)—a $1.7\times$ speedup. The wider, shallower models are faster because fewer sequential layers dominate inference latency on modern GPUs, even when total FLOPs are higher. Per-token latency is flat across sequence length, confirming that KV caching amortizes attention cost to $O(T)$ per token. Full timing details in Appendix C.9.

KV cache savings. Fewer layers also reduce KV cache memory ($C\times D$ per token). Despite being wider, $D{=}5$ at $C{=}1120$ uses $1.6\times$ less cache than $N{=}12$ at $C{=}768$; $D{=}1$ at $C{=}2048$ uses $3.2\times$ less than $N{=}6$ at $C{=}1088$.

Proposition 1(a) predicts that $D{=}1$ may match deeper transformers when the task is dominated by context propagation and sufficient $K$ and width are provided. We test $D{=}1$, $C{=}2048$ (149M FLOPs/tok) against roformer $N{=}6$, $C{=}1088$ (120M FLOPs/tok) at block size 1024 using three training strategies: fixed $K{=}10$, random-depth $K{=}10$ with $k_{\min}{=}2$, and fine-tuning from a pretrained $N{=}1$ roformer (batch 16).

Table 4 compares the three strategies. In these $D{=}1$ experiments, larger training depth $K$ substantially improves the model’s ability to exploit the available recurrent depth. With fixed $K{=}10$, $D{=}1$ surpasses $N{=}6$ at ${\sim}65$K iterations and reaches 34.40 at 100K ($\Delta=-0.97$, still improving) at $2\times$ per-iteration cost.

Random-depth training ($k_{\min}{=}2$) incurs a modest penalty of 1.88 PPL at 100K relative to fixed depth, but converges to the same quality ${\sim}25$K later (both reach 33.63). The penalty yields $\hat{L}$ values consistent with contraction ($\hat{L}\approx 0.55$ vs. $\hat{L}>1.0$ for fixed $K$).

Fine-tuning from a pretrained $N{=}1$ roformer ($K{=}1$) with a zero-initialized correction FFN at $K{=}10$ reaches the best final PPL: 31.35 at 215K total iterations (the $N{=}6$ baseline is at 100K, so total compute differs).

The gap between $D{=}1$ and $N{=}6$ shrinks with block size ($+1.19$ at 256, $+0.46$ at 512, $+0.31$ at 1024 with $K{=}5$), consistent with longer contexts providing more sequential steps. Full block-size scaling in Appendix C.5.

An $N$-layer transformer can compose at most $N$ sequential reasoning steps in a single forward pass (Merrill and Sabharwal, 2023). To test whether the context-ready architecture can exceed this depth limit, we design a pointer-chasing task. The input contains a base table that maps keys to values, followed by $H{=}10$ levels of index tables, each of which maps new keys to keys at the previous level. Answering a query at level $\ell$ therefore requires chaining $\ell$ sequential lookups through the tables, and we use windowed causal attention (window $=38$) to prevent the model from bypassing these chains by attending directly to the base table. A full task specification with a worked example is given in Appendix B.

Table 5 shows a staircase-like depth dependence: deeper transformers solve more levels, while shallow transformers fail well before full depth. The context-ready $D{=}1$ model, trained with BPTT to exploit the full sequential depth of the recurrent correction chain, solves all 11 levels in ${\sim}16$K iterations and scales to 20 hops (all 21 levels).

Any standard transformer can be converted to context-ready by adding a zero-initialized correction FFN and fine-tuning. To isolate the effect of conversion from additional training, we compare against a same-iteration control: the original roformer trained for the same total number of iterations without conversion (per-iteration compute differs because fine-tuning runs the block $K$ times). At $C{=}1408$, $N{=}12$ reaches 29.92 PPL at 200K iterations; continued training to 400K yields 27.20. Converting at 200K and fine-tuning to 400K total iterations yields 26.14—a gain of $-1.06$ PPL over the continued-training baseline at matched iterations. At $C{=}1024$, converting $N{=}24$ to $D{=}24$ improves from 29.42 to 28.99 ($-0.43$) in 18K fine-tuning iterations. The zero-initialized correction ensures no disruption at conversion (the model is function-preserving). During fine-tuning, $D{=}12$ sees a transient $+0.11$ PPL increase before recovering, while $D{=}24$ shows no transient increase.

We compare against alternative architectures that attempt the same goal. Among the variants tested, the gain depends on a dedicated correction network with its own weights.

Convergence and sequential exactness. Table 6 shows the full depth progression. Convergence is geometric: at $D{=}1$ (block size 1024), $K{=}2$ closes 91% of the $K{=}1$-to-$K{=}5$ gap and $K{=}3$ closes 98%, consistent with Theorem 3. Sequential $K{=}1$ matches parallel $K{=}10$ to within 0.01 PPL at every configuration, confirming Theorem 2. The correction contribution (“Corr.” column) grows as $D$ shrinks: 38–55 PPL at $D{=}1$ vs. 3.9 at $D{=}23$.

Contraction. With $k_{\min}=2$, empirical $\hat{L}\in[0.50,0.72]$ (measured as $\hat{L}=\max_{t}\|c_{K,t}-c_{K-1,t}\|/\|c_{K-1,t}-c_{K-2,t}\|$); without $k_{\min}$, $\hat{L}\in[0.88,1.20]$. $\hat{L}$ is a trajectory-local diagnostic, not the global $L$ in Theorem 3.

Correction FFN is essential. Using the block’s own residual as the correction (“block_head”) gives no improvement (27.32 vs. 27.19 for a standard transformer). With a dedicated correction FFN, context-ready beats the FLOP-matched baseline by 1.8 PPL. Tying correction FFN weights to the block FFN collapses performance. The add variant ($\texttt{LN}(z_{t-1}+e_{t})$) matches or beats token-blind at all depths.

Training efficiency. $K{=}5$ suffices at $D\geq 5$; $K{=}2$ with torch.compile achieves $1.7\times$ faster training at 1.07 PPL cost. Full ablation tables in Appendix C.2.