Status: 施工中🏗️

0. Contents

1. Titans: Learning to Memorize at Test Time

1. Titans: 测试时记忆

Abstract

摘要部分提到：

We present a new neural long-term memory module that learns to memorize historical context and helps an attention to attend to the current context while utilizing long past information, …, has the advantage of a fast parallelizable training while maintaining a fast inference.

Titans提出的长程记忆模块解决Attention在做注意力计算时捕捉不到很早期的历史信息（scale to larger than 2M context window size）。从记忆形成的视角，作者认为Attention扮演着短期记忆的功能：上下文短但能精确建模依赖关系；而Titans的记忆模块扮演着长程记忆的功能：能够记住很长时间之前的信息。

Intro

Attention曾在NeurIPS24的一篇文章：Birth of a Transformer: A Memory Viewpoint中认为起到associative memory（联想记忆）的功能。因为Attention捕获的K-V关系，在输入（Q）到来时与K对比，从而得到有效输出。

具体内容这篇文章zb还没有看，mark一下以后有时间看看

Modern Methods in Associative Memory ICML 2025 Tutorial

Birth of a Transformer: A Memory Viewpoint, NeurIPS 2024

联想记忆（AM）貌似是一个比较老的研究热点，但是又和在线学习、Hebb规则、Hopfield网络有一定联系，这篇博客中总结了很多关于AM的相关内容。

接上段，正是因为Attention的KV匹配机制，每一次得到的输出很明显是和当前的输入ctx窗内的内容有关系。它的时间和空间复杂度都是相对于ctx length是二次的，因此无法做ctx len极长的任务。

一些线性注意力模型的诞生，能够将内存使用压缩到线性，但是一个极长的ctx的信息无法被有效压缩成一个matrix-valued states。

抛开高效性，从远古的Hopfield网络，到LSTM、Transformer，模型的泛化性、长度外延、（复杂）推理能力都是真实场景中复杂任务需要的。一些从所谓“脑启发”得到的idea需要再被审视：

学习过程中，将以往信息注入到当前输入的关键模块
这些组件如何构成可独立运行的互连系统
能够主动从数据中学习并记住过去历史的抽象概念

Memory Perspective

在生物学里，记忆关乎基本的反射和刻板行为形成。大多数现有神经网络架构将记忆视为由输入引起的神经更新（Neural Update，我真不知道怎么翻译…），并将学习定义为：在给定目标的情况下，获取有效且有用记忆的过程。。我们能得到以下的表格：

模型	记忆模块	记忆的更新方式	检索记忆的方式（输出）
RNN	向量表示的记忆模块（vector-valued memory module） $\mathcal{M} \in \mathbb{R}^{n}$	输入 $x_t$ , $\mathcal{F}(x_{t}, \mathcal{M}_{t-1})$	$\mathcal{G}(x_{t}, \mathcal{M}_{t})$
Transformer	K、V矩阵表示的记忆模块 $K,V \in \mathbb{R}^{T \times d}$	输入 $x_t$ , $(K_t,V_t)=f(x_t)$ ， $\text{KV}_t=\text{KV}_{t-1}\cup(K_t,V_t)$	基于 Query 的注意力检索： $\text{Attn}(Q_t, K_{1:t}, V_{1:t})$
Linear Attention	矩阵表示的记忆模块（matrix-valued memory module） $\mathcal{M}_t \in \mathbb{R}^{d_k \times d_v}$	特征映射后做外积累积（压缩/汇总）： $\;\phi(K_t)=\phi(f_K(x_t)),\; V_t=f_V(x_t)$ ， $\;\mathcal{M}_t=\mathcal{M}_{t-1}+\phi(K_t)V_t^\top$	线性读出（内容寻址）： $Q_t=f_Q(x_t)$ ， $\;\text{out}_t=\phi(Q_t)^\top \mathcal{M}_t$

Transformer和Linear Attention的本质区别在于，前者的记忆模块.是.append()的，会不断增大的，；而者是通过运算压缩成一个固定长度的matrix。

但记忆不是一个单一的功能模块，而是由好多个模块构成的复杂系统（短期记忆、工作记忆、长期记忆…(至少已现在人类对于记忆机制的探索是这些)）。

Titans文章的贡献：

