Attention机制竟有bug？Softmax是罪魁祸首，影响所有Transformer

公众号新闻

2023-07-28 18:07

©作者 | 机器之心编辑部

来源 | 机器之心

「大模型开发者，你们错了。」

「我发现注意力公式里有个 bug，八年了都没有人发现。所有 Transformer 模型包括 GPT、LLaMA 都受到了影响。」

昨天，一位名叫 Evan Miller 的统计工程师的话在 AI 领域掀起了轩然大波。

我们知道，机器学习中注意力公式是这样的：

自 2017 年 Transformer 问世，这个公式已被广泛使用，但现在，Evan Miller 发现这个公式是错的，有 bug！

Evan Miller 的这篇博客解释了当前流行的 AI 模型如何在关键位置出现错误，并使得所有 Transformer 模型都难以压缩和部署。

总结而言，Evan Miller 引入了一种新函数 Quiet Attention，也叫 Softmax_1，这是对传统 softmax 函数的创新调整。

有网友对该博客总结出了一个「太长不看版」。博客作者建议在注意力机制使用的 softmax 公式分母上加 1（不是最终输出 softmax）。注意力单元中的 softmax 使其可以将键 / 查询匹配作为概率；这些概率支持一个键 - 值查找的连续值版本（我们得到的权重不是一个查找的 1/0 输出，而是高权重 = 所需的键 - 值查找）。

分母上加 1 将改变注意力单元，不再使用真实的权重概率向量，而是使用加起来小于 1 的权重。其动机是该网络可以学习提供高权重，这样调整后的 softmax 非常接近概率向量。同时有一个新的选项来提供 all-low 权重（它们提供 all-low 输出权重），这意味着它可以选择不对任何事情具有高置信度。

有人甚至猜测「这就是微软 RetNet 比 transformer 性能更优的原因？」

还有网友表示，这项研究可以促进 LLM 的改进，从而极大对权重进行压缩，使得较小的模型媲美较大的模型：

Miller 表示：你可以像使用传统的 softmax 函数一样使用 Softmax_1 函数，示例如下。

import torchfrom softmax_one.softmax_one import softmax_onex = torch.randn(5)y = softmax_one(x, dim=0)

基于这样的修改，Miller 还做了实验，结果如下：

接下来我们看看 Miller 到底发现了什么错误。

异常值

Evan Miller 是在阅读关于量化的论文时发现了这个 bug。当前，内存和存储已经成为限制人工智能发展的重要因素。人们一直在努力压缩模型，并尝试在云端、在边缘设备上运行大型语言模型（LLM）。

在计算机中，信息是用二进制数据流来存储的。如果数据流是高度可预测的，例如总是包含在有限的范围内，那么我们就可以用相对较少的位（bit）来存储它们。反之，如果一串数字是不可预测的，可能是千载难逢的巨大数字，我们就需要更多的二进制数字来编码和存储。而 Transformer 模型包含一些异常值权重。

在高通 AI Research 6 月发表的一篇论文《Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing》中，研究团队将这些异常值的存在追溯到注意力机制的 softmax 函数。

论文链接：

https://arxiv.org/abs/2306.12929

这听起来令人意外，但 Evan Miller 认为这是对的，并进一步发现 softmax 函数存在一个错误。

我们来看下 Evan Miller 是如何说明 softmax 函数在注意力机制方面并不是一个合适的工具的。

Softmax引出的问题

为什么说 softmax 不适合注意力机制，这还得从注意力机制可以做什么开始。

一般来讲，数值错误一般是由程序错误引起的，然而，当程序没有错误，这时就需要从修复复杂的数学公式入手，耗费大量时间。

Evan Miller 大概阅读了 50 篇 arXiV 论文才有点眉目。Miller 从「输入嵌入」入手，我们可以将「输入嵌入」理解为一个浮点向量，代表输入字符串中的一个单词。

举例来说，Meta 最近推出的 LLaMA 2 模型使用了一个长度为 3204 的嵌入向量，以半精度浮点数表示，这仅仅是为了表示词汇表中的一个单词，而词汇表通常包含 30000 到 50000 个条目（entry）。意味着一个单词的嵌入向量占用 6KB + 的存储空间。随着技术的发展，「输入嵌入」的长度逐渐增加，所占存储空间也随之增加。

如果你是一个对存储占用非常敏感的 C 程序员，你可能接受不了这一数字，明明是 2 字节就能存储的东西，为什么偏偏要用 6KB？如果按照 2 字节来计算，假如词汇量少于 2^16=65384，那么我们只需要 16 位来表示一个条目。

但是，实际上 Transformer 的工作原理是这样的：它将输入向量转换为大小相同的输出向量，最终的 6KB 输出向量用来预测下一个 token。运行中，Transformer 每一层的工作流都将信息添加到原始的单词向量中。在这其中，还用到了残差连接：所有的注意力机制都在为原始的两个字节的信息添加补充材料，从而是的 LLM 能够分析更长的上下文。

