「我發現注意力公式里有個 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?cnotallow=5d0e431f4edf1d8cccea47871e82fbc4