基于人类长期记忆提出了一个Test-time memory module（测试时记忆模块），能够对expectations产生violation的/感到surprised的输入记得更牢（more memorable），这个过程依赖于associative memory loss（关联记忆损失函数）来衡量surprise程度，本质是神经网络对输入的梯度。
鉴于记忆容量的有限性，施加衰减机制（decay mechanism）用于考虑现有data surprise和propotion of the memory size。其本质也和很多网络中的forgetting gate没有区别。这种机制等价于优化一个具有小批量梯度下降、动量和权重衰减的元神经网络。
为了高效快速训练这个记忆模块，提出了一种依赖增加矩阵运算的训练算法，灵感来源于原始的TTT（埋坑ing…）
提出Titans系列的模型架构，将上面的记忆模块整合进现有模型: 包含three hyper-heads）：

hyper-heads为什么要叫这个名字我也不知道…

组件	实现方式
Core（核心）/Short-term memory（短期记忆模块）	attention with limited window size（带窗口的注意力），是主要的数据流动的组件
Long-term Memory（长期记忆模块）	储存和记忆很长时间之前的东西
Persistent Memory（永久记忆）	a set of learnable but date-independent parameters （一组可学习但是输入依赖的参数）

三种Titans的变体：把记忆整合成：

(i) a context【上下文】,
(ii) a layer 【层】,
(iii) a gated branch 【门控分枝】.

实验：包括language modeling, commonsense reasoning, recall-intensive, needle in haystack, time series forecasting, and DNA modeling tasks

Preliminaries

Notations

Notation	Descr.
$x \in \mathbb{R}^{n \times d_{in}}$	模型接受的输入
$\mathcal{M}$	神经网络模块/记忆模块
$Q,K,V$	自注意力机制的查询（query）、键（key）、值（value）
$M$	自注意力的attention mask（掩码）
$S^{(i)}$	将一个序列分段的第 $i$ 段子序列
$S^{(i)}_{j}$	第 $i$ 段子序列中的第 $j$ 个token
$\mathcal{N}(x)$	输入 $x$ 进入神经网络 $\mathcal{N}$ 中做带权重修正（with weight adjustment）的前向传播
$\mathcal{N}^*(x)$	输入 $x$ 进入神经网络 $\mathcal{N}$ 中做不带权重修正（without weight adjustment）的前向传播
$\mathcal{N}^{(j)}$	神经网络 $\mathcal{N}$ 中的第 $j$ 层

回顾：自注意力与Transformer

给定输入$x \in \mathbb{R}^{n \times d_{in}}, $ 其中 $n$ 为序列长度、 $d_{in}$ 为输入特征维度. 因果自注意力的输出为 $y \in \mathbb{R}^{n \times d_{in}}$ , 自注意力的计算如下：

$Q=xW_Q, \space K=xW_k, \space V=xW_v.\quad(W_{\square} \in \mathbb{R^{d_{in} \times d_{in}}})$

$y_i = \sum_{j=1}^{i} \frac{\exp(\frac{Q_i^\top K_j}{\sqrt{d_{in}}})V_j}{\sum_{\ell=1}^{j}\exp(\frac{Q_i^\top K_\ell}{\sqrt{d_{in}}})}$

高效自注意力的策略：

1. 注意力矩阵稀疏化

2. $\text{Softmax()}$ 近似

3. Kernel-based (linear) attentions

Kernel function $\phi(\cdot)$ 需满足：

$\phi(x,y)=\phi(x)\phi(y)$

自注意力可以改写为：

$y_i = \sum_{j=1}^{i}\frac{\phi(Q_i,K_j)V_j}{\sum_{\ell=1}^{i}\phi(Q_i,K_\ell)}=\sum_{j=1}^{i}\frac{\phi(Q_i)^\top\phi(K_j)V_j}{\sum_{\ell=1}^{i}\phi(Q_i)^\top\phi(K_\ell)}=\frac{\phi(Q_i)^\top\sum_{j=1}^{i}\phi(K_j)V_j}{\phi(Q_i)^\top\sum_{\ell=1}^{i}\phi(K_\ell)}$

上式中，项 $\sum_{j=1}^{i}\phi(K_j)$ 和 $\sum_{\ell=1}^{i}\phi(K_\ell)$ 不用每次做计算，只需要每次把新的K加到原本的部分和（partial sum）即可。

当 $\phi(\cdot)$ 是identity function时，即满足 $\phi(x)=x$ , 那么上述自注意力可以写为：

$y_i = \frac{Q_i^\top\sum_{j=1}^{i}K_j^\top V_j}{Q_i^\top\sum_{\ell=1}^{i}K_\ell}$

令 $\mathcal{M}_t=\sum_{j=1}^{t}K_t^\top V_t$ , 则有：

$\mathcal{M}_t=\mathcal{M}_{t-1} + K_t^\top V_t$

