Attention機制竟有bug,Softmax是罪魁禍首,影響所有Transformer

「大模型開發者,你們錯了。」

「我發現注意力公式裡有個 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

相關文章

萬字長文全面了解 GPT-4

萬字長文全面了解 GPT-4

今天凌晨1點,openai發佈了全網期待已久的GPT-4。 GPT-4的看點都有哪些呢? 邏輯推理,強到可以代替你考律師啦 角色扮演,cos...

ChatGPT 的「神功」,是如何煉成的?

ChatGPT 的「神功」,是如何煉成的?

最強對話式 AI,站在巨人的肩膀上。 AIGC 在這幾個月成了一個大熱的話題。 頗有些風水輪流轉的感覺,如同年初大火特火的 web3 一樣,...