DLO: Dynamic Layer Operation for Efficient Vertical Scaling of LLMs

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.

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.

To facilitate dynamic depth adjustment, we introduce a group-based layer expansion strategy. Suppose the LLM has $R$ transformer layers, which we group into $P$ groups with $Q$ layers each, such that $R=P\times Q$. Each group is expanded to $Q^{\prime}=Q+q$ layers, where $q$ is the number of additional layers introduced per group. The resulting number of layers will be $R^{\prime}=P\times Q^{\prime}$.

Let $\mathcal{G}_{i}$ denote the $i$-th group with layers $\mathcal{L}_{i1},\mathcal{L}_{i2},\ldots,\mathcal{L}_{iQ}$. The expanded group $\mathcal{G}_{i}^{\prime}$ will contain layers $\mathcal{L}_{i1},\mathcal{L}_{i2},\ldots,\mathcal{L}_{i(Q+q)}$. The expanded layers are initialized using a policy $\Pi$, and we consider several initialization strategies:

• •

  Random Initialization ($\Pi_{\text{rand}}$): Initialize the new layers’ weights $\theta^{\prime}_{ij}$ using Xavier initialization Glorot and Bengio (2010).

  $$\theta^{\prime}_{ij}\sim\mathcal{U}\left(-\sqrt{\frac{6}{n_{in}+n_{out}}},\sqrt{\frac{6}{n_{in}+n_{out}}}\right),$$

  where $\forall j\in\{Q+1,Q+2,\ldots,Q+q\}$, $\mathcal{U}$ denotes the uniform distribution, $n_{in}$ is the number of input units, and $n_{out}$ is the number of output units in the layer.

• •

  Copy from Previous Layer ($\Pi_{\text{copy}}$): Copy the parameters from the preceding layer.

  $$\theta^{\prime}_{ij}=\theta_{i(Q+q-1)},\forall j\in\{Q+1,\ldots,Q+q\}.$$

• •

  Identity Initialization ($\Pi_{\text{Identity}}$) Wu et al. (2024): Copy from the preceding layer but set the output linear matrix of the multi-head self-attention (MHSA) to zero.

  1. a.

     Copy the parameters of the previous layer:

     $$\theta^{\prime}_{ij}=\theta_{i(Q+q-1)},\forall j\in\{Q+1,\ldots,Q+q\}.\vspace{-2mm}$$

  2. b.

     Set the weights of the output linear layer $W^{\prime}_{\text{out}}$ in the MHSA to zero: $W^{\prime}_{\text{out}}=0.$

  In this way, the output of the expanded layers will preserve the features from the original layers.

• •

  Linear Merge ($\Pi_{\text{linear}}$): Merge from the preceding $\tau$ layers using a linear function.

  $$\theta^{\prime}_{ij}=\sum_{k=1}^{\tau}\alpha_{k}\theta_{i(Q+q-k)},\sum_{k=1}^{\tau}\alpha_{k}=1.\vspace{-2mm}$$

• •

  Spherical Linear Interpolation (SLERP) ($\Pi_{\text{slerp}}$) Shoemake (1985): Merge from the preceding $\tau$ 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 $\mathbf{u}$ and $\mathbf{v}$ is defined as:

  $$\text{SLERP}(\mathbf{u},\mathbf{v},\alpha)=\frac{\sin((1-\alpha)\Omega)}{\sin(\Omega)}\mathbf{u}+\frac{\sin(\alpha\Omega)}{\sin(\Omega)}\mathbf{v},$$

  where $\Omega$ is the angle between $\mathbf{u}$ and $\mathbf{v}$:

  $$\Omega=\arccos\left(\frac{\mathbf{u}\cdot\mathbf{v}}{\|\mathbf{u}\|\|\mathbf{v}\|}\right),$$

  and $\alpha\in[0,1]$ is the interpolation parameter.

  In our context, for the new layer $j$, the parameters are initialized by interpolating between the weights of the previous layers $\theta_{i(Q+q-1)}$ and $\theta_{i(Q+q-\tau)}$:

  $$\theta^{\prime}_{ij}=\text{SLERP}(\theta_{i(Q+q-1)},\theta_{i(Q+q-\tau)},\alpha),$$

  where $\alpha$ 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 $\Pi$ in Section 4.3.

DLO dynamically skips the multi-layer perceptron (MLP) module within the transformer layer $\mathcal{L}_{i}$ 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 $S$ for a given layer $\mathcal{L}_{i}$, that is $\mathbf{h}_{i}=\{\mathbf{h}^{s}_{i}|s\in\mathbb{N}^{*},s\leq S\}$, where in following contents we omit the superscript $s$ for better readability. Considering feature redundancy, we use the router weights $W_{i}$ to process the input token $\mathbf{h}_{i}$ and obtain the decision score $r_{i}$, which is given by:

where $\sigma$ is the sigmoid function, $\beta$ and $\gamma$ are hyperparameters controlling the output range. During the inference stage, this score $r_{i}$ determines whether layer $\mathcal{L}_{i}$ is active: the layer is activated if and only if $\mathcal{L}_{i}\geq\frac{\beta}{2}$, otherwise it’s skipped. The final activated output for the layer is:

where $\mathcal{M}_{i}$, $\mathcal{A}_{i}$ are the MLP and the attention modules within layer $\mathcal{L}_{i}$, 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 $W_{i}$ as zeros and set $\beta=2.0,\gamma=0.05$, so that $r_{i}=1.0$ on the first step and $r_{i}\in(0.975,1.025)$ during training. This ensures benign initialization for activated tokens and avoids excessive disturbance on activated outputs brought by the decision scores.

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 $\hat{\rho}_{i}$ will be determined completely by the router based on Equation (1)-(2). The adaptive FLOPs are computed as:

where $\text{FLOPs}_{\text{full}}$ represents the FLOPs for a fully active layer. Since $\hat{\rho}_{i}$ is predicetd based on each specific token, DLO acheive adaptive FLOPs that entails better generalizability.

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 $R=32$ original transformer layers, which we group into $P=4$ clusters, each containing $Q=8$ layers. For layer expansion, we increase the group size to $Q^{\prime}=10$ 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.

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:

We conduct further analyses of key components of DLO in the subsequent subsections.

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.