Source-linked AI summary

LLaMA Pro: Progressive LLaMA with Block Expansion

Chengyue Wu, Yukang Gan, Yixiao Ge, Zeyu Lu, Jiahao Wang, Ye Feng, Ying Shan, Ping Luo

arXiv:2401.02415v2cs.CL

TL;DR

LLMs can acquire domain-specific capabilities while losing some of their original general abilities. The paper adds Transformer blocks and tunes only the added blocks on new corpus; LLaMA Pro models are reported to perform strongly across general, code, math, and agent-oriented evaluations.

  • Problem

    Domain-specific post-pretraining can diminish LLMs’ original general abilities, despite the need for specialized capabilities in areas such as programming and mathematics.

    Source: §1 · p. 1
  • Method

    Block expansion adds copied Transformer blocks, tunes the new blocks on domain-specific corpus while freezing the original blocks, and can be followed by instruction tuning.

    Source: §1 · p. 2, §3.2 · p. 4
  • Results

    LLaMA Pro-Instruct is reported to reach state-of-the-art performance across general, code, and math tasks and to outperform other LLaMA-family models on benchmarks and practical applications.

    Source: §1 · p. 2, §6 · p. 9
  • Takeaways & Limitations

    The authors conclude that block expansion can balance general and domain-specific performance in LLaMA Pro, initialized from LLaMA2-7B with eight added blocks.

    Source: §6 · p. 9

    The study is limited to the language modality, especially programming and English; extending the method to other domains remains future work.

    Source: §5 · p. 9

Abstract

from arXiv · show

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 without catastrophic forgetting. In this paper, we experiment on the corpus of code and math, yielding LLaMA Pro-8.3B, a versatile foundation model initialized from LLaMA2-7B, 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.

1 Introduction

LLMs can lack capabilities in specialized domains, while domain-specific post-pretraining can compromise general abilities. The paper proposes block expansion to add domain knowledge while preserving the original blocks and evaluates the resulting models across general, code, math, and agent-oriented tasks.

1 Introduction

  • The method addresses the risk that domain-specific post-pretraining can reduce original general abilities, a challenge for developing language agents for broader applications.Source: §1 · p. 1
  • Block expansion adds copied Transformer blocks, zero-initializes their linear layers, and tunes only the new blocks on domain-specific data while freezing the original blocks.The resulting extended model is described as excelling in both general and domain-specific tasks.Source: §1 · p. 2
  • LLaMA Pro extends LLaMA2-7B with eight blocks, trains them on 80B code and math tokens, and uses approximately 80M tokens for subsequent instruction tuning.The paper reports 2,830 GPU hours for expanded-block pretraining and says the resulting model is compatible with standard supervised instruction tuning without specific modifications.Source: §1 · p. 2
  • LLaMA Pro-Instruct reaches state-of-the-art performance across general, code, and math tasks and outperforms other LLaMA-family models on benchmarks and practical applications.The evaluation also covers agent scenarios, including tool use and grounding in environmental and human feedback.Source: Figure 1 · p. 1, §1 · p. 2
  • The authors report benchmarking the LLaMA Pro family on extensive traditional and agent-oriented datasets, with superiority and potential in complex applications.Source: §1 · p. 2

2 Related Work

Prior work includes generalist-model development, domain-specific post-pretraining, and methods for growing model capacity. This paper places block expansion within that landscape as a depth-growth strategy intended to specialize models while maintaining general capabilities.

2 Related Work

  • Prior advances in model and data scale produced strong performance across tasks, while generalist models and open-source releases such as LLaMA and CodeLLaMA supported diverse problem-solving and adaptation.Source: §2 · p. 2
  • Post-pretraining commonly follows general-domain pretraining, while later fine-tuning can target instruction following, human-preference alignment, parameter-efficient adaptation, or continual learning.Source: §2 · p. 2
  • Unlike adaptor and LoRA methods generally applied during instruction tuning, this work increases model depth during continued pretraining while seeking to maintain general capabilities.Source: §2 · p. 2
  • Earlier progressive-learning methods expanded depth, feed-forward networks, hidden size, or multiple growth dimensions; this work uses depth growth for domain adaptation while maintaining general performance.Source: §2 · p. 3