Transformer 的最后一步是将这个输出向量与一个矩形矩阵相乘，并将得到的词汇长度向量压缩到一个 softmax 函数中，将这些指数化的输出视为下一个 token 的概率。这是合理的，但众所周知，这并不完全正确，因为我们不能确定这些输出概率是正确的。相反，每个 Transformer 实现和其衍生版本都使用采样机制来隐藏 softmax 过度表示概率较低的事实。

接下里，Miller 介绍了 softmax 的发展史。softmax 最初出现在统计学中，最早作为一种基于能级预测状态分布的方法，其形式如下：

之后经济学家又将其修改为

这一修改，softmax 才拥有了多项逻辑函数。由于 Miller 对 softmax 函数的研究颇深，因而，他能识别出 softmax 使用不恰当的地方。

Softmax 应用广泛，在物理学中，它非常有效；在经济学中，它可能不那么准确；但将其应用到机器学习领域时，只要涉及离散选择，它似乎总是有效的：

Miller 进一步表示，softmax 的关键在于，如果你不想保留一些项，必须对 softmax 进行修改，否则结果就会产生扭曲。

举例来说，在 LLM 上下文中，扭曲产生的原因是对非语义 token（逗号等）进行大量加权导致的，这些较高的权重成为难以压缩的异常值，使得研究变得更加困难。来自高通的 AI 研究员也发现了这一现象，在 LLM 中，97% 以上的异常激活发生在空格和标点符号位置上。

接下来，Miller 介绍了 softmax 是如何在注意力中使用的，从而发现问题到底出现在哪里：

对上述公式进行分解，在仅解码器模型中，𝑄、𝐾和𝑉源自相同的输入序列。它们又不完全相同，即投影方式不同。但在每一层中，它们都以相同的注释嵌入向量开始。

𝑄𝐾^𝑇项用于寻找不同位置 token 向量之间的相关性，实质上构建了一个相关性矩阵（点积按缩放），其中每一列和每一行对应一个 token 位置。然后，对这个方阵的每一行进行 softmax 操作，得到的概率用作𝑉矩阵中值向量的混合函数。概率混合后的𝑉与输入向量相加，将求和结果传递给神经网络进行进一步处理。

多头注意力每层并行执行多次上述过程。从本质上讲，这种方法划分了嵌入向量，每个头使用整个向量中的信息来注释输出向量的一个（非重叠）片段。这就是原始 Transformer 论文中的串联操作。

使用 softmax 的问题在于，它强制每个注意力头进行注释，即使没有信息可添加到输出向量中。

Softmax_1和QuietAttention

来了，在这里你将看到 Softmax Super-Mod 点燃了 LLM 黑客频道。

有点失望，对吧？Miller 所做的只是在分母上加 1。如果想要的话，这可以让该向量作为一个趋于 0 的整体。否则只会将值缩小一点，并且缩小的值会在归一化过程中得到补偿，这在注意力之后发生。

当 𝑥 中的条目显著小于零并且模型试图完全避免注释时，主要的区别在于负值限制。将如下原始 softmax 的限制行为

与新的、改进后的 softmax_1 相比较。

Vanilla softmax 将始终释出相同的总权重；softmax_1 看起来大部分相同，但在负象限中有一个「逃出口」（escape hatch）。需要明确的是，这里的核心问题在本质上是数学而非数值问题。额外的精度并不能拯救 softmax，所有的 Transformers 都会受到影响。

你还可以观察到关于 softmax_1 的其他一些事项。导数是正的，所以总是有一个非零梯度，并且它的和介于 0 和 1 之间，所以输出不会失控。该函数保持以下属性

即输出向量中的相对值不变。

最开始 Miller 打算将这个函数称为 ghostmax，这是因为你可以认为中有一个额外的零值条目，并且 V 矩阵中有一个能够衰减结果的零向量。

尽管 softmax_1 表面上看起来很无聊，但 Miller 99.44% 确信它将解决异常值反馈循环，使量化成为级联研究的主题。Miller 表示，如果你想进行一些实验来证明他是对的，可以联系他。他将撰写一篇论文。

改进后的机制可以被称为 QuietAttention，它允许注意力头保持「沉默」。

Miller 认为很快可以整合一项测试：如果你在每个输入上下文的前面加上一个零向量，并确保你选择的神经网络不添加任何偏差（包括位置编码），那么零在通过时不会改变，并对每个后续的 softmax 分母添加 unity 产生影响。这样你不会因为处理梯度代码失去理智。Miller 认为这可以通过使用固定嵌入和特殊前缀 token 的 LLaMA 模型来完成。

你仍然需要重新训练模型，因此暂时不要在树莓派（RPi）上尝试此操作。但 Miller 想知道这些权重峰度和激活无穷范数在运行几次后是什么样子的。他认为这会成为有影响力的研究，无论是高通 AI Research 团队的论文，还是 LLM 黑客频道有人计算出 biblatex，但自己最先发现的。

项目地址：

https://github.com/kyegomez/AttentionIsOFFByOne

博客链接：

https://www.evanmiller.org/attention-is-off-by-one.html?continueFlag=5d0e431f4edf1d8cccea47871e82fbc4

更多阅读