$y_t = \frac{Q_t\mathcal{M}_t}{Q_i^\top \sum_{\ell=1}^{i}K_\ell^{}}$

现代循环模型以及与记忆系统的关联

如果把RNN中的hidden states当成网络的“记忆储存”，那么这个记忆系统的读和写过程就是：

$\mathcal{M}_t = f(\mathcal{M}_{t-1}, x_t), \text{(writing process)}$

$y_t = g(\mathcal{M}_{t}, x_t), \text{(reading process)}$

直观上来看，linear attention单纯将记忆更新的过程建模成一个纯加性的，势必会导致记忆溢出（memory overflow）。因此，后来的研究集中在记忆更新时 $f(\cdot)$ 和 $g(\cdot)$ 的优化上，包括但不限于：

1. 添加遗忘机制（forgetting mechanism）：GLA、LRU、Griffin、xLSTM、Mamba2
1. Delta Rule：在添加记忆（即键值对）之前，模型首先会删除其过去的值。并对这个过程加入chunkwise的高效训练。

记忆模型

将线性层视为键值（关联）记忆（Key-value associative memory）系统的想法可以追溯到快速权重程序（fast weight programs），其中动态快速程序（dynamic fast
programs）被纳入循环神经网络中，用作可写存储器（JH Schmidhuber 1992）。

FWP和Linear attention的形式真的是一模一样…

Fast weights的训练形式也有两种基本的学习规则来主导：Hebbian Rule和Delta Rule。

Hebbian Rule：“一起激发的神经元，会建立更强的连接（Cells that fire together, wire together）。”

但是，上面的所有研究忽略了两个重要的问题：tokens在网络中的传播过程 以及 大多都缺乏遗忘门

❓原文：“most of them lacks a forgetting gate, resulting in a poor memory management.”

Methods：

长期记忆模块的设计

原理：

LLM 已被证明能够记住其训练数据，受此启发，我们只需要训练一个神经网络去做这种事情不就好了。但记忆本身的泛化性、隐私问题导致的Test-time表现很差。并且，记住训练数据对于test-time任务提升并不突出。因此，我们需要一个在线元模型（online meta-model），该模型能够学习如何在测试时记忆/遗忘数据。还要要求能够记忆数据的同时，不会过度拟合。

转化为一个online learning problem（在线学习的问题）

一个令“神经网络 $\mathcal{M}$ ”更“意外”的输入更容易被记住，用网络相对于输入的梯度来量度这一意外的程度，训练一个神经网络用于压缩前面所有时刻的输入 $x_1, ..., x_t$ .因此，记忆的更新过程为：

$M_t = M_{t-1} - \theta_t \, \underbrace{\nabla_ \ell\!\left(M_{t-1}; x_t\right)}_{\text{Surprise}}$

随之而来的问题是，当一个有记忆价值的moment到来，假如很大，相当于一步学习的步子迈得很大，于是存在容易陷入局部最优（local minima）的可能。从人类记忆的角度而言，一个被认为有记忆价值的瞬间，过了一阵子可能就不那么令人意外了。也就是说，时间长了之后，这个Surprise也就不够surprise了。于是我们将上式拆成以下形式：

$\begin{aligned} M_t &= M_{t-1} + S_t, \\ S_t &= \eta_t \, \underbrace{S_{t-1}}_{\text{Past Surprise}} \;-\; \theta_t \, \underbrace{\nabla_\ell\!\left(M_{t-1}; x_t\right)}_{\text{Momentary Surprise}} . \end{aligned}$

有趣的是，这种公式类似于带动量的梯度下降法，其中 $S_t$ 是动量元素。因此，这里的动量可以作为时间（序列长度）上的意外记忆。

动量梯度下降法（Momentum Gradient Descent）的标准数学形式:

$\begin{aligned} v_t &= \mu v_{t-1} - \eta \nabla f(\theta_{t-1}) \\ \theta_t &= \theta_{t-1} + v_t \end{aligned}$

其中， $\theta_t$ 为第t步的参数， $v_t$ 为动量， $\mu$ 为动量系数, $\eta$ 为学习率。当前更新方向 = 过去更新方向的加权平均 + 当前梯度的修正。
改进版本：Nesterov 动量：梯度在 $\theta_{t-1} + \mu v_{t-1}$ 处评估，而不是仅仅在 $\theta_{t-1}$ 处。即

$v_t = \mu v_{t-1} - \eta \nabla f(\theta_{t-1}+\mu v_{t-1})$

这预见了运动，尤其是在病态问题中，可能提供更准确的梯度方向信息。

至于到底能否解决提到的过度surprise的问题：

