LL A MA P RO: Progressive LLaMA with Block Expansion

Chengyue Wu1,2 Yukang Gan2 Yixiao Ge2 *

3
Zeyu Lu Jiahao Wang1 Ye Feng4
Ying Shan2 Ping Luo1
1 2
The University of Hong Kong ARC Lab, Tencent PCG
3 4
Shanghai Jiao Tong University Beijing Language and Culture University

[Link]

Abstract
Humans generally acquire new skills with-
arXiv:2401.02415v2 [[Link]] 30 May 2024

out compromising the old; however, the op-

posite 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 Trans-
former blocks. We tune the expanded blocks
using only new corpus, efficiently and ef-
fectively improving the model’s knowledge
while mitigating forgetting. In this paper,
we experiment on the corpus of code and
math, yielding LL A MA P RO-8.3B, a versatile
foundation model initialized from LLaMA2-
7B, excelling in general tasks, programming,
Figure 1: LL A MA P RO - I NSTRUCT delivers state-
and mathematics. LL A MA P RO and its
of-the-art performance across a wide variety of tasks,
instruction-following counterpart (LL A MA
ranging from general language to specific domains, su-
P RO - I NSTRUCT) achieve advanced perfor-
perior to existing models from the LLaMA series.
mance among various benchmarks, demonstrat-
ing superiority over existing open models in
the LLaMA family and the immense potential
of reasoning and addressing diverse tasks as an lored data recipes. While feasible, they require sub-
intelligent agent. Our findings provide valuable stantial computational resources and vast amounts
insights into integrating natural and program- of data, which poses a challenge to the democra-
ming languages, laying a solid foundation for tization of LLM research. Consequently, another
developing advanced language agents that op- line of research, known as domain-adaptive pre-
erate effectively in various environments. training, focuses on post-pretraining with domain-
specific corpora (Gururangan et al., 2020). These
1 Introduction approaches have demonstrated efficacy in adapting
various LLMs to specific domains (Roziere et al.,
The advent of Large Language Models (LLMs) 2023; Azerbayev et al., 2023; Wu et al., 2023b; Xu
has revolutionized the field of natural language et al., 2023b), resulting in enhanced performance
processing, exhibiting remarkable proficiency in a on downstream domain-specific tasks at a reduced
variety of real-world tasks (OpenAI, 2023; Chowd- computational cost.
hery et al., 2023). Despite the versatility, LLMs Nonetheless, a considerable obstacle emerges
still fall short in certain domains, for example, pro- in catastrophic forgetting (De Lange et al., 2021).
gramming, mathematics, biomedical, or finance. Post-pretraining often leads to a decline in the
This limitation impedes the progress of developing model’s original general abilities, inhibiting the
generic language agents for broader applications. fine-tuned performance of the model on diverse
Existing works (Liu et al., 2023; Li et al., 2023a; tasks (Cheng et al., 2023; Dong et al., 2023).
Wu et al., 2023b) attempted to improve the multi- This necessitates a method that can inject domain-
faceted capabilities of pre-trained LLMs with tai- specific knowledge into LLMs while preserving
* Correspondence to yixiaoge@[Link]. their general abilities, thereby enhancing their com-
prehensive capabilities. excelling in general tasks, programming, and
Towards this end, we introduce a simple yet ef- mathematics.
fective post-pretraining method, termed block ex-
pansion. We expand the off-the-shelf pre-trained • We benchmark the family of LL A MA P RO
LLM using copied Transformer blocks, as illus- on extensive datasets, including both tradi-
trated in Figure 2. The newly added blocks, whose tional and agent-oriented tasks, demonstrating
linear layers are zero-initialized to enable iden- its superiority and great potential in broader
tity mapping, are further tuned with only domain- complex applications.
specific corpus while the remaining blocks are
frozen. After tuning, the extended pre-trained 2 Related Work
model excels in both general and domain-specific
tasks. Advancements in Large Language Models. Re-
In practice, we extend the pre-trained LLaMA2- cent advancements in large language models have
7B (Touvron et al., 2023) by eight more blocks, led to significant progress, with model and data
yielding LL A MA P RO, a foundation model with scale growth driving state-of-the-art performance
8.3B parameters, and enhanced performance in pro- across various tasks (Hoffmann et al., 2022; Ka-
gramming, coding, and reasoning. We pre-train plan et al., 2020; Chowdhery et al., 2023). The
LL A MA P RO’s expanded blocks on 80B tokens development of generalist models has enabled
using open-source code and math data for 2830 addressing diverse problems and rapid adapta-
GPU Hours (16 NVIDIA H800 GPUs for about tion to new tasks (Radford et al., 2019; Brown
7 days). We further perform supervised instruc- et al., 2020). The open-source community has
tion tuning (fully fine-tuning of all the blocks, aka further contributed by releasing powerful models
SFT) on LL A MA P RO with approximately 80M like LLaMA (Touvron et al., 2023) and CodeL-
tokens, yielding LL A MA P RO - I NSTRUCT. It LaMA (Roziere et al., 2023). Our work builds
is noted that pre-trained models produced by our upon these developments, providing a method for
block expansion method are well-compatible with specializing LLMs in the code domain, fostering
the subsequent SFT techniques without specific future research and applications.
modification.
As shown in Figure 1, LL A MA P RO - Post-pretraining. Language model applications
I NSTRUCT reaches state-of-the-art performance typically involve a two-step process: general-
across a broad range of general, code (i.e., Hu- domain pretraining followed by domain-specific
manEval), and math (i.e., GSM8K) tasks. Fur- training (Roziere et al., 2023; Azerbayev et al.,
thermore, we assess the capabilities of LL A MA 2023). Fine-tuning often aims to enhance
P RO - I NSTRUCT as a language agent across vari- instruction-following abilities (Sanh et al., 2021;
ous scenarios (i.e., MINT-Bench), with a focus on Wei et al., 2021; Wang et al., 2023d) or align model
the tool usage abilities and the capacity to ground outputs with human preferences (Ziegler et al.,
in environmental and human feedback. We also 2019; Ouyang et al., 2022; Bai et al., 2022). Some
employ GPT-4 (OpenAI, 2023) automatic evalua- research explores parameter-efficient fine-tuning
tion to assess LL A MA P RO’s ability to serve as an methods for adapting pretrained models to new
effective assistant (i.e., MT-Bench). Comprehen- domains (Houlsby et al., 2019; Hu et al., 2021;
sive experimental results indicate the superiority of Wu et al., 2023a), while others focus on contin-
LL A MA P RO - I NSTRUCT over other models from ual learning post-pretraining (Wang et al., 2023b;
the LLaMA family on both benchmarks and practi- Gupta et al., 2023; Scialom et al., 2022). Parameter-
cal applications. Our contributions are three-fold: efficient tuning methods like adaptor and LoRA
are generally applied during the instruction tuning
• We propose a novel post-pretraining method
phase rather than the pretraining phase. In contrast,
for LLMs, termed block expansion, enabling
our focus is on enhancing the capacity of LLMs by
the injection of new knowledge while preserv-
increasing their depth during continued pretraining.
ing the initial capabilities.
Our work proposes an adaptation strategy that com-
• We introduce LL A MA P RO and LL A MA bines continued training with general capability
P RO - I NSTRUCT, versatile LLMs that well maintenance, allowing LLMs to specialize without
integrate natural and programming languages, sacrificing overall performance.
 Language Model Language Model with Block-expansion
(LLaMA2) (LLaMA Pro)

Output Tokens Output Tokens Identity

Copy
Expand 🔥 Decoder Block ×P
×N ×N
🔥 Decoder Block ×M ❄ Decoder Block ×M

Input Tokens Input Tokens

Training
Data
Huge Unlabeled Corpus Aspect Corpus

(a) Pre-Training (b) Block-expansion Post-pretraining

