跳到正文
Hacker News· speckx·· 1 天前AI 评分28

从零理解 Transformer:线性层、MLP 与自注意力如何充当记忆库

You could have invented transformers from memory

AI 导读

一篇技术文章把线性层、MLP 和自注意力统一看作联想记忆库,用干扰信号比衡量召回质量。以 64 维 key 为例,三者可存储的记忆量从约 64 条跃升到约 2,700 条,再到超过一百万条。文章推导出线性层容量约为 key 维度 d_k,并说明扩大 d_k 受硬件限制。

正文 · AI 翻译
从记忆中检索一个字符。随着记忆数量的增长,回忆效果会变差。

如今的大型语言模型能够回忆起世界各个角落的冷门事实。 最近,我参加朋友的桌游之夜要迟到了,于是我让 Astra 提醒我《沙丘:帝国》的规则。 在我开车前往的路上,它愉快地为我讲述了规则,帮我记得足够清楚,让我能够再次上手游玩。 我们太习惯于向模型询问任何事情,以至于把它们广博的知识视为理所当然。

模型是如何记住这么多信息的?而它们又能存储多少呢?

在这篇文章中,我将介绍一种直观的方式,把线性映射、MLP 和自注意力看作不同类型的记忆存储。从这个直觉出发,我们可以推导出一些关于每一层能存储多少记忆的启发式方法。


衡量一个“好的”记忆存储的关键特性是它回忆所存内容的能力。毕竟,即使我们能够存储无限多的事实,如果需要在时无法回忆起来,那也没什么用。 我们存储的记忆越多,当我们试图回忆其中一条时,它们之间的干扰就越大。因此,在整篇文章中,我们将聚焦于一个量——干扰信号比,它衡量在回忆过程中其他记忆造成了多大的妨碍。一个存储在其干扰增长到与信号一样大之前所能容纳的记忆数量,就是它的容量。

我们将展开这一视角,展示如何从一个线性层走向更好的东西——MLP,它同时提升了容量和回忆能力,然后一路走到自注意力,它具有非常好的回忆能力,并且能够存储数量可能无界的记忆。以 64 维的键来说,这是从大约 64 条记忆,到大约 2,700 条,再到超过一百万条的跃升。

以下是我们将在全文中使用的术语。一个记忆存储是一个函数 $f$,它持有 $M$ 条记忆,其中 $i = 1, \ldots, M$。一条记忆是一个键值对:一个键 $k_i \in \reals^{d_k}$,用于基于相似度的查找,以及一个值 $v_i \in \reals^{d_v}$,即我们想要检索的内容。要使用 $f$,我们用一个输入 $x$ 对它进行查询,它会回忆出一个输出 $f(x)$。一个好的记忆存储在我们用 $k_i$ 或另一个与 $k_i$ 相似的向量查询它时,会回忆出 $v_i$。在文献中,通过相似度进行回忆的记忆存储被称为联想记忆。

外积记忆存储#

最基本的记忆存储是线性的:它简单地通过外积累积记忆,

$$W = \sum_{i=1}^{M} v_i k_i\transpose = V\transpose K$$

其中 $K \in \reals^{M\times d_k}$ 和 $V \in \reals^{M\times d_v}$ 的行分别是键 $k_i\transpose$ 和值 $v_i\transpose$。1

要存储一条新记忆 $(k, v)$,我们只需将 $vk\transpose$ 加到现有权重上。2

附注:一步梯度下降会写入一条记忆

考虑网络内部的一个线性层 $y = Wx$,其标量损失为 $L$,并令 $\bar y = \grad L(y)$ 为损失对该层输出的梯度。关于 $W$ 的梯度是 $\bar y x\transpose$,因此学习率为 $\eta$ 的一步是

$$W \leftarrow W - \eta\, \bar y x\transpose.$$

这个递推式是一次记忆更新:它用一个新记忆来更新权重,该记忆的键为 $x$,值为 $-\eta\bar y$。 关于这一观点的更多内容,参见 Irie、Csordás & Schmidhuber,他们将通过梯度下降训练的线性层表示为对其所有训练输入的注意力。

