用于序列建模 (Sequence Modeling) 的结构化状态空间 (Structured State Spaces)

本仓库提供了与 S4 相关模型的官方实现和实验， 包括 HiPPO、LSSL、SaShiMi、 DSS、HTTYH、S4D 和 S4ND。

每个模型的项目特定信息，包括源代码概述和具体实验复现， 可以在 models/ 下找到。

目录

设置环境并将 S4 移植到外部代码库：

• 设置
• S4 入门

使用本仓库训练模型：

• 训练
• 生成
• 仓库结构
• 引用

变更日志

参见 CHANGELOG.md

路线图

• 更多关于使用本仓库从头开始训练的文档
• S4 资源和实现的汇编
• pip 包

设置

要求

本仓库需要 Python 3.9+ 和 Pytorch 1.10+。 已测试至 Pytorch 1.13.1。 其他包列在 requirements.txt 中。 需要注意使某些库版本兼容，特别是 torch/torchvision/torchaudio/torchtext。

安装示例：

conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia
pip install -r requirements.txt

结构化核函数 (Structured Kernels)

S4 的核心操作是 论文 中描述的柯西 (Cauchy) 和范德蒙德 (Vandermonde) 核函数 (kernels)。 这些是非常简单的矩阵乘法；这些操作的朴素实现可以在 standalone 中的 cauchy_naive 和 log_vandermonde_naive 函数中找到。 然而，正如论文所述，这需要次优的内存使用，目前需要自定义核函数 (kernel) 来在 PyTorch 中克服。

支持两种更高效的方法。代码会自动检测是否安装了其中任何一种并调用相应的核函数。

自定义 CUDA 核函数 (Custom CUDA Kernel)

此版本更快，但需要为每个机器环境手动编译。 从 extensions/kernels/ 目录运行 python setup.py install。

Pykeops

此版本由 pykeops 库 提供。 安装通常只需 pip install pykeops cmake 即可开箱即用，这些也列在 requirements 文件中。

S4 入门

S4 模块

S4 层及其变体的独立文件可以在 models/s4/ 中找到， 其中包括调用模块的说明。

参见 notebooks/ 以获取解释 HiPPO 和 S4 背后一些概念的可视化内容。

示例训练脚本（外部使用）

example.py 是一个用于 MNIST 和 CIFAR 的独立训练脚本，导入独立的 S4 文件。默认设置 python example.py 在顺序 CIFAR 上达到 88% 的准确率，使用的是一个非常简单的 200k 参数的 S4D 模型。 此脚本可用作在外部仓库中使用 S4 变体的示例。

使用本仓库训练（内部使用）

本仓库旨在提供一个非常灵活的框架来训练序列模型。支持许多模型和数据集。

基本入口点是 python -m train，或等价于

python -m train pipeline=mnist model=s4

这在 Permuted MNIST 数据集上训练一个 S4 模型。 这应该在 1 个 epoch (训练轮次) 后达到约 90%，根据 GPU 不同需要 1-3 分钟。

更多使用本仓库的示例在整个文档中都有记录。参见 训练 以获得概述。

优化器超参数 (Optimizer Hyperparameters)

此代码库的一个重要功能是支持需要不同优化器超参数 (optimizer hyperparameters) 的参数。 特别是，SSM 核函数 (kernel) 对 $(A, B)$（有时是 $\Delta$ 参数）特别敏感， 因此这些参数的学习率有时会降低，并且权重衰减 (weight decay) 始终设置为 $0$。

参见模型中的 register 方法（例如 s4d.py）和训练脚本中的 setup_optimizer 函数（例如 example.py），以了解如何在外部仓库中实现此功能的示例。

训练

本仓库的核心训练基础设施基于 Pytorch-Lightning，配置方案基于 Hydra。

主入口点是 train.py，配置文件在 configs/ 中。

数据

基本数据集会自动下载，包括 MNIST、CIFAR 和 Speech Commands。 创建和加载数据集的所有逻辑都在 src/dataloaders 目录中。 此子目录内的 README 记录了如何下载和组织其他数据集。

模型

模型定义在 src/models 中。参见此子目录中的 README 以获得概述。