3 Method

Block expansion adds interleaved identity copies to a pretrained Transformer, preserving its output at initialization. Only the added blocks are then tuned on domain-specific data while the original blocks remain frozen.

3.1 Preliminaries: The LLaMA Block

  • A LLaMA block combines multi-head self-attention with a position-wise feed-forward network using SwiGLU and residual connections.The feed-forward network’s SwiGLU operation uses element-wise multiplication and SiLU activation.Source: §3.1 · p. 3

3.2 Block Expansion

  • Interleaved expansion partitions the original blocks into groups and stacks identity copies of each group’s top P blocks, increasing depth while preserving output behavior.The arrangement follows the model’s structural prior that deeper blocks encode more complex information.Source: §3.2 · p. 4
  • Zero-initializing RMSNorm scale parameters would give them zero loss gradient and prevent RMSNorm training, so the method instead zeros W_O and W3.This is why the cited prior identity-block initialization may not work for the LLaMA block.Source: §3.2 · p. 4
  • Zeroing the LLaMA block’s W_O and W3 matrices makes the added block an identity at initialization, so only the residual flows through and the base model’s output is preserved.The block has residual connections and no bias terms.Source: §3.2 · p. 4, Figure 3 · p. 4
  • Tuning only newly added blocks on domain-specific corpora while freezing original blocks is intended to add domain knowledge while preserving general abilities.The method operates in the post-pretraining stage.Source: §3.2 · p. 4

4 Experiments

Experiments test block expansion on code and math pretraining, instruction tuning, a second base model, and key design choices. LLaMA Pro-Instruct performs strongly across general, code, and math tasks, while ablations examine capability retention, domain adaptation, and training cost.

4.1 Experimental Settings

  • LLaMA Pro-8B tunes 1B parameters on 80B tokens, using less training overhead than CodeLLaMA-7B’s 7B parameters on 500B tokens.The evaluation includes six general benchmarks, HumanEval and MBPP pass@1, and GSM8K; training adds inference cost relative to the initial model.Source: §4.1 · p. 5, §4.2 · p. 6

4.2 Pretrain Results

  • LLaMA Pro preserves LLaMA2-7B’s general performance, surpasses its average on general tasks, and delivers code performance on par with code-oriented LLMs.The comparison includes general-purpose and code-oriented models trained around the same period and at similar sizes.Source: §4.2 · p. 6

4.3 SFT Results

  • LLaMA Pro-Instruct attains state-of-the-art performance across general, code, and math tasks, including comparisons with WizardCoder and WizardMath.It also boosts code and math performance while maintaining reliable general performance.Source: §4.3 · p. 7
  • LLaMA Pro-Instruct reports average-performance gains of 13.81% over LLaMA2-7B-chat and 14.50% over CodeLLaMA-7B-instruct.The evaluation also reports that it surpasses LLaMA-community chatbots in GPT-4-scored MT-Bench and reaches state-of-the-art performance among similar-size models on multi-turn tool use.Source: §4.3 · p. 7

4.4 Mistral-Pro Results

  • Mistral-Pro’s instruction-tuned scores are 78.4 on GSM8k and 30.3 on MATH, versus Mistral’s 77.7 and 28.2, respectively.The paper also reports superior pretrained benchmark outcomes, competitiveness with Gemma, and significantly lower training overhead.Source: §4.4 · p. 7, Table 4 · p. 8

