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

2023-07-31 14:47 670 阅读 ID:1292
机器之心
机器之心

「我发现注意力公式里有个 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 torch
from softmax_one.softmax_one import softmax_one
x = 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
免责声明:作者保留权利,不代表本站立场。如想了解更多和作者有关的信息可以查看页面右侧作者信息卡片。
反馈
to-top--btn