配置和超参数

提供了复现论文中端到端实验的预定义配置，位于 models/ 中的项目特定信息下，例如 原始 S4 论文。

配置也可以通过命令行轻松修改。 一个实验示例是

python -m train pipeline=mnist dataset.permute=True model=s4 model.n_layers=3 model.d_model=128 model.norm=batch model.prenorm=True wandb=null

这使用 Permuted MNIST 任务，带有指定层数、主干维度 (backbone dimension) 和归一化类型 (normalization type) 的 S4 模型。

参见 configs/README.md 以获取有关配置的更详细文档。

Hydra

建议阅读 Hydra 文档 以充分了解配置框架。如需帮助启动特定实验，请提交 issue。

恢复训练

每个实验都将记录到其自己的目录中（由 Hydra (配置管理工具) 生成），格式为 ./outputs/<date>/<time>/。检查点（Checkpoints）将保存于此文件夹内，并在创建新检查点时打印到控制台。 要恢复训练，只需指向所需的 .ckpt 文件（PyTorch Lightning (深度学习框架) 检查点，例如 ./outputs/<date>/<time>/checkpoints/val/loss.ckpt），并将标志 train.ckpt=<path>/<to>/<checkpoint>.ckpt 附加到原始训练命令后。

PyTorch Lightning Trainer

PTL Trainer 类控制整体训练循环，并提供许多有用的预定义标志。下面解释了一些有用的示例。 允许的标志完整列表可在 PTL 文档以及我们的 trainer configs 中找到。参见默认训练器配置 configs/trainer/default.yaml 以获取最有用的选项。

多 GPU 训练

只需传入 trainer.gpus=2 即可使用 2 个 GPU (图形处理器) 进行训练。

检查模型层

trainer.weights_summary=full 打印出模型的每一层及其参数计数。用于调试模型内部结构。

数据子采样

trainer.limit_{train,val}_batches={10,0.1} 仅在 10 个批次（所有批次的 0.1 比例）上进行训练（验证）。用于在不遍历所有数据的情况下测试训练循环。

WandB

使用 WandB (实验跟踪工具) 记录日志的功能已内置到此仓库中。 为了使用此功能，只需设置 WANDB_API_KEY 环境变量，并更改 configs/config.yaml 的 wandb.project 属性（或在命令行上传递，例如 python -m train .... wandb.project=s4）。

设置 wandb=null 以关闭 WandB 日志记录。

生成

自回归（Autoregressive）生成可以使用 generate.py 脚本执行。 在使用此代码库训练模型后，该脚本可以通过两种方式使用。

选项 1：检查点路径

更灵活的选项需要训练好的 PyTorch Lightning 模型的检查点路径。 生成脚本接受与训练脚本相同的配置选项，还有一些额外的标志记录在 configs/generate.yaml 中。 使用 python -m train <train flags> 训练后，生成使用

python -m generate <train flags> checkpoint_path=<path/to/model.ckpt> <generation flags>

配置中的任何标志都可以被覆盖。

注意：此选项可与 .ckpt 检查点（PyTorch Lightning，包含 Trainer 信息）或 .pt 检查点（PyTorch，仅是模型状态字典）一起使用。

选项 2：实验路径

生成的第二个选项不需要再次传递训练标志，而是从 Hydra 实验文件夹读取配置，以及实验文件夹内的 PyTorch Lightning 检查点。

示例 1 (语言)

下载 WikiText-103 模型检查点，例如到 ./checkpoints/s4-wt103.pt。 此模型是使用命令 python -m train experiment=lm/s4-wt103 训练的。注意从配置中我们可以看到模型是使用长度为 8192 的感受野（Receptive field）训练的。

要生成，运行

python -m generate experiment=lm/s4-wt103 checkpoint_path=checkpoints/s4-wt103.pt n_samples=1 l_sample=16384 l_prefix=8192 decode=text

这将生成一个长度为 16384 的样本，以长度为 8192 的前缀为条件。

示例 2 (音频)

让我们在 SC09 数据集上训练一个小型 SaShiMi 模型。我们还可以减少训练和验证批次的数量以更快地获得检查点：

