用于序列建模的卷积

本仓库提供了以下论文的实现与实验，同时也对早期工作（如 S4）进行了简化呈现。

请参阅以下说明，了解如何下载权重并运行我们的预训练模型：

• H3（125m–2.7B）
• Hyena（小型，150M）

Hyena

Hyena 层次结构：迈向更大的卷积语言模型
Michael Poli*, Stefano Massaroli*, Eric Nguyen*, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, Christopher Ré
ICML 2023。口头报告。
论文

Long Convs

用于序列建模的简单且硬件高效的长卷积
Daniel Y. Fu*, Elliot L. Epstein*, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, Christopher Ré
ICML 2023。
论文

Hungry Hungry Hippos (H3)

Hungry Hungry Hippos：迈向基于状态空间模型的语言建模
Daniel Y. Fu*, Tri Dao*, Khaled K. Saab, Armin W. Thomas, Atri Rudra, Christopher Ré
ICLR 2023。引人注目的前25%（亮点展示）。
论文

路线图

• 将 H3、LLM 训练和合成数据纳入本仓库
• 迁移快速卷积代码
• 添加 Hyena 的实现与实验
• pip 包

更改日志

请参阅 CHANGELOG.md

设置

需求

本仓库需要 Python 3.8 及以上版本，以及 PyTorch 1.10 及以上版本。 其他依赖包列于 requirements.txt 中。

快速上手

最简单的入门方式是运行 standalone_cifar.py 脚本。 该脚本会在 CIFAR-10 数据集上训练一个简单的长卷积模型：

python -m standalone_cifar

更多内容请参阅 experiments.md 页面：

• 来自 Long Convs 论文的 LRA 实验
• H3 实验（语言模型、合成数据）
• H3 + 长卷积实验
• Hyena 的语言与视觉实验

资源

我们乐于分享关于本仓库中方法的独立复现及解读文章。

Hyena：

• irhum 的 JAX 复现及与 nanoGPT 的对比
• expz 的复现及其解读文章

引用

如果您使用了本代码库，或以其他方式认为我们的工作有价值，您可以按如下方式引用我们：

@article{poli2023hyena,
  title={Hyena 层次结构：迈向更大的卷积语言模型},
  author={Poli, Michael 和 Massaroli, Stefano 以及 Nguyen, Eric 等},
  journal={arXiv 预印本 arXiv:2302.10866},
  year={2023}
}

@article{fu2023simple,
  title={用于序列建模的简单且硬件高效的长卷积},
  author={Fu, Daniel Y. 等},
  journal={国际机器学习会议},
  year={2023}
}

@inproceedings{fu2023hungry,
  title={Hungry Hungry Hippos：迈向基于状态空间模型的语言建模},
  author={Fu, Daniel Y. 等},
  booktitle={国际表征学习会议},
  year={2023}
}

致谢

本仓库基于 Albert Gu 的 state spaces 仓库分叉而来，并沿用了其结构。 此外，它还包含了来自 FlashAttention 训练脚本的代码。

Safari 快速上手指南

Safari 是一个专注于序列建模的卷积模型库，实现了 Hyena、Long Convs 和 H3 (Hungry Hungry Hippos) 等前沿算法。本指南将帮助你快速搭建环境并运行基础示例。

环境准备

在开始之前，请确保你的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python 版本: 3.8 或更高
• PyTorch 版本: 1.10 或更高
• 硬件: 建议使用支持 CUDA 的 NVIDIA GPU 以获得最佳性能

安装步骤

1. 克隆仓库

首先从 GitHub 克隆项目代码：

git clone https://github.com/HazyResearch/safari.git
cd safari

2. 安装依赖

创建虚拟环境（推荐）并安装所需依赖包。为了加快下载速度，国内用户可配置使用清华或阿里镜像源：

# 创建虚拟环境
python -m venv venv
source venv/bin/activate  # Windows 用户使用: venv\Scripts\activate

# 使用国内镜像源安装依赖
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：如果 requirements.txt 中未明确指定 PyTorch 版本，请根据你的 CUDA 版本手动安装对应的 PyTorch。例如（CUDA 11.8）：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

基本使用

安装完成后，最简单的入门方式是运行官方提供的独立脚本，该脚本将在 CIFAR-10 数据集上训练一个基础的长卷积（Long Convolution）模型。

运行示例

在项目根目录下执行以下命令：

python -m standalone_cifar

运行成功后，你将看到模型的训练进度和损失值输出。

进阶实验

如需复现论文中的其他实验（如 LRA 基准测试、H3 语言模型训练、Hyena 视觉/语言任务等），请参考项目目录下的 experiments.md 文件获取具体配置和运行指令。

# 查看实验文档
cat experiments.md