DLO: Dynamic Layer Operation for Efficient Vertical Scaling of LLMs
Abstract
In this paper, we introduce Dynamic Layer Operations (DLO), a novel approach for vertically scaling transformer-based Large Language Models (LLMs) by dynamically expanding, activating, or skipping layers using a sophisticated routing policy based on layerwise feature similarity. Unlike traditional Mixture-of-Experts (MoE) methods that focus on extending the model width, our approach targets model depth, addressing the redundancy observed across layer representations for various input samples. Our framework is integrated with the Supervised Fine-Tuning (SFT) stage, eliminating the need for resource-intensive Continual Pre-Training (CPT). Experimental results demonstrate that DLO not only outperforms the original unscaled models but also achieves comparable results to densely expanded models with significantly improved efficiency. Our work offers a promising direction for building efficient yet powerful LLMs. We will release our implementation and model weights upon acceptance.
[ linewidth=2pt, roundcorner=10pt, linecolor=black, backgroundcolor=gray!10, skipabove=5pt, skipbelow=5pt ]custombox
DLO: Dynamic Layer Operation for Efficient Vertical Scaling of LLMs
Zhen Tan♠††thanks: Equal contribution. Daize Dong♥11footnotemark: 1 Xinyu Zhao♠ Jie Peng❧ Yu Cheng♣ Tianlong Chen♠ ♠School of Computing, and Augmented Intelligence, Arizona State University ♥Shanghai Artificial Intelligence Laboratory ♣Department of Computer Science, Chinese University of Hong Kong ❧School of Artificial Intelligence and Data Science, University of Science and Technology of China ♠Department of Computer Science, University of North Carolina at Chapel Hill [email protected] [email protected] [email protected] {xinyu,tianlong}@cs.unc.edu [email protected]
1 Introduction
![Refer to caption](x1.png)
Large Language Models (LLMs) Achiam et al. (2023); Team et al. (2023) have shown remarkable success across various natural language processing (NLP) tasks Hadi et al. (2023); Tan et al. (2024); Li et al. (2024b, a), leveraging their vast capacity to capture complex patterns in data. Traditional scaling of these models has predominantly focused on horizontal expansion, as seen in Mixture-of-Experts (MoE) architectures Shazeer et al. (2017); Fedus et al. (2022b); Lepikhin et al. (2020), where the width of the model is increased by adding more experts. This approach primarily optimizes parameter usage and computational cost by activating a fixed portion of parameters conditioned on the given input Fedus et al. (2022a).
However, the potential for vertical expansion remains underexplored. Inspired by how the human brain allocates more neurons for complex tasks and forms deeper neural chains Baddeley (1992); Koechlin et al. (2003), we propose focusing on vertical scaling. Our method dynamically expands, activates, or skips layers to optimize model depth and reduce redundancy, as shown in Figure 1 (a).
There are three critical challenges in vertically scaling LLMs: ❶ Optimization Complexity. Dynamically adding or pruning layers making the process hard to optimize. Obtaining optimal such operations have been proved to be a NP-hard problem Glorot and Bengio (2010); Hestness et al. (2017), while an approximation method Wang et al. (2023a) has shown compromised improvement. ❷ Computation Cost. The inherent computational cost is associated with processing deeper networks. Each additional layer contributes to the overall latency and resource consumption. ❸ Feature Collapse. Our analysis in Figure 1 (b) reveals that for a significant number of inputs, the representations across consecutive layers exhibit substantial similarity, suggesting that many layers may be redundant for certain samples.
To address these challenges, in this paper, we propose Dynamic Layer Operation (DLO), that consists of three operations: (i) expansion, (ii) activation, and (iii) skipping, for dynamic vertical scaling of LLMs without a proportional increase in computational cost. Our specific designs are as follows: ❶ Expansion: Additional layers are dynamically expanded from existing ones, easing optimization complexity. ❷ Activation & Skip: Feature Similarity guides the activation and skipping of layers. We propose similarity-induced labels to train the router that controls these operations. ❸ Adaptive FLOPs: Sparsity settings vary for layers facilitate adaptive FLOPs for different tokens, maintaining efficiency. ❹ Enhanced Generalizability: Layer-specific learning rates, based on sparsity, further improve the model’s ability to generalize across tasks. Note that all modules are trained during the Supervised Fine-Tuning (SFT) stage, eliminating the need for Continual Pre-training (CPT) and simplifying the training process. Our primary contributions are as follows:
-
•
Method. We introduce a novel method, DLO, for dynamically scaling LLMs vertically by dynamically expanding, activating, or skipping layers.
-
•
Performance & Efficiency. Through rigoerous experiments, we demonstrate that DLO not only surpasses the performance of the original unscaled models but also achieves comparable results to densely expanded models with significantly enhanced efficiency.
-
•
Applicability. Fine-tuned on language understanding, math, and coding tasks, we manifest DLO’s effectiveness across multiple NLP tasks.
2 Related Work
2.1 Mixture-of-Experts (MoE)
MoE architectures have emerged as a promising approach for enhancing the efficiency and scalability of LLMs Shazeer et al. (2017). Traditional neural networks activate all parameters for every input, leading to significant computational overhead, particularly as models scale up. In contrast, MoE models activate only a subset of parameters for each input, optimizing computational resource usage and enabling models to scale to billions of parameters without a corresponding increase in computational cost per input Lepikhin et al. (2020); Fedus et al. (2022b); Zoph et al. (2022); Team (2023a); jiang2024mixtral,zhu2024llama. This selective activation makes MoE highly efficient for both training and inference by focusing on horizontal expansion and adding more experts. However, MoE’s primary aim is to optimize width, potentially leaving layer redundancy unaddressed. Our Dynamic Layer Operation (DLO) approach complements MoE by focusing on vertical scaling through dynamic layer expansion and activation, targeting depth scalability and reducing potential feature redundancy.
2.2 Efficient Model Stacking
Model stacking is a common ensemble learning technique that improves predictive performance by combining multiple models to leverage their complementary strengths Ting and Witten (1997); Chen et al. (2015). In the context of LLMs, stacking can involve integrating various models into a hierarchical structure, where outputs from one model serve as inputs to another, capturing a broader range of features and patterns Dabre and Fujita (2019); Chen et al. (2021a); Wang et al. (2023b); Kim et al. (2023).
Recent advancements have focused on progressively stacking pre-trained transformer or self-attention layers to create composite language models Gong et al. (2019); Gu et al. (2020); Shen et al. (2022); Evci et al. (2022); Yao et al. (2023); Du et al. (2024); Wu et al. (2024). This approach reduce training costs by reusing pre-trained components. However, the increased depth and complexity of stacked models lead to high inference latency.
To mitigate this issue, layer-skipping methods have been developed, allowing models to “early exit” using additional layer-wise classifiers, thereby reducing the number of layers processed during inference Wang et al. (2022); Chen et al. (2023); Zhang et al. . More recently, conditional computation techniques have been proposed to dynamically skip layers based on token-specific conditions, further enhancing efficiency Ainslie et al. (2023); Raposo et al. (2024). However, these methods often require modifications during the pre-training stage, adding computation complexity and limiting their application to existing pre-trained LLMs. In contrast, our DLO method focuses on efficiency and scalability through dynamic vertical scaling within a single model during the SFT stage. It provides a comprehensive, high-performance solution to scaling LLMs without the extensive computational demands associated with stacked ensembles.
3 Methodology
![Refer to caption](x2.png)
In this section, we introduce the Dynamic Layer Operation (DLO) framework for efficienct vertical scaling of LLMs. DLO consists of three key operations: expansion, activation, and skipping. These operations dynamically adjust the model structure during the Supervised Fine-Tuning (SFT) phase to optimize computational efficiency and improve performance. A pseudo code style description is included in Appendix A.
3.1 Layer Expansion
To facilitate dynamic depth adjustment, we introduce a group-based layer expansion strategy. Suppose the LLM has transformer layers, which we group into groups with layers each, such that . Each group is expanded to layers, where is the number of additional layers introduced per group. The resulting number of layers will be .
Let denote the -th group with layers . The expanded group will contain layers . The expanded layers are initialized using a policy , and we consider several initialization strategies:
-
•
Random Initialization (): Initialize the new layers’ weights using Xavier initialization Glorot and Bengio (2010).
where , denotes the uniform distribution, is the number of input units, and is the number of output units in the layer.
-
•
Copy from Previous Layer (): Copy the parameters from the preceding layer.
-
•
Identity Initialization () Wu et al. (2024): Copy from the preceding layer but set the output linear matrix of the multi-head self-attention (MHSA) to zero.
-
a.
Copy the parameters of the previous layer:
-
b.
Set the weights of the output linear layer in the MHSA to zero:
In this way, the output of the expanded layers will preserve the features from the original layers.
-
a.
-
•
Linear Merge (): Merge from the preceding layers using a linear function.
-
•
Spherical Linear Interpolation (SLERP) () Shoemake (1985): Merge from the preceding layers using SLERP. The SLERP method smoothly interpolates between two weight vectors on a unit sphere, maintaining constant velocity. The interpolation between two weight vectors and is defined as:
where is the angle between and :
and is the interpolation parameter.
In our context, for the new layer , the parameters are initialized by interpolating between the weights of the previous layers and :
where controls the interpolation. This ensures a smooth transition between layers, aiding in gradient flow and stable training.
We conduct comprehensive experiments on the choice of the policy in Section 4.3.
3.2 Layer Activation & Skipping
DLO dynamically skips the multi-layer perceptron (MLP) module within the transformer layer for input tokens. To achieve this, we uses a linear router to determine the activation of layers. Suppose we have the set of token embeddings in a sequence of length for a given layer , that is , where in following contents we omit the superscript for better readability. Considering feature redundancy, we use the router weights to process the input token and obtain the decision score , which is given by:
(1) |
where is the sigmoid function, and are hyperparameters controlling the output range. During the inference stage, this score determines whether layer is active: the layer is activated if and only if , otherwise it’s skipped. The final activated output for the layer is:
(2) |
where , are the MLP and the attention modules within layer , respectively. This activation & skipping mechanism aims to encourage the utilization of the most relevant layers, thus reducing unnecessary computation. In this paper we initialize the router weights as zeros and set , so that on the first step and during training. This ensures benign initialization for activated tokens and avoids excessive disturbance on activated outputs brought by the decision scores.
3.3 Training and Integration
![Refer to caption](x3.png)
Similarity-induced Label & Router Skip Loss.
Given a pre-defined overall spasity , we define a sparsity factor for each layer that controls the layer-wise actived tokens. To determine the status of token , we utilize predicted router labels , which is obtained through the decision scores :
(3) |
where indicates layer is predicted to be activated, and vice versa. To train the routers, we utilize the supervised router labels to guide the learning of sparsity, i.e., to skip or not. The router labels of layer at training step is determined by the following procedures:
-
1.
The cosine feature similarity of layer is calculated through features across MLP :
(4) -
2.
The similarity are sorted over all the layers, and the similarity-induced label is given as follows:
(5) where indicates the labels for tokens with the least portion of cosine similarity are set to s, which are expected to be activated.
-
3.
A skip loss based on the Binary-Cross-Entropy loss is incorporated to guide the learning of the router attached to each layer:
(6)
Given the task-specific loss , the overall loss function for DLO training is:
(7) |
Skip Rate Dynamics.
The redundancy exhibits an imbalanced distribution across layers, as is evidenced in Figure 1 (b). To this end, we adjust the next-step for each training step over the total steps, where the initial skip rate . The layer-wise sparsity factor is calculated using the router labels as follows:
(8) |
where is the supervised router label of the -th token in layer at step . Additionally, we employ an annealing technique on the skip rate to ensure the warm start. During training, the overall skipping rate gradually increases from an initial low value to the target sparsity level over a predefined number of steps . The overall skip rate at step is given by:
(9) |
where we set . This annealing process helps the model to progressively adapt to higher sparsity levels with smoother changes, leading to more stable training and better convergence.
Layer-Wise Learning Rates.
DLO also employs layer-wise learning rates , adjusted based on sparsity to promote generalizability. The learning rate for each layer is defined as:
(10) |
where is the base learning rate. It is noteworthy that all DLO components are trained during the Supervised Fine-Tuning (SFT) stage in an end-to-end manner, eliminating the need for Continual Pre-Training (CPT). By integrating DLO, we achieve dynamic vertical scaling, optimizing model depth, and maintaining high performance with reduced computational demands.
3.4 Adaptive Inference-Time FLOPs
During inference time, DLO uses layer-specific sparsity settings to maintain computational efficiency and ensure adaptive floating-point operations (FLOPs) for different tokens. In other words, the predicted sparsity will be determined completely by the router based on Equation (1)-(2). The adaptive FLOPs are computed as:
(11) |
where represents the FLOPs for a fully active layer. Since is predicetd based on each specific token, DLO acheive adaptive FLOPs that entails better generalizability.
4 Experiments
In this section, we present an empirical evaluation of the proposed DLO framework, detailing the experimental settings, results, and analysis.
4.1 Experimental Settings
Model Selection. We utilize LLaMA2-7B Touvron et al. (2023) as the primary backbone due to its open-source availability and extensive usage. It consists of original transformer layers, which we group into clusters, each containing layers. For layer expansion, we increase the group size to layers, resulting in a dense model, LLaMA-DLO, with a total of 40 layers and 8 billion parameters. For comparison, we also employ LLaMA-Pro-8B Wu et al. (2024), a competitive model trained with Continual Pre-Training (CPT) on specialized datasets. We demonstrate DLO achieves an optimal balance between performance and computational cost in Section 4.4.
Fine-tuning Details. Following common practices Wu et al. (2024), we fine-tune using a mixture of five instruction tuning datasets: ShareGPT Team (2023b), EvolInstruct Luo et al. (2023), SlimOrca Team (2023c), MetaMath Yu et al. (2023), and Evol-CodeAlpaca Team (2022), with ShareGPT replicated three times, totaling approximately 1.44 million instances. We use a batch size of 128 and a maximum sequence length of 4,096 tokens. The learning rate is set to with a warmup ratio of 0.03 and cosine scheduling, and we utilize AdamW Loshchilov and Hutter (2017) as the optimizer. Flash Attention Dao et al. and bfloat16 mixed-precision training are adopted to accelerate training. Fine-tuning LLaMA-DLO under different skip ratios yields the sparse models, with each training run taking approximately 36 hours on eight NVIDIA A100 GPUs.
Evaluation Benchmarks. We assess the fine-tuned models using the EleutherAI LM Harness Gao et al. (2023) and BigCode Harness Ben Allal et al. (2022) across three domains: ❶ Language [ARC-C Clark et al. (2018), GLUE Wang et al. (2018), MMLU Hendrycks et al. (2020), OBQA Mihaylov et al. (2018), PIQA Bisk et al. (2020), SQuAD Rajpurkar et al. (2016), TruthfulQA Lin et al. (2021), WinoGrande Sakaguchi et al. (2021)], ❷ Math [GSM8K Cobbe et al. (2021), MathQA Amini et al. (2019)], and ❸ Code [HumanEval Chen et al. (2021b), MBPP Austin et al. (2021)]. Detailed metrics are in Appendix B.
Sparsity | Model | FLOPs | Sprache | Math | Code | Avg. | |||||||||
ARC-C | GLUE | MMLU | OBQA | PIQA | SQuAD | TruthfulQA | WinoGrande | GSM8K | MathQA | HumanEval | MBPP | ||||
0% | LLaMA2-7B | 29.3T | 53.1 | 40.6 | 46.9 | 44.2 | 79.0 | 26.4 | 38.8 | 74.0 | 14.5 | 28.3 | 21.8 | 29.0 | 41.38 |
LLaMA2-7B | 54.0 | 72.4 | 53.0 | 44.4 | 78.8 | 22.9 | 40.8 | 74.2 | 56.6 | 30.8 | 57.3 | 30.5 | 51.31 | ||
\cdashline2-16 | LLaMA-Pro-8B | 36.4T | 54.1 | 40.7 | 47.9 | 41.6 | 78.2 | 14.2 | 39.0 | 74.0 | 17.9 | 29.5 | 28.7 | 33.2 | 41.58 |
LLaMA-Pro-8B | 51.0 | 71.0 | 53.0 | 45.0 | 79.0 | 15.1 | 38.0 | 73.6 | 58.6 | 30.8 | 58.4 | 30.5 | 50.33 | ||
\cdashline2-16 | LLaMA-DLO-8B | 36.5T | 53.2 | 75.5 | 53.2 | 43.7 | 79.0 | 22.0 | 38.7 | 74.0 | 57.4 | 31.0 | 57.0 | 30.2 | 51.24 |
10% | LLaMA2-7B | 27.5T | 51.0 | 72.5 | 51.1 | 40.2 | 78.0 | 21.0 | 39.3 | 71.0 | 53.4 | 29.6 | 49.7 | 28.3 | 48.76 |
LLaMA-DLO-8B | 34.2T | 52.5 | 75.4 | 51.4 | 43.6 | 78.7 | 21.7 | 41.0 | 73.4 | 55.0 | 31.4 | 55.3 | 29.1 | 50.71 | |
20% | LLaMA2-7B | 25.8T | 33.0 | 69.9 | 50.4 | 35.0 | 65.6 | 16.4 | 36.9 | 54.2 | 1.0 | 24.5 | 0.0 | 0.0 | 32.24 |
LLaMA-DLO-8B | 32.0T | 51.2 | 73.3 | 50.8 | 43.2 | 78.2 | 20.4 | 38.9 | 73.2 | 50.1 | 30.5 | 57.6 | 28.2 | 49.63 | |
30% | LLaMA2-7B | 24.1T | 28.1 | 2.8 | 47.1 | 35.0 | 53.8 | 13.9 | 37.9 | 52.2 | 0.0 | 21.6 | 0.0 | 0.0 | 24.40 |
LLaMA-DLO-8B | 29.8T | 44.6 | 73.0 | 50.1 | 41.2 | 77.1 | 21.7 | 37.1 | 63.2 | 46.9 | 28.1 | 31.2 | 6.0 | 43.35 |
4.2 Overall Performance
Table 1 summarizes the performance of the DLO framework across various datasets using LLaMA2-7B as the backbone. From the results, we draw several key observations are as follows:
❶ Dense Models’ Superiority: Dense models, indicated by , generally outperform their sparse counterparts across most datasets. For instance, LLaMA-DLO models consistently achieve high average scores, such as 51.31 for LLaMA2-7B+SFT, compared to the baseline LLaMA2-7B’s 41.38. This indicates that our DLO-expansion method enhances model performance significantly while leveraging additional parameters effectively. ❷ Efficiency of Sparse Models: Sparse models, marked with , show a notable reduction in inference-time FLOPs while maintaining competitive accuracy. At 10% sparsity, the LLaMA-DLO model achieves an average score of 50.71 with 34.2T FLOPs, compared to the dense model’s 51.31 with 36.5T FLOPs. This demonstrates the efficiency of DLO’s activation and skipping operations in optimizing computational resources without significantly sacrificing performance. ❸ DLO-Expansion Advantages: Models expanded using DLO expansion with up to 8B parameters outperform the original LLaMA2-7B model across multiple metrics. For example, dense LLaMA-DLO-8B+SFT achieves a higher average score of 51.24 compared to LLaMA2-7B’s 41.38, highlighting the effectiveness of vertical scaling through layer expansion in improving model capacity and performance. On the other hand ❹ Balanced Performance of Sparse Models: Sparse models with DLO’s dynamic activation and skipping () provide a well-balanced trade-off between performance and computational efficiency. At 30% sparsity, LLaMA-DLO models maintain strong performance on tasks like GLUE (51.0 vs. 28.1) and HumanEval (50.5 vs. 21.8), while significantly reducing FLOPs. This makes them suitable for scenarios requiring computational efficiency without substantial performance loss. ❺ Effective Inference Optimization: DLO demonstrates effective inference optimization. For instance, the dense LLaMA-DLO model with 8B parameters achieves lower inference-time FLOPs compared to LLaMA-Pro-8B (36.5T vs. 36.4T) while maintaining a competitive average performance (51.24 vs. 50.33). This highlights DLO’s capability to enhance model efficiency without compromising accuracy. ❻ General Observations: Overall, the DLO framework successfully balances performance and efficiency across various datasets and tasks. The adoption of both expansion and skipping strategies enables LLaMA-DLO to achieve robust performance improvements while maintaining lower computational costs, suggesting that DLO is a viable approach for scalable and efficient LLM deployment.
We conduct further analyses of key components of DLO in the subsequent subsections.
4.3 Ablation Studies
Method | Sprache | Math | Code | Avg. | |||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|
ARC-C | GLUE | MMLU | OBQA | PIQA | SQuAD | TruthfulQA | WinoGrande | GSM8K | MathQA | HumanEval | MBPP | ||
Random | 25.1 | 41.4 | 24.4 | 26.8 | 50.3 | 42.2 | 37.7 | 49.3 | 0.3 | 20.3 | 0.0 | 1.3 | 26.6 |
Identity | 52.5 | 75.4 | 51.4 | 43.6 | 78.7 | 21.7 | 41.0 | 73.4 | 55.0 | 31.4 | 55.3 | 29.1 | 50.7 |
Kopieren Sie | 52.0 | 71.3 | 52.3 | 43.2 | 78.6 | 23.3 | 40.8 | 73.0 | 52.2 | 29.0 | 41.1 | 26.7 | 48.6 |
Linear | 48.6 | 70.7 | 53.4 | 39.4 | 72.9 | 18.7 | 38.7 | 57.5 | 29.7 | 25.7 | 24.0 | 0.7 | 40.0 |
Slerp | 42.2 | 71.6 | 49.6 | 39.8 | 76.4 | 22.3 | 36.2 | 63.0 | 43.0 | 27.1 | 40.4 | 4.9 | 43.0 |
Method | Sprache | Math | Code | Avg. | |||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|
ARC-C | GLUE | MMLU | OBQA | PIQA | SQuAD | TruthfulQA | WinoGrande | GSM8K | MathQA | HumanEval | MBPP | ||
DLO | 52.5 | 75.4 | 51.4 | 43.6 | 78.7 | 21.7 | 41.0 | 73.4 | 55.0 | 31.4 | 55.3 | 29.1 | 50.7 |
w/o Zero Init | 51.9 | 73.0 | 52.1 | 44.6 | 77.7 | 20.8 | 40.3 | 74.3 | 55.0 | 30.5 | 56.1 | 28.2 | 50.4 |
w/o Rescaling | 47.4 | 71.3 | 49.4 | 35.8 | 75.3 | 21.1 | 39.6 | 67.5 | 35.0 | 24.5 | 55.6 | 28.2 | 45.9 |
Method | ARC-C | MMLU | TruthfulQA | WinoGrande | GSM8K | Avg. |
---|---|---|---|---|---|---|
SOLAR | 24.8 | 24.8 | 38.8 | 50.7 | 2.2 | 28.3 |
SD-Stack | 23.5 | 23.4 | 36.0 | 51.1 | 2.5 | 27.3 |
\hdashlineDLO | 52.5 | 51.4 | 41.0 | 73.4 | 55.0 | 54.7 |
Initialization Strategies for Expanded Layers.
We explored the effectiveness of various layer initialization strategies for the expanded layers, as detailed in Section 3.1. Table 2 evaluates the impact of different initialization strategies on the performance of LLaMA-DLO models with 10% sparsity. The results highlight that the choice of initialization plays a critical role in determining model performance across various tasks.
❶ The identity and copy initialization strategies demonstrate the most consistent and high-performing results, suggesting that leveraging existing layer information is beneficial for stabilizing and enhancing model performance. These methods help maintain coherence in the model’s internal representations, leading to robust results across a wide range of tasks, including GLUE and HumanEval.
❷ Interestingly, while linear and SLERP initializations were expected to offer smoother transitions and potentially enhance performance, their results were only moderately effective. This indicates that while sophisticated initialization techniques can offer benefits, they may not always outperform simpler strategies like identity and copy initialization, which directly utilize pre-existing model structures.
❸ Random initialization yields the lowest performance. The variability in task performance with this method highlights the challenges of using non-specific weights, which can lead to unstable and suboptimal model behavior, particularly in complex tasks like math and coding.
Overall, the findings emphasize that initialization strategies that leverage prior information from existing layers tend to provide a better foundation for training expanded models, leading to improved performance. We thus choose as the default initialization strategy.
Zeros Router Initialization & Score Rescaling.
In this experiment, we investigate the impact of zeros router initialization and score rescaling on mitigating performance degradation.
Zeros Router Initialization. Initializing the router parameters to zero aims to start the model from a neutral state, avoiding any initial bias towards layer activation or skipping. This method allows the model to learn activation patterns from scratch without being influenced by predefined weights. Results in Table 3 indicate that this approach helps maintain balanced training dynamics and mitigates premature convergence, as reflected in the performance stability observed across tasks.
Score Rescaling. Score rescaling adjusts the routing scores to maintain them within a specific range, typically 0 to 1. This adjustment is intended to preserve gradient flow and prevent extreme activations, ensuring that the model remains responsive to training signals. Our findings suggest that score rescaling helps avoid over-activation of layers, leading to more efficient use of the model’s capacity.
The combined use of zeros router initialization and score rescaling appears to prevent performance degradation effectively. As shown in Table 3, models with these techniques generally achieve more consistent accuracy and efficiency across various tasks. These results suggest that careful initialization and rescaling strategies are beneficial for maintaining robust performance during adaptation.
Method | Sprache | Math | Code | Avg. | |||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|
ARC-C | GLUE | MMLU | OBQA | PIQA | SQuAD | TruthfulQA | WinoGrande | GSM8K | MathQA | HumanEval | MBPP | ||
Uniform | 36.0 | 61.7 | 48.7 | 36.6 | 58.4 | 47.9 | 36.8 | 68.2 | 25.4 | 21.7 | 23.7 | 2.2 | 38.9 |
Layer-Wise | 52.5 | 75.4 | 51.4 | 43.6 | 78.7 | 21.7 | 41.0 | 73.4 | 55.0 | 31.4 | 55.3 | 29.1 | 50.7 |
![Refer to caption](x4.png)
![Refer to caption](x5.png)
Efficient Expansion.
In addition to the high-cost LLaMA-Pro approach (studied in Table 1), we compare our expansion method with two state-of-the-art efficient vertical expansion baselines: SOLAR Kim et al. (2023) and Self-Duplicate Stack (SD-Stack) Team (2024). These two methods duplicate blocks of transformer layers and stack them together in a training-free manner. As shown in Table 4, the proposed DLO-expansion significantly outperforms both SOLAR and SD-Stack by a considerable margin. This highlights the critical role of Supervised Fine-Tuning (SFT) in adapting the expanded layers effectively. Unlike training-free approaches, DLO-expansion achieves a superior balance between training cost and performance, demonstrating the importance of fine-tuning in maximizing the effectiveness of layer expansion.
Layer-Wise Skip Rates & Sparsity Allocation.
This experiment evaluates the impact of using layer-specific skip rates on sparsity allocation.
Layer-Wise Skip Rates. Adjusting skip rates for each layer aims to selectively activate or skip layers based on their contribution to task performance, which is measured by the layer-wise similarity. This method helps focus computational resources on more critical layers. Results in Table 5 suggest that this approach can lead to more efficient sparsity allocation with less impact on model performance. Figure 4 also show that DLO can skip layers that have high layer-wise similarity.
Sparsity Allocation. Tailoring skip rates by layer helps distribute sparsity more effectively, potentially reducing computational overhead. As indicated in Table 5, models with layer-wise skip rates tend to maintain performance while achieving better computational efficiency.
4.4 Scalability
The proposed LLaMA-DLO model surpasses the performance of the original dense LLaMA, while also achieving comparable results to the dense LLaMA-Pro. Notably, it does so at a significantly lower training cost by eliminating the need for expensive CPT. Additionally, LLaMA-DLO facilitates efficient inference through adaptively reduced FLOPs, making it a cost-effective choice for both training and deployment.
Figure 5 illustrates the trade-off between model performance and both training and inference costs. LLaMA-DLO emerges as the optimal solution, achieving the best balance across these metrics. This demonstrates the model’s scalability, ensuring that high performance is maintained while keeping computational costs manageable.
5 Conclusion
This paper presents LLaMA-DLO, a framework for efficient vertical scaling of LLMs that dynamically expands, activates, and skips layers to optimize computational resources. Our experiments demonstrate that LLaMA-DLO achieves performance on par with expensive dense expansion model like LLaMA-Pro, while significantly reducing training costs and enhancing inference efficiency. These results highlight LLaMA-DLO’s potential as a cost-effective solution for scaling LLMs in various NLP tasks, offering a balanced approach between model performance and resource management.
Limitation Discussions & Future Work
Disentanglement of Routing Decisions and Rescaling Scores.
Currently, the routing decisions and rescaling scores in our framework are interdependent, which may impact the model’s accuracy. For instance, the skipped outputs are optimized to match the original outputs primarily through cosine similarity, which does not account for the difference in L2 magnitude. This discrepancy could potentially be mitigated by applying an additional rescaling factor to the skipped outputs, ensuring a better match in magnitude and improving the overall performance.
Improved Supervision for Router Labels.
The current method relies on cosine similarity for supervising router labels, which may not be the most effective approach. Exploring alternative supervision methods, such as task-specific metrics or direct gradients, could lead to more accurate router decisions. For example, drawing inspiration from works like Jiang et al. (2023) or employing gradient-based techniques similar to those used in network pruning, could enhance the router’s ability to prioritize important tokens and improve performance.
Skipping Attention Modules.
This work primarily focuses on skipping MLP layers due to the observed instability in decoding when skipping attention modules, such as excessive repetition and degraded accuracy (e.g., achieving accuracy on GSM8K). Future work could explore strategies to stabilize the skipping of attention layers, potentially improving model efficiency without compromising output quality significantly.
Ethical Statement
The development and deployment of large language models, including the LLaMA-DLO framework presented in this work, can raise important ethical considerations. Our research aims to enhance the efficiency and scalability of LLMs while maintaining high standards of responsibility and ethical practice. We recognize the potential impact of our work on various stakeholders and are committed to the following ethical principles:
Fairness and Bias Mitigation.
We are aware that language models can inadvertently learn and propagate biases present in training data. Efforts have been made to ensure that LLaMA-DLO is trained on diverse and representative datasets to minimize the risk of bias. Effective ways to mitigate bias is a pressing problem worth further study.
Transparency and Accountability.
We strive to maintain transparency in our research and development processes. Detailed documentation and open access to our methodologies and results will be provided to allow for scrutiny and reproducibility. Accountability mechanisms are in place to ensure that any adverse effects of our technology are promptly identified and addressed.
Social Impact.
The potential societal impact of LLaMA-DLO is carefully evaluated to prevent misuse or harm. We are committed to the ethical deployment of our technology, ensuring it is used for beneficial purposes such as advancing research, improving accessibility, and enhancing communication. We actively discourage and take steps to prevent the use of our models for malicious activities, misinformation, or any application that could harm individuals or society.
Continuous Ethical Review.
The ethical implications of our work are continually assessed to adapt to evolving norms and expectations. We engage with interdisciplinary experts and stakeholders to identify and address ethical concerns, ensuring that our research and its applications remain aligned with societal values and ethical standards.
By adhering to these principles, we aim to contribute positively to the field of artificial intelligence and ensure that the benefits of our research are realized in an ethical and responsible manner.
References
- Achiam et al. (2023) Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.
- Ainslie et al. (2023) Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón, Siddhartha Brahma, Yury Zemlyanskiy, David Uthus, Mandy Guo, James Lee-Thorp, Yi Tay, et al. 2023. Colt5: Faster long-range transformers with conditional computation. arXiv preprint arXiv:2303.09752.
- Amini et al. (2019) Aida Amini, Saadia Gabriel, Peter Lin, Rik Koncel-Kedziorski, Yejin Choi, and Hannaneh Hajishirzi. 2019. Mathqa: Towards interpretable math word problem solving with operation-based formalisms. arXiv preprint arXiv:1905.13319.
- Austin et al. (2021) Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021. Program synthesis with large language models. arXiv preprint arXiv:2108.07732.
- Baddeley (1992) Alan Baddeley. 1992. Working memory. Science, 255(5044):556–559.
- Ben Allal et al. (2022) Loubna Ben Allal, Niklas Muennighoff, Logesh Kumar Umapathi, Ben Lipkin, and Leandro von Werra. 2022. A framework for the evaluation of code generation models. https://github.com/bigcode-project/bigcode-evaluation-harness.
- Bisk et al. (2020) Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. 2020. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432–7439.
- Chen et al. (2021a) Cheng Chen, Yichun Yin, Lifeng Shang, Xin Jiang, Yujia Qin, Fengyu Wang, Zhi Wang, Xiao Chen, Zhiyuan Liu, and Qun Liu. 2021a. bert2bert: Towards reusable pretrained language models. arXiv preprint arXiv:2110.07143.
- Chen et al. (2021b) Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. 2021b. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374.
- Chen et al. (2015) Tianqi Chen, Ian Goodfellow, and Jonathon Shlens. 2015. Net2net: Accelerating learning via knowledge transfer. arXiv preprint arXiv:1511.05641.
- Chen et al. (2023) Yanxi Chen, Xuchen Pan, Yaliang Li, Bolin Ding, and Jingren Zhou. 2023. Ee-llm: Large-scale training and inference of early-exit large language models with 3d parallelism. arXiv preprint arXiv:2312.04916.
- Clark et al. (2018) Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457.
- Cobbe et al. (2021) Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
- Dabre and Fujita (2019) Raj Dabre and Atsushi Fujita. 2019. Recurrent stacking of layers for compact neural machine translation models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 6292–6299.
- (15) T Dao, DY Fu, S Ermon, A Rudra, and C Flashattention Ré. Fast and memory-efficient exact attention with io-awareness, 2022. URL https://arxiv. org/abs/2205.14135.
- Du et al. (2024) Wenyu Du, Tongxu Luo, Zihan Qiu, Zeyu Huang, Yikang Shen, Reynold Cheng, Yike Guo, and Jie Fu. 2024. Stacking your transformers: A closer look at model growth for efficient llm pre-training. arXiv preprint arXiv:2405.15319.
- Evci et al. (2022) Utku Evci, Bart van Merrienboer, Thomas Unterthiner, Max Vladymyrov, and Fabian Pedregosa. 2022. Gradmax: Growing neural networks using gradient information. arXiv preprint arXiv:2201.05125.
- Fedus et al. (2022a) William Fedus, Jeff Dean, and Barret Zoph. 2022a. A review of sparse expert models in deep learning. arXiv preprint arXiv:2209.01667.
- Fedus et al. (2022b) William Fedus, Barret Zoph, and Noam Shazeer. 2022b. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39.
- Gao et al. (2023) Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, d Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2023. A framework for few-shot language model evaluation.
- Glorot and Bengio (2010) Xavier Glorot and Yoshua Bengio. 2010. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249–256. JMLR Workshop and Conference Proceedings.
- Gong et al. (2019) Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. 2019. Efficient training of bert by progressively stacking. In International conference on machine learning, pages 2337–2346. PMLR.
- Gu et al. (2020) Xiaotao Gu, Liyuan Liu, Hongkun Yu, Jing Li, Chen Chen, and Jiawei Han. 2020. On the transformer growth for progressive bert training. arXiv preprint arXiv:2010.12562.
- Hadi et al. (2023) Muhammad Usman Hadi, Rizwan Qureshi, Abbas Shah, Muhammad Irfan, Anas Zafar, Muhammad Bilal Shaikh, Naveed Akhtar, Jia Wu, Seyedali Mirjalili, et al. 2023. A survey on large language models: Applications, challenges, limitations, and practical usage. Authorea Preprints.
- Hendrycks et al. (2020) Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.
- Hestness et al. (2017) Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. 2017. Deep learning scaling is predictable, empirically. arXiv preprint arXiv:1712.00409.
- Jiang et al. (2023) Huiqiang Jiang, Qianhui Wu, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2023. Llmlingua: Compressing prompts for accelerated inference of large language models. arXiv preprint arXiv:2310.05736.
- Kim et al. (2023) Dahyun Kim, Chanjun Park, Sanghoon Kim, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeonwoo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, et al. 2023. Solar 10.7 b: Scaling large language models with simple yet effective depth up-scaling. arXiv preprint arXiv:2312.15166.
- Koechlin et al. (2003) Etienne Koechlin, Chrystele Ody, and Frédérique Kouneiher. 2003. The architecture of cognitive control in the human prefrontal cortex. Science, 302(5648):1181–1185.
- Lepikhin et al. (2020) Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2020. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668.
- Li et al. (2024a) Dawei Li, Shu Yang, Zhen Tan, Jae Young Baik, Sunkwon Yun, Joseph Lee, Aaron Chacko, Bojian Hou, Duy Duong-Tran, Ying Ding, et al. 2024a. Dalk: Dynamic co-augmentation of llms and kg to answer alzheimer’s disease questions with scientific literature. arXiv preprint arXiv:2405.04819.
- Li et al. (2024b) Yifan Li, Anh Dao, Wentao Bao, Zhen Tan, Tianlong Chen, Huan Liu, and Yu Kong. 2024b. Facial affective behavior analysis with instruction tuning. arXiv preprint arXiv:2404.05052.
- Lin et al. (2021) Stephanie Lin, Jacob Hilton, and Owain Evans. 2021. Truthfulqa: Measuring how models mimic human falsehoods. arXiv preprint arXiv:2109.07958.
- Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter. 2017. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101.
- Luo et al. (2023) Ziyang Luo, Can Xu, Pu Zhao, Qingfeng Sun, Xiubo Geng, Wenxiang Hu, Chongyang Tao, Jing Ma, Qingwei Lin, and Daxin Jiang. 2023. Wizardcoder: Empowering code large language models with evol-instruct. arXiv preprint arXiv:2306.08568.
- Mihaylov et al. (2018) Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. 2018. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789.
- Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250.
- Raposo et al. (2024) David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. 2024. Mixture-of-depths: Dynamically allocating compute in transformer-based language models. arXiv preprint arXiv:2404.02258.
- Sakaguchi et al. (2021) Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2021. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106.
- Shazeer et al. (2017) Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538.
- Shen et al. (2022) Sheng Shen, Pete Walsh, Kurt Keutzer, Jesse Dodge, Matthew Peters, and Iz Beltagy. 2022. Staged training for transformer language models. In International Conference on Machine Learning, pages 19893–19908. PMLR.
- Shoemake (1985) Ken Shoemake. 1985. Animating rotation with quaternion curves. In Proceedings of the 12th Annual Conference on Computer Graphics and Interactive Techniques, pages 245–254. ACM.
- Tan et al. (2024) Zhen Tan, Alimohammad Beigi, Song Wang, Ruocheng Guo, Amrita Bhattacharjee, Bohan Jiang, Mansooreh Karami, Jundong Li, Lu Cheng, and Huan Liu. 2024. Large language models for data annotation: A survey. arXiv preprint arXiv:2402.13446.
- Team (2022) EvolCodeAlpaca Team. 2022. Evolcodealpaca dataset.
- Team et al. (2023) Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. 2023. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805.
- Team (2023a) LLaMA-MoE Team. 2023a. Llama-moe: Building mixture-of-experts from llama with continual pre-training.
- Team (2024) Self-Duplicate Stack Team. 2024. Self-duplicate stack.
- Team (2023b) ShareGPT Team. 2023b. Sharegpt dataset.
- Team (2023c) SlimOrca Team. 2023c. Slimorca dataset.
- Ting and Witten (1997) Kai Ming Ting and Ian H Witten. 1997. Stacking bagged and dagged models.
- Touvron et al. (2023) Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.
- Wang et al. (2018) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. 2018. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461.
- Wang et al. (2022) Jue Wang, Ke Chen, Gang Chen, Lidan Shou, and Julian McAuley. 2022. Skipbert: Efficient inference with shallow layer skipping. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7287–7301.
- Wang et al. (2023a) Peihao Wang, Rameswar Panda, Lucas Torroba Hennigen, Philip Greengard, Leonid Karlinsky, Rogerio Feris, David Daniel Cox, Zhangyang Wang, and Yoon Kim. 2023a. Learning to grow pretrained models for efficient transformer training. arXiv preprint arXiv:2303.00980.
- Wang et al. (2023b) Yite Wang, Jiahao Su, Hanlin Lu, Cong Xie, Tianyi Liu, Jianbo Yuan, Haibin Lin, Ruoyu Sun, and Hongxia Yang. 2023b. Lemon: Lossless model expansion. arXiv preprint arXiv:2310.07999.
- Wu et al. (2024) Chengyue Wu, Yukang Gan, Yixiao Ge, Zeyu Lu, Jiahao Wang, Ye Feng, Ping Luo, and Ying Shan. 2024. Llama pro: Progressive llama with block expansion. arXiv preprint arXiv:2401.02415.
- Yao et al. (2023) Yiqun Yao, Zheng Zhang, Jing Li, and Yequan Wang. 2023. Masked structural growth for 2x faster language model pre-training. In The Twelfth International Conference on Learning Representations.
- Yu et al. (2023) Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James T Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. 2023. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309.12284.
- (59) Xingjian Zhang, Jiaxi Tang, Yang Liu, Xinyang Yi, Li Wei, Lichan Hong, Qiaozhu Mei, and Ed H Chi. Conditional transformer fine-tuning by adaptive layer skipping. In 5th Workshop on practical ML for limited/low resource settings.
- Zoph et al. (2022) Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08906.
Appendix A Pseudo Code style Description of Dynamic Layer Operation (DLO)
Appendix B Evaluation Details
We list the number of shots and the metric used for each dataset as follows:
-
•
ARC-C: 25 shots, normalized accuracy.
-
•
GLUE: 0 shot, accuracy.
-
•
MMLU: 5 shots, normalized accuracy.
-
•
PIQA: 0 shot, normalized accuracy.
-
•
OBQA: 0 shot, normalized accuracy.
-
•
SQuAD: 0 shot, F1 Score.
-
•
TruthfulQA: 0 shot, accuracy.
-
•
WinoGrande: 5 shots, accuracy.
-
•
GSM8K: 5 shots, accuracy.
-
•
MathQA: 0 shot, normalized accuracy.
-
•
HumanEval: 200 rounds, pass@100.
-
•
MBPP: 15 rounds, pass@10.
Appendix C Acknowledgment of AI Assistance in Writing and Revision
We utilized ChatGPT-4 for revising and enhancing sections of this paper.