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

Authors: Sangmin Bae, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Seungyeon Kim, Tal Schuster
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: Efficient LLM, Parameter Sharing, Relaxed Recursive Transformer, Continuous Depth-wise Batching, Early-Exiting

Abstract:
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-3$\times$) gains in inference throughput.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730405073492

thread_type:
official_review

<thread_object>
creator: reviewer_lz2e

<content>
summary:
This paper revisits the concept of "layer tying" as a form of parameter sharing in Transformers for developing smaller models. These models utilize a single block of unique layers repeated in a loop derived from pre-trained models to reduce model size while maintaining high performance. The paper presents Relaxed Recursive Transformers incorporating depth-wise LoRA modules to enhance these recursive models. These modules add flexibility to the layer-tying process but maintain the model's compactness. Furthermore, the paper proposes a new inference technique, Continuous Depth-wise Batching, which leverages the recursive nature of these models combined with early exiting strategies to increase throughput by 2-3 times.

strengths:
1. The introduction of Recursive Transformers with a single repeated layer block offers a novel approach to reducing model size and complexity without substantial loss in performance.

2. The incorporation of Relaxed Recursive Transformers with LoRA modules successfully improves flexibility and performance.

3. The proposed Continuous Depth-wise Batching method, which takes advantage of the recursive model structure and early exiting, increases inference throughput.

weaknesses:
1. The approach of uptraining a full-size pre-trained model is questionable, especially given that the goal is to achieve comparable performance to the original model. In practical applications, deployment considerations often prioritize model size rather than training data size or consumption. Using a smaller, lower-quality dataset for these experiments may unfairly favor smaller (1B) models, potentially skewing the results and making the claimed improvements less persuasive. 

2. Furthermore, this uptraining step could be prohibitive for settings where on-device retraining is required. The reliance on extensive uptraining, which consumes 60 billion tokens, presents challenges for practical deployment, particularly when computational resources are constrained.

3. The current methodology appears to train reduced models from scratch, using the same number of tokens, while recursive transformers are initialized with pre-trained layers. This discrepancy in training procedures introduces a potential bias in favor of the recursive transformers.

4. The paper lacks a clear rationale for determining the number of initial layers to be recursively repeated and the frequency of such recursions. Furthermore, an automated approach for selecting the optimal recursion size would greatly increase the novelty and technical contribution of the proposed method.

5. All experiments are conducted with models around the 1B parameter mark, yet hardware advancements make models larger than 1B feasible for many applications. Testing the proposed approach on larger models is essential to validate its scalability and effectiveness beyond the 1B parameter range, ensuring the method’s relevance for state-of-the-art large models. 

6. The selection of layers for initializing looped layers, as indicated in Figure 5, appears to depend heavily on the specific model used. A clear explanation of how to automate this layer selection process would substantially improve usability and consistency across different models. Without this, the method’s applicability may be limited by model-specific manual tuning. 

7. The impact of using different base models (other than a 2B model) for recursive transformer initialization has not been evaluated. Testing a range of base model sizes would provide insight into how the choice of initial model affects recursive initialization outcomes, offering a more thorough assessment of the method’s robustness.

8. The experiments focus exclusively on few-shot tasks, but in real-world applications, 1B models are often fine-tuned for specific tasks. Testing the proposed method on downstream QA or classification tasks with full training on task-specific data would provide a more realistic evaluation of its performance and applicability across different use cases.

questions:
Please refer to the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730562290278

thread_type:
official_review

<thread_object>
creator: reviewer_ud6h

<content>
summary:
This paper introduces a novel framework for parameter sharing in LLMs, namely relaxed recursive transformers. The method revist the way of layer tying and further convert LLMs into recursive transformers by sharing parameters across layers. Moreover, the performance is further improved by relaxing the layer tying constraint. The method is demonstrated to be well performed on various LLMs, it also get significant throughput gains by paring with early exiting, namely continuous depth-wise batching.

strengths:
1. The paper is well motivated on a hot topic in the field of LLMs.
2. The paper is well organized and the insight of the method is clearly stated.
3. The experimental results are convincing. I appreciate the way the authors present the takeaways for relaxed recursive transformers.

weaknesses:
1, It lacks evaluation on larger scales, which makes the scalability of the method remain questionalbe. For instance, I expect results on LLMs with 7B, 70B or larger scales.

questions:
1. I wonder the performance of relaxed recursive transformers on larger scales.
2. How does the authors compare the proposed parameter sharing methods with other model deployment optimization methods, such as quantization, sparcification，pruning or distillation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711991019

thread_type:
official_review

<thread_object>
creator: reviewer_9eth

<content>
summary:
The paper presents an innovative approach to parameter efficiency in Large Language Models through Recursive Transformers, by iteratively reusing a single block. The concept is extended to Relaxed Recursive Transformers by incorporating layer-specific LoRA modules, introducing flexibility to the parameter-sharing scheme. The authors demonstrate the conversion of pre-trained LLMs into more compact recursive versions while maintaining performance and achieving improved throughput through depth-wise batching.

strengths:
- The core concept of recursive layer sharing combined with LoRA-based relaxation is straightforward yet powerful.
- The paper provides extensive empirical validation across multiple model architectures, recursive layer initialization strategies, LoRA initialization strategies, and evaluation metrics, including thorough ablation studies and practical deployment.
- The authors not only show the benchmarking results of Relaxed Recursive Transformers but also explore hypothetical training speedup with continuous depth-wise batching strategy and early-exiting.