4.5 Ablation Study

  • Training loss decreases regardless of added-block count and falls faster for larger models; adding blocks at the top preserves initial performance but underperforms interleaving on domain-specific tasks.Adding blocks at the bottom produces poor evaluation performance; the authors suggest it disrupts the model’s foundation and propagates errors.Source: Figure 5 · p. 7, §4.5 · pp. 8–9
  • Eight added blocks provide optimal performance at minimal cost; all expanded models preserve general capabilities, while larger models perform better on the law-specific task.MoE is comparable to block expansion with four added blocks; LoRA preserves general ability but struggles on the new-domain distribution, while full fine-tuning causes a larger general-performance drop.Source: §4.5 · p. 8, Table 5 · p. 8
  • With the same instruction dataset, LLaMA Pro outperforms LLaMA2-7B across all tasks, and its token-distribution shift on general Alpaca questions is described as subtle.The study also evaluates law adaptation using the Pile’s freelaw subset and UNFAIR-ToS, though the supplied passages do not report its outcome.Source: §4.5 · pp. 7, 9, Figure 6 · p. 8, Figure 7 · p. 9

6 Conclusion

Block expansion is presented as a post-pretraining way to improve domain-specific abilities while preserving general capabilities. LLaMA Pro, initialized from LLaMA2-7B with eight added blocks, outperformed other LLaMA-series models on comprehensive benchmarks.

6 Conclusion

  • Block expansion aims to improve domain-specific abilities while preserving a model’s original general capabilities.The conclusion identifies LLaMA Pro as initialized from LLaMA2-7B with eight added blocks.Source: §6 · p. 9
  • LLaMA Pro outperformed other LLaMA-series models on comprehensive benchmarks.Source: §6 · p. 9
  • The paper presents the method as a post-pretraining approach for large language models.Source: §6 · p. 9

7 Ethical Statement

LLaMA Pro and LLaMA Pro-Instruct target general language tasks, programming, and mathematics. The paper cautions that specialized domains may remain challenging and calls for responsible use mindful of bias and the LLaMA-2 license.

7 Ethical Statement

  • The models are designed for general language tasks, programming, and mathematics.The stated focus is scenarios that integrate natural and programming languages.Source: §7 · p. 9
  • Highly specialized domains or tasks may still pose challenges for LLaMA Pro.The statement frames this as a possibility, not a categorical limitation.Source: §7 · p. 9
  • Users are urged to consider potential model biases, responsible use, and the LLaMA-2 license.Source: §7 · p. 9

A Gradient Derivation

The derivation discusses RMSNorm’s forward pass and the chain-rule gradient for its weight. It also establishes that the FFN derivative with respect to its input is zero at input zero.

A Gradient Derivation

  • The RMSNorm forward pass uses input x, weight w, variance across the last dimension, and epsilon for numerical stability.Source: A Gradient Derivation · p. 14
  • The loss gradient with respect to the RMSNorm weight is computed using the chain rule through the FFN output y.The setup denotes the loss as L and the FFN output as y, then considers the partial derivative of y with respect to w.Source: A Gradient Derivation · p. 14
  • At t = 0, the FFN derivative with respect to t is zero.The derivation identifies the zero gradient at this input as its conclusion.Source: A Gradient Derivation · p. 14

B Dataset Details

The appendix describes the code-and-math pretraining corpus and the instruction datasets used for LLaMA Pro-Instruct, then outlines continued pretraining and SFT for Mistral-Pro.

B Dataset Details

  • The pretraining corpus comprises approximately 80 billion tokens from math and code corpora.Table 6 is described as reporting pretraining sources, token counts, and mixture weights.Source: B Dataset Details · p. 15
  • LLaMA Pro-Instruct SFT blends general, math, and code instruction datasets, including ShareGPT3, WizardLM evolution, CodeAlpaca, MetaMath, and SlimOrca.The account characterizes SlimOrca as a curated OpenOrca subset and notes its approximately 500,000 GPT-4 completions.Source: B Dataset Details · p. 15

C Mistal-Pro Details

  • Mistral-Pro adds Transformer blocks to Mistral and continues pretraining on code and math datasets, with SFT using Mistral-Pro as the base model.The listed continued-pretraining sources include automath from Cosmopedia, proof-pile-2, and the Python subset of Stack; its SFT approach follows MetaMath-Mistral.Source: C Mistal-Pro Details · p. 15

D Evaluation Benchmark

The evaluation section lists general reasoning, knowledge, truthfulness, and commonsense benchmarks alongside math and code tasks. It also identifies a table comparing GSM8k and MATH Pass@1 performance.