Figure 2: (a) We begin with a large language model (LLM) pre-trained on a massive unlabeled corpus, resulting in a
model with strong general capabilities. Here we select the off-the-shelf LLaMA2 for convenience. (b) We employ
backbone expansion and fine-tune the expanded identity blocks using the aspect corpus while freezing the blocks
inherited from the base model. The model after post-pretraining can be used for instruction tuning as usual.

Progressive Learning. Progressive training has output y has the same dimension as the input x.
gained attention for accelerating large-scale model The MHSA operation is a crucial component of the
training in computer vision (Zhang et al., 2023) and transformer, defined as:
NLP research (Yao et al., 2023; Li et al., 2023b).
Gong et al. (2019) proposed a stacking method MHSA(Q, K, V ) = Concat(head1 , . . . , headh )W O
doubling model depth successively. Compound- (2)
Grow(Gu et al., 2020) extends stacking with Feed- where Q, K, and V are the query, key, and value
Forward Network expansion in schedule design. matrices, respectively, and W O is the output weight
Shen et al. (2022) introduced a staged method sup- matrix without bias . Each head is computed as:
porting hidden size expansion. Bert2BERT(Chen
et al., 2021a) and LiGO (Wang et al., 2023a) ac- headi = Attention(xWiQ , xWiK , xWiV )
Qi KiT
 
commodate all growth dimensions. Our method
Attention(Qi , Ki , Vi ) = Softmax √ Vi
utilizes depth growth to maintain general perfor- dk
mance while adapting to specific domains. (3)

3 Method with WiQ , WiK , and WiV being the corresponding

weight matrices for the i-th head.
3.1 Preliminaries: The LLaMA Block
The FFN block in the LLaMA model utilizes the
The LLaMA block consists of a multi-head self- SwiGLU activation function, which is defined as:
attention (MHSA) mechanism followed by a
position-wise feed-forward network (FFN) with SwiGLU(x, W, V ) = SiLU(xW ) ⊗ (xV )
(4)
residual connections and a Swish-Gated Linear FFN(x) = SwiGLU(x, W1 , W2 )W3
Unit (SwiGLU) operation as Figure 3 shows. Given
an input x, the LLaMA block produces an output y where ⊗ denotes element-wise multiplication, W1 ,
as described by the following equations: W2 , and W3 are the weight matrices without bias,
SiLU(x) = x ⊗ σ(x).
x′ = x + MHSA(RMSNorm(x))
(1) 3.2 Block Expansion
y = x′ + FFN(RMSNorm(x′ ))
Given a model with blocks (ϕ0 , ϕ1 , ..., ϕL ), the
The input x has a dimension of n × d, where n is block expansion incorporates an identity block ϕid
the sequence length and d is the hidden size. The after each block in the original model, ensuring
 ent of the loss function L with respect to the RM-
Tokens Tokens
SNorm weight w during backpropagation would
be zero. This would prevent the training of RM-
Linear Zero-Linear SNorm, implying that when RMSNorm(x′ ) = 0,
SwiGLU SwiGLU the following condition will hold:
SiLU SiLU ∂L ∂L ∂FFN(RMSNorm(x′ )) ∂RMSNorm(x′ )
= = 0.
Linear Linear Linear Linear
∂w ∂y ∂RMSNorm(x′ ) ∂w
(5)
This equation signifies that the gradient of the
RMSNorm RMSNorm
loss function with respect to the weight of RM-
MHSA MHSA SNorm is zero, which would hinder the training of
Linear Zero-Linear
the RMSNorm module. This is further explained
in Appendix A. Referring to the LLaMA block
Scaled Dot-Product Scaled Dot-Product formulation in Equation 1, the identity can be
Attention Attention
achieved as long as MHSA(RMSNorm(x)) = 0
Linear Linear and FFN(RMSNorm(x′ )) = 0. We initialize the
W O and W3 weight matrices in the identity blocks
RMSNorm RMSNorm to zero. Due to the presence of residual connec-
tions and the absence of bias terms in the LLaMA
Tokens Tokens block, only the residual flows through the identity
block. As a result, the entire block is reduced to
(a) LLaMA Block (b) LLaMA Block after Identity Copy
an identity block at initialization, preserving the
Figure 3: (a) An overview of the LLaMA Block, com- output from the initial model.
prising an MHSA mechanism followed by the FFN with The entire training pipeline is depicted in Fig-
SwiGLU activation. (b) The Identity LLaMA block af- ure 2. Our method concentrates on the post-
ter an identity copy, achieved by initializing the output pretraining stage, targeting specific domain cor-
linear matrix to zero in order to preserve the output from pora. We begin by initializing our model with large
the base LLaMA model.
language models trained on extensive unlabeled
general corpora, where all blocks will be fine-tuned.
that the expanded model maintains the same output To enhance the model’s capacity for accommodat-
after expansion. The identity block is defined as ing additional domain knowledge while retaining
ϕid (x) = x so the input and output are identical. its general knowledge, we employ block expansion
Suppose we have an initial model with L blocks to increase the number of blocks in the LLM. Dur-
that needs to be expanded to L′ blocks. First, we ing this process, we only fine-tune the newly added
partition the original L blocks into N groups, with blocks while freezing the original blocks, thereby
L preserving the general abilities of the model.
each group containing N blocks. For each group,
we create identity copies of the top P blocks and
stack them on top of each group, as depicted in 4 Experiments
Figure 3. We arrange these blocks in an inter- This section presents our key experimental findings.
leaved manner to maintain the structural character- We begin with experimental settings (described
istic of the transformer model, whose prior is that in Sec. 4.1), and then verify the effectiveness of
deeper blocks encode more complex information block expanded tuning after pretraining (described
(Van Aken et al., 2019; Tenney et al., 2019). This in Sec. 4.2). Next, we give the supervised finetun-
process leads to an increased depth in the model ing (SFT) results (described in Sec. 4.3). Finally,
while maintaining its output behavior. ablation studies of the key design choices are pre-
Shen et al. (Shen et al., 2022) proposed the ini- sented (described in Sec. 4.5).
tialization of scale parameters in the Norm modules
within the identity blocks to zero for the construc- 4.1 Experimental Settings
tion of the identity block. However, this approach Pretrain details. We construct a dataset that con-
may not be effective when applied to the LLaMA centrates on code and math. For the code compo-
block. The reason lies in the fact that the gradi- nent, we rely on the Stack-dedup dataset, which
is a compilation of permissively licensed source 54

codes from GitHub. Among all the programming LLaMA Pro

languages available in Stack-dedup, we specifically LLaMA2-7B
50
utilize the Python split. As for the math component,
we opt for the Proof-pile-2 dataset (Azerbayev CrystalCoder
LLaMA-7B
46
et al., 2023), a 55-billion-token amalgamation of

Language Tasks Avg.

scientific papers, web data containing mathemat- Falcon-7B OpenLLaMA-v2-7B

ical content, and mathematical code. The details 42

can be found in Appendix B. CodeLLaMA-7B

We initialize our base model with LLaMA2-7B 38
and expand the number of blocks from 32 to 40 us-
StarCoder-15B
ing an interleaved approach. In the block expansion Tokens

34
process, we configure the parameters as P = 1,
1.0T 1.5T 2.0T 2.5T
M = 4, and N = 8, resulting in 8 groups where
each group expands from 4 blocks to 5 blocks. For 30
6 14 22 30 38
the code and math corpus pretraining, we employ Code Tasks Avg.

a batch size of 1024, a sequence length of 4096, a

