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 的记忆检索功能。