LLAMA PRO: Progressive LLaMA with Block Expansion (2024)

Abstract: Humans generally acquire new skills without compromising the old; however, the opposite holds for Large Language Models (LLMs), e.g., from LLaMA to CodeLLaMA. To this end, we propose a new post-pretraining method for LLMs with an expansion of Transformer blocks. We tune the expanded blocks using only new corpus, efficiently and effectively improving the model’s knowledge while mitigating forgetting. In this paper, we experiment on the corpus of code and math, yielding LLAMA PRO-8.3B, a versatile foundation model initialized from LLaMA27B, excelling in general tasks, programming, and mathematics. LLAMA PRO and its instruction-following counterpart (LLAMA PRO INSTRUCT) achieve advanced performance among various benchmarks, demonstrating superiority over existing open models in the LLaMA family and the immense potential of reasoning and addressing diverse tasks as an intelligent agent. Our findings provide valuable insights into integrating natural and programming languages, laying a solid foundation for developing advanced language agents that operate effectively in various environments.

heading: 1 Introduction
text: The advent of Large Language Models (LLMs) has revolutionized the field of natural language processing, exhibiting remarkable proficiency in a variety of real-world tasks (OpenAI, 2023; Chowdhery et al., 2023). Despite the versatility, LLMs still fall short in certain domains, for example, programming, mathematics, biomedical, or finance. This limitation impedes the progress of developing generic language agents for broader applications.
Existing works (Liu et al., 2023; Li et al., 2023a; Wu et al., 2023b) attempted to improve the multifaceted capabilities of pre-trained LLMs with tai-
*Correspondence to yixiaoge@tencent.com.
lored data recipes. While feasible, they require substantial computational resources and vast amounts of data, which poses a challenge to the democratization of LLM research. Consequently, another line of research, known as domain-adaptive pretraining, focuses on post-pretraining with domainspecific corpora (Gururangan et al., 2020). These approaches have demonstrated efficacy in adapting various LLMs to specific domains (Roziere et al., 2023; Azerbayev et al., 2023; Wu et al., 2023b; Xu et al., 2023b), resulting in enhanced performance on downstream domain-specific tasks at a reduced computational cost.
Nonetheless, a considerable obstacle emerges in catastrophic forgetting (De Lange et al., 2021). Post-pretraining often leads to a decline in the model’s original general abilities, inhibiting the fine-tuned performance of the model on diverse tasks (Cheng et al., 2023; Dong et al., 2023). This necessitates a method that can inject domainspecific knowledge into LLMs while preserving their general abilities, thereby enhancing their com-
6518
prehensive capabilities. Towards this end, we introduce a simple yet effective post-pretraining method, termed block expansion. We expand the off-the-shelf pre-trained LLM using copied Transformer blocks, as illustrated in Figure 2. The newly added blocks, whose linear layers are zero-initialized to enable identity mapping, are further tuned with only domainspecific corpus while the remaining blocks are frozen. After tuning, the extended pre-trained model excels in both general and domain-specific tasks.
In practice, we extend the pre-trained LLaMA27B (Touvron et al., 2023) by eight more blocks, yielding LLAMA PRO, a foundation model with 8.3B parameters, and enhanced performance in programming, coding, and reasoning. We pre-train LLAMA PRO’s expanded blocks on 80B tokens using open-source code and math data for 2830 GPU Hours (16 NVIDIA H800 GPUs for about 7 days). We further perform supervised instruction tuning (fully fine-tuning of all the blocks, aka SFT) on LLAMA PRO with approximately 80M tokens, yielding LLAMA PRO - INSTRUCT. It is noted that pre-trained models produced by our block expansion method are well-compatible with the subsequent SFT techniques without specific modification.
As shown in Figure 1, LLAMA PRO - INSTRUCT reaches state-of-the-art performance across a broad range of general, code (i.e., HumanEval), and math (i.e., GSM8K) tasks. Furthermore, we assess the capabilities of LLAMA PRO - INSTRUCT as a language agent across various scenarios (i.e., MINT-Bench), with a focus on the tool usage abilities and the capacity to ground in environmental and human feedback. We also employ GPT-4 (OpenAI, 2023) automatic evaluation to assess LLAMA PRO’s ability to serve as an effective assistant (i.e., MT-Bench). Comprehensive experimental results indicate the superiority of LLAMA PRO - INSTRUCT over other models from the LLaMA family on both benchmarks and practical applications. Our contributions are three-fold:
• We propose a novel post-pretraining method for LLMs, termed block expansion, enabling the injection of new knowledge while preserving the initial capabilities.
• We introduce LLAMA PRO and LLAMA PRO - INSTRUCT, versatile LLMs that well integrate natural and programming languages,
excelling in general tasks, programming, and mathematics.
• We benchmark the family of LLAMA PRO on extensive datasets, including both traditional and agent-oriented tasks, demonstrating its superiority and great potential in broader complex applications.

