GAIR-NLP / Entropy-ABF

Official implementation for 'Extending LLMs’ Context Window with 100 Samples'

Extending LLMs' Context Window with 100 Samples

This is the official repo for "Extending LLMs' Context Window with 100 Samples". Preprint

Introduction

We introduce 'Entropy-Aware ABF' that supports efficient context window extension of RoPE-based LLMs with only 100 samples. The repository contains code and data to reproduce our model.

Model and Data

We release long-context Llama-2-7b-chat extended with our method trained with different data amounts on 🤗Hugging Face:

Data	Link
0.1k	🤗eabf-llama2-7b-chat-0.1k
1k	🤗eabf-llama2-7b-chat-1k
3.5k	🤗eabf-llama2-7b-chat-3.5k

We also release our training data on 🤗Hugging Face Datasets.

Quick Guide

Use Entropy-Aware ABF

To use our code, your transformers library should be version 4.31 or higher.

We adopt the paper summarization test proposed in NTK-Aware scaling's blog to serve as a sanity check.

In short, to load the LLaMA model with our method, you should first import the required packages:

from transformers.models.llama.modeling_llama import LlamaForCausalLM
import patch.eabf as eabf

Then, you can load the model by using the right rope_scaling argument and our monkey patching function:

model = LlamaForCausalLM.from_pretrained(MODEL_NAME_OR_PATH, ..., rope_scaling={"type": "eabf", "factor": 4})
eabf.apply_eabf(model)

Reproduce Observation of Attention Scores

Other RoPE-based LLMs might or might not follow the same attention scores pattern as Llama-2-7b-chat, we release our code for retrieving attention scores and computing the 'attention entropy' so that users can apply our method tailored to their model.