xlstm

xLSTM 项目介绍

xLSTM 是一种最新的循环神经网络架构，它基于原始 LSTM（长短时记忆网络）的理念，进行了扩展和改进。通过指数门控结合合适的规范化和稳定化技术，以及全新的矩阵记忆机制，它可以克服原始 LSTM 的限制，并在语言建模领域相对于 Transformer 和状态空间模型展现出有潜力的性能。

项目背景

xLSTM 的设计旨在解决 LSTM 在处理长时间序列数据时的一些固有问题。LSTM 虽然在很多任务中表现出色，但其门控机制和存储容量限制了其在某些复杂任务上的应用。xLSTM 引入了更为复杂的门控和存储方式，例如指数门控和矩阵记忆，来增强其在语言建模等任务中的表现。

安装指南

为了开始使用 xLSTM，用户可以选择在 conda 环境中安装，推荐使用environment_pt220cu121.yaml文件进行环境搭建。安装该模型的代码可以通过以下方式实现：

使用 pip 安装：
```
pip install xlstm
```

从 GitHub 克隆仓库并安装：

git clone https://github.com/NX-AI/xlstm.git
cd xlstm
pip install -e .

环境要求

xLSTM 是基于 PyTorch 构建的，并在 PyTorch 版本不低于 1.8 的环境下进行测试。若需要使用 CUDA 版本的 sLSTM，用户的计算设备需具备计算能力不低于 8.0。推荐使用 environment_pt220cu121.yaml 以确保环境的兼容性和稳定性。

使用方法

xLSTM 提供了灵活的模块化设计，适用于不同的应用场景：

xLSTM 块堆叠

xLSTMBlockStack 可以作为现有项目中的替代骨干使用。它类似于 Transformer 块的叠加，但每个块采用 xLSTM 的设计：

import torch
from xlstm import xLSTMBlockStack, xLSTMBlockStackConfig, ...

cfg = xLSTMBlockStackConfig(...)
xlstm_stack = xLSTMBlockStack(cfg)

x = torch.randn(4, 256, 128).to("cuda")
xlstm_stack = xlstm_stack.to("cuda")
y = xlstm_stack(x)

xLSTM 语言模型

xLSTMLMModel 在 xLSTMBlockStack 的基础上增加了词嵌入和语言模型头，用于语言建模：

from xlstm import xLSTMLMModel, xLSTMLMModelConfig, ...

xlstm_cfg = """ 
vocab_size: 50304
mlstm_block:
  ...
"""
cfg = OmegaConf.create(xlstm_cfg)
cfg = from_dict(data_class=xLSTMLMModelConfig, ...)
xlstm_stack = xLSTMLMModel(cfg)

x = torch.randint(0, 50304, size=(4, 256)).to("cuda")
xlstm_stack = xlstm_stack.to("cuda")
y = xlstm_stack(x)

实验

xLSTM 在多项实验中展示了其性能与应用前景。在任务如奇偶校验任务和多查询关联记忆任务中，其不同构件（sLSTM 和 mLSTM）分别展现了出色的状态跟踪和记忆能力。在运行以上实验时，可以通过执行 main.py 文件进行：

python experiments/main.py --config experiments/parity_xLSTM01.yaml  # 仅sLSTM
python experiments/main.py --config experiments/parity_xLSTM10.yaml  # 仅mLSTM
python experiments/main.py --config experiments/parity_xLSTM11.yaml  # sLSTM与mLSTM结合

请注意，目前的训练循环还未包含早停或测试评估机制。

引用

如果使用 xLSTM 的代码库或其研究成果，请引用相关论文：

@article{xlstm,
  title={xLSTM: Extended Long Short-Term Memory},
  author={Beck, Maximilian and P{\"o}ppel, Korbinian and ...},
  journal={arXiv preprint arXiv:2405.04517},
  year={2024}
}