PaperlayerView exact revision on arXiv

Source-linked AI summary

RecurrentGPT: Expressive Depth through Recurrent Modulation in Transformers

Amr Hegazy, Amr Alanwar, Mostafa Elhoushi

arXiv:2608.15062v1cs.CLcs.LG

TL;DR

Transformers couple depth to parameter memory, while uniform weight sharing can reduce the functional diversity that makes depth valuable. RecurrentGPT uses gated, context-conditioned reuse of a shared core, achieving competitive quality with 36–37% of dense-baseline parameters and improving validation loss at matched parameter counts.

  • Problem

    Transformer depth increases memory through unique layer weights, while standard sharing applies identical transformations despite changing hidden states, leaving the efficiency–functional-diversity tradeoff unresolved.

  • Method

    RecurrentGPT brackets a shared recurrent core with fixed prelude and coda blocks, using per-element gates conditioned on hidden state, context, and stochastic noise.

  • Results

    Under isoFLOPS, RECURRENTGPT stores 36–37% of dense-baseline parameters while beating MoR and heavy-tail Poisson sampling; under isoPARAMS, deeper recurrence improves validation loss.

  • Takeaways & Limitations

    At fixed memory, recurrent depth provides a supported way to trade inference FLOPs for improved model expressivity and validation loss.

  • Takeaways & Limitations

    Recurrence depth is fixed at inference, without per-token halting.

Abstract

from arXiv · show

Scaling transformer language models creates an inherent tension between expressivity and memory efficiency. While unique weights across layers preserve functional specialization---from input-grounding to abstract refinement---they incur a substantial memory footprint. Conversely, standard depth-sharing enforces uniform transformations that collapse representational diversity and degrade modeling quality. We introduce RecurrentGPT, a recurrent depth transformer where fixed-depth prelude and coda blocks bracket a single shared core iterated R times. Inspired by gated recurrent neural networks, we employ a lightweight projection and an elementwise update gate---conditioned on the hidden state, the fixed prelude output, and noise resampled at every step---to modulate the recurrent update. This allows the model to specialize the input to the same few layers across recurrences, rather than requiring many unique layers to achieve functional diversity. Under an isoFLOPS constraint, a 3-layer RecurrentGPT matches the accuracy of a 12-layer GPT-2 Small baseline with similar training and inference FLOPs, and leads MoR and heavy-tail depth sampling in all nine scale-by-budget cells; at medium and large scale it approaches dense quality at the standard token budget and overtakes it at medium scale once that budget is doubled. Under an isoPARAMS constraint, deeper recurrence achieves a 2.76 validation loss versus 2.84 for a non-recurrent counterpart at matched parameter and data budget. Our results demonstrate that adaptive depth reuse is a principled strategy for trading parameters for quality: at large scale, 63% fewer parameters and 59% less peak decoding memory for a 10% increase in compiled generation latency.

1 Introduction

RecurrentGPT addresses the memory bottleneck caused by coupling transformer depth to unique parameters by iterating a shared transformation in hidden states. Its gated, stochastic modulation lets shared layers specialize across recurrences, and the paper evaluates this approach under isoFLOPS and isoPARAMS regimes.

  • Motivation: Standard transformers couple every added layer to fresh weights, creating a memory bottleneck that limits effective depth under fixed hardware budgets.Scaling also increases parameter count and training compute.
  • Motivation: Recurrent depth iterates a shared transformation over input representations, enabling deeper computation in hidden states without extra output tokens or parameters.This provides a complementary realization of iterative computation to chain-of-thought in the output space.
  • Motivation: Weight sharing grows effective depth without inflating parameter count, but identical transformations can create a tension between memory efficiency and representational diversity.Prior work used recurrent depth both to match quality with fewer parameters and to improve accuracy at matched parameter counts.
  • Method: RECURRENTGPT uses per-element gates conditioned on the hidden state, fixed prelude representation, and stochastic noise so one shared tensor behaves as multiple specialized layers.Gates initially keep the residual stream nearly unchanged across recurrences, while training enables distinct context-grounded inputs at each step.
  • Contributions: RECURRENTGPT combines per-element gating and stochastic perturbations to prevent representational collapse and is evaluated under isoFLOPS and isoPARAMS at three scales.The stated evaluation includes comparisons with depth-sampling and recurrent baselines across scale and budget conditions.