Figure 4: We compare LL A MA P RO’s general perfor-
warmup ratio of 6%, a learning rate of 2e-4, and a mance and code performance to a set of models trained
Cosine learning rate scheduler. We also use bf16 around the same time, spanning from general LLMs to
mixed precision, a weight decay of 0.1, and gradi- code-oriented LLMs. The size of the blobs is propor-
ent clipping at 1.0. To speed up the training process, tional to the number of tokens trained. Mistral-7B is not
we apply the flash-attention mechanism. included here, as the number of tokens is not reported
Our experiment is conducted on 16 NVIDIA in its paper.
H800 GPUs. LL A MA P RO 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 LL A MA P RO - I NSTRUCT as shown in Table 7.
resources to achieve a better performance of the The final sft dataset consists of approximately 1M
domain specific tasks in the inference. samples. To fine-tune the basic models, we employ
Training stage cost: Our approach requires specific configurations, including a batch size of
fewer computational resources since only the newly 128, a sequence length of 4096, 0.03 warmup ra-
added blocks are tuned during training. As illus- tio, a learning rate of 2e-5, a Cosine learning rate
trated in Figure 4, LLaMA Pro-8B (1B parameters scheduler, and bf16 mixed precision.
tuned for 80B tokens) incurs less training over-
head 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 Evaluation details. We conduct a comparative
and domain-specific performance, offering a more analysis of LL A MA P RO with the latest state-of-
cost-effective solution. the-art (SOTA) Large Language Models (LLMs).
Inference stage cost: Although our method re- The evaluation is performed on six key general
quires more resources during inference than the ini- benchmarks using the Eleuther AI Language Model
tial LLM, it strikes a balance between performance Evaluation Harness1 , a unified framework designed
and efficiency. LLaMA Pro-8B outperforms larger to test generative language models across a vast
models like LLaMA2-13B and LLaMA2-34B in array of evaluation tasks. For code-related tasks,
the code domain while demanding significantly we employ the BigCode Evaluation Harness2 to
fewer resources during training and inference. evaluate HumanEval and MBPP, and we report the
pass@1 rate of code tasks with greedy decoding.
SFT details. During the instruction fine-tuning The evaluation details can be found in Appendix
phase, we combine five data sources to create D.
 Model Language Tasks Math Tasks Code Tasks Avg.
ARC HellaSwag MMLU TruthfulQA Winogrande GSM8K GSM8K-PoT HumanEval MBPP
Pretrained comparison
LL A MA P RO (8B) 54.10 77.94 47.88 39.04 73.95 17.89 25.42 28.66 33.20 44.23
CrystalCoder (7B) 47.01 71.97 48.78 35.91 67.17 10.77 24.96 28.38 36.38 41.26
LLaMA2-7B 53.07 78.59 46.87 38.76 74.03 14.48 17.68 13.05 20.09 39.62
CodeLLaMA-7B 39.93 60.80 31.12 37.82 64.01 5.16 25.20 33.50 41.40 37.66
StarCoder-15B 30.38 47.93 29.96 41.28 56.12 9.48 25.09 33.63 43.28 35.24
LLaMA-7B 50.94 77.81 35.69 34.33 71.43 8.04 10.46 10.61 17.04 35.15
OpenLLaMA-v2-7B 43.69 72.20 41.29 35.54 69.38 3.49 5.46 15.32 12.69 33.23
Falcon-7B 47.87 78.13 27.79 34.26 72.38 4.62 4.32 9.42 13.39 32.46
SFT comparison
LL A MA P RO - I NSTRUCT 52.30 76.88 52.57 48.80 72.53 43.59 55.61 44.51 37.88 53.85
LLaMA2-7B-Chat 52.90 78.55 48.32 45.57 71.74 7.35 19.73 14.63 21.60 40.04
CodeLLaMA-7B-Instruct 36.52 55.44 34.54 41.25 64.56 7.96 34.67 34.80 44.4 39.35
WizardCoder-Python-7B 41.81 65.06 32.29 36.32 61.72 4.70 17.60 42.07 47.20 38.75
WizardMath-7B 54.10 79.55 45.97 43.65 72.69 2.73 25.57 12.20 18.00 39.38

Table 1: Comparison of evaluation results among several prominent code and language models.

4.2 Pretrain Results Model MT Bench

We evaluate LL A MA P RO’s performance with Alpaca-13B 4.53
benchmark datasets from the Open LLM Leader- CodeLLaMA-7B-Instruct 5.71
board. Furthermore, we incorporate coding bench- Vicuna-7B 6.17
mark datasets, including HumanEval pass@1 and LLaMA2-7B-Chat 6.27
MBPP pass@1, as well as the math benchmark LL A MA P RO - I NSTRUCT 6.32
GSM8K, to provide a comprehensive evaluation.
We compare the performance of LL A MA P RO with Table 2: GPT-4 automatic evaluation of Chatbot models.
a selection of state-of-the-art pretrained models LL A MA P RO - I NSTRUCT outperforms widely used
that were trained around the same period with sim- LLaMA community chatbots.
ilar size. This includes general-purpose pretrained
models like LLaMA2 and code-oriented pretrained dergone fine-tuning with an additional 80B tokens
models like CodeLLaMA. The results are presented in conjunction with LLaMA2, which more than
in Table 1. doubles the code tasks average performance. In
The results highlight that LL A MA P RO effec- contrast, CodeLLaMA is fine-tuned with 500B to-
tively balances natural language processing and kens. LL A MA P RO excels in general performance
coding capabilities. It not only preserves the gen- while maintaining code performance that is com-
eral performance of its base model, LLaMA2-7B, petitive with code-oriented LLMs, whether they are
but also surpasses it in the average performance of trained from scratch, such as StarCoder-15B and
general language tasks. Conversely, CodeLLaMA- CrystalCoder, or fine-tuned like CodeLLaMA-7B.
7B sacrifices general performance. We attribute
4.3 SFT Results
this improvement to our expansion design, which
freezes the initial LLaMA blocks to maintain their Modern LLMs typically undergo supervised fine-
capabilities and increases the blocks to accommo- tuning or instruction tuning after pretraining on
date domain-specific knowledge. vast amounts of unlabeled data. In this section, we
As depicted in Figure 4, LL A MA P RO shows aim to demonstrate that our expansion strategy can
robust general performance alongside code perfor- adapt to this widely used training pipeline, just as
mance that is on par with code-oriented LLMs. Sit- traditional LLMs do.
uated on the Pareto frontier, LL A MA P RO has un- Table 1 presents a comparison of evaluation
results among several prominent supervised fine-
1
[Link] tuning (SFT) LLMs from the LLaMA community,
lm-evaluation-harness
2
[Link] across general tasks, math tasks, and code tasks
bigcode-evaluation-harness benchmarks. As a singular SFT model, LL A MA
 Model Interaction Turns Avg. Add 1 Block
1 2 3 4 5 1.4 Add 2 Block
AgentLM-7B 0.0 4.44 5.29 6.48 7.34 4.71
Add 4 Block
1.3 Add 8 Block
CodeLLaMA-7B-Instruct 0.34 7.85 10.24 9.73 8.70 7.37 Add 16 Block
LLaMA2-7B-Chat 1.02 4.27 6.66 6.48 7.34 5.77 Add 32 Block
1.2 MoE

Training loss
Mistral-Instruct-v0.1 1.54 12.12 13.31 14.16 13.99 11.02
LoRA
LL A MA P RO - I NSTRUCT 0.68 12.63 11.95 11.95 14.68 10.38 1.1 Finetune

Table 3: : In the tool-augmented reasoning assessments, 1.0

we evaluate the model’s proficiency in integrating tools 0.9
into its reasoning workflow. The model’s effectiveness
is measured by its success rate across various stages of 0.8
interaction.
0 1000 2000 3000 4000 5000
Training steps

Figure 5: The training loss is analyzed with respect to

