论文

重塑和循环:通过输入重塑和深度循环改进 SSM

Reshape and Recur: Improving SSMs with Input Reshaping and Depth Recurrence

摘要

状态空间模型 (SSM) 越来越多地部署在边缘,因为与 大语言模型 (LLM) 相比,它们在相当的性能下提供更小的内存/训练/推理占用空间。这三个优点是 SSM 架构中固有的时间循环的直接结果。在这里,我们通过积极回答两个先前未充分探索的正交问题来进一步改进这种循环架构:(1)我们是否可以通过也采用深度循环来减少 SSM 内存占用而不造成任何性能损失? (2) 我们能否通过在所有任务中使用固定且一致的时间粒度来提高 SSM 性能?鉴于 SSM 已经反复出现,第一个问题有点出乎意料。然而,正交深度递归进一步减少了 SSM 的内存占用。我们证明,具有 $k$ 参数的循环 SSM 自适应迭代 $M$ 次,其性能与具有 $k \cdot L$ 独立参数的标准 SSM 相当,其中 $M \leq L$。考虑到 SSM 架构的时间循环性质,第二个问题也是出乎意料的。然而,它对于在整个输入序列上进行 SSM 的时间并行训练非常有意义。我们表明,连接低维序列元素的时间步长,或者对高维序列元素展平和重新分块联合特征时间维度,可以通过增强向模型呈现信息的方式来改进基线。我们对这两种扩展的结果在四种代表性 SSM 架构中带来了一致的好处:LRU、S5、LinOSS、LrcSSM。