2 Related Work

Recurrent depth develops from weight-tied iterative networks and recurrent sequence models, while transformer approaches differ in whether they share entire stacks or selectively recur middle layers. Prior work frames recurrent depth around isoPARAMS and isoFLOPS tradeoffs, dynamic computation, and equilibrium methods, motivating gains under both regimes.

  • Origins: Weight sharing traces to the 1986 synchronous iterative net and later recurrent architectures, including RNNs, LSTMs, and GRUs.These models apply shared weights across repeated computational steps or sequence positions.
  • Transformer recurrent depth: ALBERT and Universal Transformers share the entire transformer stack, whereas later work selectively recurs middle layers while keeping other layers fixed.Universal Transformers additionally use per-token adaptive halting; RecurrentGPT follows selective recurrence of middle layers.
  • Evaluation regimes: In the isoPARAMS regime, prior studies increase FLOPs by recurring layers at fixed parameter count and report consistent quality gains.This regime explicitly trades additional computation for improved accuracy.
  • Dynamic computation: Adaptive Computation Time, Mixture-of-Depths, and Mixture-of-Recursions dynamically vary computation across positions, tokens, or layer subsets rather than fixing it uniformly.These methods adapt computation through learned halting or routing mechanisms.
  • Equilibrium models: Deep equilibrium models iterate tied weights toward a fixed point, while RecurrentGPT unrolls fixed discrete steps and trains each step as an exit.RecurrentGPT therefore shares the weight-tied update but not the equilibrium objective or implicit-gradient training approach.
  • Scaling regimes: Kaplan et al. found recurrent transformers better at matched parameters but worse at matched compute; RecurrentGPT targets gains in both scaling regimes.The paper presents its architecture as improving loss scaling with parameter count and training FLOPs as model size and data scale.

3 Recurrent Depth Reuse for Language Modeling

RecurrentGPT reuses a shared transformer core between fixed prelude and coda blocks, combining recurrent state updates with context-conditioned gating and noise. This reduces unique blocks while preserving recurrent depth and enables stochastic-depth training with early exits.

  • Architecture: RecurrentGPT partitions blocks into a prelude, a shared core repeated R times, and a coda, yielding more stable training than uniform weight sharing.The prelude provides a fixed context-aware signal, while the shared core iteratively processes the recurrent state conditioned on that anchor.
  • Efficiency: 2.6× fewer unique blocks: a 2+5×4+2 model executes 24 blocks while storing weights for 9, matched to the 24-layer GPT-2 medium baseline.The configuration visits 2 + 5 · 4 + 2 = 24 block executions per forward pass, while storing only 2 + 5 + 2 = 9 blocks.
  • Adaptive gating: A learned elementwise gate selectively writes shared-core outputs into the residual state using the current state, fixed prelude output, projection, and injected noise.The gate is produced by a feed-forward network, while state noise perturbs the recurrent projection input and gate noise prevents collapse toward constant gating.
  • Adaptive gating: Conditioning every recurrence on h(pre) provides a constant view of the original input while modulating it with the recurrent state and stochasticity.This offers a memory-efficient alternative to Per-Layer Embeddings, which require L unique embedding layers for layer-specific context.
  • Training: Uniformly sampling r from 1 to R trains every exit point for early inference exits and supplies stochastic-depth regularization.The procedure avoids auxiliary loss terms and is reported to improve final validation loss over fixed-depth schedules.

4 Experiments & Results