heading: 2 Related Work
text: Advancements in Large Language Models. Recent advancements in large language models have led to significant progress, with model and data scale growth driving state-of-the-art performance across various tasks (Hoffmann et al., 2022; Kaplan et al., 2020; Chowdhery et al., 2023). The development of generalist models has enabled addressing diverse problems and rapid adaptation to new tasks (Radford et al., 2019; Brown et al., 2020). The open-source community has further contributed by releasing powerful models like LLaMA (Touvron et al., 2023) and CodeLLaMA (Roziere et al., 2023). Our work builds upon these developments, providing a method for specializing LLMs in the code domain, fostering future research and applications.
Post-pretraining. Language model applications typically involve a two-step process: generaldomain pretraining followed by domain-specific training (Roziere et al., 2023; Azerbayev et al., 2023). Fine-tuning often aims to enhance instruction-following abilities (Sanh et al., 2021; Wei et al., 2021; Wang et al., 2023d) or align model outputs with human preferences (Ziegler et al., 2019; Ouyang et al., 2022; Bai et al., 2022). Some research explores parameter-efficient fine-tuning methods for adapting pretrained models to new domains (Houlsby et al., 2019; Hu et al., 2021; Wu et al., 2023a), while others focus on continual learning post-pretraining (Wang et al., 2023b; Gupta et al., 2023; Scialom et al., 2022). Parameterefficient tuning methods like adaptor and LoRA are generally applied during the instruction tuning phase rather than the pretraining phase. In contrast, our focus is on enhancing the capacity of LLMs by increasing their depth during continued pretraining. Our work proposes an adaptation strategy that combines continued training with general capability maintenance, allowing LLMs to specialize without sacrificing overall performance.
Progressive Learning. Progressive training has gained attention for accelerating large-scale model training in computer vision (Zhang et al., 2023) and NLP research (Yao et al., 2023; Li et al., 2023b). Gong et al. (2019) proposed a stacking method doubling model depth successively. CompoundGrow(Gu et al., 2020) extends stacking with FeedForward Network expansion in schedule design. Shen et al. (2022) introduced a staged method supporting hidden size expansion. Bert2BERT(Chen et al., 2021a) and LiGO (Wang et al., 2023a) accommodate all growth dimensions. Our method utilizes depth growth to maintain general performance while adapting to specific domains.

heading: 3 Method

heading: 3.1 Preliminaries: The LLaMA Block
text: The LLaMA block consists of a multi-head selfattention (MHSA) mechanism followed by a position-wise feed-forward network (FFN) with residual connections and a Swish-Gated Linear Unit (SwiGLU) operation as Figure 3 shows. Given an input x, the LLaMA block produces an output y as described by the following equations:
x′ = x+ MHSA(RMSNorm(x))
y = x′ + FFN(RMSNorm(x′)) (1)
The input x has a dimension of n× d, where n is the sequence length and d is the hidden size. The
output y has the same dimension as the input x. The MHSA operation is a crucial component of the transformer, defined as:
MHSA(Q,K, V ) = Concat(head1, . . . , headh)WO
(2) where Q, K, and V are the query, key, and value matrices, respectively, and WO is the output weight matrix without bias . Each head is computed as:
headi = Attention(xW Q i , xW K i , xW V i )
Attention(Qi,Ki, Vi) = Softmax ( QiK T i√
dk
) Vi
(3)
with WQi , W K i , and W V i being the corresponding weight matrices for the i-th head. The FFN block in the LLaMA model utilizes the SwiGLU activation function, which is defined as:
SwiGLU(x,W, V ) = SiLU(xW )⊗ (xV ) FFN(x) = SwiGLU(x,W1,W2)W3 (4)
where ⊗ denotes element-wise multiplication, W1, W2, and W3 are the weight matrices without bias, SiLU(x) = x⊗ σ(x).

