访问官网
Github
论文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}
}