Experiments show that RecurrentGPT preserves or improves quality with substantially fewer unique parameters, especially when inference FLOPs are traded for recurrent depth. It also provides adaptive compute behavior and meaningful decoding-memory savings, though recurrent KV caching remains an important systems cost.

  • Main Results: 36–37% of baseline parameters suffice for competitive isoFLOPS quality, with RecurrentGPT outperforming the dense baseline at small scale but trailing at medium and large standard budgets.At small scale, validation loss is 3.145 ± 0.004 versus 3.188 ± 0.056; the dense baseline leads by 0.05 and 0.06 nats at medium and large scale.
  • Main Results: 0.08 nat and 0.06 nat improvements result from increasing recurrence at fixed parameters at medium and large scale, respectively.The results support trading inference FLOPs for recurrent depth to improve expressivity under a fixed memory footprint.
  • Downstream Evaluation: 42.08 vs. 42.05 average accuracy shows isoFLOPS parity with the dense model, while isoPARAMS outperforms the dense baseline on eight of nine benchmarks.At large scale, the isoPARAMS model exceeds the dense baseline by +2.10 average points and the isoFLOPS model uses 37% of its parameters.
  • Early Exit Analysis: Fewer recurrences can yield better matched-FLOPs loss than earlier dense exits, producing a continuous compute–quality dial across model scales.Uniform depth sampling trains intermediate recurrent states to predict final losses, enabling early-exit evaluation.
  • Wall-Clock Latency and Decoding Memory: 63% fewer parameters and 59% less peak memory cost RecurrentGPT-Large only +10% generation latency under torch.compile.At medium scale, compiled latency overhead is +11%; in eager mode, both overheads are +23%.
  • KV Cache Sharing: 0.39× dense decoding memory is achieved at batch size 32 by averaging K/V across recurrence steps, versus 0.91× with the naive recurrent cache.The naive R× cache is 0.55× dense memory at batch size 1 but becomes costly when cache memory dominates.

5 Ablations

The ablations show that recurrence benefits substantially from architectural mechanisms that differentiate processing steps, especially the elementwise gate. Across training-data budgets, RecurrentGPT maintains a superior small-scale Pareto frontier, while larger-scale behavior exhibits a data-dependent crossover with the dense baseline.

  • Component Ablation: 0.107 nats higher validation loss results from recurrence alone, as shared weights conflate early and late processing.This ablation uses the small model at the full 20,000-step budget.
  • Component Ablation: 0.035 nats are recovered by adding prelude and coda boundaries, while state noise contributes 0.018 more.The noise result isolates injecting noise at every recurrence rather than only at initialisation.
  • Component Ablation: −0.048 nats makes the elementwise gate the single largest contributor in the sequential ablation.The ablations use validation loss as the primary metric.
  • Data Efficiency: RecurrentGPT consistently maintains a superior Pareto frontier to the dense baseline at smaller model scales across observed data volumes.The comparison varies the training data budget and uses validation loss versus training data budget.
  • Data Efficiency: At larger scales, the dense model initially leads in low-data regimes, but the gap narrows significantly as the token budget increases.This larger-scale behavior is described as a more complex crossover dynamic.

6 Conclusion

RecurrentGPT decouples model depth from parameter count through a learned per-element gate conditioned on hidden state, fixed prelude output, and injected stochastic noise. Under isoFLOPS and isoPARAMS constraints, it reduces unique parameters, outperforms depth alternatives, and gains from deeper recurrence.

  • Conclusion: A learned per-element gate conditions on the current hidden state, fixed prelude output, and injected stochastic noise to decouple depth from parameter count.The gate is presented as a principled mechanism for recurrent depth reuse.
  • Conclusion: 36–37% of the dense baseline’s unique parameters are stored under an isoFLOPS constraint.RecurrentGPT achieves this reduction while beating both MoR and heavy-tail Poisson sampling at every scale.
  • Conclusion: Deeper recurrence yields consistent validation-loss gains under an isoPARAMS constraint.The passage reports these gains without specifying their numerical magnitude.