heading: 3.2 Block Expansion
text: Given a model with blocks (ϕ0, ϕ1, ..., ϕL), the block expansion incorporates an identity block ϕid
after each block in the original model, ensuring that the expanded model maintains the same output after expansion. The identity block is defined as ϕid(x) = x so the input and output are identical.
Suppose we have an initial model with L blocks that needs to be expanded to L′ blocks. First, we partition the original L blocks into N groups, with each group containing LN blocks. For each group, we create identity copies of the top P blocks and stack them on top of each group, as depicted in Figure 3. We arrange these blocks in an interleaved manner to maintain the structural characteristic of the transformer model, whose prior is that deeper blocks encode more complex information (Van Aken et al., 2019; Tenney et al., 2019). This process leads to an increased depth in the model while maintaining its output behavior.
Shen et al. (Shen et al., 2022) proposed the initialization of scale parameters in the Norm modules within the identity blocks to zero for the construction of the identity block. However, this approach may not be effective when applied to the LLaMA
block. The reason lies in the fact that the gradient of the loss function L with respect to the RMSNorm weight w during backpropagation would be zero. This would prevent the training of RMSNorm, implying that when RMSNorm(x′) = 0, the following condition will hold:
∂L ∂w = ∂L ∂y ∂FFN(RMSNorm(x′)) ∂RMSNorm(x′) ∂RMSNorm(x′) ∂w = 0 (5)
This equation signifies that the gradient of the loss function with respect to the weight of RMSNorm is zero, which would hinder the training of the RMSNorm module. This is further explained in Appendix A. Referring to the LLaMA block formulation in Equation 1, the identity can be achieved as long as MHSA(RMSNorm(x)) = 0 and FFN(RMSNorm(x′)) = 0. We initialize the WO and W3 weight matrices in the identity blocks to zero. Due to the presence of residual connections and the absence of bias terms in the LLaMA block, only the residual flows through the identity block. As a result, the entire block is reduced to an identity block at initialization, preserving the output from the initial model.
The entire training pipeline is depicted in Figure 2. Our method concentrates on the postpretraining stage, targeting specific domain corpora. We begin by initializing our model with large language models trained on extensive unlabeled general corpora, where all blocks will be fine-tuned. To enhance the model’s capacity for accommodating additional domain knowledge while retaining its general knowledge, we employ block expansion to increase the number of blocks in the LLM. During this process, we only fine-tune the newly added blocks while freezing the original blocks, thereby preserving the general abilities of the model.

heading: 4 Experiments
text: This section presents our key experimental findings. We begin with experimental settings (described in Sec. 4.1), and then verify the effectiveness of block expanded tuning after pretraining (described in Sec. 4.2). Next, we give the supervised finetuning (SFT) results (described in Sec. 4.3). Finally, ablation studies of the key design choices are presented (described in Sec. 4.5).