P RO - I NSTRUCT attains state-of-the-art perfor-
the addition of varying blocks and mixture-of-expert
mance, even when compared to specifically tuned (MoE) expansion, in conjunction with traditional train-
models such as WizardCoder and WizardMath. ing strategies such as finetuning and LoRA.
This demonstrates its more comprehensive capabil-
ities. 4.4 Mistral-Pro Results
As seen in Figure 1, LL A MA P RO - I NSTRUCT We experimented with block expansion on Mistral-
boosts both code and math tasks to SOTA perfor- 7B (Jiang et al., 2023), training it on code and
mances while maintaining reliable general perfor- mathematics datasets. The resulting pretrained per-
mance. We enhance the average performance of formance is detailed in Table 4, highlighting supe-
LLaMA2-7B-chat and CodeLLaMA-7B-instruct rior outcomes across various benchmarks, partic-
by 13.81% and 14.50% respectively, which high- ularly in the domains of code and math. Notably,
lights the benefits of balancing textual and coding it demonstrates competitive results compared to
abilities. the new open-source model Gemma (Team et al.,
2024), while incurring significantly lower training
To assess the comprehensive conversational per- overhead. We further utilized the MetaMath dataset
formance of the LL A MA P RO - I NSTRUCT assis- (Yu et al., 2023) for supervised fine-tuning. Our ap-
tant, we evaluate it using the MT-Bench with GPT-4 proach yielded scores of 78.4 for GSM8k and 30.3
automatic scoring, as proposed by Vicuna (Zheng for MATH, surpassing Mistral’s scores of 77.7 and
et al., 2023). As depicted in Table 2, LL A MA P RO 28.2, respectively. Additional details are provided
- I NSTRUCT surpasses widely used chatbots from in Appendix C.
the LLaMA community. This indicates its potential
as a chatbot capable of providing helpful responses, 4.5 Ablation Study
in addition to its impressive performance in tradi-
Apart from the aspect of code corpus, we explore
tional benchmarks. The details of MT-Bench can
our method on another domain: law, with the
be found in the Appendix F.
freelaw subset of Pile dataset as our pretrain corpus
We use MINT-Bench (Wang et al., 2023c) to (Gao et al., 2020). We evaluate on UNFAIR-ToS
evaluate our model’s ability to solve multi-turn in- (Lippi et al., 2019) of the LexGLUE benchmark
teractions by using tools. MINT-Bench tests LLMs’ (Chalkidis et al., 2021).
ability to use tools by generating and executing In our experiment, we assess the scalability of
Python code, focusing on tool-augmented task- our block expansion method in terms of training
solving and leveraging natural language feedback. loss and downstream task performance as we in-
MINT includes eight datasets covering reasoning, crease the number of added blocks. We also com-
code generation, and decision-making. The details pare our method with the Mixture-of-Expert (MoE)
of MINT can be found in the Appendix E. The expansion method (Fedus et al., 2022) and tradi-
results are shown in Table ??. LL A MA P RO - I N - tional training strategies, such as fine-tuning and
STRUCT achieves SOTA performance compared to LoRA (Hu et al., 2021). The details can be found
similar size models in multi-turn interactions with in Appendix H.
the use of tools. We analyze the training loss with varying added
 Model ARC Hellaswag MMLU TruthfulQA Winogrande GSM8K HumanEval
Gemma-7B 61.9 82.2 64.6 44.8 79.0 50.9 32.3
Mistral-7B 60.8 83.3 62.7 42.6 78.0 39.2 28.7
Mistral-Pro (Ours) 63.2 82.6 60.6 48.3 78.9 50.6 32.9

Table 4: Comparison between the original Mistral-7B (Jiang et al., 2023), Gemma-7B (Team et al., 2024), and our
Mistral-Pro with the Open LLM leaderboard metrics.

Method Language Tasks Law Task Avg.

ARC HellaSwag MMLU TruthfulQA Winogrand Avg. Unfair-ToS
Add 1 Block 52.30 77.92 38.62 37.80 73.16 55.96 67.45 61.71
Add 2 Block 53.16 77.91 39.62 38.92 73.01 56.52 69.57 63.05
Add 4 Block 52.39 76.92 37.30 40.53 72.22 55.87 71.31 63.59
Add 8 Block 52.90 76.63 41.74 39.83 72.38 56.70 75.11 65.91
Add 16 Block 51.88 76.59 41.35 40.13 71.82 56.35 75.17 65.76
Add 32 Block 50.77 76.72 40.68 41.66 72.77 56.52 73.93 65.23
Mixture-of-Expert (MoE) 51.45 76.51 42.47 40.13 72.23 56.56 67.27 61.92
Fine-tuning 48.81 74.49 41.13 41.49 69.14 55.01 70.63 62.82
LoRA 53.50 78.12 44.30 40.96 73.88 58.15 65.34 61.75
Prefix Stacking (8 Block) 27.82 26.12 23.12 22.52 47.20 29.36 0.81 15.08
Suffix Stacking (8 Block) 52.56 77.89 39.10 39.03 72.38 56.19 60.98 58.59

Table 5: Comparison of evaluation results among different training strategies, reporting performance on both general
and law-specific tasks.

not necessarily guarantee superior performance

LLaMA-2-7B LLaMA-PRO-8B on domain-specific tasks. Therefore, we evaluate
Unshifted (92.6%) Marginal (6.8%) Shifted (0.7%) 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 domain-
'\n', '•', 'the', 'The', '3', '8', 'C', 'Y', 'S', '1', 'a', '4', specific task, larger models achieve better perfor-
'9', '2', '.', 'This', 'in', 'In', '#', ',', 'A', 'can', 'size', mance. We find that adding eight blocks provides
'6', 'of', '7', 'o', 'and', 'P', 'F', 'Frank', 'you', 'find', optimal performance with minimal cost compared
to larger models, hence we adopt this as our default
'p', '$', 'Le', '5', 'B', '-', '(', '0', 'f', ':', 'for', 'Am', 'I',
strategy. The performance of MoE is comparable
'c', 'av', ’g‘, …
to our method with four added blocks. Figure 9
illustrates the differences between traditional train-
Figure 6: Token distribution shift after block expan- ing strategies such as fine-tuning and LoRA, and
sion compared to the initial LLaMA-2-7B. The propor- our proposed method. We observe that while LoRA
tions of unshifted, marginally shifted, and significantly effectively preserves the general ability, it struggles
shifted tokens are color-coded and presented as percent- to model the distribution of a new domain, as also
ages. Frequently shifted tokens are displayed below. evidenced by the training loss depicted in Figure
5. In contrast, full fine-tuning results in a more
blocks (Figure 5). The loss consistently decreases significant drop in general performance. Here we
during training, regardless of the number of added use a rank of 1024 for LoRA, resulting in a number
blocks, and decreases more rapidly with larger of trainable parameters comparable to our method.
models. These findings indicate that our method In line with the approach of Lin et al. (2023), we
demonstrates strong scalability with larger models analyze the token distribution between the original
and more data. LLaMA and LL A MA P RO to assess the similarity
However, a lower overall training loss does in their behavior when answering general ques-
 63.61 64.67 LLaMA2-7B 5 Scope and Limitations
60 LLaMA Pro
Although our study presents a promising method
50 47.56 for balancing general and domain-specific capabili-
41.4 ties in LLMs, its scope is limited to the language
40 39
modality, especially programming language and
Score

31.2
30 English. Future research could explore extending
20 19.6 21.2 the application of our block expansion method to
other domains, such as maintaining original lan-
10 guage ability in multimodal large language mod-
els(Ge et al., 2023; Bai et al., 2023), and multi-
0
GSM8K MATH HumanEval MBPP lingual domains.
Tasks

Figure 7: By fine-tuning both LLaMA2-7B and