A Implementation Details … B.1 Run-to-Run Variance Across Seeds

RecurrentGPT combines a fixed prelude with a shared recurrent core whose updates are modulated by state-dependent gates and noise, while training samples recurrence depth uniformly. Across three seeds, it outperforms the dense small isoFLOPS counterpart with substantially lower run-to-run variance.

  • A Implementation Details: The forward pass anchors the hidden state after prelude blocks before iterating a shared core for the sampled recurrence depth.At inference, depth is set to R; during training, r is sampled between r_min and R.
  • A Implementation Details: Each recurrent update projects the noisy current state concatenated with the frozen prelude output before applying the shared block.State noise is resampled at every recurrence step.
  • A Implementation Details: The update gate uses normalized current and anchor states, temperature scaling, and gate noise to interpolate between the prior state and shared-core output.The gate is elementwise and lies in [0,1]^(S × d).
  • A.1 Hyperparameter Table: Training configurations define p, b, R, c, d, and h for prelude, shared, recurrence, coda, embedding, and attention dimensions.All runs use AdamW, weight decay 0.1, gradient clipping 1.0, bfloat16 mixed precision, and cosine learning-rate scheduling.
  • A.1 Hyperparameter Table: All recurrent runs initialize gate bias at +4, corresponding to g ≈0.98, with σx = 0.1, σg = 0.1, and τ = 1.0.These settings are used together with uniform recurrence-depth sampling during training.
  • B.1 Run-to-Run Variance Across Seeds: 3.145 ± 0.004 was RecurrentGPT’s mean result across three seeds, versus 3.188 ± 0.056 for the dense baseline.The comparison uses the small isoFLOPS configuration at the full 20,000-step budget.
  • B.1 Run-to-Run Variance Across Seeds: 3.148 was RecurrentGPT’s worst seed, still below the dense baseline’s best seed of 3.154.Seeds were 1337, 42, and 123; seed 1337 is reported throughout the main paper.

B.2 Component Ablation at Short Horizon … B.5 Noise Magnitude Sensitivity

The ablations show that structural components contribute early, while the learned gate accrues benefits over training; comparisons and sensitivity sweeps identify the model’s distinctive design and preferred hyperparameters. Short-horizon results, method properties, gate temperature, and state-noise magnitude are evaluated across small and medium configurations.

  • B.2 Component Ablation at Short Horizon: At 2,000 steps, the sequential medium-configuration ablation reverses the ordering of the final two rows relative to 20,000 steps.Table 7 uses RECURRENTGPT-medium and a standard 24-layer GPT dense baseline trained identically.
  • B.2 Component Ablation at Short Horizon: Table 6 reports three-seed validation loss for the small configuration after 20,000 steps under the isoFLOPS setting of Table 1.Lower validation loss is better.
  • B.2 Component Ablation at Short Horizon: Prelude re-injection contributes −0.198 nats, compared with −0.115 for the gate.Structural components land early, whereas the learned gate’s contribution accrues over training.
  • B.3 Design Comparison of Recurrent-Depth Methods: RECURRENTGPT uniquely combines full sharing, an input-dependent update, and variable inference depth among the compared recurrent-depth methods.MoR and Ouro gate updates, Ouro supervises every iteration, MoR requires a router, heavy-tail Poisson updates unconditionally, and RRT is input-independent and depth-fixed.
  • B.3 Design Comparison of Recurrent-Depth Methods: Small-scale isoFLOPS comparisons use 20,000-step training times, but recurrent-method timings are not directly comparable with the optimistic compiled dense baseline.The dense model used torch.compile, while recurrent methods—including RECURRENTGPT—ran without compilation because compilation failed on the hardware.
  • B.4 Gate Temperature Sensitivity: Validation loss is 3.62, 3.60, and 3.63 for gate temperatures τ = 0.5, 1.0, and 2.0, respectively.The sweep uses RECURRENTGPT-medium at 5,000 steps, making τ = 1.0 optimal.
  • B.5 Noise Magnitude Sensitivity: Removing state noise at σx = 0.0 increases validation loss by 0.031 nats relative to σx = 0.1, while σx = 0.2 degrades by 0.018 nats.The optimal range is σx ∈[0.05, 0.1], and σx = 0.1 is used throughout.