heading: 4.1 Experimental Settings
text: Pretrain details. We construct a dataset that concentrates on code and math. For the code component, we rely on the Stack-dedup dataset, which
is a compilation of permissively licensed source codes from GitHub. Among all the programming languages available in Stack-dedup, we specifically utilize the Python split. As for the math component, we opt for the Proof-pile-2 dataset (Azerbayev et al., 2023), a 55-billion-token amalgamation of scientific papers, web data containing mathematical content, and mathematical code. The details can be found in Appendix B.
We initialize our base model with LLaMA2-7B and expand the number of blocks from 32 to 40 using an interleaved approach. In the block expansion process, we configure the parameters as P = 1, M = 4, and N = 8, resulting in 8 groups where each group expands from 4 blocks to 5 blocks. For the code and math corpus pretraining, we employ a batch size of 1024, a sequence length of 4096, a warmup ratio of 6%, a learning rate of 2e-4, and a Cosine learning rate scheduler. We also use bf16 mixed precision, a weight decay of 0.1, and gradient clipping at 1.0. To speed up the training process, we apply the flash-attention mechanism.
Our experiment is conducted on 16 NVIDIA H800 GPUs. LLAMA PRO is trained for a total of 15,900 steps. This training process corresponds to approximately 2830 H800 GPU hours.
We want to highlight that our approach does not incur higher training costs, and it is worth the extra resources to achieve a better performance of the domain specific tasks in the inference.
Training stage cost: Our approach requires fewer computational resources since only the newly added blocks are tuned during training. As illustrated in Figure 4, LLaMA Pro-8B (1B parameters tuned for 80B tokens) incurs less training overhead compared to CodeLLaMA-7B (7B parameters tuned for 500B tokens). It also uses fewer resources than training domain-specific models from scratch, such as StarCoder and CrystalCoder. Despite this, our method achieves a better balance of general and domain-specific performance, offering a more cost-effective solution.
Inference stage cost: Although our method requires more resources during inference than the initial LLM, it strikes a balance between performance and efficiency. LLaMA Pro-8B outperforms larger models like LLaMA2-13B and LLaMA2-34B in the code domain while demanding significantly fewer resources during training and inference.
SFT details. During the instruction fine-tuning phase, we combine five data sources to create
LLAMA PRO - INSTRUCT as shown in Table 7. The final sft dataset consists of approximately 1M samples. To fine-tune the basic models, we employ specific configurations, including a batch size of 128, a sequence length of 4096, 0.03 warmup ratio, a learning rate of 2e-5, a Cosine learning rate scheduler, and bf16 mixed precision.
Evaluation details. We conduct a comparative analysis of LLAMA PRO with the latest state-ofthe-art (SOTA) Large Language Models (LLMs). The evaluation is performed on six key general benchmarks using the Eleuther AI Language Model Evaluation Harness1, a unified framework designed to test generative language models across a vast array of evaluation tasks. For code-related tasks, we employ the BigCode Evaluation Harness2 to evaluate HumanEval and MBPP, and we report the pass@1 rate of code tasks with greedy decoding. The evaluation details can be found in Appendix D.

heading: 4.2 Pretrain Results
text: We evaluate LLAMA PRO’s performance with benchmark datasets from the Open LLM Leader-
1https://github.com/EleutherAI/ lm-evaluation-harness
2https://github.com/bigcode-project/ bigcode-evaluation-harness
board. Furthermore, we incorporate coding benchmark datasets, including HumanEval pass@1 and MBPP pass@1, as well as the math benchmark GSM8K, to provide a comprehensive evaluation. We compare the performance of LLAMA PRO with a selection of state-of-the-art pretrained models that were trained around the same period with similar size. This includes general-purpose pretrained models like LLaMA2 and code-oriented pretrained models like CodeLLaMA. The results are presented in Table 1.
The results highlight that LLAMA PRO effectively balances natural language processing and coding capabilities. It not only preserves the general performance of its base model, LLaMA2-7B, but also surpasses it in the average performance of general language tasks. Conversely, CodeLLaMA7B sacrifices general performance. We attribute this improvement to our expansion design, which freezes the initial LLaMA blocks to maintain their capabilities and increases the blocks to accommodate domain-specific knowledge.
As depicted in Figure 4, LLAMA PRO shows robust general performance alongside code performance that is on par with code-oriented LLMs. Situated on the Pareto frontier, LLAMA PRO has undergone fine-tuning with an additional 80B tokens in conjunction with LLaMA2, which more than doubles the code tasks average performance. In contrast, CodeLLaMA is fine-tuned with 500B tokens. LLAMA PRO excels in general performance while maintaining code performance that is competitive with code-oriented LLMs, whether they are
trained from scratch, such as StarCoder-15B and CrystalCoder, or fine-tuned like CodeLLaMA-7B.

heading: 4.3 SFT Results
text: Modern LLMs typically undergo supervised finetuning or instruction tuning after pretraining on vast amounts of unlabeled data. In this section, we aim to demonstrate that our expansion strategy can adapt to this widely used training pipeline, just as traditional LLMs do.
Table 1 presents a comparison of evaluation results among several prominent supervised finetuning (SFT) LLMs from the LLaMA community, across general tasks, math tasks, and code tasks benchmarks. As a singular SFT model, LLAMA PRO - INSTRUCT attains state-of-the-art performance, even when compared to specifically tuned models such as WizardCoder and WizardMath. This demonstrates its more comprehensive capabilities.
As seen in Figure 1, LLAMA PRO - INSTRUCT boosts both code and math tasks to SOTA perfor-
mances while maintaining reliable general performance. We enhance the average performance of LLaMA2-7B-chat and CodeLLaMA-7B-instruct by 13.81% and 14.50% respectively, which highlights the benefits of balancing textual and coding abilities.
To assess the comprehensive conversational performance of the LLAMA PRO - INSTRUCT assistant, we evaluate it using the MT-Bench with GPT-4 automatic scoring, as proposed by Vicuna (Zheng et al., 2023). As depicted in Table 2, LLAMA PRO - INSTRUCT surpasses widely used chatbots from the LLaMA community. This indicates its potential as a chatbot capable of providing helpful responses, in addition to its impressive performance in traditional benchmarks. The details of MT-Bench can be found in the Appendix F.
We use MINT-Bench (Wang et al., 2023c) to evaluate our model’s ability to solve multi-turn interactions by using tools. MINT-Bench tests LLMs’ ability to use tools by generating and executing Python code, focusing on tool-augmented tasksolving and leveraging natural language feedback. MINT includes eight datasets covering reasoning, code generation, and decision-making. The details of MINT can be found in the Appendix E. The results are shown in Table ??. LLAMA PRO - INSTRUCT achieves SOTA performance compared to similar size models in multi-turn interactions with the use of tools.

