Temporal Superposition and Feature Geometry of RNNs under Memory Demands¶
会议: ICLR 2026
OpenReview: https://openreview.net/forum?id=7cMzTpbJHC
代码: https://github.com/kashparty/iclr-rnn-superposition
领域: 可解释性 / 表示几何 / 循环网络
关键词: 时间叠加, 特征几何, RNN, 叠加假设, 工作记忆
一句话总结¶
本文把前馈网络里的"特征叠加"(superposition)概念扩展到时间维度,提出时间叠加(temporal superposition):在 \(k\)-延迟回忆任务上训练线性/非线性 RNN,通过把损失解析分解为四项、识别出 ReLU 诱导的"无干扰空间"、以及密集→稀疏两个机制之间的相变,机制性地解释了 RNN 在记忆压力下为何、如何选择不同的表示几何。
研究背景与动机¶
领域现状:在可解释性领域,"叠加假设"(Elhage et al., 2022) 已经成为理解神经网络多义性(polysemanticity)的主流框架。当数据的特征数多于神经元数、且特征稀疏(不常同时出现)时,网络会把特征非正交地塞进激活空间——用更少的维度表示更多的特征,代价是特征之间产生干扰。这种"空间叠加"在前馈/Transformer 里被反复验证,机制可解释性的大量工作都在研究如何从叠加里把单个特征抽取出来。
现有痛点:以往对叠加的研究几乎只关注空间维度(特征数 vs 神经元数),完全忽略了时间(记忆)这个额外的容量约束。在前馈和 Transformer 里时间不是问题,但循环结构天然带来记忆压力:RNN 的隐状态是一个固定大小的瓶颈,要同时容纳"现在的输入"和"过去若干步还需要保留的输入"。记忆窗口一旦超过隐状态维度,RNN 就被迫要么遗忘、要么把多个时刻的特征压进同一个低维空间。这种由时间引起的叠加,此前没有被刻画过。
核心矛盾:隐状态维度 \(N_h\) 固定,而任务要求保留的"任务相关特征"数量随记忆长度 \(k\) 线性增长。当 \(k+1 > N_h\) 时,RNN 必须在"多保留特征"和"减少特征间干扰"之间权衡——这正是叠加的根源,但带上了时间轴后干扰的形态变得更复杂。
本文目标:(1) 形式化"时间叠加"概念;(2) 搞清楚数据稀疏性、任务记忆需求 \(k\)、网络维度三者如何共同决定 RNN 学到的表示几何;(3) 解释这些几何"为什么"是最优的。
切入角度:作者用一个极简但可解析的设定——线性循环 RNN、标量输入输出、二维隐状态、\(k\)-延迟回忆任务——把问题做到能手推损失闭式解的程度,再验证这些洞察能推广到非线性 RNN 和高维情形。
核心 idea:把每个输入特征看成不仅有"是什么"(空间方向),还有"在多久之前出现"(时间分量)的二重身份;隐状态里同时叠加着输入历史的多个特征方向 \(w_s\),而 RNN 学到的几何就是在最小化由这些方向相互投影/组合产生的干扰。
方法详解¶
整体框架¶
本文不是提出一个新模型,而是建立一套理论分析框架来解释 RNN 在记忆压力下的表示几何。研究对象是一个线性循环 RNN:
其中 \(\sigma\) 可以是恒等(线性)或 ReLU。当 \(\sigma_h\) 线性时,隐状态展开成历史输入的加权和。作者引入"经历了多少次递归"的量 \(s := t-i\),定义 \(W_s := W_h^s W_x\),把隐状态重写为:
这个改写是整套分析的支点:它说明在时刻 \(t\),每个历史输入 \(x_{t-s}\) 都被独立、线性地表示在方向 \(w_s\) 上(称为"特征方向")。读出向量 \(w_y\) 只对那些投影到它身上的方向有响应。任务是 \(k\)-延迟回忆:要求 \(\hat{y}_t = x_{t-k}\),即在 \(k\) 步之后重放输入。于是 \(w_{s=k}\) 是"输出特征方向"(应当对齐 \(w_y\)),\(w_s\ (0\le s<k)\) 是仍需保留的"中间特征",而 \(s>k\) 的是应被遗忘的"历史无关特征"。
分析沿四步推进:① 写出线性情形损失的闭式期望并分解成四项,读出 RNN 学到的几何策略与学习动力学;② 给读出加 ReLU,在高稀疏极限下近似损失,揭示"无干扰空间";③ 扫描稀疏度,观察密集↔稀疏的相变;④ 把标量输入换成向量、改变 \(k\),研究空间叠加与时间叠加如何相互权衡。
关键设计¶
1. 时间叠加与特征方向:把"何时出现"编码进几何
本文最核心的概念创新是把叠加从空间维度推广到时间维度。空间叠加说的是"特征多于神经元时非正交压缩";时间叠加则指出:在 RNN 里,同一个输入 \(A\) 在不同时刻被表示成不同的方向——一个 \(A\) 的脉冲输入后,随着它"变老"(\(s\) 增大),它在隐状态里会沿着一串特征方向 \(w_{s=0}, w_{s=1}, \dots\) 移动,直到不再与任务相关。表示因此由"何时"(when)和"是什么"(what)共同决定。
由于隐状态维度固定,记忆窗口越长(\(k\) 越大),被迫压进隐状态的特征方向越多,瓶颈就越紧。本文据此把时间叠加明确区分于空间叠加:时间叠加 = 在更长时间跨度上(更高记忆需求)、于更低维激活空间里表示特征;空间叠加 = 在同一时刻把更多输入特征塞进低维空间。这个区分让后面所有几何现象都能挂到"记忆压力"这根轴上。
2. 损失的四项分解:把几何策略拆成可解释的竞争激励
为了搞清楚学习问题到底在"奖励"什么几何,作者在数据时间独立、稀疏的假设下,推导出线性 RNN 损失期望的闭式形式,恰好分解为四项:
其中 \(p\) 控制时间稀疏度(\(p\) 越小越稀疏),\(\mu, \nu\) 是输入分布的均值与方差。这四项各自对应一种激励:任务收益奖励把输出特征 \(w_{s=k}\) 对齐到读出 \(w_y\);均值修正说明 RNN 会利用投影干扰来抵消非零均值(若输出加 bias 或 \(\mu=0\) 则此项消失);投影干扰代价惩罚在错误时刻(\(s\ne k\))投影到读出的特征方向;组合干扰惩罚 \(w_s\) 之间的正相关、奖励负相关,因此鼓励特征方向尽量散开、理想情况下形成对跖(antipodal)对。
这个分解不仅与实测损失曲线高度吻合,还能解释学习动力学:训练初期所有 \(w_s\) 都先对齐到读出(任务误差下降、其它三项上升),之后特征方向按时间顺序在激活空间里散开,呈现"阶梯式"损失(对应奇异值的逐级学习、saddle-to-saddle 动力学)。线性 RNN 最终收敛到一个螺旋汇(spiral sink):每次递归把特征下缩并旋转,老特征螺旋进原点,实现"平滑遗忘"——在 2D 情形这已被证明是 \(k\)-延迟任务的最优解。
3. 无干扰空间:ReLU 读出腾出半空间,把中间特征全塞进去
给读出加上 ReLU(\(\sigma_y = \text{ReLU}\))后解析解更难求,作者在高时间稀疏极限下近似损失:
关键转折在于:ReLU 只对正投影产生输出,因此所有落在读出 \(w_y\) 对侧半空间的特征方向,其投影被 ReLU 截成 0,完全不贡献投影干扰。在极稀疏(组合干扰可忽略)时,这个半空间就成了"无干扰空间"。这给了模型一个强烈激励:把尽可能多的中间特征方向 \(w_{s\ne k}\) 打包进这个半空间,只让输出特征 \(w_{s=k}\) 留在外面对齐读出。
不同模型对这个空间的利用程度不同,正好暴露了表达力的边界:线性 RNN 根本没有无干扰空间(只能螺旋遗忘);线性递归+ReLU 读出的 SSM 因为递归仍是线性的、\(w_s\) 只能沿椭圆螺旋等距排列,所以只能近似地把最大的几个特征方向塞进无干扰空间;而非线性 RNN(\(\sigma_h\) 也用 ReLU)表达力足够,能把全部 \(k\) 个中间特征精确打包进无干扰空间(常常压进单个象限),并实现锐利遗忘——直接把特征送进 \(h_t\) 的负象限被 ReLU 清零,而不是线性 RNN 那种逐步缩小的平滑遗忘。
4. 密集↔稀疏相变与空间×时间叠加的 all-or-none 权衡
扫描时间稀疏度,SSM 会在两种离散机制间发生相变。密集机制(低稀疏)下,由于组合干扰更易发生(把大特征方向塞进负的无干扰空间后,它若与小的正特征方向相加,ReLU 会把输出截成 0),SSM 退化成类似线性 RNN 的螺旋汇,任务相关特征方向挤在约 \(90°\) 的小锥里;稀疏机制下,SSM 充分利用无干扰空间,任务相关特征方向铺开到约 \(270°\)、横跨整个无干扰空间才够到读出。最直接的相变指标是任务相关特征张成的角度 \(k\theta\)(从 \(w_{s=0}\) 到 \(w_{s=k}\) 转过的角),它随稀疏度突变,同时伴随谱半径 \(\rho = \|w_{s=0}\|\) 下降(密集机制需要更大的特征方向来压过投影干扰)。
把标量输入换成向量、同时变化 \(k\),就能看到空间叠加与时间叠加的相互权衡。\(k=0\) 时回到纯空间叠加(5 个特征排成五边形)。随 \(k\) 增大,非线性 RNN 采取"全或无"(all-or-none)策略:要从表示某个特征中获益,必须能把它在全部 \(k+1\) 个时刻都维持住,否则它对降低损失毫无贡献,干脆完全不表示。于是最重要的特征 \(A\) 在所有 \(k\) 下都保留 \(k+1\) 个特征方向,而次要特征只在 \(k\) 较小、容量富余时才出现。这揭示了一个清晰的权衡:为多个特征各保留短时间 vs 为单个重要特征保留全程,RNN 在容量受限时坚定地选择后者。
损失函数 / 训练策略¶
任务损失就是 \(k\)-延迟回忆的平方误差 \(L = \sum_t \|y_t - \hat{y}_t\|^2\),其中 \(y_t = x_{t-k}\)(\(t\le k\) 时 \(y_t=0\))。\(k\) 作为控制参数指定输入必须维持多久;\(p\) 控制时间稀疏度。理论分析在"时间独立 + 稀疏"假设下推导损失期望,并分别在线性(精确)与 ReLU 读出(高稀疏近似)两种情形给出闭式/近似表达,再用实测损失曲线验证。
实验关键数据¶
本文是理论+机制分析为主,"实验"主要是用小型 RNN 验证理论预测的几何策略,而非刷 benchmark。
主结果:理论预测 vs 训练所得几何¶
| 模型 | 递归/读出 | 稀疏机制下的几何策略 | 遗忘方式 |
|---|---|---|---|
| 线性 RNN | 线性 / 线性 | 无无干扰空间,老特征螺旋进原点(螺旋汇) | 平滑遗忘 |
| SSM | 线性 / ReLU | 近似把最大特征方向塞进无干扰空间 | 介于两者之间 |
| 非线性 RNN | ReLU / ReLU | 完整把 \(k\) 个中间特征打包进无干扰空间,仅 \(w_{s=k}\) 在外 | 锐利遗忘 |
损失实测曲线与四项分解的期望损失 \(\mathbb{E}[L]=\)(i)+(ii)+(iii)+(iv) 高度吻合,且呈现"先读出对齐、后按时间顺序分离"的阶梯式学习动力学。
相变与高维验证¶
| 分析 | 设置 | 关键发现 |
|---|---|---|
| 密集→稀疏相变 | 扫描稀疏度 \(1-p\) | 角度 \(k\theta\) 从约 \(90°\) 突变到约 \(270°\),谱半径 \(\rho\) 同步下降 |
| 高维隐状态 | \(N_x=10\), \(N_h\in\{2,5,10\}\), 2-延迟 | \(W_y^\top W_{s=2}\) 对角线为正(正确时刻正确输出),其余及 \(W_y^\top W_{s\ne 2}\) 为负/0(落在无干扰空间);隐状态越大对角线捕获特征越多 |
| 大规模验证 | \(N_h=100\), 75 特征, 2-延迟 | 最优模型把最大特征方向分组进无干扰空间、并在恰当时刻把输出特征投影到读出,与预测一致 |
关键发现¶
- 四项损失分解是全文枢纽:它既精确预测实测损失,又把"对齐读出 / 散开特征 / 抵消均值 / 避免组合干扰"这些几何动机一一对应,让"RNN 为什么学这种几何"变得可解释。
- 表达力决定能否吃到无干扰空间:线性 RNN 吃不到、SSM 近似吃、非线性 RNN 完整吃,并对应平滑 vs 锐利两种遗忘——同一套理论解释了三类架构的差异。
- 容量受限时 RNN 是"全或无"的:宁可完整保留一个重要特征 \(k+1\) 步,也不愿为多个特征各保留半截,这与工作记忆里的"slot"和"resource"两种理论同时吻合。
亮点与洞察¶
- 把"时间"作为容量约束引入叠加研究:之前叠加只谈空间(特征数 vs 神经元数),本文指出记忆长度同样是压缩压力,且会随时间线性消耗容量——这个视角直接迁移到对长程依赖 RNN/SSM/Mamba 的可解释性分析。
- 无干扰空间是个很"干净"的机制解释:ReLU 截掉对侧半空间的投影,于是模型有动机把无关特征全藏进去——一个非线性带来的"免费午餐",也解释了非线性 RNN 为何能锐利遗忘。
- 可迁移的方法论:把隐状态展开成 \(\sum_s w_s x_{t-s}\)、再按 \(w_s\) 对读出的投影来定义投影/组合两类干扰,这套"特征方向 + 干扰分解"的分析范式可以套到其它带时间轴的表示几何问题上。
局限与展望¶
- 作者承认的局限:理论假设特征时间独立、只研究小型 RNN;空间稀疏假设较合理,但时间稀疏假设较强、是否普遍成立取决于任务,仍是开放问题。
- 任务单一:只研究 \(k\)-延迟回忆(纯重放),没有涉及需要对输入做操作/变换、或记忆需求变化的任务,这类任务的几何与行为有待刻画。
- 线性表示假设:整套分析建立在"特征是激活空间里的方向"这一线性表示假设上;对高度过参数化的现代大模型在多大程度上适用,作者也坦承存疑(但指出叠加已在 LLM 中被观测到,且递归与记忆只会加剧它)。
- 改进思路:把分析推广到带门控/复值参数化的 SSM、变延迟任务、以及高维高容量设定,验证无干扰空间与相变是否依然成立。
相关工作与启发¶
- vs Elhage et al. (2022)(空间叠加 toy model):他们在前馈网络里形式化了空间叠加(\(k=0\) 恰好退化到他们的设定),本文把它扩展到时间维度,新增了投影/组合干扰区分、无干扰空间和密集↔稀疏相变。
- vs 低秩 RNN 理论(Mastrogiuseppe & Ostojic 等):低秩连接同样是一种容量约束、可视化相图,但很少显式研究"记忆引起的容量约束"下的特征几何,本文正好补这一块。
- vs 工作记忆的 slot vs resource 理论:本文发现 RNN 同时体现两者——把不同时刻特征排进各自"槽位"(方向),但在叠加下槽位非正交;同时隐空间又是连续的有限资源,需求(记忆长度/特征数)越大越易干扰。
- vs François et al. (2025):他们在频域研究欠参数化线性 RNN 的 \(k\)-延迟任务但聚焦密集(线性)机制;本文进一步研究非线性 RNN 的稀疏机制,识别出新的行为。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次把叠加从空间扩展到时间维度,提出时间叠加并给出完整机制解释
- 实验充分度: ⭐⭐⭐⭐ 理论与小规模实证吻合扎实,但任务单一(仅 \(k\)-延迟回忆)、未触及大模型
- 写作质量: ⭐⭐⭐⭐⭐ 概念清晰、损失分解优雅、图文对照把抽象几何讲得很直观
- 价值: ⭐⭐⭐⭐ 为 RNN/SSM 可解释性和神经科学工作记忆建模提供了可迁移的几何分析框架