B.6 Gate Bias and State Noise at Full Training Horizon

At the full 20,000-step horizon, gate-bias and state-noise sweeps remain shallow and non-monotonic rather than sharpening at convergence. No tested setting caused a loss spike or divergence.

  • Full Training Horizon: 20,000 steps: the small-configuration response surface does not sharpen at convergence.The sweep was repeated at the full training budget to test convergence behavior.
  • Gate Bias: 0.019 nats: moving gate bias from the default +4 to 0 increases loss; −2 recovers to 3.151.Gate-bias sensitivity is shallow and non-monotonic rather than peaked.
  • State Noise: 0.018 nats: removing state noise increases loss, while doubling it costs 0.019 nats.State-noise sensitivity follows the same shallow pattern as gate bias.
  • Robustness: No setting produced a loss spike or divergence.This held across the full-horizon gate-bias and state-noise sweeps.

C Additional Qualitative Results · C.1 Per-Token Gate Activations Across Recurrence Steps

Per-token gate activations reveal recurrence-dependent writing and copying patterns across prompts. Gates are primarily driven by token identity and content, producing implicit routing without an explicit per-token mechanism.

  • C Additional Qualitative Results: Table 10 reports state-noise sensitivity for the small configuration at 20,000 steps.
  • C.1 Per-Token Gate Activations Across Recurrence Steps: Gate visualizations average the elementwise gate over model dimensions for each token position and recurrence step across four example prompts.
  • C.1 Per-Token Gate Activations Across Recurrence Steps: 0.70–0.82 marks write-heavy positions where the gate is open and substantially absorbs the shared block’s proposal.
  • C.1 Per-Token Gate Activations Across Recurrence Steps: 0.95–1.00 marks copy-heavy positions where the hidden state passes through largely unchanged.
  • C.1 Per-Token Gate Activations Across Recurrence Steps: ∼0.87–0.89 at step 1 falls to a minimum around step 2, then recovers toward ∼0.88–0.91 at step 6.The minimum corresponds to a write-heavy phase, while step 6 is copy-heavy.
  • C.1 Per-Token Gate Activations Across Recurrence Steps: Token positions that light up at step 1 tend to remain the same at step 6, while their writes progressively dim over recurrence steps.This indicates token identity, rather than recurrence depth, drives the gate’s primary signal.
  • C.1 Per-Token Gate Activations Across Recurrence Steps: Content words and syntactically load-bearing tokens receive lower gate values than function words, including “man” and “sat” versus “the” and “by.”
  • C.1 Per-Token Gate Activations Across Recurrence Steps: In code, operators and identifiers are write-heavy while keywords such as “def” and “for” are copy-heavy, yielding implicit routing without explicit per-token control.The routing arises purely from conditioning the gate on the evolving hidden state.

D Broader Impact · E Recurrence Dynamics Analysis

