对角状态空间模型中的门控位置:输入门控与状态门控的表达能力研究
Where Should a Gate Look? Input vs. State Gating in Diagonal State Space Models
这篇论文回答了一个架构设计里常被忽略的问题:门控该看输入还是看状态。结论很实用——做上下文检索类任务(induction head)时状态门控有理论优势,做 selective copying 时输入门控更强,两者结合最好。调 Mamba 类模型架构的人值得把实验部分仔细看看。
一篇 arXiv 论文研究了门控机制在对角状态空间模型中的设计位置问题,比较了仅依赖当前输入的门控、依赖循环状态的门控以及两者结合的三种方案。理论层面,论文给出了严格的表达能力分离证明:在标量模型中,状态门控可以解决任意词表大小和序列长度下的 induction-head 任务,而输入门控在状态转移矩阵满足标准稳定性假设的条件下,在一个有限实例上就会失败。
实验部分验证了理论结论:状态门控在高维模型的 induction-head 任务上保持优势,而输入门控在 selective copying 任务上表现更好。核心发现是,即使是对角形式的简单状态反馈门控,也能解锁输入控制型模型无法实现的能力;同时两类门控各有互补优势,将输入门控与状态门控结合可以获得最强的整体性能。
对开发者而言,这项工作为 Mamba 类对角状态空间模型的架构微调提供了设计依据:如果目标能力偏向上下文检索类任务,应优先引入状态反馈门控,而非只调输入门控。
Where Should a Gate Look? Input vs. State Gating in Diagonal State Space Models
Modern state-space models owe much of their computational efficiency to their linear, often diagonal, dynamics. Classical recurrent neural networks instead allow nonlinear state-to-state interactions, which can be beneficial when memory updates need to depend on the past context. In this paper, we ask whether state feedback through a simple gate regulating memory updates can enhance the expressivity of diagonal state-space models. To investigate this question, we compare gates that depend on the current input, the recurrent state, or both. We establish a strict expressivity separation: for scalar models, state gating solves induction-head tasks for arbitrary vocabulary size and sequence length, whereas input gating fails on a finite instance under standard stability assumptions on the state transition matrix. Experiments show that state gating retains its advantage on induction-head tasks in higher-dimensional models, whereas input gating performs better on selective copying tasks. Together, our results show that even a simple, diagonal form of state-feedback gating can unlock capabilities unavailable to input-controlled models, while input and state gating offer complementary strengths and achieve the strongest overall performance when combined.
This report is categorized under AITOP's "Papers" desk. Key entities and technical topics under continuous tracking include 状态空间模型, Mamba, 门控机制, induction head.
Source attribution: curated from arXiv cs.AI (@Riccardo Zattra, Giacomo Baggio, Leonardo De Marchi, Augusto Ferrante, Francesco Ticozzi, Marco Del Prà, Antonio Orvieto).