「我發(fā)現(xiàn)注意力公式里有個 bug，八年了都沒有人發(fā)現(xiàn)。所有 Transformer 模型包括 GPT、LLaMA 都受到了影響。」

昨天，一位名叫 Evan Miller 的統(tǒng)計工程師的話在 AI 領(lǐng)域掀起了軒然大波。

我們知道，機(jī)器學(xué)習(xí)中注意力公式是這樣的：

圖片

自 2017 年 Transformer 問世，這個公式已被廣泛使用，但現(xiàn)在，Evan Miller 發(fā)現(xiàn)這個公式是錯的，有 bug！

Evan Miller 的這篇博客解釋了當(dāng)前流行的 AI 模型如何在關(guān)鍵位置出現(xiàn)錯誤，并使得所有 Transformer 模型都難以壓縮和部署。

總結(jié)而言，Evan Miller 引入了一種新函數(shù) Quiet Attention，也叫 Softmax_1，這是對傳統(tǒng) softmax 函數(shù)的創(chuàng)新調(diào)整。

有網(wǎng)友對該博客總結(jié)出了一個「太長不看版」。博客作者建議在注意力機(jī)制使用的 softmax 公式分母上加 1（不是最終輸出 softmax）。注意力單元中的 softmax 使其可以將鍵 / 查詢匹配作為概率；這些概率支持一個鍵 - 值查找的連續(xù)值版本（我們得到的權(quán)重不是一個查找的 1/0 輸出，而是高權(quán)重 = 所需的鍵 - 值查找）。

分母上加 1 將改變注意力單元，不再使用真實(shí)的權(quán)重概率向量，而是使用加起來小于 1 的權(quán)重。其動機(jī)是該網(wǎng)絡(luò)可以學(xué)習(xí)提供高權(quán)重，這樣調(diào)整后的 softmax 非常接近概率向量。同時有一個新的選項(xiàng)來提供 all-low 權(quán)重（它們提供 all-low 輸出權(quán)重），這意味著它可以選擇不對任何事情具有高置信度。

有人甚至猜測「這就是微軟 RetNet 比 transformer 性能更優(yōu)的原因？」

還有網(wǎng)友表示，這項(xiàng)研究可以促進(jìn) LLM 的改進(jìn)，從而極大對權(quán)重進(jìn)行壓縮，使得較小的模型媲美較大的模型：

Miller 表示：你可以像使用傳統(tǒng)的 softmax 函數(shù)一樣使用 Softmax_1 函數(shù)，示例如下。

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

基于這樣的修改，Miller 還做了實(shí)驗(yàn)，結(jié)果如下：

接下來我們看看 Miller 到底發(fā)現(xiàn)了什么錯誤。

異常值

Evan Miller 是在閱讀關(guān)于量化的論文時發(fā)現(xiàn)了這個 bug。當(dāng)前，內(nèi)存和存儲已經(jīng)成為限制人工智能發(fā)展的重要因素。人們一直在努力壓縮模型，并嘗試在云端、在邊緣設(shè)備上運(yùn)行大型語言模型（LLM）。

在計算機(jī)中，信息是用二進(jìn)制數(shù)據(jù)流來存儲的。如果數(shù)據(jù)流是高度可預(yù)測的，例如總是包含在有限的范圍內(nèi)，那么我們就可以用相對較少的位（bit）來存儲它們。反之，如果一串?dāng)?shù)字是不可預(yù)測的，可能是千載難逢的巨大數(shù)字，我們就需要更多的二進(jìn)制數(shù)字來編碼和存儲。而 Transformer 模型包含一些異常值權(quán)重。

在高通 AI Research 6 月發(fā)表的一篇論文《Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing》中，研究團(tuán)隊(duì)將這些異常值的存在追溯到注意力機(jī)制的 softmax 函數(shù)。

高通論文地址：https://arxiv.org/abs/2306.12929

這聽起來令人意外，但 Evan Miller 認(rèn)為這是對的，并進(jìn)一步發(fā)現(xiàn) softmax 函數(shù)存在一個錯誤。

我們來看下 Evan Miller 是如何說明 softmax 函數(shù)在注意力機(jī)制方面并不是一個合適的工具的。

Softmax 引出的問題

為什么說 softmax 不適合注意力機(jī)制，這還得從注意力機(jī)制可以做什么開始。