RecurrentGPT offers environmental and accessibility benefits through reduced serving memory and energy use, while preserving an inference-time compute–quality tradeoff. Recurrence dynamics analysis examines six steps with diagnostics showing a transition from writing to copying and lower gate values for structurally informative tokens.

  • D Broader Impact: 37% of a baseline’s parameter count is required by a RecurrentGPT checkpoint, reducing serving GPU memory and energy consumption per inference step.This can enable deployment on hardware unable to accommodate the full dense model.
  • D Broader Impact: Recurrent depth provides an inference-time compute–quality tradeoff through the early-exit capability described in Section 4.The supplied passage identifies early exit as a capability of a single trained model.
  • E Recurrence Dynamics Analysis: The per-step mean gate activation rises from step 2 onward, reflecting a transition from writing to copying across recurrence steps.Bright colors indicate write-heavy updates, while dark purple indicates copy-heavy hidden-state passthrough.
  • E Recurrence Dynamics Analysis: Structurally informative tokens consistently receive lower gate values than function words and punctuation across recurrence steps.The gate interpretation identifies lower values with more copy-heavy behavior.
  • D Broader Impact: RECURRENTGPT inherits standard dual-use risks of capable web-trained language models, including potential misinformation or harmful content generation.The passage states that standard safeguards include content filtering, output monitoring, and responsible release practices.
  • E Recurrence Dynamics Analysis: Six recurrence steps are probed in the large 1+5×6+5 checkpoint using ten diagnostics, grouped into behavioural and mechanistic analyses.Inference noise is disabled and evaluation uses a validation split.

E.1 Rapid Convergence of Loss Across Recurrence Steps … E.11 Unified Summary

Across recurrence, RecurrentGPT rapidly restructures representations in early steps, then applies increasingly selective, aligned corrections toward a fixed point. The learned gate, contrastive state–anchor projection, and stable attention roles explain parameter-efficient adaptive refinement across tokens.

  • E.1 Rapid Convergence of Loss Across Recurrence Steps: 77% of loss reduction occurs within the first two recurrence steps, reaching 2.68 final validation loss; output distributions are essentially committed by step 4.Extending inference beyond six trained steps marginally worsens loss to 2.69, 2.72, and 2.74 at R = 8, 10, and 12.
  • E.2 Gate Behaviour: Write-Heavy Early, Copy-Heavy Late: The gate shifts from write-heavy early steps (mean ≈0.82) to copy-heavy late steps (mean ≈0.87), while effective openness declines from 0.18 to 0.07.Copy-saturated dimensions increase from 19.7% at step 1 to 28.8% at step 6, whereas write-saturated dimensions remain below 10^-4.
  • E.3 Update Dynamics: Diminishing Magnitude, Increasing Alignment: Applied update norms decline from 17.6 to 6.5 while raw proposals remain 93–99, and applied-update cosine similarity rises from 0.94 to 0.99.Proposal cosine similarity increases from 0.19 to 0.51, while per-step gain peaks at 10.4 at step 2 and falls to 4.7 at step 6.
  • E.4 Representational Geometry: Two-Phase Structure: A representational phase transition separates steps 1 and 2: steps 0–1 align with baseline layer 2, whereas steps 2–6 cluster around layer 11.Within-model CKA is 0.996 between steps 0 and 1, while steps 2–6 have pairwise CKA ≥0.907; recurrent steps do not strongly align with layers > 20.
  • E.6 Fixed-Point Convergence and Effective Rank: Effective rank rises from 0.546 at the prelude to 0.635 by step 3, while fixed-point update magnitude declines from 0.324 to 0.112 by step 6.The late-delta subspace rank is 0.69, exceeding representation rank and indicating diverse feature refinement.
  • E.7 Gate Attribution: Current State Gradually Supersedes the Anchor; E.8 Projection Weight Analysis: State and Anchor as Complements; E.11 Unified Summary: Gate attribution shifts from 59% current-state versus 41% anchor at step 1 to 74% versus 26% at step 6, supported by Wx’s 44% larger leading singular value.The projection is contrastive: Wx and Wh have mean row cosine −0.189, so changed state directions are emphasized rather than averaged.
  • E.9 Attention Entropy: Stable Head Specialisations, Gradual Sharpening; E.10 Recurrence Preferentially Serves Two Token Regimes; E.11 Unified Summary: Across steps, attention heads retain stable specializations while entropy falls from 3.00 to 2.77 nats before rebounding to 2.82; byte/punctuation tokens resolve 49% of prior loss at step 1.Rare tokens improve more gradually, with recurrence resolving < 4% of preceding loss for all frequency bands beyond step 3.
Loading 2608.15062v1…