6 Conclusion
LL A MA P RO using the same instruction dataset,
In this study, we introduced a novel block
LL A MA P RO consistently outperforms LLaMA2-7B
across all tasks. This result highlights the effectiveness
expansion method for Large Language Mod-
of our method, as it demonstrates that LL A MA P RO els (LLMs) post-pretraining, aiming to enhance
successfully encodes more domain knowledge during domain-specific abilities while preserving the orig-
the pretraining process. inal general capabilities. Our approach effectively
balances the model’s performance across both gen-
eral and domain-specific tasks. We demonstrated
tions from the Alpaca dataset (Taori et al., 2023). the effectiveness of our method through LL A MA
As depicted in Figure 6, the token distribution shift P RO, an LLM initialized from LLaMA2-7B with 8
between LLaMA and LL A MA P RO is subtle. De- added blocks, which outperformed other LLaMA-
tailed information can be found in Appendix G. series models on comprehensive benchmarks.
We also analyze the impact of the position where
the identity blocks are added, either at the bottom
7 Ethical Statement
or the top of the model, compared to adding them LL A MA P RO and LL A MA P RO - I NSTRUCT are
interleaved, as shown in Table 5. We observe that designed for a wide range of NLP tasks, with a
adding blocks at the bottom results in poor eval- focus on programming, mathematics, and general
uation performance, likely because it disrupts the language tasks. It suits scenarios requiring integra-
model’s foundation, causing errors to propagate tion of natural and programming languages. While
throughout the model. Adding blocks at the top LLaMA-Pro addresses some limitations of previ-
of the model (Gong et al., 2019) preserves the ini- ous models in the series, it may still encounter
tial model’s performance, but its performance on challenges specific to highly specialized domains
domain-specific tasks is lower than when adding or tasks. Users should be aware of potential biases
blocks interleaved. in the model and use it responsibly, considering its
As highlighted in the LIMA study (Zhou et al., impact on various applications with the LLaMA-2
2023), the majority of knowledge in large language license.
models is acquired during pretraining, with only a
limited amount of instruction tuning data required
References
to generate high-quality output. To investigate
the extent of knowledge encoded during pretrain- Jacob Austin, Augustus Odena, Maxwell Nye, Maarten
ing, we conducted a comparative analysis between Bosma, Henryk Michalewski, David Dohan, Ellen
Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. 2021.
LLaMA2-7B and LL A MA P RO using the same Program synthesis with large language models. arXiv
instruction dataset, as illustrated in Figure 7. Our preprint arXiv:2108.07732.
results showed that LL A MA P RO consistently out-
performs LLaMA2-7B across all tasks, indicating Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster,
Marco Dos Santos, Stephen McAleer, Albert Q Jiang,
that our method effectively enables LL A MA P RO Jia Deng, Stella Biderman, and Sean Welleck. 2023.
to encode more domain-specific knowledge during Llemma: An open language model for mathematics.
the pretraining phase. arXiv preprint arXiv:2310.10631.
Jinze Bai, Shuai Bai, Shusheng Yang, Shijie Wang, Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian,
Sinan Tan, Peng Wang, Junyang Lin, Chang Zhou, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias
and Jingren Zhou. 2023. Qwen-vl: A frontier large Plappert, Jerry Tworek, Jacob Hilton, Reiichiro
vision-language model with versatile abilities. arXiv Nakano, et al. 2021. Training verifiers to solve math
preprint arXiv:2308.12966. word problems. arXiv preprint arXiv:2110.14168.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Matthias De Lange, Rahaf Aljundi, Marc Masana, Sarah
Askell, Anna Chen, Nova DasSarma, Dawn Drain, Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh,
Stanislav Fort, Deep Ganguli, Tom Henighan, et al. and Tinne Tuytelaars. 2021. A continual learning
2022. Training a helpful and harmless assistant with survey: Defying forgetting in classification tasks.
reinforcement learning from human feedback. arXiv IEEE transactions on pattern analysis and machine
preprint arXiv:2204.05862. intelligence, 44(7):3366–3385.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Guanting Dong, Hongyi Yuan, Keming Lu, Cheng-
Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind peng Li, Mingfeng Xue, Dayiheng Liu, Wei Wang,
Neelakantan, Pranav Shyam, Girish Sastry, Amanda Zheng Yuan, Chang Zhou, and Jingren Zhou. 2023.
Askell, et al. 2020. Language models are few-shot How abilities in large language models are affected
learners. Advances in neural information processing by supervised fine-tuning data composition. arXiv
systems, 33:1877–1901. preprint arXiv:2310.05492.
Ilias Chalkidis, Abhik Jana, Dirk Hartung, Michael
Bommarito, Ion Androutsopoulos, Daniel Martin William Fedus, Barret Zoph, and Noam Shazeer. 2022.
Katz, and Nikolaos Aletras. 2021. Lexglue: A bench- Switch transformers: Scaling to trillion parame-
mark dataset for legal language understanding in en- ter models with simple and efficient sparsity. The
glish. arXiv preprint arXiv:2110.00976. Journal of Machine Learning Research, 23(1):5232–
5270.
Cheng Chen, Yichun Yin, Lifeng Shang, Xin Jiang,
Yujia Qin, Fengyu Wang, Zhi Wang, Xiao Chen, Leo Gao, Stella Biderman, Sid Black, Laurence Gold-
Zhiyuan Liu, and Qun Liu. 2021a. bert2bert: To- ing, Travis Hoppe, Charles Foster, Jason Phang, Ho-
wards reusable pretrained language models. arXiv race He, Anish Thite, Noa Nabeshima, et al. 2020.
preprint arXiv:2110.07143. The pile: An 800gb dataset of diverse text for lan-
guage modeling. arXiv preprint arXiv:2101.00027.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan,
Henrique Ponde de Oliveira Pinto, Jared Kaplan, Yuying Ge, Sijie Zhao, Ziyun Zeng, Yixiao Ge, Chen Li,
Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Xintao Wang, and Ying Shan. 2023. Making llama
Brockman, et al. 2021b. Evaluating large lan- see and draw with seed tokenizer. arXiv preprint
guage models trained on code. arXiv preprint arXiv:2310.01218.
arXiv:2107.03374.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei
Wenhu Chen, Xueguang Ma, Xinyi Wang, and Wang, and Tieyan Liu. 2019. Efficient training
William W. Cohen. 2023a. Program of thoughts of bert by progressively stacking. In International
prompting: Disentangling computation from reason- conference on machine learning, pages 2337–2346.
ing for numerical reasoning tasks. Transactions on PMLR.
Machine Learning Research.
Xiaotao Gu, Liyuan Liu, Hongkun Yu, Jing Li, Chen
Wenhu Chen, Ming Yin, Max Ku, Pan Lu, Yixin Wan, Chen, and Jiawei Han. 2020. On the transformer
Xueguang Ma, Jianyu Xu, Xinyi Wang, and Tony Xia. growth for progressive bert training. arXiv preprint
2023b. Theoremqa: A theorem-driven question an- arXiv:2010.12562.
swering dataset. arXiv preprint arXiv:2305.12524.
Daixuan Cheng, Shaohan Huang, and Furu Wei. 2023. Kshitij Gupta, Benjamin Thérien, Adam Ibrahim,
Adapting large language models via reading compre- Mats L Richter, Quentin Anthony, Eugene Belilovsky,
hension. arXiv preprint arXiv:2309.09530. Irina Rish, and Timothée Lesort. 2023. Continual pre-
training of large language models: How to (re) warm
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, your model? arXiv preprint arXiv:2308.04014.
Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul
Barham, Hyung Won Chung, Charles Sutton, Se- Suchin Gururangan, Ana Marasović, Swabha
bastian Gehrmann, et al. 2023. Palm: Scaling lan- Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey,
guage modeling with pathways. Journal of Machine and Noah A Smith. 2020. Don’t stop pretraining:
Learning Research, 24(240):1–113. Adapt language models to domains and tasks. arXiv
preprint arXiv:2004.10964.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot,
Ashish Sabharwal, Carissa Schoenick, and Oyvind Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou,
Tafjord. 2018. Think you have solved question an- Mantas Mazeika, Dawn Song, and Jacob Steinhardt.
swering? try arc, the ai2 reasoning challenge. arXiv 2020. Measuring massive multitask language under-
preprint arXiv:1803.05457. standing. arXiv preprint arXiv:2009.03300.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Stephanie Lin, Jacob Hilton, and Owain Evans. 2021.
Arora, Steven Basart, Eric Tang, Dawn Song, and Truthfulqa: Measuring how models mimic human
Jacob Steinhardt. 2021. Measuring mathematical falsehoods. arXiv preprint arXiv:2109.07958.
problem solving with the math dataset. NeurIPS.
Marco Lippi, Przemysław Pałka, Giuseppe Contissa,
Jordan Hoffmann, Sebastian Borgeaud, Arthur Men- Francesca Lagioia, Hans-Wolfgang Micklitz, Gio-
sch, Elena Buchatskaya, Trevor Cai, Eliza Ruther- vanni Sartor, and Paolo Torroni. 2019. Claudette: an
ford, Diego de Las Casas, Lisa Anne Hendricks, automated detector of potentially unfair clauses in
Johannes Welbl, Aidan Clark, et al. 2022. Train- online terms of service. Artificial Intelligence and
ing compute-optimal large language models. arXiv Law, 27:117–139.
preprint arXiv:2203.15556.
Zhengzhong Liu, Aurick Qiao, Willie Neiswanger,
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Hongyi Wang, Bowen Tan, Tianhua Tao, Junbo Li,
Bruna Morrone, Quentin De Laroussilhe, Andrea Yuqi Wang, Suqi Sun, Omkar Pangarkar, et al. 2023.
Gesmundo, Mona Attariyan, and Sylvain Gelly. Llm360: Towards fully transparent open-source llms.
2019. Parameter-efficient transfer learning for nlp. arXiv preprint arXiv:2312.06550.
In International Conference on Machine Learning,
pages 2790–2799. PMLR. Ziyang Luo, Can Xu, Pu Zhao, Qingfeng Sun, Xi-
ubo Geng, Wenxiang Hu, Chongyang Tao, Jing Ma,
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Qingwei Lin, and Daxin Jiang. 2023. Wizardcoder:
Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Empowering code large language models with evol-
and Weizhu Chen. 2021. Lora: Low-rank adap- instruct.
tation of large language models. arXiv preprint
OpenAI. 2023. Gpt-4 technical report. arXiv preprint
arXiv:2106.09685.
arXiv:2303.08774.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Men- Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida,
sch, Chris Bamford, Devendra Singh Chaplot, Diego Carroll Wainwright, Pamela Mishkin, Chong Zhang,
de las Casas, Florian Bressand, Gianna Lengyel, Guil- Sandhini Agarwal, Katarina Slama, Alex Ray, et al.
laume Lample, Lucile Saulnier, et al. 2023. Mistral 2022. Training language models to follow instruc-
7b. arXiv preprint arXiv:2310.06825. tions with human feedback. Advances in Neural
Albert Q Jiang, Alexandre Sablayrolles, Antoine Information Processing Systems, 35:27730–27744.
Roux, Arthur Mensch, Blanche Savary, Chris Alec Radford, Jeffrey Wu, Rewon Child, David Luan,
Bamford, Devendra Singh Chaplot, Diego de las Dario Amodei, Ilya Sutskever, et al. 2019. Language
Casas, Emma Bou Hanna, Florian Bressand, et al. models are unsupervised multitask learners. OpenAI
2024. Mixtral of experts. arXiv preprint blog, 1(8):9.
arXiv:2401.04088.
Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi,
Brown, Benjamin Chess, Rewon Child, Scott Gray, Jingyu Liu, Tal Remez, Jérémy Rapin, et al. 2023.
Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Code llama: Open foundation models for code. arXiv
Scaling laws for neural language models. arXiv preprint arXiv:2308.12950.
preprint arXiv:2001.08361.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhaga-
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas vatula, and Yejin Choi. 2021. Winogrande: An
Muennighoff, Denis Kocetkov, Chenghao Mou, Marc adversarial winograd schema challenge at scale.
Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Communications of the ACM, 64(9):99–106.
2023a. Starcoder: may the source be with you!
arXiv preprint arXiv:2305.06161. Victor Sanh, Albert Webson, Colin Raffel, Stephen H
Bach, Lintang Sutawika, Zaid Alyafeai, Antoine
Xiang Li, Yiqun Yao, Xin Jiang, Xuezhi Fang, Xuying Chaffin, Arnaud Stiegler, Teven Le Scao, Arun
Meng, Siqi Fan, Peng Han, Jing Li, Li Du, Bowen Raja, et al. 2021. Multitask prompted training en-
Qin, et al. 2023b. Flm-101b: An open llm and ables zero-shot task generalization. arXiv preprint
how to train it with $100 k budget. arXiv preprint arXiv:2110.08207.
arXiv:2309.03852.
Thomas Scialom, Tuhin Chakrabarty, and Smaranda
Wing Lian, Guan Wang, Bleys Goodson, Eugene Pent- Muresan. 2022. Fine-tuned language models
land, Austin Cook, Chanvichet Vong, and "Teknium". are continual learners. In Proceedings of the
2023. Slimorca: An open dataset of gpt-4 augmented 2022 Conference on Empirical Methods in Natural
flan reasoning traces, with verification. Language Processing, pages 6107–6122.
Bill Yuchen Lin, Abhilasha Ravichander, Ximing Lu, Sheng Shen, Pete Walsh, Kurt Keutzer, Jesse Dodge,
Nouha Dziri, Melanie Sclar, Khyathi Chandu, Chan- Matthew Peters, and Iz Beltagy. 2022. Staged
dra Bhagavatula, and Yejin Choi. 2023. The unlock- training for transformer language models. In
ing spell on base llms: Rethinking alignment via in- International Conference on Machine Learning,
context learning. arXiv preprint arXiv:2312.01552. pages 19893–19908. PMLR.
Mohit Shridhar, Xingdi Yuan, Marc-Alexandre Côté, Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin
Yonatan Bisk, Adam Trischler, and Matthew Guu, Adams Wei Yu, Brian Lester, Nan Du, An-
Hausknecht. 2020. Alfworld: Aligning text and em- drew M Dai, and Quoc V Le. 2021. Finetuned lan-
bodied environments for interactive learning. arXiv guage models are zero-shot learners. arXiv preprint
preprint arXiv:2010.03768. arXiv:2109.01652.

Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Chengyue Wu, Teng Wang, Yixiao Ge, Zeyu Lu,
Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, Ruisong Zhou, Ying Shan, and Ping Luo. 2023a. π-
and Tatsunori B. Hashimoto. 2023. Stanford alpaca: tuning: Transferring multimodal foundation models
An instruction-following llama model. https:// with optimal multi-task interpolation. In Proceedings
[Link]/tatsu-lab/stanford_alpaca. of the 40th International Conference on Machine
Learning, volume 202 of Proceedings of Machine
Learning Research, pages 37713–37727. PMLR.
Gemma Team, Thomas Mesnard, Cassidy Hardin,
Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Shijie Wu, Ozan Irsoy, Steven Lu, Vadim Dabravolski,
Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Mark Dredze, Sebastian Gehrmann, Prabhanjan Kam-
Juliette Love, et al. 2024. Gemma: Open models badur, David Rosenberg, and Gideon Mann. 2023b.
based on gemini research and technology. arXiv Bloomberggpt: A large language model for finance.
preprint arXiv:2403.08295. arXiv preprint arXiv:2303.17564.

