LongLM - 通过Self-Extend方法扩展大语言模型的上下文窗口

LongLM 项目介绍

LongLM项目是为了解决大语言模型（LLM）在长上下文处理上的困难而设计的一个项目。通常情况下，大语言模型在训练时的序列长度是有限的，这限制了它们在推理时处理长输入序列的能力。LongLM项目提出了一种名为“Self-Extend”的方法，通过自我扩展机制来最大化LLM在长上下文处理上的固有能力，不需要对模型进行再训练。

项目更新

2024年5月31日: SelfExtend能力在谷歌I/O会议中被介绍。
2024年5月1日: SelfExtend被ICML 2024大会接受。
最近的更新包括对Llama-3模型的支持、增加了一些超参数搜索结果，以及在不同的环境下进行了一系列的改进和优化。

SelfExtend的使用方式

1. 概述

Self-Extend的方法依赖于模型中原有的自注意力机制，以组内和邻居两级生成注意力信息，不需要额外的训练。通过这种方法，能够刺激和利用LLM处理长上下文的潜能。

2. 如何使用SelfExtend

设置

需要使用的Python包有 transformers==4.38.2 及 flash_attn==2.5.6。推荐使用提供的docker镜像以避免环境问题。

安装与运行

从仓库克隆代码，并在Python中加载你的模型。使用 SelfExtend.apply 方法来应用扩展功能。可以通过一个简单的代码示例了解使用方式：

import SelfExtend

# 加载模型，例如 loaded_model = AutoModelForCausalLM.from_pretrained(model_path)

# 设置群大小和邻居窗口

SelfExtend.apply(loaded_model, group_size, window_size, enable_flash_attention=False)

# 推理，例如 loaded_model.generate(...)

运行示例代码：

python example.py

3. 如何选择群大小和邻居窗口

对于Llama-2这样的基础模型，合理的群大小为2到64，而邻居窗口大小可在512到1536之间选择。具体的选择需要根据输入序列长度和期望的窗口尺寸来确定。由于不同模型的预训练程度有差异，选择偏大或者偏小的群大小都可能在不同条件下提供更好的效果。

实验与启示

通过实验发现，SelfExtend在超参数选择上不敏感，使用预定或启发式的参数值往往能获得令人满意的性能。可以利用简单的任务来确定合适的超参数，或者参考经验不等式来进行选择。

贡献与许可

LongLM项目欢迎来自研究界的贡献以提高SelfExtend的效率。在项目的GitHub页面上，研究人员可以提交问题或拉取请求。所有的代码都以MIT许可发布，任何有兴趣的人都可以免费使用和修改。

引用格式

如果您觉得我们的工作对您有帮助，请使用下面的引用格式：

@misc{jin2024llm,
      title={LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning}, 
      author={Hongye Jin and Xiaotian Han and Jingfeng Yang and Zhimeng Jiang and Zirui Liu and Chia-Yuan Chang and Huiyuan Chen and Xia Hu},
      year={2024},
      eprint={2401.01325},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

通过这些努力，LongLM项目致力于在无需对模型进行复杂调试的情况下提升长文本上下文的处理能力，提高大语言模型在实际应用中的效能。