heading: 4.4 Mistral-Pro Results
text: We experimented with block expansion on Mistral7B (Jiang et al., 2023), training it on code and mathematics datasets. The resulting pretrained performance is detailed in Table 4, highlighting superior outcomes across various benchmarks, particularly in the domains of code and math. Notably, it demonstrates competitive results compared to the new open-source model Gemma (Team et al.,
2024), while incurring significantly lower training overhead. We further utilized the MetaMath dataset (Yu et al., 2023) for supervised fine-tuning. Our approach yielded scores of 78.4 for GSM8k and 30.3 for MATH, surpassing Mistral’s scores of 77.7 and 28.2, respectively. Additional details are provided in Appendix C.

heading: 4.5 Ablation Study
text: Apart from the aspect of code corpus, we explore our method on another domain: law, with the freelaw subset of Pile dataset as our pretrain corpus (Gao et al., 2020). We evaluate on UNFAIR-ToS (Lippi et al., 2019) of the LexGLUE benchmark (Chalkidis et al., 2021).
In our experiment, we assess the scalability of
our block expansion method in terms of training loss and downstream task performance as we increase the number of added blocks. We also compare our method with the Mixture-of-Expert (MoE) expansion method (Fedus et al., 2022) and traditional training strategies, such as fine-tuning and LoRA (Hu et al., 2021). The details can be found in Appendix H.
We analyze the training loss with varying added blocks (Figure 5). The loss consistently decreases during training, regardless of the number of added blocks, and decreases more rapidly with larger models. These findings indicate that our method demonstrates strong scalability with larger models and more data.
However, a lower overall training loss does not necessarily guarantee superior performance on domain-specific tasks. Therefore, we evaluate models of different sizes on both general language tasks and Unfair-ToS, as shown in Table 5. All the expanded models effectively preserve the general capabilities of the initial model. For the domainspecific task, larger models achieve better performance. We find that adding eight blocks provides optimal performance with minimal cost compared to larger models, hence we adopt this as our default
strategy. The performance of MoE is comparable to our method with four added blocks. Figure 9 illustrates the differences between traditional training strategies such as fine-tuning and LoRA, and our proposed method. We observe that while LoRA effectively preserves the general ability, it struggles to model the distribution of a new domain, as also evidenced by the training loss depicted in Figure 5. In contrast, full fine-tuning results in a more significant drop in general performance. Here we use a rank of 1024 for LoRA, resulting in a number of trainable parameters comparable to our method.
In line with the approach of Lin et al. (2023), we analyze the token distribution between the original LLaMA and LLAMA PRO to assess the similarity in their behavior when answering general questions from the Alpaca dataset (Taori et al., 2023). As depicted in Figure 6, the token distribution shift between LLaMA and LLAMA PRO is subtle. Detailed information can be found in Appendix G.
We also analyze the impact of the position where the identity blocks are added, either at the bottom or the top of the model, compared to adding them interleaved, as shown in Table 5. We observe that adding blocks at the bottom results in poor evaluation performance, likely because it disrupts the
model’s foundation, causing errors to propagate throughout the model. Adding blocks at the top of the model (Gong et al., 2019) preserves the initial model’s performance, but its performance on domain-specific tasks is lower than when adding blocks interleaved.
As highlighted in the LIMA study (Zhou et al., 2023), the majority of knowledge in large language models is acquired during pretraining, with only a limited amount of instruction tuning data required to generate high-quality output. To investigate the extent of knowledge encoded during pretraining, we conducted a comparative analysis between LLaMA2-7B and LLAMA PRO using the same instruction dataset, as illustrated in Figure 7. Our results showed that LLAMA PRO consistently outperforms LLaMA2-7B across all tasks, indicating that our method effectively enables LLAMA PRO to encode more domain-specific knowledge during the pretraining phase.

