Github
论文项目简介
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及相关工作。引用格式参见上文的代码示例。
