Flash STU:快速谱变换单元
机器学习
2026-01-21 v6 人工智能
摘要
状态空间模型架构的最新进展在高效序列建模方面展现出巨大潜力,但在平衡计算效率与模型表达能力方面仍存在挑战。我们提出了Flash STU架构,这是一种将谱状态空间模型层与滑动窗口注意力交替排列的混合模型,在保持近线性时间复杂度的同时,实现了语言建模向数十亿参数规模的可扩展性。我们在多种序列预测任务上评估了Flash STU及其变体,包括线性动力系统、机器人控制和语言建模。我们发现,在给定固定参数预算的情况下,Flash STU架构始终优于Transformer以及其他领先的状态空间模型,如S4和Mamba-2。
引用
@article{arxiv.2409.10489,
title = {Flash STU: Fast Spectral Transform Units},
author = {Y. Isabel Liu and Windsor Nguyen and Yagiz Devre and Evan Dogariu and Anirudha Majumdar and Elad Hazan},
journal= {arXiv preprint arXiv:2409.10489},
year = {2026}
}