heading: 5 Scope and Limitations
text: Although our study presents a promising method for balancing general and domain-specific capabilities in LLMs, its scope is limited to the language modality, especially programming language and English. Future research could explore extending the application of our block expansion method to other domains, such as maintaining original language ability in multimodal large language models(Ge et al., 2023; Bai et al., 2023), and multi-
lingual domains.

heading: 6 Conclusion
text: In this study, we introduced a novel block expansion method for Large Language Models (LLMs) post-pretraining, aiming to enhance domain-specific abilities while preserving the original general capabilities. Our approach effectively balances the model’s performance across both general and domain-specific tasks. We demonstrated the effectiveness of our method through LLAMA PRO, an LLM initialized from LLaMA2-7B with 8 added blocks, which outperformed other LLaMAseries models on comprehensive benchmarks.

heading: 7 Ethical Statement
text: LLAMA PRO and LLAMA PRO - INSTRUCT are designed for a wide range of NLP tasks, with a focus on programming, mathematics, and general language tasks. It suits scenarios requiring integration of natural and programming languages. While LLaMA-Pro addresses some limitations of previous models in the series, it may still encounter challenges specific to highly specialized domains or tasks. Users should be aware of potential biases in the model and use it responsibly, considering its impact on various applications with the LLaMA-2 license.

heading: A Gradient Derivation
text: To calculate the gradient of the RMSNorm weight during backpropagation, we first need to consider the forward pass equation for the Llama RMSNorm:
RMSNorm(x) = w ⊙ x√
Var(x) + ϵ (6)
where x is the input tensor, w is the weight parameter, Var(x) is the variance of x across the last dimension, and ϵ is a small constant for numerical stability.
Now, let’s consider the chain rule for the gradient of the loss function with respect to the RMSNorm weight during backpropagation. Denote the loss function as L, and the output of the FFN as y. We have:
∂L ∂w = ∂L ∂y ∂y ∂w (7)
To compute the gradient, we need to find the partial derivative ∂y∂w . From the FFN equation, we have:
y = x′ + FFN(RMSNorm(x′)) (8)
Taking the derivative with respect to w, we get:
∂y ∂w = ∂FFN(RMSNorm(x′)) ∂w
(9)
Now, let’s differentiate the RMSNorm function with respect to w:
∂RMSNorm(x) ∂w = x√ Var(x) + ϵ (10)
Using the chain rule, we can compute the gradient of the loss function with respect to the RMSNorm weight:
∂L ∂w = ∂L ∂y ∂FFN(RMSNorm(x′)) ∂RMSNorm(x′) ∂RMSNorm(x′) ∂w
(11)
Given that RMSNorm(x′) = t, we need to find the derivative of the FFN with respect to t. Recall the FFN equation:
FFN(t) = SwiGLU(t,W1,W2)W3 (12)
Now we want to find the partial derivative of the FFN with respect to t. Recall the SwiGLU activation function:
SwiGLU(t,W1,W2) = SiLU(tW1)⊗ (tW2) (13) Taking the derivative of the SwiGLU function with respect to t, we get:
∂SwiGLU(t,W1,W2) ∂t =
( ∂SiLU(tW1)
∂t
) ⊗ (tW2) + SiLU(tW1)⊗ ( ∂(tW2)
∂t
) (14)
Now, recall the SiLU activation function:
SiLU(x) = x⊗ σ(x) (15) Thus, the gradient of the FFN with respect to t when t = 0 is also zero:
∂FFN(t) ∂t = 0 (16)
In conclusion, when t = 0, the gradient of the FFN with respect to t is zero, which demonstrates that the gradient is zero when the input to the FFN is zero.

