s4

项目简介

Structured State Spaces for Sequence Modeling（简称S4）项目是一个专注于序列建模的开源代码库，提供了多种模型的官方实现和实验。该项目聚集了多篇论文中的模型，包括HiPPO、LSSL、SaShiMi、DSS、HTTYH、S4D以及S4ND等。项目的详细信息及实验再现方法可在models目录下找到。

环境设置

要求

S4项目需要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

结构化核函数

S4的核心运算是Cauchy与Vandermonde核，与相关论文中描述的简单矩阵乘法类似。虽然有一种简单的实现方式，但可能在内存使用上较低效，因此需要自定义的核来优化内存使用。项目中提供了两种高效的方法来处理这些计算，并能自动识别安装情况以调用相应的核。

自定义CUDA核

这种方式速度较快，但需要针对每个计算机环境手动编译。在extensions/kernels/目录下运行python setup.py install即可。

Pykeops实现

Pykeops库提供了另一种实现方式，通常通过pip install pykeops cmake即可安装使用。

使用S4项目

S4模块

项目中的S4层及其变种在models/s4/目录中都有提供独立的文件，并附带调用该模块的说明。notebooks/目录中提供了相关概念的可视化讲解。

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

example.py是一个针对MNIST和CIFAR的训练脚本，默认设置下，通过简单的200k参数的S4D模型在顺序CIFAR上达到88%的准确率。此脚本可作为外部项目中使用S4变种的示例。

使用本项目进行训练（内部使用）

该项目旨在提供一个灵活的框架，用于训练序列模型，支持多种模型和数据集。基本的训练入口是python -m train，例如：

python -m train pipeline=mnist model=s4

这将在Permuted MNIST数据集上训练一个S4模型，大约需要1-3分钟，以获得约90%的准确率。

优化器超参数

代码库提供了支持不同优化器超参数的特性，尤其是对SSM核的$(A, B)$等参数，需要调整学习率和权重衰减。

训练

项目训练基础设施基于Pytorch-Lightning，配置使用Hydra。主要的训练入口是train.py，配置文件位于configs/目录下。

数据与模型

项目支持自动下载和加载基本数据集如MNIST、CIFAR。数据加载逻辑在src/dataloaders目录中。模型定义在src/models目录中。预定义的配置用于再现论文中的实验结果。

继续训练和多GPU

实验日志存储在由Hydra生成的目录中，可以从指定的.ckpt文件中继续训练。多GPU训练只需设置trainer.gpus=2即可。

生成

项目可以通过generate.py脚本进行自回归生成，此脚本可以在训练了模型后使用。

选项1：检查点路径

需要已经训练好的模型检查点路径，生成脚本接受与训练脚本相同的配置选项。

选项2：实验路径

无需再次传入训练标志，而是从Hydra实验文件夹读取配置。

整体目录结构

• configs/: 模型、数据管道、训练相关的配置文件
• data/: 原始数据的默认位置
• extensions/: CUDA扩展
• src/: 主代码模块
• models/: 模型特定的信息和资料
• example.py: 外部使用S4的例子
• train.py: 项目的训练入口
• generate.py: 自回归生成脚本

如何引用

如果您使用该代码库或认为我们的工作有价值，欢迎引用S4及相关工作。引用格式参见上文的代码示例。