要用输入 $x$ 查询记忆存储,我们将输入乘以矩阵:

$$f(x) = V\transpose K x = \sum_{i=1}^{M} v_i\,(k_i\transpose x).$$

因此回忆就是通过与存储键的相似度进行查找:我们将查询与每个键进行比较,并返回按这些相似度加权的值之和。

回忆会引入干扰#

所以存储记忆很容易:只需不断将外积加到权重上。但当我们尝试回忆其中一个时会发生什么?

为简单起见,假设每个键都被归一化为长度为 1,即 $k_j\transpose k_j = 1$。那么

$$f(k_j) = \underbrace{v_j}_{\text{信号}} + \underbrace{\sum_{i\neq j} v_i\,(k_i\transpose k_j)}_{\text{干扰}},$$

这意味着我们取回了目标记忆的值,即我们用其键进行查询的那个记忆,再加上来自其他每个键与查询相似的记忆的干扰。

如果所有键都是正交归一的,干扰将为零,回忆将是精确的。但 $\reals^{d_k}$ 只有 $d_k$ 个正交方向,因此一旦我们存储的记忆数超过键维度:$M > d_k$,就无法获得完美的回忆。

外积存储的容量是多少?这里有一个直观的论证。有 $M-1$ 个干扰项。如果键是随机单位向量,每个相似度 $k_i\transpose k_j$ 约为 $\pm 1/\sqrt{d_k}$。如果值也是随机的,且长度相似,那么 $M-1$ 个干扰项会像随机游走一样累加,因此干扰的平方长度约为信号的 $(M-1)/d_k$ 倍,并且

$$\frac{\norm{\text{干扰}}}{\norm{\text{信号}}} \sim \sqrt{\frac{M}{d_k}}.$$

将此比值设为 1 可得到容量

$$M \approx d_k,$$

一旦我们存储的记忆数与键维度一样多,干扰就会与信号一样大,将其淹没。

下面的每个条形代表一次回忆中一个存储记忆的权重,相对于目标记忆的权重(红色)。每个条形都会将其值的一部分添加到回忆中。将 $M$ 拖过 $d_k = 64$,观察小条形不断累积,直到回忆失败。其他按钮切换到我们将在本文后面构建的记忆存储。

目标记忆其他每个存储记忆,相对于目标记忆的权重

干扰 / 信号 = 回忆失败

如何减少干扰?#

即使在其约 $d_k$ 个记忆的容量以下,外积存储在回忆过程中也会产生显著的干扰。3 为了提高容量,或在容量以下更干净地回忆,我们需要减少干扰。

暴力方法是增大键维度 $d_k$,因为线性存储的容量随 $d_k$ 增长。$d_k$ 能有多大受硬件约束的限制。

或者,我们可以进行算法上的改进。我们可以更谨慎地存储记忆:外积存储会逐个添加每个记忆而不考虑其他记忆,因此它无法抵消它们之间的干扰。或者我们可以更有选择性地回忆:线性存储中的每个记忆都会造成干扰,即使其键与查询的相似度很低,因此我们可以通过抑制相似度较低的键来改善回忆。

在本文的剩余部分,我们将看到,更谨慎地存储外积记忆能给我们带来最优线性记忆存储,其容量大约翻倍。通过非线性进行选择性回忆则能给我们带来 MLP 和自注意力。

最优线性记忆存储#

外积存储写入每个记忆时都不考虑其他记忆。相反,我们可以选择权重 $W$,使回忆能同时拟合所有已存储的记忆。这是一个线性回归问题:4

$$\operatorname*{minimize}_W \sum_{i=1}^M \norm{W k_i - v_i}_2^2.$$

当我们存储的记忆数量超过键的维度,即 $M > d_k$ 时,最小二乘解为

$$W = \parens*{\sum_i v_i k_i\transpose}\parens*{\sum_i k_i k_i\transpose}\inverse = V\transpose K\,(K\transpose K)\inverse.$$

这就是外积存储 $V\transpose K$,只做了一处改动:在将查询与键进行比较之前,我们将其乘以 $(K\transpose K)\inverse$,这抵消了键的协方差 $K\transpose K$。这白化了键,因此存储不再过度加权许多键共有的方向。