heading: B Dataset Details
text: In this section, we provide detailed information about the dataset used for both pretraining and Supervised Fine-Tuning (SFT). Table 6 outlines the composition of our pretraining dataset, which comprises approximately 80 billion tokens from both math and code corpora. The specifics of the SFT data are delineated in Table 7.
For our proposed LLAMA PRO - INSTRUCT, we employ a blend of multiple instruction datasets spanning general instruction, math, and code for the SFT process. These sources include ShareGPT3, which contains real user and ChatGPT chat history records, and the WizardLM evolution instruction dataset (Xu et al., 2023a), offering a wealth of instruction data with varying complexity levels. We also incorporate the evolution CodeAlpaca dataset (Luo et al., 2023), which includes complex coding tasks generated by ChatGPT and their corresponding solutions. Additionally, we use MetaMath (Yu et al., 2023), which reframes questions from multiple perspectives, and SlimOrca (Lian et al., 2023), a curated subset of our OpenOrca data. SlimOrca provides an efficient route to achieve performance comparable to using larger data slices, while only incorporating approximately 500,000 GPT-4 completions.

heading: C Mistal-Pro Details
text: Mistral-Pro is an advanced version of the original Mistral model (Jiang et al., 2023), enhanced through the addition of Transformer blocks. This version excels in combining general language understanding with domain-specific knowledge, particularly in programming and mathematics. It employs the same methodology for creating additional blocks as LLaMA-Pro but utilizes only 110 of LLaMA Pro’s learning rate, as recommended by MetaMath-Mistral 4. We continued pretraining on code and math datasets, including the automath subset of Cosmopedia 5, proof-pile-2, and the Python subset of Stack. The supervised fine-tuning (SFT) approach remains consistent with MetaMath-Mistral, except that we switch the base model to our Mistral-Pro. The detailed results of GSM8k and MATH can be found in Table 8.

heading: D Evaluation Benchmark
text: The benchmarks used for evaluation include: 3https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered 4https://huggingface.co/spaces/TencentARC/MetaMath-Mistral-Pro 5https://huggingface.co/datasets/HuggingFaceTB/cosmopedia
• AI2 Reasoning Challenge (Clark et al., 2018) (25-shot): a set of grade-school science questions.
• HellaSwag (10-shot) (Zellers et al., 2019): a test of commonsense inference, which is easy for humans (approximately 95%) but challenging for SOTA models.
• MMLU (5-shot) (Hendrycks et al., 2020): a test to measure a text model’s multitask accuracy. The test covers 57 tasks including elementary mathematics, US history, computer science, law, and more.
• TruthfulQA (0-shot) (Lin et al., 2021): a test to measure a model’s propensity to reproduce falsehoods commonly found online.
• Winogrande (5-shot) (Sakaguchi et al., 2021): an adversarial and difficult Winograd benchmark at scale, for commonsense reasoning.
• GSM8k (5-shot) (Cobbe et al., 2021): diverse grade school math word problems to measure a model’s ability to solve multi-step mathematical reasoning problems. Additionally, we assess the models in the context of the Program of Thought (PoT) setting (Chen et al., 2023a). The PoT setting utilizes Python code to solve mathematical problems, which serves to evaluate the code generation capabilities of the models.
• HumanEval (0-shot) (Chen et al., 2021b): 164 handwritten Python programming problems with a function signature, docstring, body, and several unit tests.
• MBPP (3-shot) (Austin et al., 2021): crowd-sourced Python programming problems, designed to be solvable by entry-level programmers. Each problem consists of a task description in English, a code solution and 3 automated test cases.

heading: E MINT-Bench
text: The MINT-Bench (Wang et al., 2023c) details are provided in this section. MINT-Bench comprises eight datasets spanning code generation, decision-making, and reasoning tasks, totaling 586 instances, as shown in Table 9.
We use the Success Rate (SR) as our evaluation metric, which measures the percentage of successful task instances. For an interaction limit of k, MINT-Bench starts from scratch and allows each LLM to
interact up to the k-th turn, measuring the corresponding SRk. Unless specified otherwise, MINT-Bench limits k ∈ [1, 5], where k = 1 indicates no interaction, and k = 5 maximizes interaction turns within the context window (4,096 tokens) of most modern LLMs.
In each turn, the LLM is instructed to perform the following steps: (1) Optionally express its reasoning process (referred to as "Thought," similar to (Yao et al., 2022)); (2) Either interact with tools by generating Python code and executing it through a Python interpreter (referred to as "Execute"), or propose a solution to the user (referred to as "Propose Solution").
Table 10 displays the success rate for each model evaluated on various task type benchmarks, as well as the micro average when k = 5. The LLAMA PRO - INSTRUCT model demonstrates robust performance across all task types compared to other models of similar size. Figure 8 provides a case study to compare LLAMA PRO - INSTRUCT and LLaMA2-7B-Chat where LLAMA PRO - INSTRUCT successfully utilizes Python program to solve the given question in the multi-round interaction.