python -m train experiment=audio/sashimi-sc09 model.n_layers=2 trainer.limit_train_batches=0.1 trainer.limit_val_batches=0.1

第一个 Epoch (训练轮次) 完成后，会打印一条消息指示检查点保存的位置。

Epoch 0, global step 96: val/loss reached 3.71754 (best 3.71754), saving model to "<repository>/outputs/<date>/<time>/checkpoints/val/loss.ckpt"

选项 1：

python -m generate experiment=audio/sashimi-sc09 model.n_layers=2 checkpoint_path=<repository>/outputs/<date>/<time>/checkpoints/val/loss.ckpt n_samples=4 l_sample=16000

此选项重新定义完整配置，以便可以构建模型和数据集。

选项 2：

python -m generate experiment_path=<repository>/outputs/<date>/<time> checkpoint_path=checkpoints/val/loss.ckpt n_samples=4 l_sample=16000

此选项只需要 Hydra 实验文件夹的路径以及其中所需的检查点。

整体仓库结构

configs/         Config files for model, data pipeline, training loop, etc.
data/            Default location of raw data
extensions/      CUDA extensions (Cauchy and Vandermonde kernels)
src/             Main source code for models, datasets, etc.
  callbacks/     Training loop utilities (e.g. checkpointing)
  dataloaders/   Dataset and dataloader definitions
  models/        Model definitions
  tasks/         Encoder/decoder modules to interface between data and model backbone
  utils/
models/          Model-specific information (code, experiments, additional resources)
example.py       Example training script for using S4 externally
train.py         Training entrypoint for this repo
generate.py      Autoregressive generation script

引用

如果您使用此代码库，或发现我们的工作有价值，请引用 S4 和 其他相关论文。

@inproceedings{gu2022efficiently,
  title={Efficiently Modeling Long Sequences with Structured State Spaces},
  author={Gu, Albert and Goel, Karan and R\'e, Christopher},
  booktitle={The International Conference on Learning Representations ({ICLR})},
  year={2022}
}

S4 快速上手指南

环境准备

在开始使用 S4 之前，请确保您的开发环境满足以下系统要求：

• Python: 3.9 或更高版本
• PyTorch: 1.10 或更高版本（本项目已测试至 1.13.1）
• 其他依赖: 详见项目根目录下的 requirements.txt
• 硬件: 推荐使用 NVIDIA GPU 以加速训练（部分内核需要 CUDA 支持）

安装步骤

1. 安装 PyTorch 及相关组件 建议使用 Conda 创建环境，并安装兼容版本的 PyTorch 家族库：

   conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia
   

2. 安装项目依赖 进入项目目录，安装所需 Python 包：

   pip install -r requirements.txt
   

3. 配置高效内核（可选但推荐） S4 的核心运算涉及 Cauchy 和 Vandermonde 内核。为了优化内存和速度，代码会自动检测以下两种实现之一：

   • 自定义 CUDA 内核（速度最快，需编译）：

     cd extensions/kernels/
     python setup.py install
     

   • Pykeops 库（通常可直接安装）：

     pip install pykeops cmake
     

基本使用

1. 运行独立示例脚本

仓库提供了自包含的训练脚本 example.py，适合快速验证或将 S4 集成到外部代码库。默认设置下在序列 CIFAR 数据集上训练：

python example.py

预期结果：使用约 200k 参数的 S4D 模型，准确率可达 88%。

2. 使用仓库框架训练

本仓库基于 PyTorch-Lightning 和 Hydra 配置系统。以下命令在 Permuted MNIST 数据集上训练 S4 模型：

python -m train pipeline=mnist model=s4

预期结果：1 个 epoch 后准确率约为 90%（耗时 1-3 分钟，取决于 GPU）。

3. 自定义超参数

您可以通过命令行直接修改模型配置。例如，指定层数、维度并关闭 WandB 日志：

python -m train pipeline=mnist dataset.permute=True model=s4 model.n_layers=3 model.d_model=128 model.norm=batch model.prenorm=True wandb=null

更多配置详情可查阅 configs/ 目录文档。