[Paper Review] XGen-7B Technical Report

Erik Nijkamp, Tian Xie+23 more|arXiv (Cornell University)|Sep 7, 2023

Topic Modeling4 citations

TL;DR

XGen-7B is a 7B-parameter large language model pre-trained on up to 1.5 trillion tokens with support for 8K context length, using a two-stage training process (2K → 4K → 8K) and diverse data mixtures including natural language and code. It achieves state-of-the-art performance on standard benchmarks and significantly outperforms 2K-sequence open-source models on long-context tasks like dialogue summarization and code generation.

ABSTRACT

Large Language Models (LLMs) have become ubiquitous across various domains, transforming the way we interact with information and conduct research. However, most high-performing LLMs remain confined behind proprietary walls, hindering scientific progress. Most open-source LLMs, on the other hand, are limited in their ability to support longer sequence lengths, which is a key requirement for many tasks that require inference over an input context. To address this, we have trained XGen, a series of 7B parameter models on up to 8K sequence length for up to 1.5T tokens. We have also finetuned the XGen models on public-domain instructional data, creating their instruction-tuned counterparts (XGen-Inst). We open-source our models for both research advancements and commercial applications. Our evaluation on standard benchmarks shows that XGen models achieve comparable or better results when compared with state-of-the-art open-source LLMs. Our targeted evaluation on long sequence modeling tasks shows the benefits of our 8K-sequence models over 2K-sequence open-source LLMs.

Motivation & Objective

• Address the limitation of open-source LLMs that are restricted to 2K context length, which hinders performance on long-context tasks such as code generation and document summarization.
• Develop a scalable, efficient 7B-parameter model that maintains strong performance while supporting longer context lengths through staged pre-training.
• Enable broader scientific and commercial use by open-sourcing both base and instruction-tuned variants of the model.
• Improve instruction-following capabilities through supervised fine-tuning on diverse public-domain instruction datasets, enhancing zero-shot and few-shot generalization.

Proposed method

• Employ a two-stage pre-training strategy: first pre-training on 2K context length with 1.2T tokens (natural language + code), then continuing on 4K and finally on 8K context length with an additional 300B tokens.
• Use a mixture of data sources including natural language from diverse public corpora, GitHub code from RedPajama, and BigCode Starcoder for code-specific pre-training with permissive licensing.
• Apply byte pair encoding (BPE) tokenization using GPT-2 as a base, with added special tokens for code and instruction tuning.
• Fine-tune the XGen-7B-8K base model on two public instruction datasets: WizardLM-196K and general instruction data (OAsst, Baize, Dolly2, ShareGPT, SCROLLS) to produce XGen-7B-Inst variants.
• Optimize training efficiency through mixed-precision training and model parallelism across TPU pods, with a total training cost of ~57 MWh and 4.5 tCO₂eq emissions.
• Use standard LLM evaluation protocols including perplexity, accuracy, and ROUGE scores on standard and long-context benchmarks.

Figure 1: Cross-entropy over training time. The model is pre-trained in three stages with a step-wise increase of the sequence length from 2K to 4K to 8K tokens. Note, the pre-training does not suffer from any loss spikes. The spikes at the transitions in sequence lengths are expected as the model a

Experimental results

Research questions

• RQ1 Can a 7B-parameter model trained on 1.5T tokens with progressive context length scaling (2K → 8K) achieve competitive performance on standard and long-context benchmarks?
• RQ2 How does the 8K context length of XGen-7B compare to 2K-sequence open-source models in long-context tasks such as dialogue summarization and code generation?
• RQ3 To what extent do instruction-tuned variants of XGen-7B outperform similarly sized or larger models on instruction-following benchmarks like AlpacaEval and MTBench?
• RQ4 What is the environmental cost of training such a model, and how does it compare to other LLMs in terms of energy consumption and carbon emissions?
• RQ5 Can a smaller, efficient model like XGen-7B match or exceed the performance of larger proprietary or open-source models in both base and instruction-tuned settings?

Key findings

• XGen-7B-8K achieves comparable or better results than state-of-the-art open-source LLMs on standard benchmarks, including strong performance on code generation tasks.
• On the AlpacaEval benchmark, XGen-7B-Inst (general) outperforms models like WizardLM-7B and MPT-7B, and even surpasses larger models such as Falcon-40B-instruct.
• On the MTBench benchmark, XGen-7B-Inst (general) achieves a score of 8.25 on the Mixture-of-Model (MoE) metric, outperforming Alpaca-13B and Falcon-40B-instruct.
• In long dialogue summarization tasks (AMI, FD, TMS), XGen-7B-Inst (general) achieves the highest ROUGE-1 score of 31.34 on AMI (5.5K tokens), significantly outperforming MPT-7B-instruct and Alpaca-7B, which struggled without input truncation.
• The model's performance on long-context tasks demonstrates the tangible benefit of 8K context length, with XGen-7B-8K showing clear superiority over 2K-sequence models in modeling long-distance dependencies.
• The carbon footprint of training XGen-7B is estimated at 4.5 tCO₂eq, with a total energy consumption of 57 MWh, which is comparable to other large models but with a focus on efficiency and open access.

Figure 2: Perplexity over sequence length. If a model can utilize the information encoded in long sequences, then, in expectation, the perplexity should decrease over the length of such sequences. That is, the information contained in previous tokens increases the certainty of the next token predict