$\prod n_j \leq 1 \rightarrow \text{old surprise} \rightarrow 0$ :“一个 moment 当下很惊讶，过一阵子就不那么 surprise 了”

与动量梯度下降不同的是，记忆更新公式中的 $\eta_t$ 是一个data-dependent function，即是一个有关输入 $x_t$ 的函数。记忆的更新率 $\theta_t$ 也是一个data-dependent的。这样设计的理由是：

虽然前一个标记的意外性可能会影响下一个标记的意外性，但只有当所有标记都相关且处于同一上下文中时，这种依赖性才基本有效。
While surprise of previous tokens might be needed to affect the surprise of the next token, it is mostly valid if all tokens are relevant and are in the same context.

记忆模块是一个元模型（meta model），它基于损失函数 $\ell(\cdot, \cdot)$ 来学习一个函数。按照Key-value关联记忆的模型来说，Key和Value矩阵就是我们要学习的关联对象，损失函数自然而然就会设计成：

$\ell(\mathcal{M}_{t-1}; x_t) = \left\| \mathcal{M}_{t-1}(\mathbf{k}_t) - \mathbf{v}_t \right\|_2^2 .$

跟RWKV-v7一模一样…

类似于元学习模型，记忆的训练发生在内循环中，因此参数 $W_K$ 和 $W_V$ 在上述损失函数中被视为超参数（hyperparameters）。在内循环中我们优化的是 $\mathcal{M}$ 的权重，而在外循环（outer-loop）中，我们则优化整个架构中的其他参数。

即便对于很大容量的记忆空间，一个记忆系统也要学会怎么遗忘。所以，进一步改进上面的记忆更新公式，加入遗忘机制：

$\begin{aligned} M_t &= \red{(1-\alpha_t)}M_{t-1} + S_t \quad (\alpha_t \in [0,1]),\\ S_t &= \eta_t S_{t-1} -\theta_t \nabla_\ell\left(M_{t-1}; x_t\right). \end{aligned}$

架构设计：

Titans为了验证基本的理论构想，使用了较为简单的大于等于1层的MLP做长期记忆模块。

当使用向量值或矩阵值记忆时，记忆模块会压缩历史数据并将其拟合成一条直线。

从元学习或在线学习的角度来看, $\mathcal{M}$ 如果是一个 $\mathbb{R}^{d_{in} \times d_{in}}$ ,记忆过程等价于优化：

$\ell(W_{t-1}; x_t) = \left\| W_{t-1}\mathbf{k}_t - \mathbf{v}_t \right\|_2^2 .$

最优解假设历史数据的潜在依赖关系是线性的。

从记忆中检索：

使用前向传播（不更新权重，即推理）来检索与查询对应的记忆。

训练优化

利用小批量梯度下降、数据相关的学习率和权重衰减来计算内循环中的权重，可以重新表述为 仅使用矩阵乘法matmul和求和(sum) 的过程。

split the sequence into chunks of size $b \geq 1$ (将序列分割成b份)

记忆更新过程展开可以重写为

$\begin{aligned} S_t&=\eta_t S_{t-1}-\theta_t \nabla_\ell\left(M_{t-1}; x_t\right) \\ &=\eta_t (\eta_{t-1} S_{t-2}-\theta_{t-1} \nabla_\ell\left(M_{t-2}; x_{t-1}\right))-\theta_t \nabla_\ell\left(M_{t-1}; x_t\right) \\ &= \eta_t (\eta_{t-1} S_{t-2}-\theta_{t-1} g_{t-1})-\theta_t g_{t} \\ &= \eta_t \eta_{t-1} S_{t-2}- \eta_t \theta_{t-1} g_{t-1}-\theta_t g_{t} \\ &= \dots \\ &= \prod_{j=1}^{t}\eta_j S_0 - \sum_{i=1}^{t} \theta_i (\prod_{\ell=i+1}^{t} \eta_{\ell}) g_i \\ M_t&=\prod_{j=1}^{t}(1-\alpha_j)M_0+\sum_{k=1}^{t}S_k \\ &=\prod_{j=1}^{t}(1-\alpha_j)M_0+\sum_{k=1}^{t}[\prod_{j=1}^{k}\eta_j S_0 - \sum_{i=1}^{k} \theta_i (\prod_{\ell=i+1}^{k} \eta_{\ell}) g_i] \\ &=\prod_{j=1}^{t}(1-\alpha_j)M_0+ S_0 \sum_{k=1}^{t} [\prod_{j=1}^{k}\eta_j] - \sum_{i=1}^{k} \theta_i [\sum_{k=1}^{t} (\prod_{\ell=i+1}^{k} \eta_{\ell})] g_i \end{aligned}$

