[Paper Review] Gemma 2: Improving Open Language Models at a Practical Size

Gemma Team, Morgane Rivière+196 more|arXiv (Cornell University)|Jul 31, 2024

Natural Language Processing Techniques128 citations

TL;DR

Gemma 2 introduces 2B, 9B, and 27B open models trained with knowledge distillation and Transformer modifications, achieving state-of-the-art performance for their sizes and competitive results versus larger models.

ABSTRACT

In this work, we introduce Gemma 2, a new addition to the Gemma family of lightweight, state-of-the-art open models, ranging in scale from 2 billion to 27 billion parameters. In this new version, we apply several known technical modifications to the Transformer architecture, such as interleaving local-global attentions (Beltagy et al., 2020a) and group-query attention (Ainslie et al., 2023). We also train the 2B and 9B models with knowledge distillation (Hinton et al., 2015) instead of next token prediction. The resulting models deliver the best performance for their size, and even offer competitive alternatives to models that are 2-3 times bigger. We release all our models to the community.

Motivation & Objective

• Improve open language model performance at small-to-mid scales without solely increasing training length.
• Evaluate the effectiveness of knowledge distillation and architectural modifications (local-global attention, GQA) in Gemma 2.
• Assess pre-training and post-training strategies, safety, and deployment considerations for practical-use LLMs.
• Compare Gemma 2 models against open models of larger size and against baselines on automated and human evaluations.

Proposed method

• Adopt decoder-only Transformer architecture with interleaved local sliding window and global attention layers (4096 local, 8192 global).
• Apply Grouped-Query Attention (GQA) with num_groups=2 to improve inference speed.
• Use RMSNorm for stability and GeGLU nonlinearity for efficiency.
• Train 2B/9B models with knowledge distillation from a large teacher model; train 27B from scratch.
• Pre-train on up to 13T tokens (27B), 8T (9B), and 2T (2B) using a 256k SentencePiece vocabulary; filter data to reduce unsafe outputs.
• Post-train via supervised fine-tuning (SFT) and reinforcement learning from human feedback (RLHF); average models from multiple hyperparameters to form final IT models.
• Evaluate safety, memorization, and multi-turn conversational capabilities alongside standard benchmarks.

Figure 1: Comparing memorization rates. We find significantly lower memorization rates across-the-board. (Left) Overall memorization across model families. (Right) Exact and approximate memorization per data source.

Experimental results

Research questions

• RQ1 Does knowledge distillation improve performance of small to mid-size open LLMs when trained on very large token counts?
• RQ2 What is the impact of architectural tweaks (local-global interleaving, GQA) on efficiency and downstream performance at practical sizes?
• RQ3 How do Gemma 2 models compare to larger open models on automated benchmarks and human evaluations, including safety and alignment measures?
• RQ4 What are the safety, memorization, and deployment implications of releasing open Gemma 2 models?
• RQ5 How do post-training (SFT/RLHF) and model merging affect instruction-following and user satisfaction in practice?

Key findings

• Distillation from a larger model significantly improves small-model performance; for example, 2B and 9B models show large gains over prior versions trained with the same token counts, with up to about 10% improvement on some benchmarks for the 9B model.
• 27B Gemma 2 achieves competitive performance relative to models more than twice its size, and in some evaluations is competitive with LLaMA-3 70B despite being trained on less data.
• In post-training, Gemma 2 IT models outperform many open models in human evaluations and achieve high Elo scores on LMSYS Chatbot Arena, surpassing several closed and open competitors of larger scale (e.g., Gemma 27B Elo 1218 vs LLaMA-3 70B Elo 1206).
• Memorization rates are significantly lower than prior open models of similar size, with verbatim memorization <0.1% and low approximate memorization across data sources.
• Gemma 2 shows improved safety metrics in external and internal evaluations, and IT models demonstrate stronger instruction following and safety performance in human assessments.