weaknesses:
**Major**
- It compares to different-sized dense baselines. However, the evaluation lacks comparison with sparse designs, particularly layer-skipping mechanisms, and layerwise parameter-sharing approaches [1-2].
- The paper does not fully investigate how recursive architectures affect model training dynamics and optimization landscapes. The relationship between shared parameters and gradient flow during the uptraining phase. Also, whether shared parameter amounts affect model learning capacity could be explored.
- The study would benefit from experiments on optimal recursion depths across varying model scales and task domains.

**Minor**
- The experiments are primarily conducted on relatively small models with 1-2B parameters. The effectiveness and scalability of the approach on larger models remain unproven. This raises questions about whether the benefits would hold at scales where current LLMs demonstrate more sophisticated capabilities.
- Typos: "achieved a 10%p performance" on page6.

[1] Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

[2] Zamba: A Compact 7B SSM Hybrid Model

questions:
- The choice of LoRA for relaxing parameter sharing appears somewhat arbitrary. And prefix-tuning results show limiting performance gain. Could you elaborate more on the relaxation technique selections and other alternatives (*e.g.* Adapter, IA3)?
- For Section 3.8, how would the throughput gains hold up when using practical confidence-based early-exiting rather than the idealized oracle approach that assumes perfect prediction matching?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730780802912

thread_type:
official_review

<thread_object>
creator: reviewer_wb6t

<content>
summary:
The paper introduces a new approach to reducing the deployment costs of large language models by implementing recursive parameter sharing. The proposed Recursive Transformer model reuses a single block of unique layers across multiple iterations, effectively minimizing model size while maintaining strong performance. To enhance flexibility, Relaxed Recursive Transformers incorporate low-rank adaptation (LoRA) modules at each layer, allowing for slight layer-specific modifications that improve performance with minimal overhead. Additionally, the authors propose a Continuous Depth-wise Batching technique, which leverages the recursive architecture to increase inference throughput by enabling multiple depths to be processed simultaneously, achieving up to a 2-3× improvement in efficiency over traditional models.

strengths:
1. Propose a novel model structure that can effectively compress pre-trained large language models for deployment.
2. LoRA adaptation makes the model structure more flexible to adapt to different depths.
3. Depth-wise Batching translates the advantages of the proposed recursive transformer block to real deployment acceleration.

weaknesses:
1. Experiment design is not so convincing, which can be further improved.
2. Authors do not push the recursive computation block’s advantage to the limit, which can be interesting to explore
3. Some highly-related papers are missing in the reference

questions:
1. In experiment design (Line 268), authors mention “As a result, we set the target performance for Gemma and Pythia models as the performance achieved by uptraining a full-size pretrained model with an equivalent number of tokens.”. This design is unreasonable, especially since your training (SlimPajama) and testing datasets (llm harness evaluation) are different. You need to provide the original LLM’s performance to help readers understand the real performance of the proposed recursive transformer instead of fine-tuning a worse model to mislead readers.
2. To test the effectiveness of your model compression method, it is highly recommended to achieve a larger compression ratio. E.g. compressing 7B model to 1B. Current experiments only compress 2B models to 1B, which lacks practical meaning as they are all models with similar sizes.
Reusing the transformer computation block has been explored before [1] without properly cited within this paper.

[1] Liu, Zechun, Changsheng Zhao, Forrest Iandola, Chen Lai, Yuandong Tian, Igor Fedorov, Yunyang Xiong et al. "Mobilellm: Optimizing sub-billion parameter language models for on-device use cases." arXiv preprint arXiv:2402.14905 (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732179456102

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Thank you all for taking the time and effort to review our paper and provide thoughtful and constructive feedback. We have carefully addressed each reviewer's comments and uploaded a revised version of the paper, which includes the following additions:

- **Clarification of baseline models:** We have further clarified the rationale behind the full-size and reduced-size non-recursive baselines in Section 3.2.
- **Discussion of sparse designs:** Appendix A now includes a discussion of how sparsity-based approaches (pruning, quantization, and layer-skipping) compare to our proposed methods and how they might be integrated with our structures.
- **Analysis of individual effects in Recursive Transformer:** We conducted additional experiments on Pythia 410M to isolate and demonstrate the individual effects of leveraging pretrained layers and recursive patterns (see Appendix O).
- **Evaluation of initialization methods in other model scales:** To confirm the robustness of our proposed initialization methods, we expanded our experiments to include two different model scales---Pythia 410M and 160M (see Appendix P).
- **Missing references:** We have added missing references related to parameter sharing and relaxation techniques.

Thank you again for your valuable feedback.  We appreciate your time and insights. If you have any further questions or concerns, please feel free to share them. We are happy to engage in further discussions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes relaxed recursive transformers, a new approach to parameter sharing in LLMs by iteratively reusing a single block of layers. The incorporation of LoRA enhances model flexibility without significant performance loss. Additionally, the paper introduces continuous depth-wise batching for improved inference throughput. Reviewers praised the innovative combination of recursive architecture and LoRA, as well as the comprehensive experiments. I recommend acceptance.