heading: F MT-Bench
text: MT-bench is a collection of demanding multi-turn open-ended questions designed for evaluating chat assistants. In order to automate the evaluation process, we employ powerful LLMs, such as GPT-4, to act as judges and assess the quality of the models’ responses. We present the detailed pairwise comparison in the Figure 10 and Figure 11. Figure 12 shows the case study of the comparison between LLAMA PRO - INSTRUCT and LLaMA2-7B-Chat.

heading: G Token Distribution
text: We assess the token distribution between LLaMA-2-7B and LLAMA PRO, employing the methodology proposed by Lin et al. (2023). Specifically, for a given user query q = {q1, q2, ...}, we input it into LLAMA PRO to obtain its output o = {o1, o2, ...} using greedy decoding. For each position t, we define a context at this position as xt = q + {o1, ..., ot−1}. We denote the aligned model’s probability distribution for predicting the next token at this position as Palign, where ot has the highest probability.
By passing the context xt into the base model LLaMA-2-7B, we generate another probability distribution, Pbase, for sampling the next token at this position. First, the aligned model with greedy decoding is
used to generate a full output o. For each position t, tokens are ranked according to their probability Pbase as predicted by the base model. The rank of ot in this sorted list is defined as the ’base rank’, denoted as η. This categorizes positions into three types: (1) unshifted positions (η = 1): ot is the top-ranked token in both Pbase and Palign, having the highest probability; (2) marginal positions (1 < η ≤ 3): although ot is not the top-ranked token in Pbase, it is still likely to be sampled for decoding, with the 2nd or 3rd highest probability; (3) shifted positions (η > 3): in this case, ot is rather unlikely to be sampled by Pbase, indicating a significant distribution shift from Pbase to Palign.
We conduct a perplexity evaluation of LLaMA-2-7B and LLAMA PRO across general and code corpora. For the general domain, we utilize two different versions of the LAMBADA dataset. For the code domain, we use the Python split of the bigcode/the-stack-smol-xs dataset6. The results, presented in Table 11, indicate that LLAMA PRO effectively retains the language modeling ability for the general corpus while enhancing its proficiency in the code domain.
6https://huggingface.co/datasets/bigcode/the-stack-smol-xs

heading: H Domain of Law
text: Table 12 shows the hyper-parameters we use to do the ablation study in the domain of law. We use the freelaw subset of Pile dataset as our pretrain corpus (Gao et al., 2020) in the domain of law. This subset has 51.2 GiB raw size and 16.7B tokens with 3.6M documents.
The Unfair-ToS dataset, which we use to evaluate the performance of law, contains Terms of Service (ToS) from online platforms (e.g., YouTube, Ebay, Facebook, etc.). The dataset has been annotated on the sentence-level with 8 types of unfair contractual terms (sentences), meaning terms that potentially violate user rights according to the European consumer law. The UNFAIR-ToS task is a multilabel classification task. To get model predictions for this task, we categorize it as a multiple-choice question as the method Cheng et al. (2023) uses. The accuracy of an individual data example is considered true if the model prediction (i.e., the option with the highest per-token likelihood) belongs to the label(s) set. We evaluate the Unfair-ToS dataset in a 4-shot scenario just like Cheng et al. (2023).
Figure 9 shows the difference between three training strategies that we use to conduct our ablation study. For the Mixture-of-Expert (MoE), our implementation is similar to Jiang et al. (2024). We use 2 experts and for each token, both experts will be activated. Specifically, We extend each FFN for all 32 layers, keep the original ‘W3‘ unchanged, learn an additional Linear layer with weights ‘Ŵ3‘, and at the same time add two new learnable parameters ‘α1, α2‘ , when forward the output of Linear corresponding to W3, Ŵ3 will be weighted and summed with softmax(α1, α2) and fed into the next block.

heading: Future Work
text: Future research could explore extending the application of our block expansion method to other domains, such as maintaining original language ability in multimodal large language models(Ge et al. , 2023; Bai et al. , 2023), and multi-
lingual domains