阅读笔记:Memory Caching:RNNs with Growing Memory
0. 发展脉络
解决 Miras 和 NL 架构面临的超长文本内存上限问题。
其实可以尝试用 NL 的视角来看待 MC 这篇文章,把 MC 当做一种跨频率的知识路由机制
0.1 分段长度 -> 频率
在 NL 中,模型不同层级是不同更新频率。在 MC 中,核心操作是将长序列切分为多个自定义长度的片段,并缓存这些片段的检查点。
- 在线记忆(Online Memory)= 最高频:在 MC 中,当前正在处理的在线记忆,每输入一个 Token 就会利用Mirsa内部目标(如梯度下降)更新一次(即频率 f=1)。它等价于 NL 中负责处理瞬时上下文的极高频神经元。
- 缓存记忆(Cached Memories)= 低频:一旦一个片段结束,它的记忆状态就会被“冻结”并缓存。这意味着这些缓存的记忆是经过了一段较长的时间周期(即片段长度 C)才最终成型的一次“低频快照”。
- 频率视角的折中:在 MC 中,Transformer 可以被看作是片段长度为1(无限高频,保留所有细节)的特例,而传统 RNN 是片段长度为无限大(极低频,全部压缩到一起)的特例。MC 的分段大小,实际上就是在为系统人为设定一个“中频”,也就是==折中==。
0.2 Log-Linear++ -> CMS
在 NL 中,CMS ,即系统应包含从高频到低频的连续渐变谱,高频记近期,低频记远期。其实 MC 中的 Log-Linear++ 就有点这个思想
用 NL 的频率视角来看,这是多频嵌套结构:
- 距离当前 Token 近的片段,切分得短(更新频率高,保留细粒度的近期记忆)
- 距离当前 Token 远的片段,切分得长(更新频率低,提取宏观长期规律)
0.3 Aggregation -> 跨频知识传递
在 NL 中,不同频率层级之间必须有通信机制,在 MC 中,作者设计的 Agg 聚合函数,其本质工作就是实现“跨频知识传递”:
- 当模型在处理当前的 Query 时,它不仅读取最高频的“在线记忆”,还通过 Agg 函数同时读取过去所有低频的“缓存记忆”。
- GRM(门控残差记忆)最高频模块在主动决定需要从哪几个低频模块中提取历史知识。
0.4 Memory Soup -> 动态生成专属频率层
NL 强调了系统应该具有自我修改的能力, MC 提出的 Memory Soup 机制针对深度非线性内存(如 Titans),将历史低频的参数矩阵(Checkpoints)进行加权平均,动态插值出一个专属于当前 Token 的新记忆网络。用频率视角来看,这是高频模块利用过去的多个低频知识库,实时合成了一个最适合当下语境的中频记忆网络来进行检索。
前面发展脉络说的有点多了,接下来开始把注意力放在 MC 这篇文章本身上
1. 现有问题
- RNN:时间复杂度 O(L) 内存尺寸固定,内存占用极低且恒定。但也是因为内存大小恒定,导致面对长上下文和经常回忆的任务时会遗忘过去信息
- Transformer:时间复杂度 O(L2) ,序列越长,处理时间平方级增长。内存随长度线性增长,易引发内存溢出,但也因此可以看到所有历史信息,获得了长上下文的能力
2. 基本逻辑
作者想提出一种算法,兼顾 Transformer 的记忆容量与 RNN 的计算效率,很显然就是计算机领域常见的==折中思想==。MC(Memory Caching)算法有3个主要的步骤
- Segmenting:最直观的,要避免 RNN 把超长序列当做一个大片段导致的内存溢出,同时避免 Transformer 把每个词当做一个片段导致的平方计算代价,当然就是要==折中==,把长序列切分为多个适中长度的片段 S(1),…,S(N),分别处理(有个小小的好处就是,切分的长度和方法可以自由设计,导致时间复杂度可以由设计者自己决定)。
- Caching:在每个片段结束时,把它压缩后的记忆状态作为检查点存入缓存,相当于一个历史“快照”。
- Aggregation:在输出时,通过 Agg 函数同时调动当前在线记忆与历史缓存记忆,结合现在的记忆和历史的记忆,就能避免遗忘历史信息。
3. 模块选择
3.1 Agg函数的几种选择
作者提出了 MC 算法的基本逻辑之后,就要着手设计用来聚合多个缓存记忆的函数(Agg函数)。作者从简单到复杂,提出了以下几种Agg函数的设计方案:
- Residual Memory(残差记忆):最简单,将当前输出直接与所有历史缓存的检索结果求和。作者模仿 Transformer 中的残差连接起了这么个名字。
- Gated Residual Memory(GRM)(门控残差记忆):从名字来看,显然就是在上面那种方法的基础上增加了一个“门控”。模型衡量当前 Q 与各个片段整体语义的关联度。若高度相关,则提高这段记忆的权重;若无关,则忽略这个“干扰”。
- Memory Soup: 这是针对 Titans 等非线性深度记忆模块的特殊优化。它使用插值参数,对历史权重 θM(i) 进行加权平均,在每一个时间步动态生成一个专属于当前语境的非线性检索函数。它在线性内存(矩阵存储)下与 GRM 是完全等价的,但在非线性深度内存下才能体现优化成果。
- Sparse Selective Caching(SSC)(稀疏选择性缓存): 模仿 MoE 的 Router 机制,计算与各片段缓存的点积相似度评分,仅加载评分最高的 k 个缓存状态,其余保持休眠。极大地降低了超长序列下的内存读取压力,只加载最需要的几个记忆片段。
3.2 分段方式的几种选择
- 等长分段:
- 复杂度:O(L2/C)。
- 表现:记忆精度均匀,但在极限长度下仍然有计算压力大的问题。
- 对数分段:
- 复杂度:O(LlogL)。
- 表现:对临近的记忆表现好(细粒度分段),但对远期的记忆表现差(粗粒度大段压缩),但时间复杂度很好。
3.3 缓存初始化的选择:
- 作为检查点 (Checkpoints):下一个片段继承上个片段的最终内存状态,强调序列学习的连续性。
- 作为独立压缩器 (Independent Compressors):每个片段的内存从零独立初始化,确保缓存是该片段纯粹的压缩代表,避免跨片段的信息干扰。
3.3 线性或深度内存模块的影响
- 线性内存:如果使用简单的线性内存模块,MC 在数学上会退化或等价于“门控全局 Softmax 注意力”。另外,作者指出目前流行的 Hybrid Models 本质上其实就是片段长度为1、采用检查点机制的内存缓存的一个特例。
- 深度内存 :如果使用 MLP 等深度内存,MC 不会像线性内存那样发生退化,而是会创造出一类与现有混合模型完全没有重叠的全新架构,检索过程高度依赖于当前的查询词。
4. 实际尝试
作者把 MC 实际插入到 Titans 和几种线性注意力变体中进行尝试,各自解决了一些问题:
- 滑动窗口线性注意力(SWLA):证明 MC 可以适配简单的线性模型。由于内存是线性的(用的矩阵存储),在这里使用 GRM(门控残差内存)和 Memory Soup 在数学上会得到完全相同的公式。
- 深度线性注意力(DLA):证明 MC 能够适配使用 Hebbian rule ,并基于负内积损失进行梯度下降优化的深度内存模块。
- Titans:证明 MC 能够在复杂的非线性内部优化器中使用。Titans 模型在内部使用了带动量和权重衰减的梯度下降来优化 L2 回归损失,MC 成功为其增加了长文本的动态内存检索能力。
- Log-Linear++:原版的 Log-Linear 模型虽然实现了对数级增长的隐藏状态,但它有位置偏置问题,检索时缺乏上下文感知能力。作者结合 MC 的 GRM 对其重构(得到Log-Linear++),成功为其增加了根据当前语境动态评估历史片段权重的能力。
5. 实验
实验得到了以下结论:
- 添加 MC 机制后,所有 RNN 模型的性能都得到了稳定、显著的提升
- 传统 RNN 在长文本中容易遗忘缺陷被弥补。不过 Transformer 确实精度最高,但 MC 已经极大地缩小了 RNN 与 Transformer 的差距
- 对数级的分段方式(Log-Linear++)不如固定长度分段的 MC (这是前面用公式证明过的一个理论)
- 消融实验证明以上这几个模块的改进都是很有用的

