PaperlayerView exact revision on arXiv

Source-linked AI summary

Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA

Sangmin Bae, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Seungyeon Kim, Tal Schuster

arXiv:2410.20672v3cs.CLcs.LG

TL;DR

Modern LLM deployment requires balancing performance against substantial memory and computation costs, while parameter sharing has had limited reported success in modern LLMs. The paper converts pretrained Transformers into recursively reused architectures, adds layer-specific LoRA for flexibility, and proposes continuous depth-wise batching with early exiting. These models perform strongly against comparable non-recursive models, can approach full-size Gemma performance, and offer simulated throughput gains of up to 2–3×.

  • Problem

    Modern LLMs are expensive to deploy, while parameter sharing has achieved limited reported success in modern LLMs.

    Source: §1 · p. 1

  • Method

    The paper converts pretrained Transformers into Recursive Transformers with repeatedly reused layer blocks, relaxes tying with layer-specific LoRA modules, and pairs recursion with continuous depth-wise batching and early exiting.

    Source: §1 · pp. 1–2

  • Results

    Recursive Gemma 1B improved few-shot accuracy by up to 13.5 absolute points, and distilled recursive Gemma matched full-size Gemma while early exiting yielded up to 2–3× simulated throughput gains.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Recursive parameter sharing can produce compact LLMs with strong comparable-model performance and potential inference-throughput benefits.

    Source: §1 · pp. 1–3

    Practical deployment still requires confidence-based early-exiting algorithms and solutions for key-value-cache computation for exited tokens.

    Source: §5 · p. 15

Abstract

from arXiv · show

Large language models (LLMs) are expensive to deploy. Parameter sharing offers a possible path towards reducing their size and cost, but its effectiveness in modern LLMs remains fairly limited. In this work, we revisit "layer tying" as form of parameter sharing in Transformers, and introduce novel methods for converting existing LLMs into smaller "Recursive Transformers" that share parameters across layers, with minimal loss of performance. Here, our Recursive Transformers are efficiently initialized from standard pretrained Transformers, but only use a single block of unique layers that is then repeated multiple times in a loop. We further improve performance by introducing Relaxed Recursive Transformers that add flexibility to the layer tying constraint via depth-wise low-rank adaptation (LoRA) modules, yet still preserve the compactness of the overall model. We show that our recursive models (e.g., recursive Gemma 1B) outperform both similar-sized vanilla pretrained models (such as TinyLlama 1.1B and Pythia 1B) and knowledge distillation baselines -- and can even recover most of the performance of the original "full-size" model (e.g., Gemma 2B with no shared parameters). Finally, we propose Continuous Depth-wise Batching, a promising new inference paradigm enabled by the Recursive Transformer when paired with early exiting. In a theoretical analysis, we show that this has the potential to lead to significant (2-3x) gains in inference throughput.

1. Introduction

The paper revisits parameter sharing for modern LLMs, converting pretrained Transformers into compact Recursive Transformers and extending them with layer-specific LoRA modules. It reports strong performance against comparable non-recursive models and potential throughput gains from continuous depth-wise batching with early exiting.

Motivation

  • Parameter sharing can reduce memory footprint and accelerator requirements, but its application to modern LLMs has previously yielded limited reported success.The paper frames its contribution as revisiting this trade-off with recursive architectures.Source: §1 · p. 1

Method

  • Recursive Transformers replace independently parameterized layers with one shared block reused across multiple loops, initialized from the original pretrained model and briefly uptrained.This design targets smaller, more efficient models while preserving performance.Source: §1 · p. 1, Figure 1 · p. 2
  • Relaxed Recursive Transformers add layer-specific low-rank LoRA modules, allowing controlled deviations between shared layers with minimal overhead.The LoRA rank controls the degree of relaxation between fully tied and vanilla architectures.Source: §1 · p. 2