Ian Tenney, Dipanjan Das, and Ellie Pavlick. 2019. Bert Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng,
rediscovers the classical nlp pipeline. arXiv preprint Pu Zhao, Jiazhan Feng, Chongyang Tao, and Daxin
arXiv:1905.05950. Jiang. 2023a. Wizardlm: Empowering large lan-
guage models to follow complex instructions. arXiv
Hugo Touvron, Louis Martin, Kevin Stone, Peter Al- preprint arXiv:2304.12244.
bert, Amjad Almahairi, Yasmine Babaei, Nikolay
Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Yiheng Xu, Hongjin Su, Chen Xing, Boyu Mi, Qian
Bhosale, et al. 2023. Llama 2: Open founda- Liu, Weijia Shi, Binyuan Hui, Fan Zhou, Yitao Liu,
tion and fine-tuned chat models. arXiv preprint Tianbao Xie, et al. 2023b. Lemur: Harmonizing
arXiv:2307.09288. natural language and code for language agents. arXiv
preprint arXiv:2310.06830.
Betty Van Aken, Benjamin Winter, Alexander Löser, Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Ben-
and Felix A Gers. 2019. How does bert an- gio, William W Cohen, Ruslan Salakhutdinov, and
swer questions? a layer-wise analysis of trans- Christopher D Manning. 2018. Hotpotqa: A dataset
former representations. In Proceedings of the 28th for diverse, explainable multi-hop question answer-
ACM international conference on information and ing. arXiv preprint arXiv:1809.09600.
knowledge management, pages 1823–1832.
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak
Peihao Wang, Rameswar Panda, Lucas Torroba Hen- Shafran, Karthik Narasimhan, and Yuan Cao. 2022.
nigen, Philip Greengard, Leonid Karlinsky, Roge- React: Synergizing reasoning and acting in language
rio Feris, David Daniel Cox, Zhangyang Wang, and models. arXiv preprint arXiv:2210.03629.
Yoon Kim. 2023a. Learning to grow pretrained mod-
els for efficient transformer training. arXiv preprint Yiqun Yao, Zheng Zhang, Jing Li, and Yequan
arXiv:2303.00980. Wang. 2023. 2x faster language model pre-training
via masked structural growth. arXiv preprint
Xiao Wang, Yuansen Zhang, Tianze Chen, Songyang arXiv:2305.02869.
Gao, Senjie Jin, Xianjun Yang, Zhiheng Xi, Rui
Zheng, Yicheng Zou, Tao Gui, et al. 2023b. Trace: Longhui Yu, Weisen Jiang, Han Shi, Jincheng Yu,
A comprehensive benchmark for continual learn- Zhengying Liu, Yu Zhang, James T Kwok, Zhen-
ing in large language models. arXiv preprint guo Li, Adrian Weller, and Weiyang Liu. 2023.
arXiv:2310.06762. Metamath: Bootstrap your own mathematical ques-
tions for large language models. arXiv preprint
arXiv:2309.12284.
Xingyao Wang, Zihan Wang, Jiateng Liu, Yangyi Chen,
Lifan Yuan, Hao Peng, and Heng Ji. 2023c. Mint: Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali
Evaluating llms in multi-turn interaction with tools Farhadi, and Yejin Choi. 2019. Hellaswag: Can a
and language feedback. machine really finish your sentence? arXiv preprint
arXiv:1905.07830.
Yizhong Wang, Hamish Ivison, Pradeep Dasigi, Jack
Hessel, Tushar Khot, Khyathi Raghavi Chandu, Lvmin Zhang, Anyi Rao, and Maneesh Agrawala. 2023.
David Wadden, Kelsey MacMillan, Noah A. Smith, Adding conditional control to text-to-image diffu-
Iz Beltagy, and Hannaneh Hajishirzi. 2023d. How sion models. In Proceedings of the IEEE/CVF
far can camels go? exploring the state of instruction International Conference on Computer Vision, pages
tuning on open resources. 3836–3847.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan
Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin,
Zhuohan Li, Dacheng Li, Eric Xing, et al. 2023.
Judging llm-as-a-judge with mt-bench and chatbot
arena. arXiv preprint arXiv:2306.05685.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao
Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu,
Lili Yu, et al. 2023. Lima: Less is more for alignment.
arXiv preprint arXiv:2305.11206.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B
Brown, Alec Radford, Dario Amodei, Paul Chris-
tiano, and Geoffrey Irving. 2019. Fine-tuning lan-
guage models from human preferences. arXiv
preprint arXiv:1909.08593.
A Gradient Derivation
To calculate the gradient of the RMSNorm weight during backpropagation, we first need to consider the
forward pass equation for the Llama RMSNorm:
w⊙x
RMSNorm(x) = p (6)
Var(x) + ϵ
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 ∂L ∂y
= (7)
∂w ∂y ∂w
∂y
To compute the gradient, we need to find the partial derivative ∂w . From the FFN equation, we have:

y = x′ + FFN(RMSNorm(x′ )) (8)
Taking the derivative with respect to w, we get:

∂y ∂FFN(RMSNorm(x′ ))
= (9)
∂w ∂w
Now, let’s differentiate the RMSNorm function with respect to w:

∂RMSNorm(x) x
=p (10)
∂w Var(x) + ϵ
Using the chain rule, we can compute the gradient of the loss function with respect to the RMSNorm
weight:

∂L ∂L ∂FFN(RMSNorm(x′ )) ∂RMSNorm(x′ )
= (11)
∂w ∂y ∂RMSNorm(x′ ) ∂w
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 ) ∂SiLU(tW1 ) ∂(tW2 )
= ⊗ (tW2 ) + SiLU(tW1 ) ⊗ (14)
∂t ∂t ∂t
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)
=0 (16)
∂t
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.
 Data source Tokens Weight
Proof-Pile-2 55B
AlgebraicStack 11B
1.00
OpenWebMath 15B
ArXiv 29B
The-Stack-Dedup
Python 22B 1.50

Table 6: Pretrain data sources, tokens, and the mixture weights of each component during training.

Datasets Query Source Response Source # Instances N̄rounds L̄prompt L̄completion

ShareGPT User prompts GPT-3.5/GPT-4 63,817 2.9 293.2 1157.1
WizardLM_evol_instruct_V2 GPT-4 GPT-4 143,000 1.0 602.6 1704.9
SlimOrca Human-written GPT-4 517,982 1.0 574.3 599.3
MetaMath Human-written/GPT-4 GPT-4 395,000 1.0 209.4 498.2
Evol-CodeAlpaca GPT-4 GPT-4 111,272 1.0 652.5 1552.0

Table 7: Instruction datasets investigated in this work. We report the average number of rounds (N̄rounds ), average
length of prompts (L̄prompt ), average length of completion (L̄completion ).

B Dataset Details
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 approxi-
mately 80 billion tokens from both math and code corpora. The specifics of the SFT data are delineated in
Table 7.
For our proposed LL A MA P RO - I NSTRUCT, 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.

C Mistal-Pro Details
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
1
methodology for creating additional blocks as LLaMA-Pro but utilizes only 10 of LLaMA Pro’s learning
4
rate, as recommended by MetaMath-Mistral . 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.

D Evaluation Benchmark
The benchmarks used for evaluation include:
3
[Link]
4
[Link]
5
[Link]
 Model GSM8k Pass@1 MATH Pass@1
MPT-7B 6.8 3.0
Falcon-7B 6.8 2.3
LLAMA-1-7B 11.0 2.9
LLAMA-2-7B 14.6 2.5
MPT-30B 15.2 3.1
LLAMA-1-13B 17.8 3.9
GPT-Neo-2.7B 19.5 –
Falcon-40B 19.6 2.5
Baichuan-chat-13B 23.9 –
Vicuna-v1.3-13B 27.6 –
LLAMA-2-13B 28.7 3.9
MetaMath-7B 66.5 19.8
MetaMath-13B 72.3 22.4
MetaMath-Mistral-7B 77.7 28.2
MetaMath-Llemma-7B 69.2 30.0
MetaMath-Mistral-Pro 78.4 30.3

Table 8: Performance of various models on GSM8k Pass@1 and MATH Pass@1

• 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.

E MINT-Bench
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
 Task Type Task Name # Instances
HumanEval (Chen et al., 2021b) 45
Code Generation
MBPP (Austin et al., 2021) 91
Decision Making ALFWorld (Shridhar et al., 2020) 134
GSM8K (Cobbe et al., 2021) 48
HotpotQA (Yang et al., 2018) 43
Reasoning MATH (Hendrycks et al., 2021) 100
MMLU (Hendrycks et al., 2020) 76
TheoremQA (Chen et al., 2023b) 49
Total 586