一般來講，數(shù)值錯誤一般是由程序錯誤引起的，然而，當(dāng)程序沒有錯誤，這時就需要從修復(fù)復(fù)雜的數(shù)學(xué)公式入手，耗費(fèi)大量時間。

Evan Miller 大概閱讀了 50 篇 arXiV 論文才有點(diǎn)眉目。Miller 從「輸入嵌入」入手，我們可以將「輸入嵌入」理解為一個浮點(diǎn)向量，代表輸入字符串中的一個單詞。

舉例來說，Meta 最近推出的 LLaMA 2 模型使用了一個長度為 3204 的嵌入向量，以半精度浮點(diǎn)數(shù)表示，這僅僅是為了表示詞匯表中的一個單詞，而詞匯表通常包含 30000 到 50000 個條目（entry）。意味著一個單詞的嵌入向量占用 6KB + 的存儲空間。隨著技術(shù)的發(fā)展，「輸入嵌入」的長度逐漸增加，所占存儲空間也隨之增加。

如果你是一個對存儲占用非常敏感的 C 程序員，你可能接受不了這一數(shù)字，明明是 2 字節(jié)就能存儲的東西，為什么偏偏要用 6KB？如果按照 2 字節(jié)來計算，假如詞匯量少于 2^16=65384，那么我們只需要 16 位來表示一個條目。

但是，實(shí)際上 Transformer 的工作原理是這樣的：它將輸入向量轉(zhuǎn)換為大小相同的輸出向量，最終的 6KB 輸出向量用來預(yù)測下一個 token。運(yùn)行中，Transformer 每一層的工作流都將信息添加到原始的單詞向量中。在這其中，還用到了殘差連接：所有的注意力機(jī)制都在為原始的兩個字節(jié)的信息添加補(bǔ)充材料，從而是的 LLM 能夠分析更長的上下文。

Transformer 的最后一步是將這個輸出向量與一個矩形矩陣相乘，并將得到的詞匯長度向量壓縮到一個 softmax 函數(shù)中，將這些指數(shù)化的輸出視為下一個 token 的概率。這是合理的，但眾所周知，這并不完全正確，因?yàn)槲覀儾荒艽_定這些輸出概率是正確的。相反，每個 Transformer 實(shí)現(xiàn)和其衍生版本都使用采樣機(jī)制來隱藏 softmax 過度表示概率較低的事實(shí)。

接下里，Miller 介紹了 softmax 的發(fā)展史。softmax 最初出現(xiàn)在統(tǒng)計學(xué)中，最早作為一種基于能級預(yù)測狀態(tài)分布的方法，其形式如下：

之后經(jīng)濟(jì)學(xué)家又將其修改為

這一修改，softmax 才擁有了多項(xiàng)邏輯函數(shù)。由于 Miller 對 softmax 函數(shù)的研究頗深，因而，他能識別出 softmax 使用不恰當(dāng)?shù)牡胤健?/span>

Softmax 應(yīng)用廣泛，在物理學(xué)中，它非常有效；在經(jīng)濟(jì)學(xué)中，它可能不那么準(zhǔn)確；但將其應(yīng)用到機(jī)器學(xué)習(xí)領(lǐng)域時，只要涉及離散選擇，它似乎總是有效的：

Miller 進(jìn)一步表示，softmax 的關(guān)鍵在于，如果你不想保留一些項(xiàng)，必須對 softmax 進(jìn)行修改，否則結(jié)果就會產(chǎn)生扭曲。

舉例來說，在 LLM 上下文中，扭曲產(chǎn)生的原因是對非語義 token（逗號等）進(jìn)行大量加權(quán)導(dǎo)致的，這些較高的權(quán)重成為難以壓縮的異常值，使得研究變得更加困難。來自高通的 AI 研究員也發(fā)現(xiàn)了這一現(xiàn)象，在 LLM 中，97% 以上的異常激活發(fā)生在空格和標(biāo)點(diǎn)符號位置上。

接下來，Miller 介紹了 softmax 是如何在注意力中使用的，從而發(fā)現(xiàn)問題到底出現(xiàn)在哪里：

對上述公式進(jìn)行分解，在僅解碼器模型中，??、??和??源自相同的輸入序列。它們又不完全相同，即投影方式不同。但在每一層中，它們都以相同的注釋嵌入向量開始。