令 $\beta_t = \prod_{j=1}^{t}(1-\alpha_j)$ , 得到 $M_t = \beta_t M_0+ S_0 \sum_{k=1}^{t} [\prod_{j=1}^{k}\eta_j] - \sum_{i=1}^{k} \theta_i [\sum_{k=1}^{t} (\prod_{\ell=i+1}^{k} \eta_{\ell})] g_i$ .

原文中其实隐含了一个自定的假设：论文里 η 和 (1−α) 不是两个独立的东西，inner-loop reformulation 里，实际上 ηₜ ≡ (1−αₜ)

原论文中表示为：

$M_t = \beta_t M_0 - \sum_{i=1}^{t} \theta_i \frac{\beta_t}{\beta_i} g_{t'}$

其中， $t' = t - \text{mod}(t,b)$ .

在一个 chunk 内（长度为b），所有样本的梯度都用 chunk 起始时刻的参数来算，即 $M_{t'}, M_{t'+1} , ... , M_{t'+b-1}$ 的梯度 $g_{t'}, g_{t'+1}, ... , g_{t'+b-1}$ 都是用同一套参数 $M_{t'}$ 求得的。所以一个 chunk 内（t = b, t′ = 0），把整个 inner loop 重写成一次线性算子： $\frac{\beta_t}{\beta_i}$ 也没有了求和符号。

⚠️：不是代数恒等变形，而是模型重参数化（reparameterization）

论文在 inner loop reformulation 时，隐含做了下面这个假设：✅ Inner loop memory is zero-initialized.所以也就不考虑了 $S_0$ .

SGD 本来是一个“时间递归系统”。chunk = 把一段时间递归，压缩成一个“矩阵 × 向量 + 求和”的静态算子。
利用chunkwise简化梯度运算：

考虑第一个chunk，即 $t = b, t'=0$ , 且 $M_t = W_t$ 是一个线性层。

$\nabla_\ell(W_{0};x_t) = (W_0x_t - x_t)x_t^\top$

$\sum_{i=1}^{t} \theta_i \frac{\beta_t}{\beta_i} g_{0} = \sum_{i=1}^{b} \theta_i \frac{\beta_b}{\beta_i}\nabla_\ell(W_{0};x_i) = \sum_{i=1}^{b} \theta_i \frac{\beta_b}{\beta_i} (W_0x_i - x_i)x_i^\top = \Theta_b \Beta_b (W_0 X - X)X^\top$

其中， $\Theta_b = \text{diag}([\theta_1, \theta_2, ..., \theta_b])$ , $\Beta_b$ 矩阵是由 $\frac{\beta_b}{\beta_i}$ 定义的。

对于所有的chunk，我们只需要存一次 $\Theta_b, \Beta_b$ 。

$S_t$ 的更新过程也可以写成chunkwise的形式：

$S_t=\eta_t S_{t-1}-\theta_t \nabla_\ell\left(M_{t'}; x_t\right)$

$u_t = \nabla_\ell\left(M_{t'}; x_t\right)$ 可以在同一时间内算出，整个S的更新是一个linear RNN过程， $u_t$ 为输入， $S_t$ 为hidden state， $\eta_t$ 是输入相关的transition value。它的并行化形式可以用parallel associative scan

Associative Scan(关联/结合扫描): 一种用在SSM的并行化计算策略。详细介绍参考这篇Blog. Associativity simply mean that we can apply the operator (denoted by $\cdot$ ) in arbitrary order, i.e. that $(a \cdot b) \cdot c = a \cdot (b \cdot c) $ . Examples of associative operators are addition and multiplication. 即满足结合律的运算。一个典型的cumsum()算子就是associative scan的例子。其实简单推一下就知道了： $f(x) = a x + b, y = f_1(w) = f_1(f_2(x))$ , $f_1 \cdot f_2 = f_1(f_2(\cdot)) = a(cx+d) + b = (ac)x + (ad+b)$ 形式一样，这就是Paral. Assoc. Scan算法。
进一步：Parameter as chunk function

把 $\alpha_t, \eta_t, \theta_t$ 都变成chunkwise的参数，不是每一个输入都有一套独立的，虽然会损失一些表达能力，但是对于高效训练来说用处很大。这样，整个过程是一个LTI系统，可以用一个global convolution来做一次运算。（也可以对这个卷积做加速：FFT+IFFT）

但titans依然保留了token-wise的形式，没有采纳chunkwise的。