当 $M \leq d_k$ 时,只要键线性无关,回归就能精确拟合每个记忆。因此,即使键并不正交,最优存储的回忆也完全没有干扰。

当 $M > d_k$ 时,我们开始出现干扰,因为记忆存储只有 $d_k$ 个方向可用。对于随机的键和值,我们的干扰信号比变为

$$\frac{\norm{\text{interference}}}{\norm{\text{signal}}} \approx \sqrt{\frac{M}{d_k} - 1}.$$

该比值在以下条件下等于 1

$$M = 2d_k,$$

因此最优存储的容量是之前的两倍!

附注:推导最优线性记忆存储

当 $M \leq d_k$ 时。如果键线性无关,键之间内积构成的 $M\times M$ 矩阵 $KK\transpose$ 可逆,且

$$W = V\transpose (KK\transpose)\inverse K$$

能精确拟合每个记忆。当 $M < d_k$ 时,问题约束不足,有无穷多个解,而这个是极小范数解。$(KK\transpose)\inverse K$ 的各行 $\tilde k_i\transpose$ 被称为 $K$ 的各行的对偶基:根据构造,$(KK\transpose)\inverse KK\transpose = I$,因此 $\tilde k_i\transpose k_j = \delta_{ij}$。于是 $f(k_j) = \sum_i v_i\, \tilde k_i\transpose k_j = v_j$,因此没有干扰,我们实现了完美回忆。

当 $M > d_k$ 时。将目标函数的梯度设为零可得 $W K\transpose K = V\transpose K$,因此 $W = V\transpose K\,(K\transpose K)\inverse$。 两个解都是 $W = V\transpose (K\transpose)\pinv$,其中 $A\pinv$ 表示 $A$ 的伪逆:当键线性无关时 $(K\transpose)\pinv = (KK\transpose)\inverse K$,当它们张成 $\reals^{d_k}$ 时 $(K\transpose)\pinv = K(K\transpose K)\inverse$。

现在我们可以分析 $f(k_j) = V\transpose K\,(K\transpose K)\inverse k_j$ 的干扰。 对于存储的键 $k_j$,每个值 $v_i$ 上的权重是 $k_i\transpose (K\transpose K)\inverse k_j$,这恰好是 $M\times M$ 矩阵 $P = K(K\transpose K)\inverse K\transpose$ 的第 $(i, j)$ 项:

$$f(k_j) = \sum_{i=1}^M v_i\, P_{ij} = \underbrace{P_{jj}\, v_j}_{\text{信号}} + \underbrace{\sum_{i\neq j} P_{ij}\, v_i}_{\text{干扰}}.$$

$P$ 是到 $K$ 的列空间($\reals^M$ 的 $d_k$ 维子空间)上的投影。关于投影的两个事实有助于我们分析干扰信号比:

  • 信号项: 由迹的循环性质, $$\sum_{j=1}^M P_{jj} = \trace{K(K\transpose K)\inverse K\transpose} = \trace{(K\transpose K)\inverse K\transpose K} = \trace{I_{d_k}} = d_k,$$ 所以 $M$ 个信号权重之和为 $d_k$。如果键是随机的,没有哪个键是特殊的,因此 $P_{jj} \approx d_k/M$。
  • 干扰项: $P$ 是对称投影矩阵,所以 $P\transpose P = P^2 = P$。于是 $$\sum_{i\neq j} P_{ij}^2 = \sum_{i=1}^M P_{ij}^2 - P_{jj}^2 = (P\transpose P)_{jj} - P_{jj}^2 = P_{jj} - P_{jj}^2 \approx \frac{d_k}{M} - \parens*{\frac{d_k}{M}}^2.$$

如果值是随机的且长度相近,

$$\frac{\norm{\text{干扰}}}{\norm{\text{信号}}} \approx \frac{\sqrt{P_{jj} - P_{jj}^2}}{P_{jj}} = \sqrt{\frac{M}{d_k} - 1}.$$