Results

  • Up to 13.5 absolute accuracy improvement, or 22% error reduction, was observed for recursive Gemma 1B versus a non-recursive Gemma 1B on few-shot tasks.The recursive model was converted from pretrained Gemma 2B and uptrained rather than pretrained from scratch.Source: §1 · p. 2
  • With knowledge distillation and 60 billion uptraining tokens, recursive Gemma achieved performance on par with full-size Gemma trained on 3 trillion tokens.This comparison is reported against the full-size model in §3.3.Source: §1 · p. 2

Inference

  • Continuous Depth-wise Batching groups different loop iterations across responses and, with early exiting, is simulated to improve throughput by up to 2–3× over a same-architecture vanilla Transformer.The recursive Gemma model is stated to theoretically achieve nearly 4× throughput relative to vanilla Pythia.Source: §1 · pp. 2–3

2. Effective Model Compression with Recursive Patterns

Recursive Transformers compress pretrained Transformers by repeatedly reusing a single block of layers, while Relaxed Recursive Transformers restore depth-specific flexibility through low-rank adapters. The same recursive structure also supports depth-wise batching and early exiting for potentially faster inference.

Recursive Transformer

  • Recursive Transformers reuse a single block of unique layers across multiple loops, tying all trainable parameters and reducing stored model capacity.For example, a 2-block recursive variant of 18-layer Gemma 2B stores weights for only the first 9 layers and loops through them twice.Source: §2.2 · p. 3

Initialization

  • Initialization methods transfer knowledge from the original pretrained model into looped layers, with the Average method performing best for Relaxed Recursive Transformers.The proposed methods include Stepwise selection of source layers and alternatives illustrated in Figure 2.Source: §2.2 · p. 4, Figure 2 · p. 4

Relaxed Recursive Transformer

  • Relaxed Recursive Transformers attach independent LoRA modules to each looped layer, adding low-rank, iteration-specific weight deltas while preserving shared base parameters.The LoRA rank controls the degree of relaxation; rank zero recovers the Recursive Transformer, while sufficiently large ranks approximate the full-size vanilla model.Source: §2.3 · pp. 4–5

LoRA initialization

  • Truncated SVD initializes LoRA factors from the largest singular components of residuals between original and tied weights, making rank a controllable bridge between tied and vanilla models.The method initializes B from U_rΣ_r and A from the transpose of V_r.Source: §2.3 · p. 5

Inference

  • Continuous Depth-wise Batching schedules different loop iterations for different samples, while early exiting immediately fills vacant batch slots and can improve throughput.Recursive layer sharing allows new requests to begin before all stages finish, and early exits further reduce synchronization waste.Source: Figure 3 · p. 5, §2.4 · p. 6

3. Experiments

The experiments show that pretrained-weight initialization makes compact Recursive Transformers competitive with larger and similar-sized baselines, while LoRA relaxation improves the size–performance trade-off. The same recursive structure also supports substantial throughput gains with early exiting.

3.3. Main Results

  • Recursive Gemma retains strong performance with half the full-size model’s parameters, and achieves a 10%p gain over a reduced-size model trained on 60 billion distillation tokens.Recursive TinyLlama also surpassed its vanilla counterpart despite the latter using a larger 105-billion-token corpus.Source: §3.3 · pp. 7–8

3.5. Relaxation of Strict Parameter Sharing via LoRA Modules

  • At rank 512, relaxed Gemma reaches 58.4% versus 58.6% for the original full-size model, despite fewer parameters and only 60 billion uptraining tokens.Increasing LoRA rank interpolates between the fully tied recursive model and the vanilla model, while SVD initialization better approximates full-rank weights.Source: §3.3 · p. 8, §3.5 · p. 10

3.4. Initialization Techniques for Looped Layers

  • Stepwise initialization gives the lowest training loss and highest average few-shot accuracy among the tested initialization methods, outperforming random initialization.With 15 billion tokens of uptraining, the resulting recursive Gemma 1B surpasses full-size pretrained TinyLlama and Pythia.Source: §3.4 · pp. 8–9

3.3. Main Results

  • Up to 13.5 percentage points (51.7% vs. 38.2%) separates recursive Gemma 1B from reduced-size baselines, while it also outperforms full-size TinyLlama 1.1B and Pythia 1B.The recursive model used only 15 billion uptraining tokens, whereas Pythia was pretrained on 300 billion tokens.Source: §3.4 · p. 9