D Evaluation Benchmark

  • Table 8 compares models on GSM8k Pass@1 and MATH Pass@1.The supplied caption specifies the metrics but does not state comparative scores.Source: Table 8 · p. 16
  • General and commonsense evaluation includes ARC, HellaSwag, MMLU, TruthfulQA, and Winogrande.The section describes these benchmarks as covering science questions, commonsense inference, multitask accuracy, falsehood propensity, and commonsense reasoning.Source: D Evaluation Benchmark · p. 16
  • Math evaluation uses GSM8k for multi-step reasoning and a Program of Thought setting that uses Python to solve problems.Source: D Evaluation Benchmark · p. 16
  • Code evaluation includes HumanEval’s handwritten Python problems and MBPP’s crowdsourced programming problems.Source: D Evaluation Benchmark · p. 16

E MINT-Bench

The section evaluates LLAMA PRO - INSTRUCT as an interactive assistant on MINT-Bench and MT-Bench, and examines its token distributions and language-modeling performance across general and code corpora.

E MINT-Bench

  • MINT-Bench comprises eight datasets and 586 instances, and measures success rate across interaction limits from one to five turns.At k = 1 there is no interaction; k = 5 maximizes interaction turns within a 4,096-token context window.Source: E MINT-Bench · pp. 16–17, Table 9 · p. 17
  • LLAMA PRO - INSTRUCT performs robustly across MINT-Bench task types versus similar-size models, and a case study shows it solving a math question with Python over multiple interaction rounds.Table 10 reports success rates by task type and a micro average at k = 5.Source: Table 10 · p. 17, E MINT-Bench · p. 17, Figure 8 · p. 18
  • Each MINT-Bench turn lets the model optionally provide a thought, execute Python through a tool, or propose a solution.The benchmarks span code generation, decision-making, and reasoning.Source: E MINT-Bench · pp. 16–17

F MT-Bench

  • MT-Bench uses GPT-4 as a judge for demanding multi-turn open-ended questions, with pairwise comparisons covering math, code, and comprehensive questions.A figure also presents a dialogue comparison with LLaMA2-7B-Chat.Source: F MT-Bench · p. 17, Figure 10 · p. 19, Figure 11 · p. 19

G Token Distribution

  • Figure 9 contrasts full fine-tuning, LoRA, and block expansion as training strategies.The caption identifies the strategies but does not report their comparative outcomes.Source: Figure 9 · p. 18
  • Token-distribution analysis ranks LLAMA PRO's greedy-decoded output tokens under LLaMA-2-7B, separating unshifted, marginal, and shifted positions.The base rank is 1 for unshifted positions, 2–3 for marginal positions, and above 3 for shifted positions.Source: G Token Distribution · pp. 17–18
  • LLAMA PRO retains general-corpus language-modeling ability while improving code-domain proficiency in perplexity evaluations.The general evaluation uses two LAMBADA versions; the code evaluation uses the Python split of bigcode/the-stack-smol-xs.Source: G Token Distribution · p. 18

H Domain of Law

The law-domain study uses the Pile's freelaw subset for pretraining and evaluates on Unfair-ToS; it also describes a two-expert MoE implementation.

H Domain of Law

  • The law-domain pretraining corpus is the Pile's freelaw subset, with 51.2 GiB of raw data, 16.7B tokens, and 3.6M documents.Table 12 is identified as listing the law-domain pretraining hyperparameters.Source: H Domain of Law · p. 19, Table 12 · p. 20
  • Unfair-ToS is an eight-label, sentence-level multilabel classification task evaluated as a four-shot multiple-choice task.An example counts as correct when its highest per-token-likelihood option belongs to the example's label set.Source: H Domain of Law · p. 19
  • The MoE implementation uses two simultaneously active experts and extends each FFN across all 32 layers.It adds a learned linear layer and combines its output with the original FFN output using softmax-weighted parameters.Source: H Domain of Law · p. 19
Loading 2401.02415v2…