与外积存储的 $\sqrt{M/d_k}$ 相比,$-1$ 就是最优存储所抵消掉的相当于 $d_k$ 个记忆的干扰。

我们可以用数值方法验证我们的预测。我们在 $d_k = 64$ 维中随机采样 $M$ 个具有单位长度键的记忆,用其中一个存储的键查询记忆存储,并测量干扰信号比:

外积存储在 $M \approx d_k = 64$ 时达到其容量。最优存储恰好能回忆到 $M = d_k$。超过该值后,其干扰从零开始像 $\sqrt{M/d_k - 1}$ 一样上升,并且存储在大约 $M = 2d_k = 128$ 时达到其容量,正如预测的那样。由于它抵消了固定数量的记忆的干扰,随着 $M$ 增长,两个存储之间的差距缩小。

不过,这个最优记忆存储对相似度仍然是线性的。它必须以相同的因子缩放小相似度和大相似度。为了做得更好,我们需要改变存储回忆的方式。

MLP 作为非线性记忆存储#

线性记忆存储回忆

$$f(x) = \sum_{i=1}^M v_i (k_i\transpose x).$$

如果我们有办法在 $k_i\transpose x$ 很小时抑制权重,那么我们就可以减少来自低相似度键的干扰量。

让我们在计算值的加权和之前,将每个相似度通过一个非线性函数 $\sigma$:

$$f(x) = \sum_{i=1}^{M} v_i\,\sigma(k_i\transpose x) = V\transpose \sigma(Kx).$$

我们刚刚从第一性原理重新发现了一个两层 MLP!第一个权重矩阵的行是键,第二个权重矩阵的列是值。Sukhbaatar 等人指出了前馈层与记忆存储之间的这种相似性,而 Geva 等人发现训练后的 transformer 中的前馈层通常以这种方式工作,键在输入中可识别的模式上激活。

理想情况下,门控函数 $\sigma$ 的选择应使得当相似度较低时($k_i\transpose x \sim 0$ 甚至 $k_i\transpose x < 0$),其在和中的对应项被抑制以减少干扰。ReLU 激活将具有负相似度的键置零,而添加偏置可以移动该阈值5。平方 ReLU 激活进一步抑制相似度接近零的键。相似度 0.1 变为 0.01,而相似度 1 保持为 1。事实上,平方 ReLU 是对 MLP 激活函数鲜为人知但已在规模上得到验证的改进之一(So et al.)。

平方 ReLU MLP 是一种 GLU MLP,其门控和输入投影共享权重,因为 $\operatorname{ReLU}(Kx)^2 = \operatorname{ReLU}(Kx) \odot Kx$。当今语言模型中的 SwiGLU MLP 为门控赋予自己的权重,并使用 Swish 激活代替 ReLU。

在实践中,MLP 以叠加方式存储事实:许多事实共享相同的权重,输出是许多记忆的混合。Garcia et al. 明确构建了这一点:GLU MLP 的隐藏单元计算输入的随机特征 $\phi(x)$,第二个权重矩阵是对这些特征的外积存储 $\sum_i v_i \phi(k_i)\transpose$,并且 MLP 存储的事实远多于其隐藏单元的数量。尽管如此,训练后的 transformer 的 ReLU MLP 对于给定输入可以仅激活 1% 的隐藏单元,因此每次回忆只利用一小部分记忆(Li et al.)。

Transformer 的作者们早已知道这一点

原始 transformer 论文也建立了这种联系,但仅在 arXiv 的 v2 和 v3 版本中。在那里,Vaswani 等人将前馈层描述为“一种注意力”,其两个权重矩阵的行作为键和值,并用 ReLU 代替 softmax。他们甚至尝试用对学习到的键值参数的注意力替换前馈层,并获得了大致相同的翻译质量:一个版本稍差,另一个稍好,尽管第一个训练速度慢了约 30%。这段话在论文的所有后续版本中都消失了。我们将在下面讨论自注意力与这一切的关系。

用更高次幂减少干扰#

