LongMem 项目使用教程

1. 项目介绍

LongMem 是一个用于增强语言模型（LLMs）长期记忆的框架，由 Weizhi Wang 等人在 NeurIPS 2023 上提出。该项目通过设计一种新颖的解耦网络架构，使得语言模型能够记忆长期历史信息。原始的骨干 LLM 被冻结为记忆编码器，而自适应残差侧网络则作为记忆检索器和阅读器。这种解耦的记忆设计可以轻松缓存和更新长期过去的上下文，而不会受到记忆陈旧的影响。通过记忆增强的适应性训练，LongMem 能够记忆和利用丰富的长期上下文信息。

2. 项目快速启动

环境设置

首先，确保你已经安装了以下依赖：

• torch: 建议使用 torch>=1.8.0。请根据你的 CUDA 驱动版本选择合适的 GPU 版本。
• Faiss-GPU: 对于 Nvidia V100 GPU，可以通过 pip install faiss-gpu 安装。对于 Nvidia A100 或 A6000 GPU，请运行 conda install faiss-gpu cudatoolkit=11.0 -c pytorch。
• fairseq: 通过以下命令安装：

  pip install --editable /path/to/fairseq
  

• 其他依赖：

  pip install -r requirements.txt
  

项目结构

• Pre-trained LLM Class: fairseq/fairseq/models/newgpt.py
• Transformer Decoder with SideNetwork: fairseq/fairseq/models/sidenet/transformer_decoder_sidenet.py
• Transformer Language Model with SideNetwork Class: fairseq/fairseq/models/transformer_lm_sidenet.py
• Memory Bank and Retrieval: fairseq/fairseq/modules/dynamic_memory_with_chunk.py
• Joint Attention for Memory Fusion: fairseq/fairseq/modules/joint_multihead_attention_sum.py

数据预处理

下载 Pile 数据集，并按照以下步骤进行预处理：

python preprocess/filter_shard_tnlg.py

训练

运行以下命令进行训练：

bash train_scripts/train_longmem.sh

评估

下载预训练的 GPT2-medium 模型和 LongMem 模型的检查点，然后运行以下命令进行评估：

# 评估 GPT2 基线模型
python eval_scripts/eval_longmem_icl.py --path /path/to/gpt2_pretrained_model

# 评估 LongMem 模型
python eval_scripts/eval_longmem_icl.py --path /path/to/longmem_model --pretrained-model-path /path/to/gpt2_pretrained_model

3. 应用案例和最佳实践

LongMem 可以应用于需要长期记忆的场景，例如：

• 对话系统: 增强对话系统对用户历史对话的记忆能力，提升对话的连贯性和用户体验。
• 文档摘要: 在处理长文档时，LongMem 可以帮助模型更好地理解和总结文档内容。
• 知识问答: 通过记忆和检索长期知识，提升问答系统的准确性和覆盖范围。

4. 典型生态项目

• fairseq: LongMem 基于 fairseq 框架开发，fairseq 是一个用于序列到序列任务的强大工具包。
• Pile 数据集: 用于训练和评估 LongMem 的数据集，提供了丰富的文本数据。
• Faiss: 用于高效的向量检索，支持 LongMem 的记忆检索功能。