「注意力公式」中存在了8年的bug，竟被國外小哥發現了？

瞬間，這個話題就在網上炸開了鍋。

現在基于Transformer打造的主流模型，GPT-4、Lalma 2、PaLM等都將受到影響。

Eppo初創公司的工程師Evan Miller今天在博客中介紹了這一重大發現，并表示：

研究人員上個月分離了bug——但是他們誤失了一個簡單的解決方案，「為什么LLM設計人員應該停止使用Softmax？」

圖片

那么，究竟是什么bug，能夠暗藏8年？

作者在博文中，引用了維特根斯坦別有蘊意的一句話，「對于無法言說之事，必須保持沉默」。

圖片

注意力是Off By One

這篇博文標題為「注意力是Off By One」。

圖片

你能看到這個公式的差一錯誤嗎？

圖片

要知道，注意力公式是現代人工智能的核心等式，但其中有一個bug在上周讓作者Evan Miller抓狂。

由此，Miller決定就這個漏洞和修復建議寫篇博文。

文章中，他解釋了當前一代AI模型是如何在一個關鍵的地方出現差一錯誤，這使得每個人的Transformer模型都難以壓縮和部署。

不過，作者強調這只是一篇觀點文章，但如果網上有人想做一些實驗來證明這是對的，可以一起合作驗證。

全與「離群值」有關

首先，先談談為什么差一錯誤很重要。ChatGPT工作得很好，有什么問題嗎？

作者第一次發現了不對勁的地方，是在忙自己的事情和閱讀量化研究論文時發現，這是一種通過LLM Edgers將大型模型壓縮到Mac Minis、Rasberry Pis，以及解開家用恒溫器的技術。

在AI領域，每個人都會受到RAM限制。

所以你使用的RAM越少，你就可以做的更多酷炫的事情，無論是在云端還是在邊緣設備上。

LLM有數十億的權重，如果我們可以讓這些權重縮小，我們可以寫出更好的十四行詩，或者剽竊更優秀的文章，又或者加速世界末日，這都取決于你使用語言的個人動機。

RAM存儲信息，這聽起來像是一種同義反復。信息是負對數概率，即我們需要多少位來存儲事物。

如果一串數字流可預測，例如始終限制在一個有限的范圍內，我們需要的比特數就會少一些。

如果一個數字流不可預測，比如偶爾出現一個超大數字，我們需要更多的二進制數字來編碼這個龐然大物。

這就是在LLM中正在發生的事情（出于目前僅能部分理解的原因）。

Transformer模型包含這些離群權重（outlier weights），并且產生了相差一個數量級的巨大激活。

但是沒有人能夠消除它們。這些megalodons（研究命令行工具）看起來對這些模型的運行至關重要。

但是它們的存在與我們在構建優秀模型之前，所了解的關于神經網絡的一切知識相矛盾。

已經有很多論文討論這些離群值（outlier），人們已經想出了各種各樣的位燃燒方案，以更少的1和0來進行編碼。

因為現在，我們使用普通的比例和偏差整數量化得到的性能退化非常嚴重。

關于所有這些的最佳分析來自高通AI研究院的一篇論文：「Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing」。

圖片

論文地址：https://arxiv.org/pdf/2306.12929.pdf

作者們將這些離群值的存在，追溯到注意力機制的softmax函數。這個看似無辜的指數函數，沒有人發現其能夠產生如此嚴重的峰度異常。

而研究人員也就差點發現這個邊界錯誤。

對此，作者表示，高通研究人員暫時還未回復自己電子郵件，但必須通過這種方式呼吁國際學者社區。如果你讀了這篇鏈接的論文，就忽略他們的建議吧。

修剪后的softmax帶有一個旋轉式的零梯度，他們的門控注意力提議雖然可行，但是為了解決這只是一個增量的失敗而引入了數百萬個新的參數。

在作者看來，這里有一個簡單而明顯的解決方案，就自己閱讀的所有內容中，還沒有人想過去嘗試。

接下來，一起談談softmax函數，以及為什么在處理注意力時，它并非最適合的工具。

Softmax出現的問題

為了解釋這個錯誤，你真的需要理解注意力機制的目標。

