Academician E Weinan leads the new work: Large models not only have RAG and parameter storage, but also a third kind of memory

2.4B of Memory³ achieves better performance than larger LLM and RAG models.

In recent years, large language models (LLMs) have received unprecedented attention due to their extraordinary performance. However, LLM is expensive to train and infer, and people have been trying to reduce the cost through various optimization methods.

In this article, researchers from Shanghai Algorithm Innovation Research Institute, Peking University and other institutions were inspired by the memory hierarchy of the human brain. They equipped LLM with explicit memory (a memory format that is cheaper than model parameters and RAG). ) to reduce this cost. Conceptually, LLMs can enjoy smaller parameter sizes, training costs, and inference costs since most of their knowledge is externalized into explicit memory. T Paper Address: https: //arxiv.org/pdf/2407.01178

Thesis Title: Memory

3
• : Language Modeling with Explicit Memory
• as a preliminary concept proof proof , the researchers trained a 2.4B LLM from scratch, which achieved better performance than larger LLM and RAG models, and achieved higher decoding speed than RAG. This model is named Memory
  3
because in LLM, explicit memory is the third form of memory after implicit memory (model parameters) and working memory (context key values).

Specifically, this paper introduces a new memory format, explicit memory, which is characterized by relatively low writing costs and relatively low reading costs. As shown in Figure 1, the model first converts the knowledge base (or any text dataset) into explicit memories implemented as sparse attention key-values, then calls these memories during inference and integrates them into the self-attention layer middle.

The new memory format defines a new memory hierarchy:

In addition, this article also introduces a memory circuit theory that supports knowledge externalization and proposes memory sparsity that can make storage tractable Mechanisms and a two-stage pretraining protocol to promote memory formation.

In summary:

Memory

3

Utilizes explicit memory during the inference process, which reduces the burden of memorizing specific knowledge for model parameters;

• Explicit memory is built from Encoded in the knowledge base, where the sparse memory format maintains true storage size;
  The researchers trained a Memory
3
• model from scratch with 2.4B non-embedded parameters, and its performance exceeded that of larger scales SOTA model. It also has better performance and faster inference than RAG;
• Additionally, Memory
  3 ^{improves factuality and mitigates hallucinations, and enables rapid adaptation to professional tasks.}
• Method introduction

Memory circuit theory helps determine what knowledge can be stored as explicit memory, and which model architecture is suitable for reading and writing explicit memory.

Researchers regard the input-output relationship as the internal mechanism of the circuit, and define knowledge as the input-output relationship and its circuit. By manipulating these circuits, one can separate much of the knowledge from the LLM while keeping its functionality intact.

Memory

3

: In terms of architecture, the goal of this article is to design an explicit memory mechanism for Transformer LLM so that its writing cost and reading cost are relatively low. In addition, this article hopes to limit the modifications to the Transformer architecture to the smallest possible scope without adding any new trainable parameters, so that most existing Transformer LLMs can be converted to Memory with almost no fine-tuning

3

models. The simple design process is as follows:

Write cost: Before inference, LLM writes each reference to explicit memory, which is saved on the drive.Memories are selected from the key vectors of the self-attention layer, so the writing process does not involve training. Each reference is processed independently, avoiding the cost of long context attention.

Read cost: During inference, explicit memory is retrieved from the drive and read by self-attention along with the usual context key values. Each memory consists of a very small number of key values ​​from a small number of attention heads, greatly reducing additional computation, GPU storage, drive storage and load time. It allows LLM to retrieve many references frequently with limited impact on decoding speed.

The reasoning process is shown in Figure 9. Whenever LLM generates 64 tokens, it discards the current memory, uses these 64 tokens as query text to retrieve 5 new memories, and continues to use these memories. decoding. Likewise, when processing cues, LLM retrieves 5 memories for every 64-token block. Each block focuses on its own memory, and the memory may vary between blocks.

Writing and reading memories: During inference, LLM can directly read retrieved explicit memories through its self-attention layer by concatenating them with contextual key values ​​(Figure 9). Specifically, for each attention head h of the l-th layer, if it is selected as a memory head, then its output Y^(l,h) will change:

In addition, this study Explicit memory uses parallel position encoding, that is, all key positions are located in the same interval of length 128, as shown in Figure 9.

Two-stage pre-training: Pre-training consists of two stages, warmup and continuous training. Only the ongoing training phase involves explicit memory, while the warmup phase uses the same format as normal pre-training.

Figure 13 plots the training loss and learning rate schedule during the warmup phase.

Figure 14 plots the training loss and learning rate schedule during the continuous training phase.

Experimental results

The researchers evaluated the Memory³ model’s general abilities (benchmark tasks), conversational abilities, professional abilities (law and medicine), and hallucinations. In addition, the researchers also measured the decoding speed of Memory³ and compared it with similar and larger SOTA LLM and RAG models.

The assessment results of general ability are shown below, and the results show that explicit memory increased the average score by 2.51%. In comparison, the score difference between Llama2-7B and 13B is 4.91%. Explicit memory can increase the "effective model size" by 2.51/4.91 ≈ 51.1%.

Next, the authors evaluated the dialogue skills of Memory³, and the results are listed in Table 18, showing that the model outperforms Vicuna-7B, Falcon-40B-Instruct and ChatGLM2-6B with fewer parameters.

Currently, LLM still faces hallucination issues. Conceptually, Memory³ should be less susceptible to hallucinations because its explicit memory corresponds directly to the reference text. To evaluate hallucinations, the researchers selected two English datasets for evaluation. The results are shown in Table 19. Memory³ achieves the highest scores on most tasks.

One benefit of using explicit memory is that LLM can easily adapt to new domains and tasks by updating its knowledge base. Simply import task-related references into Memory³'s knowledge base and optionally convert them to explicit memory in the event of a warm start.The model can then leverage this new knowledge for inference, skipping the more costly and potentially lossy fine-tuning process, and run faster than RAG. This cost reduction has been demonstrated in Figure 4 and could facilitate the rapid deployment of LLM in various industries.

The table below shows that Memory³ performs better than most models.

Finally, the researchers evaluated the decoding speed or throughput of Memory³ by the number of tokens generated per second.

For more information, please refer to the original paper.

The above is the detailed content of Academician E Weinan leads the new work: Large models not only have RAG and parameter storage, but also a third kind of memory. For more information, please follow other related articles on the PHP Chinese website!