New work by Chen Danqi's team: Llama-2 context is expanded to 128k, 10 times throughput only requires 1/6 of the memory

The Chen Danqi team has just released a new LLMContext window extensionMethod:

It only uses 8k token documents for training, and can Llama-2 Window extended to 128k.

The most important thing is that in this process, the model only requires 1/6 of the original memory, and the model obtains 10 times the throughput.

In addition, it can also greatly reduce the training cost:

Use this method to train 7B alpaca 2 For transformation, you only need a piece of A100 to complete it.

The team expressed:

We hope this method will be useful and easy to use, and provide

cheap and effective long context capabilities for future LLMs.

Currently, the model and code have been released on HuggingFace and GitHub.

Just add two components

This method is called

CEPE, the full name is "Parallel Encoding Context Extension（Context Expansion with Parallel Encoding）”.

As a lightweight framework, it can be used to extend the context window of any

pre-trained and directive fine-tuning model.

For any pretrained decoder-only language model, CEPE extends it by adding two small components:

One is a small encoder for long The context is block-encoded;

One is the cross-attention module, which is inserted into each layer of the decoder to focus on the encoder representation.

The complete architecture is as follows:

In this diagram, the encoder model encodes 3 extra blocks of context in parallel and concatenates with the final hidden representation, It is then used as input to the decoder cross-attention layer.

Here, the cross-attention layer mainly focuses on the encoder representation between the self-attention layer and the feed-forward layer in the decoder model.

By carefully selecting training data that does not need to be labeled, CEPE helps the model have long context capabilities and is also good at document retrieval.

The author introduces that such CEPE mainly contains three major advantages:

(1) The length can be generalized

because it is not subject to positional encoding A constraint, instead, has its context encoded in segments, each segment having its own positional encoding.

(2) High efficiencyUsing small encoders and parallel encoding to process context can reduce computational costs.

At the same time, since cross-attention only focuses on the representation of the last layer of the encoder, and the language model that only uses the decoder needs to cache the key-value pairs of each token in each layer, so in comparison, CEPE The memory required is greatly reduced.

(3) Reduce training cost

Unlike full fine-tuning methods, CEPE only adjusts the encoder and cross-attention while keeping the large decoder model frozen.

The author introduces that by expanding the 7B decoder into a model

with a 400M encoder and cross-attention layer (a total of 1.4 billion parameters), it can be completed with an 80GB A100 GPU.

Perplexity continues to decrease

The team applies CEPE to Llama-2 and trains on the filtered version of RedPajama with 20 billion tokens

(only the Llama-2 pre-training budget 1%).

First, compared to two fully fine-tuned models, LLAMA2-32K and YARN-64K, CEPE achieves lower or comparable

perplexity## on all datasets. #, with both lower memory usage and higher throughput.

When the context is increased to 128k

(far beyond its 8k training length)

, CEPE’s perplexity continues to decrease while remaining low memory status. In contrast, Llama-2-32K and YARN-64K not only fail to generalize beyond their training length, but are also accompanied by a significant increase in memory cost.

Secondly,

retrieval capability is enhanced. As shown in the following table:

By using the retrieved context, CEPE can effectively improve model perplexity, and its performance is better than RePlug.

It is worth noting that even if paragraph k=50 (training is 60), CEPE will continue to improve the perplexity.

This shows that CEPE transfers well to the retrieval enhancement setting, whereas the full-context decoder model degrades in this ability.

Third, open domain question and answer capabilitiesare significantly surpassed.

As shown in the figure below, CEPE is significantly better than other models in all data sets and paragraph k parameters, and unlike other models, the performance drops significantly as the k value becomes larger and larger.

This also shows that CEPE is not sensitive to a large number of redundant or irrelevant paragraphs.

So to summarize, CEPE outperforms on all the above tasks with much lower memory and computational cost compared to most other solutions.

Finally, based on these, the author proposed CEPE-Distilled (CEPED) specifically for the instruction tuning model.

It uses only unlabeled data to expand the context window of the model, distilling the behavior of the original instruction-tuned model into a new architecture through assisted KL divergence loss, thereby eliminating the need to manage expensive long context instruction tracking data.

Ultimately, CEPED can expand the context window of Llama-2 and improve the long text performance of the model while retaining the ability to understand instructions.

Team Introduction

CEPE has a total of 3 authors.

Yan Heguang（Howard Yen） is a master's student in computer science at Princeton University.

The second person is Gao Tianyu, a doctoral student at the same school and a bachelor's degree graduate from Tsinghua University.

They are all students of the corresponding author Chen Danqi.

Original paper: https://arxiv.org/abs/2402.16617
Reference link: https://twitter. com/HowardYen1/status/1762474556101661158