对相似度取平方之所以有效,是因为它使低相似度权重比高相似度权重缩小得更多。我们可以用 $\sigma(s) = \operatorname{ReLU}(s)^p$ 进一步推进这一点。信号保持不变,因为 $1^p = 1$,而相似度为 $1/\sqrt{d_k}$ 的干扰项缩小为 $d_k^{-p/2}$。$M$ 个干扰项像随机游走一样累加,因此干扰与信号的比率为

$$\frac{\norm{\text{interference}}}{\norm{\text{signal}}} \approx \sqrt{\frac{M}{d_k^p}}.$$

这意味着容量大约为

$$M \approx d_k^p$$

个记忆,几何级地多于之前!

为了验证这一点,我们对五种 $\sigma$ 选择运行与之前相同的实验:

由于干扰以 $\sqrt{M}$ 增长,在双对数图上这些线几乎平行。只有它们的高度变化,随之容量也变化,即每条线与 1 的交点:

  • 线性:容量为 $M \approx d_k = 64$。
  • ReLU:将负相似度置零使干扰缩小 $\sqrt 2$,因此容量增长到 $M \approx 2d_k = 128$。
  • 平方 ReLU:每个干扰项从约 $1/\sqrt{d_k}$ 缩小到约 $1/d_k$,因此容量增长到 $M \approx 2d_k^2/3 \approx 2700$,约为线性存储的 40 倍。
  • SwiGLU MLP,其门控与输入投影共享权重,如同平方 ReLU 中那样,$\sigma(s) = \operatorname{Swish}(s)\cdot s = s^2/(1 + e^{-s})$:其容量几乎与平方 ReLU 相当。
  • 立方 ReLU:每个干扰项缩小到约 $d_k^{-3/2}$,因此容量增长到 $M \approx 2d_k^3/15 \approx 35000$。

自注意力作为非参数记忆存储#

将权重提升到幂次可将容量增加到约 $d_k^p$。但如果我们能将其推到极限呢?与其选择一个幂次,我们可以使用无限多个幂次。事实上,指数函数是每个幂次的加权和,

$$e^{\beta s} = 1 + \beta s + \frac{(\beta s)^2}{2} + \frac{(\beta s)^3}{6} + \cdots = \sum_{p=0}^{\infty} \frac{(\beta s)^p}{p!},$$

其中锐度 $\beta$ 设定了更高幂次的权重。以下是使用 $\beta = \sqrt{d_k}$(自注意力使用的值)6 对指数函数及其级数截断至 2、3 和 4 项的相同实验:

使用 2 项时,容量约为 40,低于线性存储,因为常数项使每个键的权重约为 1。使用 3 项时约为 450,使用 4 项时约为 3,100,接近平方 ReLU。完整的指数函数远低于虚线,因此其容量超出了图表的边缘。

为了找到指数函数的容量,我们用一个存储的键 $k_j$ 进行查询,因此目标记忆的相似度为 1,权重为 $w_j = e^{\beta}$。其他每个记忆的权重为 $w_i = e^{\beta s_i}$,其中对于随机单位键,相似度 $s_i$ 大致服从均值为 0、方差为 $1/d_k$ 的高斯分布。对于长度相似的随机值,

$$ \begin{aligned} \frac{\norm{\text{interference}}}{\norm{\text{signal}}} &\approx \sqrt{\sum_{i\neq j} \parens*{\frac{w_i}{w_j}}^2} && \text{干扰项像随机游走一样累加} \\ &= \sqrt{\sum_{i\neq j} e^{2\beta (s_i - 1)}} && w_i / w_j = e^{\beta s_i} / e^{\beta} \\ &\approx \sqrt{M\, e^{-2\beta}\, \ev{e^{2\beta s}}} && \text{对 } M \text{ 个其他记忆取平均} \\ &= \sqrt{M\, e^{-2\beta}\, e^{2\beta^2/d_k}} && \ev{e^{ts}} = e^{t^2/(2d_k)} \text{ 对于 } s \sim \mathcal{N}(0, 1/d_k) \\ &= \sqrt{M}\, e^{\beta^2/d_k - \beta}. \end{aligned} $$

将比值设为 1 得到容量为

$$M \approx e^{2\beta(1 - \beta/d_k)},$$