Table 9: Dataset statistics of MINT-Bench.

Model Code Generation Decision Making Reasoning Micro Avg.

AgentLM-7B 1.47 9.70 8.86 7.34

CodeLLaMA-7B-Instruct 2.21 17.16 7.91 8.70
LLaMA2-7B-Chat 0.00 0.00 13.61 7.34
Mistral-Instruct-v0.1 6.62 34.33 8.54 13.99
LL A MA P RO - I NSTRUCT 11.76 29.10 9.81 14.68

Table 10: The success rates of each model evaluated on different task type benchmarks, as well as the micro average
when k = 5.

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 LL A MA P RO - I NSTRUCT model demonstrates robust performance
across all task types compared to other models of similar size. Figure 8 provides a case study to compare
LL A MA P RO - I NSTRUCT and LLaMA2-7B-Chat where LL A MA P RO - I NSTRUCT successfully utilizes
Python program to solve the given question in the multi-round interaction.

F MT-Bench
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 LL A MA P RO -
I NSTRUCT and LLaMA2-7B-Chat.

G Token Distribution
We assess the token distribution between LLaMA-2-7B and LL A MA P RO, employing the methodology
proposed by Lin et al. (2023). Specifically, for a given user query q = {q1 , q2 , ...}, we input it into
LL A MA P RO 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 distribu-
tion, Pbase , for sampling the next token at this position. First, the aligned model with greedy decoding is
 Questions: Maria invited 4 of her friends over for a water balloon fight in the backyard. At the start of the game, Maria gave each of her friends 2 water balloons. She had one
water balloon for herself. Then her mom came out and gave each person 3 more balloons. How many total balloons did the girls have? Solution output format: an integer.
LLaMA Pro- Instruct: LLaMA2-7B-Chat:

Interaction Turn # 1 Interaction Turn # 1

Thought: Let's first determine the initial number of Thought: Great, let's solve this problem together! To start, we know that Maria had 2 water balloons
balloons Maria and her friends had. initially, and her mom gave each person 3 more balloons. So, the total number of balloons Maria's friends
Execute:
had after the mom's gift is: 2 + 3(4) = 2 + 12 = 14Therefore, the total number of balloons the girls had was
Maria_balloons = 1
friends_balloons = 4 * 2
14. Now it's your turn! Please provide the next task or question, and I'll be happy to help you solve it.
total_balloons = Maria_balloons + friends_balloons
print(total_balloons)
Observation: I don't understand your input. If you want to execute code, please use
<execute> YOUR_CODE_HERE </execute>.
Observation: Out[1]: 9. If you want to give me an answer, please use <solution> YOUR_SOLUTION_HERE </solution>. For example:
You have 4 steps left and 2 chances to propose The answer to the question is <solution> 42 </solution>.
solution left. You have 4 steps left and 2 chances to propose solution left.

Interaction Turn # 2 Interaction Turn # 2

Thought: Let's now add the additional balloons Thought: Ah, my apologies! Thank you for pointing that out. To answer the question, the total number of
given by Maria's mom. balloons the girls had was 14. So, the solution is <solution> 14 </solution>. Now, please provide the next
Execute: task or question, and I'll be happy to help you with it!
additional_balloons = 3 * 5 Propose solution: The answer is 14.
total_balloons = total_balloons + additional_balloons
print(total_balloons) Observation: Your answer is wrong. You have 3 steps left and 1 chances to propose solution left.
Observation: Out[2]: 24. Interaction Turn # 3
You have 3 steps left and 2 chances to propose
solution left. Thought: Oh, I see! Thank you for letting me know. I apologize for my mistake earlier. To solve the problem,
we need to find the total number of balloons the girls had after Maria's mom gave each person 3 more
Interaction Turn # 3 balloons. So, the total number of balloons the girls had is:2 + 3(4) = 2 + 12 = 14
Therefore, the answer is <solution> 14 </solution>.
Propose solution: The answer is 24.
Propose solution: The answer is 14.
Good job! You have successfully solved the task!
[stop interaction] Your answer is wrong! [stop interaction]

Figure 8: A case study of multi-turn interactions by using tools and environment feedback to solve math questions
with LL A MA P RO - I NSTRUCT and LLaMA2-7B-Chat.

Output Tokens Output Tokens Output Tokens

LoRA Identical
🔥 Decoder Block ×P Copy
×N
🔥 Decoder Block ×M ×N 🔥 Decoder Block ×M ❄ ×N ❄ Decoder Block ×M

Input Tokens Input Tokens Input Tokens

(a) Fully Fine-tuning (b) LoRA-tuning (c) Block-Expansion

(LLaMA Pro)

Figure 9: The difference of three training strategies, fully fine-tuning, LoRA, and our proposed block expansion.

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 LL A MA P RO 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 LL A MA P RO effectively retains the language modeling ability for the general corpus while
enhancing its proficiency in the code domain.

6
[Link]
 LLaMA Pro-Instruct Wins Tie LLaMA Pro-Instruct Loses

LLaMA2-7B-Chat 23.3% 71.7% 5.0%

CodeLlama-7b-Instruct-hf 26.7% 63.3% 10.0%

WizardCoder-Python-7B-V1.0 43.3% 43.3% 13.3%

0% 20% 40% 60% 80% 100%

Figure 10: MT-Bench pairwise comparison between LL A MA P RO - I NSTRUCT and widely used LLaMA community
models in math and code questions.
LLaMA Pro-Instruct Wins Tie LLaMA Pro-Instruct Loses

LLaMA2-7B-Chat 21.9% 46.9% 31.2%

CodeLlama-7b-Instruct-hf 35.0% 42.5% 22.5%

WizardCoder-Python-7B-V1.0 61.9% 28.8% 9.3%

0% 20% 40% 60% 80% 100%

Figure 11: MT-Bench pairwise comparison between LL A MA P RO - I NSTRUCT and widely used LLaMA community
models in comprehensive questions.

H Domain of Law
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 ‘Ŵ3 ‘, and at the
same time add two new learnable parameters ‘α1 , α2 ‘ , when forward the output of Linear corresponding
to W3 , Ŵ3 will be weighted and summed with softmax(α1 , α2 ) and fed into the next block.
 Model General Domain Perplexity Code Domain Perplexity
lambada openai lambada standard stack
LLaMA-2-7B 3.39 4.13 9.46
LL A MA P RO 3.46 4.30 5.25

Table 11: The perplexity of LLaMA and LL A MA P RO evaluated across general domain and code domain.

Hyperparameter Assignment
Batch size 1024
Maximum sequence length 2,048
Maximum learning rate 2e-4
Optimizer Adam
Adam beta weights 0.9, 0.95
Learning rate scheduler cosine
Warmup ratio 0.06
Gradient clipping 1.0

Table 12: Hyper-parameters of pretraining on the domain of law.

 Questions: Develop a Python program that reads all the text files under a
directory and returns top-5 words with the most number of occurrences.

LLaMA Pro-Instruct: LLaMA2-7B-Chat:

Follow-up Questions: Can you parallelize it?

LLaMA Pro-Instruct: LLaMA2-7B-Chat:

GPT-4 Judgment: 🎉 LLaMA Pro Wins !! 🎉

LLaMA Pro-Instruct’s LLaMA2-7B-chat’s
response to both question is correct and helpful response to the first question is correct . But its
. The assistant provides a parallelized version response to second question is incorrect. The assistant
of the program using the `ProcessPoolExecutor` mentions using the Parallel library to parallelize the
from the `[Link]` module. The assis program, but the provided code does not include
tant also explains how the parallelized version any parallelization
works, which is helpful for the user.

Figure 12: Multi-turn dialogues between a user and two Al assistants—LL A MA P RO - I NSTRUCT and LLaMA2-
7B-Chat.