這么做個類比吧，大多數數值錯誤都是程序員錯誤地實現方程。

然而，當你處理的不是錯誤的代碼，而是錯誤的數學時，你需要理解這個等式來自哪里，以及你應該怎么做，才有可能修復它。

對此，作者不得不閱讀了大約50篇arXiV論文來理解所有這些。

首先，從輸入嵌入開始理解，這是一個浮點向量，它表示輸入字符串中的一個單詞。

這個向量似乎每年都在變高，比如，最近的LLaMA 2模型從Meta使用了一個長度為3,204嵌入向量。

半精度浮點數計算為6KB+，僅僅是為了表示詞匯表中的一個單詞，而詞匯表通常包含30,000——50,000個條目。

現在，如果你是節省內存的C程序員，你可能會想，為什么這些AI goober要使用6KB，來表示應該只需要2字節就能搞定的事情？

如果他們的詞匯表小于，我們只需要16位就能表示一個條目，對吧？

這正是Transformer實際在做的事情：它將輸入向量轉換為相同大小的輸出向量，這個最終的6KB輸出向量需要編碼絕對一切，以預測當前詞語之后的詞語。

每一層Transformer的工作就是，實實在在地向原始的單詞向量添加信息。

這就是殘差（née skip）連接的作用：所有的注意力機制只是向原始的兩個字節信息添加補充材料，分析更大的上下文以指示。

例如，單詞「pupil」指的是學生，而不是你的瞳孔。重復幾十次注意力機制，你就掌握了英語和所有豐富的內容。

現在，Transformer 的最后一步是將這個輸出向量與一個矩形矩陣相乘，并將結果的詞匯長度向量塞入softmax，將那些指數輸出視為下一個詞的概率。

這是合理的，但每個人都知道它并非是完全正確的。

因為沒有模型將那些輸出概率視為正確，與之相反，每個實現和其他模型都使用采樣機制來掩蓋softmax過度表示低概率的事實。

這一切都很好，也可行。

在輸出步驟中的softmax為詞匯表中的每個詞提供了梯度，這是一個合理的選擇，直到有更好的詞出現。

但作者想要辯論的是，Transformer的輸出softmax與注意力機制的內部softmax有著不同的目的，我們都應該去除后者，或者至少用一些方便的東西支撐起它的分母。

那么什么是softmax？

softmax最初起源于統計力學中，用于基于能級預測狀態分布：

然后經濟學家意識到，如果人們的線性效用函數中的噪聲項恰好遵循Gumbel分布，那么某人選擇某個項目的概率將與效用輸入的指數成比例：

而這也使得softmax在多項式邏輯函數中有了用武之地。

可以說，softmax是一種將實數映射為總和為1的概率的「作弊代碼」。

在物理學中，它效果很好；在經濟學中，它有點虛假，但是一旦它進入機器學習領域，每當涉及到離散選擇時，它似乎就成為一種行之有效的東西。

這就是softmax的核心機制：它強制在競爭的替代方案中進行選擇，無論是粒子選擇能級狀態，還是消費者選擇汽車。

也就是說，如果softmax機制根本不想做出任何選擇，softmax將需要進行修改，否則我們預期softmax在遇到實際數據時會產生扭曲。

就LLM而言，其中一個扭曲是對非語義token（逗號等）進行重點加權，而那些權重也就變成了難以壓縮的異常值。

對此，高通AI研究人員發現，LLM中97%以上的異常激活發生在空白和標點位置。

哪里會出錯？

接下來，讓我們深入研究softmax在注意力中的使用，并看看它在哪里出錯了：

分解一下：在僅解碼器模型中（即ChatGPT之后的所有模型），??、??和??都來自同一輸入序列。

雖然它們并不相同，因為它們在途中被以不同的方式投影，但在每一層中，它們都始于相同的已注釋（已添加到）嵌入向量。

現在：????^??正在尋找不同位置的token（嵌入）向量之間的相關性，實際上正在構建一個相關性（點積按1/√??縮放）值的方陣，其中每列和行對應一個token位置。