3.7. Early-Exit Training Strategy for Recursive Transformer

  • Intermediate-output training uses an aggressive coefficient and self-distillation strategy, but applying the strategy derived for non-relaxed models may leave relaxed models’ early exits suboptimal.The authors identify a more tailored training approach as a possible improvement for intermediate loop outputs.Source: §3.7 · p. 12

3.8. Continuous Depth-wise Batching

  • Continuous depth-wise batching with early exiting yields up to 2.66× measured speedup and a theoretical 2–3× gain over vanilla Transformers.The recursive Gemma model can theoretically approach a nearly 4× throughput increase, and the method can outperform continuous sequence-wise batching.Source: §3.8 · p. 13

4. Related Work

Prior work established parameter sharing and recurrent computation in smaller models and other architectures, but this paper targets pretrained LLMs. It also connects recursive Transformers to early-exit decoding without vanilla Transformers’ batch synchronization problem.

Cross-layer parameter sharing

  • Earlier parameter-sharing research spans RNNs, CNNs, and Transformers, including Universal Transformers that achieved strong results with fewer parameters.The cited prior work primarily concerns cross-layer sharing or recurrent self-attention in smaller settings.Source: §4 · p. 13

Recurrent reasoning and computation

  • Previous studies of recurrent networks explored extrapolating learned reasoning strategies through additional inference recurrences, including looped Transformers for basic computation.These works broaden recurrence beyond standard language-model parameter efficiency.Source: §4 · p. 13

Distinction from prior work

  • This work differs by applying parameter sharing to LLMs pretrained from existing weights and by introducing a generalized framework that relaxes tying with layer-specific parameters.The paper positions its initialization strategies as the mechanism for leveraging knowledge already embedded in pretrained LLMs.Source: §4 · p. 13

Early-exit decoding

  • Recursive structure supports individual-response early exiting within serving batches, avoiding vanilla Transformers’ synchronization requirement that all tokens follow the deepest path.This property motivates the paper’s continuous depth-wise batching paradigm.Source: §4 · p. 14

5. Conclusion and Future Work

The paper compresses LLMs through recursive parameter sharing, adds layer-specific LoRA relaxation, and proposes continuous depth-wise batching for efficient serving. Future work addresses sparse or larger recursive models, latent reasoning, and practical early-exit and multi-LoRA serving constraints.

Contributions

  • Recursive Transformers compress LLMs by sharing parameters across recursively looped layer blocks, while LoRA modules provide low-rank deltas between repeated layers.The approach uses layer-specific LoRA modules to relax fully tied weights while retaining compactness.Source: §5 · p. 14
  • Oracle early exiting with continuous depth-wise batching theoretically reaches up to 2-3× throughput speedup.The paradigm is designed specifically for recursive models and early-exit inference.Source: §5 · p. 14

Future work

  • The dense recursive design is compatible with pruning, quantization, layer skipping, mixture-of-experts, activation skipping, and SSM components, but these combinations remain future work.The paper focuses on recursive dense Transformer layers rather than integrating these alternatives.Source: §5 · p. 14
  • Scaling the approach to 7B-and-larger models may require substantially higher uptraining costs, and its effect on achievable batch sizes remains unclear.Larger hidden dimensions may offset memory-footprint benefits when determining batch size.Source: §5 · p. 15

Limitations

  • Practical throughput remains unresolved because oracle exiting assumes matching intermediate predictions, while deployment must handle confidence calibration and key-value-cache computation.Continuous depth-wise batching and key-value-cache sharing are proposed as possible solutions.Source: §5 · p. 15
  • Relaxed models require distinct LoRA computations during batched inference, hindering parallelism despite concatenated-weight optimization.Optimized CUDA kernels and accelerator parallelization are suggested directions.Source: §5 · p. 15

A. Components in Transformer Architecture