????^??項(xiàng)用于尋找不同位置 token 向量之間的相關(guān)性，實(shí)質(zhì)上構(gòu)建了一個相關(guān)性矩陣（點(diǎn)積按縮放），其中每一列和每一行對應(yīng)一個 token 位置。然后，對這個方陣的每一行進(jìn)行 softmax 操作，得到的概率用作??矩陣中值向量的混合函數(shù)。概率混合后的??與輸入向量相加，將求和結(jié)果傳遞給神經(jīng)網(wǎng)絡(luò)進(jìn)行進(jìn)一步處理。

多頭注意力每層并行執(zhí)行多次上述過程。從本質(zhì)上講，這種方法劃分了嵌入向量，每個頭使用整個向量中的信息來注釋輸出向量的一個（非重疊）片段。這就是原始 Transformer 論文中的串聯(lián)操作。

使用 softmax 的問題在于，它強(qiáng)制每個注意力頭進(jìn)行注釋，即使沒有信息可添加到輸出向量中。

Softmax_1 和 QuietAttention

來了，在這里你將看到 Softmax Super-Mod 點(diǎn)燃了 LLM 頻道。

有點(diǎn)失望，對吧？Miller 所做的只是在分母上加 1。如果想要的話，這可以讓該向量作為一個趨于 0 的整體。否則只會將值縮小一點(diǎn)，并且縮小的值會在歸一化過程中得到補(bǔ)償，這在注意力之后發(fā)生。

當(dāng) ?? 中的條目顯著小于零并且模型試圖完全避免注釋時，主要的區(qū)別在于負(fù)值限制。將如下原始 softmax 的限制行為

與新的、改進(jìn)后的 softmax_1 相比較。

Vanilla softmax 將始終釋出相同的總權(quán)重；softmax_1 看起來大部分相同，但在負(fù)象限中有一個「逃出口」（escape hatch）。需要明確的是，這里的核心問題在本質(zhì)上是數(shù)學(xué)而非數(shù)值問題。額外的精度并不能拯救 softmax，所有的 Transformers 都會受到影響。

你還可以觀察到關(guān)于 softmax_1 的其他一些事項(xiàng)。導(dǎo)數(shù)是正的，所以總是有一個非零梯度，并且它的和介于 0 和 1 之間，所以輸出不會失控。該函數(shù)保持以下屬性

即輸出向量中的相對值不變。

最開始 Miller 打算將這個函數(shù)稱為 ghostmax，這是因?yàn)槟憧梢哉J(rèn)為中有一個額外的零值條目，并且 V 矩陣中有一個能夠衰減結(jié)果的零向量。

盡管 softmax_1 表面上看起來很無聊，但 Miller 99.44% 確信它將解決異常值反饋循環(huán)，使量化成為級聯(lián)研究的主題。Miller 表示，如果你想進(jìn)行一些實(shí)驗(yàn)來證明他是對的，可以聯(lián)系他。他將撰寫一篇論文。

改進(jìn)后的機(jī)制可以被稱為 QuietAttention，它允許注意力頭保持「沉默」。

Miller 認(rèn)為很快可以整合一項(xiàng)測試：如果你在每個輸入上下文的前面加上一個零向量，并確保你選擇的神經(jīng)網(wǎng)絡(luò)不添加任何偏差（包括位置編碼），那么零在通過時不會改變，并對每個后續(xù)的 softmax 分母添加 unity 產(chǎn)生影響。這樣你不會因?yàn)樘幚硖荻却a失去理智。Miller 認(rèn)為這可以通過使用固定嵌入和特殊前綴 token 的 LLaMA 模型來完成。

你仍然需要重新訓(xùn)練模型，因此暫時不要在樹莓派（RPi）上嘗試此操作。但 Miller 想知道這些權(quán)重峰度和激活無窮范數(shù)在運(yùn)行幾次后是什么樣子的。他認(rèn)為這會成為有影響力的研究，無論是高通 AI Research 團(tuán)隊(duì)的論文，還是 LLM 頻道有人計算出 biblatex，但自己最先發(fā)現(xiàn)的。

? 項(xiàng)目地址：https://github.com/kyegomez/AttentionIsOFFByOne

? 博客鏈接：https://www.evanmiller.org/attention-is-off-by-one.html?cnotallow=5d0e431f4edf1d8cccea47871e82fbc4