然后，這個方陣的每一行都經過softmax處理，得到的概率用作??矩陣中的值向量的混合函數。概率混合后的??矩陣被加到輸入向量中，并將其傳遞到神經網絡中進行進一步處理。

多頭注意力在每個層中同時經過這個過程，進行多次處理。它基本上將嵌入向量劃分成多個部分，每個頭使用整個向量中的信息來注釋輸出向量的一個（不重疊的）段。

如果你對原始Transformer論文中的Concatenation操作感到困惑，那就是在發生的事情：頭1向段1添加信息，頭2向段2添加信息，依此類推。

使用softmax的問題在于，它迫使每個注意力頭都要進行注釋，即使它沒有任何信息可以添加到輸出向量中。

在離散選擇之間使用softmax是很好的；但作為可選注釋（即輸入到加法中）使用它，就有點不太好。其中，多頭注意力則會加劇這個問題，因為專門的頭比通用的頭更有可能想要「通過」。

現在，可能應該全面替換softmax，但它在大部分情況下效果還不錯，除了一個小問題，它阻止了注意力頭發出空白注釋。

因此，我提出了一個非常小的調整，我愿意將所有未來的互聯網聲明都寄托在這個正確性上。

這個調整是如此小，又是如此明顯，自從注意力被發明（2014年）以來一直在大家的眼皮底下。

Softmax1和Quiet Attention

現在，經過改造的Softmax Super-Mod公式來了！

不過，實際上只是在分母上加了一個「1」。

作者表示，如果愿意的話，這可以讓整個向量趨向于零，但除此之外，就只是將數值縮小了一些，而這將會在歸一化過程中得到補償。其中，歸一化過程會在注意力之后進行。

關鍵的區別在于負極限，當??中的條目明顯小于零且模型試圖完全避免一個注釋時。

比較原始softmax的極限行為：

與新的改進softmax1的極限行為：

可以看到，原始的softmax總是會產生相同的總權重；softmax1雖然看起來大部分相同，但在負半軸中有一個逃生通道。

此外，softmax1還有其他一些特點。比如，它的導數是正數，因此我們始終有非零梯度；它的和在0和1之間，因此輸出不會失控。

同時，softmax1還會保持如下函數性質，即輸出向量中的相對值保持不變。

而原始的softmax即便采用更高的精度，也無法解決這些問題。也就是說，所有的Transformers都受到影響。

盡管softmax1表面上看起來相當普通，但作者有99.44%的把握，它可以解決量化的離群反饋循環問題。

對于改進后的機制，作者稱之為——QuietAttention，因為它允許注意力頭保持安靜：

基于此，作者認為可以很快地編寫一個測試：

「如果在每個輸入上下文前加上一個零向量，并確保選擇的神經網絡種不會增加任何偏差（包括位置編碼），那么零向量應該會原封不動地通過，并且在每個后續的softmax分母中都添加一個單位。這樣，也就不必糾結于梯度代碼了。」

此外，作者還認為可以使用一個使用固定嵌入和特殊前綴token的LLaMA模型來完成這項工作。

不過，由于仍然需要重新訓練模型，所以暫時不要在Raspberry Pi上進行嘗試。

順便，如果你真的進行測試了的話，記得把結果發給這位作者——他想在即將發表的arXiV論文中制作一張漂亮的表格。

作者介紹

本文作者Evan Miller其實說起來并不是那么有名。

他的履歷和出身和一些科學大牛相比，確實是比不過。但并不影響他能做出本文所講的重大發現。

Miller本科、碩士和博士三個階段其實都沒有主修計算機科學。

本科階段，他在威廉姆斯學院學的是物理。后來又在芝加哥大學攻讀了經濟學的博士。

圖片

而除了學業生涯，Miller寫過很有名的排名算法，目前在很多網站上都有應用。

他設計的統計軟件還曾被頂級醫學期刊引用。

目前，Miller在一家名為 Eppo 的初創公司擔任統計工程師。

圖片

而Miller在他自己的網站上所展示的內容，可以看出這哥們快是個全才了。

開源項目，做過7個。

圖片

各個專業的博客、文章、論文啥的，多的數不過來。

有編程類的，應用數學類的，甚至還有生意經。

圖片