This section reviews Transformer components and layer-sharing strategies before describing recursive conversion, LoRA initialization, and the evaluated pretrained models. The proposed setup uses looped layers with initialization derived from pretrained weights.

Components in Transformer Architecture

  • The Transformer block combines multi-head attention, which captures relationships across sequence elements, with a feed-forward network.Gemma’s feed-forward network uses three learned linear matrices and a GeGLU activation.Source: A. Components in Transformer… · p. 25

Parameter Sharing

  • CYCLE repeatedly stacks one block of unique layers, whereas SEQUENCE shares parameters across consecutive layers and CYCLE (REV) reverses lower layers in later repetitions.The paper selects CYCLE because it is compatible with early exiting.Source: B. Parameter Sharing Strategy · p. 25, Figure B.1 · p. 25

LoRA Initialization

  • Relaxed Recursive Transformers use SVD initialization for LoRA modules so base-plus-LoRA weights are steered toward pretrained weights at corresponding depths.When looped weights already match pretrained weights, the corresponding LoRA module receives standard zero initialization.Source: C. Illustrative Examples of SVD… · p. 26, Figure C.1 · p. 26

Models and Evaluation

  • The experiments convert Gemma 2B, TinyLlama 1.1B, and Pythia 1B into Recursive Transformers and evaluate their few-shot performance.Few-shot accuracy covers seven benchmarks, including LAMBADA, HellaSwag, PIQA, WinoGrande, ARC-easy, ARC-challenge, and OpenBookQA.Source: D. Overview of Three Pretrained LLMs · p. 26, Table D.1 · p. 26
  • The model selection enables comparisons by trainable-parameter count and tests whether leveraging pretrained weights benefits recursive models relative to comparable pretrained baselines.The set includes models with different training statuses and scales.Source: D. Overview of Three Pretrained LLMs · p. 27

E. Experimental Setup

The experiments uptrain converted models on SlimPajama, train intermediate-loop exits, and evaluate language modeling, few-shot accuracy, and simulated throughput. Throughput analysis uses oracle exits and measured component-level generation times.

Uptraining setting

  • Recursive models receive additional uptraining on either 15 billion or 60 billion SlimPajama tokens.SlimPajama is a cleaned and deduplicated RedPajama-derived corpus containing web, code, books, Arxiv, Wikipedia, and StackExchange data.Source: E. Experimental Setup · p. 27

Early-exit training setting

  • Early-exit training teaches models to predict next tokens at intermediate loops, using two exit points for two-block models and three for three-block models.The default procedure adds 15 billion SlimPajama tokens after recursive uptraining.Source: E. Experimental Setup · p. 27

Evaluation setting

  • Evaluation measures perplexity on SlimPajama, RedPajama, and PG19 plus accuracy on seven few-shot benchmarks.Few-shot evaluation uses the Language Model Evaluation Harness.Source: E. Experimental Setup · p. 27

Throughput measurement setting

  • Throughput comparisons measure per-token generation time for model components and base final comparisons on time spent inside Transformer blocks.Measurements use dummy weights and inputs across model configurations.Source: E. Experimental Setup · p. 28
  • Exit trajectories come from an oracle-exiting simulation in which tokens exit when intermediate predictions match final-loop predictions.The simulation uses SlimPajama, RedPajama, and PG19 and assumes a queue of 20K samples.Source: E. Experimental Setup · p. 28

F. Expanded Results of Initialization Methods for Looped Layers

The expanded results show that initialization from pretrained layers strongly affects recursive-model quality: Stepwise initialization is consistently strongest, while normalization handling can be architecture-sensitive. Comparisons across scales and baselines further assess the roles of pretrained weights and looping.

Ablation study of Stepwise method

Ablation study of Average method

  • Zero initialization of normalization layers caused a large performance drop for TinyLlama and Pythia, whereas averaging and single-layer selection performed similarly.The findings indicate that either averaging or selecting one normalization layer was sufficient for maintaining performance in these comparisons.Source: F. Expanded Results of Initialization… · p. 28, Figure F.1 · p. 29

Overall comparison of few-shot performance

