Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA (2410.20672v3)

Abstract: 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.

• The paper presents a Recursive Transformer architecture that shares parameters via a cycle strategy, significantly reducing model complexity.
• The paper introduces Relaxed Recursive Transformers which integrate depth-wise LoRA modules to inject layer-specific flexibility without compromising compactness.
• The paper details a Continuous Depth-wise Batching technique that boosts inference throughput by 2-3x, enabling efficient deployment of LLMs.

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

This paper explores innovative methodologies for enhancing the efficiency of LLMs by employing parameter sharing through the development of Recursive Transformers. Despite the proven efficacy of larger models in delivering superior performance, their extensive computational and memory demands pose significant challenges. Therefore, this research revisits the concept of "layer tying" and introduces Recursive Transformers as a potential solution for maintaining model performance while reducing costs.

Key Contributions and Methodologies

1. Recursive Transformer Architecture: The paper introduces Recursive Transformers, reducing model parameters by sharing them across layers. This involves repeating a single block of unique layers multiple times, forming a looped structure known as "CYCLE strategy." Importantly, this mechanism allows for initialization from standard pretrained Transformers, maintaining performance while slimming down model size.
2. Relaxed Recursive Transformers: To mitigate performance loss from parameter sharing, the researchers propose Relaxed Recursive Transformers. By incorporating depth-wise low-rank adaptation (LoRA) modules, the approach injects layer-specific flexibility without compromising the model's compactness. These modules capture low-rank deltas between shared layers, ensuring robustness through efficient adjustments.
3. Continuous Depth-wise Batching: A novel inference paradigm, Continuous Depth-wise Batching, is suggested. This technique leverages the recursive nature of Recursive Transformers, enabling dynamic grouping across layers and promoting significant gains in inference throughput (up to 2-3 times).
4. Initialization Techniques: The paper proposes several innovative initialization strategies, including the Stepwise, Average, and Lower methods, for effectively priming the shared layer blocks using pretrained model weights. Initialization via averaged weights using truncated Singular Value Decomposition (SVD) significantly enhances performance.
5. Theoretical Analysis: Comprehensive analysis suggests that these models can recover much of the performance of a full-size model, such as the case where a recursive Gemma 1B outperforms its non-recursive counterparts.
6. Training Techniques: Demonstrating the benefits of extended uptraining and knowledge distillation, the Recursive Transformers can achieve high performance with fewer resources.

Implications and Future Directions

This research contributes to both theoretical and practical advancements in model efficiency. By introducing a method to convert existing LLMs into Recursive Transformers, it paves the way for potentially scalable solutions in deploying deep learning models in resource-constrained environments. Practical implications involve reducing the computational costs and enabling larger batch sizes in serving environments.

Future research directions could involve scaling Recursive Transformers to larger LLM frameworks, further exploring batched inference paradigms, and refining the adaptation process through LoRA to enhance performance further. Additionally, significant potential exists in developing robust confidence-based early-exiting strategies to optimize throughput gains.

Overall, this paper establishes a compelling foundation for effectively managing the memory and computational constraints of LLMs, while still achieving competitive performance, marking a notable step forward in model efficiency practices within the domain of artificial intelligence.