只要 $\beta$ 远低于 $d_k$,容量就随 $\beta$ 指数增长,而不是像之前那样随 $d_k$ 多项式增长!当 $\beta = \sqrt{d_k} = 8$ 时,容量约为 $e^{14} \approx 1.2$ 百万个记忆,几乎是线性存储的 20,000 倍。同一实验的更长运行也将容量定在约 1.2 百万。

在实践中,我们可能不想在实现中天真地对数字取指数,因为小数可能下溢,大数可能上溢。幸运的是,如上面的推导所示,干扰仅取决于权重之间的比值,因此我们可以将每个权重除以它们的和,而不改变干扰与信号的比值,同时确保数值稳定性:

$$f(x) = \sum_{i=1}^{M} v_i\,\frac{\exp(\beta\,k_i\transpose x)}{\sum_j \exp(\beta\,k_j\transpose x)} = V\transpose \operatorname{softmax}(\beta K x).$$

这就是自注意力。7

与 MLP 不同——其键和值是训练得到的权重——自注意力从其上下文中的 token 计算键和值。每个 token 存储一个记忆,因此记忆数量 $M$ 就是上下文长度。这正是自注意力非参数化的原因:它的记忆不是参数,且存储随上下文增长。

再次给出柱状图,从 softmax 开始。在 $\beta = 8$ 时,在滑块范围内召回从不失败,因此将 $\beta$ 降至 2 并提高 $M$:在几十个记忆后召回失败。$M$ 每翻倍就需要稍大一点的 $\beta$ 来保持召回工作,这就是下一节背后的想法。

目标记忆其他每个已存储记忆,相对于目标记忆的权重

干扰 / 信号 = 召回失败

在长上下文中处理大量键值对#

由于 $\beta$ 决定了我们抑制不太相似键的力度,我们可以随着键值对数量的增加而增大 $\beta$。为什么不直接把 $\beta$ 设得很大?当 $\beta \to \infty$ 时,softmax 变成只选择最相似键的硬 argmax,这不可微,而且非常大的 $\beta$ 已经会使训练效果很差,因为 softmax 饱和,到其他所有键的梯度消失。因此我们只按 $M$ 所需的程度提高 $\beta$。用 $\log M$ 缩放注意力 logits 是长上下文的一种技巧,可以做到这一点。它有几个版本,每个版本都相对于固定参考长度 $L$(通常是训练长度)设置 $\beta$,使得当 $M = L$ 时 $\beta$ 保持其训练值:

  • Su (2021) 从在 $M$ 增长时保持注意力熵恒定推导出 $\log M$ 的缩放,并归一化以在 $M = L$ 时匹配标准注意力。在每个长度下,logits 都乘以 $\log M / \log L$。Scalable-Softmax 具有相同形式,将 logits 乘以 $s \log M$,其中 $s$ 可学习。
  • 原始 Qwen 模型使用 LogN-Scaling,仅在上下文长于 $L$ 时应用相同的因子。
  • Qwen2.5-1M 使用 YaRN 的注意力缩放,将 logits 乘以 $(0.1 \log(M/L) + 1)^2$,因此它像 $(\log M)^2$ 一样增长。

以下是在我们的实验中,每条规则的表现,其中在 $L = 1024$ 个记忆时 $\beta = \sqrt{d_k} = 8$:

Qwen 的规则和 Su 的规则在超过参考长度 $L$ 后一致,此时干扰与信号之比随存储增长而下降,在 200 万个记忆时降至约 0.004。YaRN 增长更快的缩放将该比率降至约 $10^{-4}$。 在每种情况下,用 $\log M$ 缩放 $\beta$ 都允许存储本质上无限数量的具有不同键的记忆。

结语#

以下是我们访问过的记忆存储,以及当我们用已存储的键进行查询且键和值随机时的容量:

记忆存储$f(x)$容量$d_k = 64$ 时的容量
外积$V\transpose K x$$d_k$64
最优线性$V\transpose K (K\transpose K)\inverse x$$2d_k$128
ReLU MLP$V\transpose \operatorname{ReLU}(Kx)$$2d_k$128
平方 ReLU(GLU MLP)$V\transpose \operatorname{ReLU}(Kx)^2$$2d_k^2/3$约 2,700
SwiGLU MLP$V\transpose \parens*{\operatorname{Swish}(Kx) \odot Kx}$约 $2d_k^2/3$约 2,900
立方 ReLU$V\transpose \operatorname{ReLU}(Kx)^3$$2d_k^3/15$约 35,000
自注意力,$\beta = \sqrt{d_k}$$V\transpose \operatorname{softmax}(\beta K x)$$e^{2\sqrt{d_k} - 2}$约 120 万
自注意力,$\beta \propto \log M$$V\transpose \operatorname{softmax}(\beta K x)$在我们的实验中未达到超过 200 万

我们从一种启发式的记忆存储——外积记忆——出发,逐步构建出如今在现代 transformer 中见到的所有标准层。 这一视角还让我们能够估计它们在容量上的差异,即作为向量值键值存储的能力。

我应该提到,这种容量估计只是审视现代神经网络架构的一种方式。我们的容量估计只计算单层在干扰主导信号之前能回忆多少随机向量。事实上,我们完全没有考虑训练数据,而训练数据会影响键和值中实际表示的内容!还有其他工作考虑了训练后的层,以比特为单位衡量存储的信息。Morris 等人估计 GPT 风格模型每个参数大约记忆 3.6 比特,并发现一旦模型填满其容量,它就开始泛化,对每个训练样本的记忆减少。Allen-Zhu & Li发现语言模型每个参数存储约 2 比特的事实知识,即使量化到 int8 也是如此。

从这里出发,还有许多同样引人入胜的相关方向:

我一直被记忆所吸引,也许是因为我切身感受到自己的记忆有多么有限。也许我内心的一部分希望,通过理解机器如何记忆,我也能对自己有所了解。


  1. 有时这种键和值的累积被称为“Hebbian学习”,因为突触(权重)的强度会因同时激活(相关)的输入(键)和输出(值)而增强。 ↩︎

  2. 这种外积记忆存储用于线性注意力的状态,也称为快速权重:每个 token 通过将 $v k\transpose$ 加到矩阵上来存储记忆 $(k, v)$,后续 token 会查询该矩阵(Katharopoulos et al.;Schlag, Irie & Schmidhuber)。 ↩︎

  3. 这是为了回忆一个干净的值向量。如果每个值只需要指向固定集合中的正确答案,比如选择最可能的 token,那么回忆只需要比错误答案更接近正确答案,而一个 $d\times d$ 矩阵可以容纳大约 $d^2$ 个事实(Nichani, Lee & Bietti)。 ↩︎

  4. 我们也可以使用这种回归目标来设计序列模型在前向传播过程中写入的短期记忆,比如线性注意力和自注意力。参见我在测试时回归方面的工作。 ↩︎

  5. 对于随机键,ReLU 将干扰减少了 $\sqrt 2$ 倍。随机键的相似度为正或负的概率相等,因此 ReLU 将大约一半的权重置零,另一半保持不变。这使干扰的方差(权重平方和)减半,因此干扰本身缩小了 $\sqrt{2}$ 倍。完整推导见Zhong et al.。 ↩︎

  6. Transformer 将相似度除以 $\sqrt{d_k}$ 而不是乘以。但它们的查询和键长度约为 $\sqrt{d_k}$,因此匹配对的 $q\transpose k \approx d_k$。我们的键和查询具有单位长度,因此匹配对的相似度为 1。因此我们按 $\beta = \sqrt{d_k}$ 进行缩放。 ↩︎

  7. 在 Transformer 的表示法中,相似度是 $k_i\transpose q$,其中 $q = W_Q x$,按 $1/\sqrt{d_k}$ 缩放。我将 $W_Q$ 合并到键中。Transformer 中的键和查询长度约为 $\sqrt{d_k}$,因此对于像我们这样的单位长度键,该缩放与 $\beta = \sqrt{d_k}$ 相匹配。 ↩︎

来源:Hacker News · alexwang.ai