Evaluation scope

  • Table F.1 reports perplexity on three language-modeling datasets and accuracy on seven few-shot benchmarks, with deltas measured against random initialization.It also records whether pretrained weights are used and the number of uptraining tokens.Source: Table F.1 · p. 30

Comparison across various base model sizes

  • Stepwise initialization remained consistently superior across recursive conversions at both 1B and 2B scales and was additionally tested at 410M and 160M.The smaller Pythia experiments use the original pretrained model as the comparison baseline because pretraining and uptraining corpora match.Source: F. Expanded Results of Initialization… · p. 32, Table F.2 · p. 32

Individual contributions

G. Expanded Results of Relaxed Recursive Transformers

SVD-initialized LoRA relaxation improves Recursive Transformer performance, with the Average initialization generally strongest and higher ranks usually helping. The results also show that rank allocation and initialization method materially affect the performance–size trade-off.

Training perplexity

  • SVD-initialized LoRA reduced training loss most clearly with Average initialization, indicating that this configuration more readily learns the difference between shared and original weights.The comparison used two-block recursive and relaxed recursive models with rank 512.Source: G. Expanded Results of Relaxed… · p. 33, Figure G.1 · p. 33

Initialization comparison

  • SVD initialization amplified LoRA-relaxation gains relative to zero initialization, except for the Stepwise method, where it reduced performance.The comparison spans three models and marks both gains and cases below the corresponding recursive baseline.Source: G. Expanded Results of Relaxed… · p. 33, Figure G.2 · p. 33

LoRA-rank ablation

  • Higher LoRA ranks improved performance for Average and Lower initialization, whereas Stepwise produced a U-shaped trend with gains at extremely low or high ranks.Intermediate Stepwise ranks yielded less accurate pretrained-weight approximations and further performance decreases.Source: G. Expanded Results of Relaxed… · p. 34, Figure G.3 · p. 34

Rank allocation

  • Performance correlated strongly with overall model size when ranks were assigned separately to query, key-value, output, and FFN linear weights.The FFN rank is especially consequential because its linear weights have a large hidden dimension.Source: G. Expanded Results of Relaxed… · p. 34, Table G.1 · p. 35

Overall performance comparison

  • The Average looping initialization with SVD-initialized LoRA consistently outperformed other configurations, and performance improved as LoRA rank increased.This comprehensive comparison covers three models, multiple looping initializations, and varying relaxation levels.Source: G. Expanded Results of Relaxed… · p. 35, Table G.2 · p. 36

H. Alternative Approaches to Relaxation of Parameter Sharing

The paper evaluates prefix tuning as a computationally cheaper alternative to LoRA-based relaxation. Although prefix tuning reduces computational overhead, its small learnable prompts provide limited gains and perform substantially worse without pretrained-weight initialization.

Prefix tuning

  • Prefix tuning offered computational advantages but produced limited performance gains because it relied on small, learnable prompts.The approach was explored as an alternative to sequentially computing multiple LoRA modules and base layers.Source: H. Alternative Approaches to… · p. 37
  • Without leveraging original pretrained weights, prefix tuning reached 47.6% versus 52.1% with the Average method in the 1.07B model.The reported comparison highlights the importance of retaining pretrained-weight information for this alternative relaxation method.Source: H. Alternative Approaches to… · p. 37
  • Prefix tuning encodes layer-specific trainable vectors that generate key-value states in each self-attention layer.Prefix length denotes the sequence length of these trainable vectors, and the added prefixes contribute to non-embedding parameter size.Source: Table H.1 · p. 37

I. Expanded Results of Extended Uptraining and Knowledge Distillation

Longer uptraining and knowledge distillation substantially improve Recursive Transformer few-shot performance across architectures and loop counts. The final configuration uses 60 billion tokens, a full-size teacher, and forward KL distillation.

Training setup

  • The study increased uptraining from 15 billion to 60 billion tokens and used each architecture’s full-size model as the distillation teacher.The teacher was itself uptrained on 15 billion SlimPajama tokens before distillation.Source: §I · p. 38

Distillation ablation

  • Forward KL was selected as the final distillation loss because it was simpler and outperformed the alternative FKL, RKL, JSD, and TVD configurations.Table I.1 reports improvements for longer training and the tested distillation objectives.Source: Table I.1 · p. 38

Overall performance

  • Longer uptraining and knowledge distillation consistently produced substantial few-shot improvements across architectures and different numbers of looping blocks.The evaluation compares the resulting models using optimal configurations for each model size and displays their size–performance Pareto frontiers.Source: Figure I.1 · p. 38, §I · p. 39

Final evaluation

  • The final evaluation uses forward KL and compares longer-uptrained recursive models against full-size Gemma and Pythia baselines also uptrained on 60 billion tokens.The baselines account for distribution shifts between SlimPajama and the models’ original pretraining datasets.Source: Table I.2 · p. 39

J. Expanded Results of Early-Exit Training

Early-exit training can preserve final-loop performance while improving intermediate outputs, especially with aggressive loss weighting and final-output distillation. However, LoRA relaxation creates a trade-off because its intermediate outputs underperform those of non-relaxed recursive models.

Early-exit training strategy

  • Post-training after uptraining preserved final-loop performance, whereas co-training intermediate outputs during uptraining significantly degraded it.Training LoRA modules on frozen classifier heads was ineffective, and overlapping versus new SlimPajama tokens made no significant difference in post-training.Source: J. Expanded Results of Early-Exit… · p. 40

Training procedure

  • The adopted procedure post-trains on new SlimPajama tokens, distills from the final loop output, and aggressively downweights intermediate-output losses.This configuration was selected to improve intermediate outputs while maintaining final-loop performance.Source: J. Expanded Results of Early-Exit… · p. 40

Final early-exit results

  • Aggressive coefficient weighting with distillation from detached final outputs yielded the best intermediate and final outputs in early-exit training.The final models used 60 billion-token uptraining followed by knowledge distillation, with comparisons across three model families.Source: J. Expanded Results of Early-Exit… · p. 40

Limitation of relaxed models

  • Intermediate outputs of LoRA-relaxed models underperformed non-relaxed recursive models, potentially reducing throughput gains because early-exit eligibility depends on intermediate performance.The paper attributes this trade-off to depth-specific LoRA specialization toward final-output quality and leaves relaxed-model strategies for future work.Source: J. Expanded Results of Early-Exit… · p. 40

Ablation design

  • Early-exit training was evaluated with static exits after the first or second iteration loops, including post-training, co-training, frozen-weight, and classifier-head LoRA variants.The static exits correspond to depths 9 and 18 in the recursive Gemma experiments.Source: J. Expanded Results of Early-Exit… · p. 40, Table J.1 · p. 41

K. Expanded Results of Hypothetical Generation Speedup

The appendix estimates generation speed and hypothetical early-exit throughput across model sizes, contexts, and memory constraints. Recursive Transformers achieve 2–3× throughput gains, while LoRA overhead reduces gains in shorter, memory-constrained settings.

Measurement setup

Model-speed comparison

  • Gemma 2B was fastest, followed by TinyLlama 1.1B and Pythia 1B, because grouped- and multi-query attention reduce key-value-cache memory access.The speed ordering was the inverse of the models’ non-embedding parameter sizes.Source: K. Expanded Results of Hypothetical… · p. 44

Hypothetical generation throughput

Relaxed-model overhead

  • Relaxed models also speed up inference, but merged LoRA computations reduce gains under shorter contexts and lower memory limits.LoRA computation occupies a larger fraction of total processing time when attention becomes less expensive at shorter sequence lengths.Source: K. Expanded Results of Hypothetical… · p. 44

Approximation errors and limitations

  • The throughput estimates remain theoretical because exit trajectories were simulated, oracle early exiting was used, and only Transformer-block computation was analyzed.Realistic confidence-based exiting, confidence-estimation costs, embedding and classifier processing, and key-value-cache handling could change realized efficiency.Source: K. Expanded Results of Hypothetical… · p. 45

Related papers

Loading 2410.20672v3…