乘法和排序也有效。
自 2017 年被提出以來,Transformer 已成為 AI 大模型的主流架構,一直穩站 C 位元。
然而,雖然所有研究者都不得不承認的是,Transformer 在算數任務中表現異常糟糕,儘管是加法,這一缺陷在很大程度上源於Transformer 無法跟踪大範圍數字中每個數字的準確位置。
為了解決這個問題,來自馬裡蘭大學、CMU等機構的研究者們向這個問題發起了挑戰。他們透過在每個數字中添加一個嵌入來解決這個問題,該嵌入編碼數字相對於開頭的位置。研究發現,只花一天時間在單一GPU上訓練20位數,就可以達到最新的表現水平,100位數數字加法問題高達99%的準確率。
論文網址:https://arxiv.org/pdf/2405.17399
專案網址:https://github.com/mcleish7/arithmetic
標題:Transformers Can Do Arithmetic with the Right Embeddings
具體而言,研究者建議對資料表顯示進行一個簡單的修改,就能解決這個缺點。他們提出了 Abacus 嵌入用於編碼每個數字符號 token 範圍內的位置。將Abacus 嵌入與標準位置嵌入結合使用後,研究觀察到Transformer 在算數任務上的準確率有顯著提高,以至於最多只訓練了20 位數操作數的模型可擴展到120 位數操作數的問題。這個數字代表了 6 倍的 SOTA 擴展因子,而前的最先進的擴展因子也只有 2.5 倍。據了解,這是迄今為止被證明的最長的學習加法序列。
除了研究優化Transformer在算術和泛化方面的表現之外,本文還探討了幾種其他方法來改善Transformer的表現。他們發現,透過在輸入註入(input injection)層和每個解碼器層之間插入跳躍連接,可以在Abacus嵌入基線上減少50%的泛化誤差。本文也發現,與嵌入結合使用的looped Transformer架構可以在加法問題上實現幾乎完美的泛化。
本文的貢獻可以總結如下:
本文提出了一個新的位置嵌入,稱為Abacus 嵌入,以更好地捕捉每個數字的重要性,從而實現近乎完美的分佈內泛化;
研究表明,當將Abacus 嵌入與輸入註入和looped transformer 相結合時,性能會進一步提高,分佈外準確率從92.9% 提高到99.1%,與單獨使用標準架構的嵌入相比,誤差降低了87%;
研究者將這些發現擴展到更複雜的問題,包括乘法和排序,在這些領域也展現了長度泛化。
實現加法的長度泛化
作者研究了一系列方法,旨在提高從頭開始訓練的語言模型在算術能力上的表現。他們主要關注兩個假設:1)數字內各位數的位置資訊正在遺失;2)循環可以提高 Transformer 架構在多步驟算術推理問題上的推理能力。在詳細描述每項改進之前,作者簡要討論了訓練和評估設定。
實驗設定
作者訓練了僅包含解碼器的因果語言模型來解決加法問題。
他們考慮了兩種標準 transformer 架構。首先,他們使用標準的自回歸 transformer 模型,多個解碼器層以前饋方式堆疊。其次,他們透過輸入註入(input injection)增強了這個標準 transformer 模型,也就是把嵌入的輸入加入到每個解碼器層的輸入中。作者在圖 20 中直觀地描述了這些架構。
Abacus 嵌入幫助對齊數字
透過先前的研究和初步實驗,作者發現,即使輸入的數字是先顯示最不重要的數字,訓練資料是分層的、豐富的(幾百萬個例子),標準transformer 也很難學習多位數加法。他們也觀察到,人類在進行長加法運算時,會先將數位相同的數字排列成列。因此,作者的第一個假設是,對於 transformer 來說,每個數字的數字並不容易表示,而且這個子問題比實際加法本身帶來的障礙更大。
為了解決 transformer 在表示位置資訊方面的局限性,作者設計了一種特殊的位置嵌入,它可以編碼每個數字相對於當前數位起始位置的位置。作者將其稱之為 Abacus 嵌入。他們將相同的位置嵌入應用於所有具有相同數字的數字,從而提供一個明確的訊號,供模型用於對齊數字,如圖 2 所示。
Abacus 嵌入解決加法問題
對於標準 transformer 架構,Abacus 嵌入可將泛化效能提高到 100 位元及以上。在圖 3(左)中,作者強調了 Abacus 嵌入與標準 transformer 架構和嵌入相比,在進行加法運算時所具有的比較優勢,取三種模型在所有情況下的平均準確度。
圖 1 也顯示了使用 FIRE 和 Abacus 訓練的標準 transformer 模型的準確度結果,這些模型經過了域內 (ID) 和域外 (OOD) 測試。
Transformer 中的循環提高了效能
在解決位置嵌入問題後,接下來作者探討了循環架構能否進一步提高transformer 執行多位數加法的能力。他們使用「循環區塊(recurrent block)」一詞來指一組具有不同權重的解碼器層,而「循環(recurrence)」則指循環塊的重複次數。作者使用有效深度(effective depth)一詞來指稱 transformer 中使用的層數,無論其權重是否唯一。除非另有說明,否則他們使用的是最大循環架構,即只循環一個唯一層來達到有效深度。他們也採用了輸入註入、 殘差連接的方式,將輸入的副本傳播到網路中的每一層。
循環的優勢
在圖3(右)中,作者比較了使用FIRE 和NoPE 嵌入對操作數多達40 位的加法進行訓練的所有架構變體。儘管參數數量僅相當於其他模型的 1/10,但可以看到,looped transformer(循環的、有輸入註入和漸進損失)在使用任何一種位置嵌入時都取得了最佳的分佈外性能。在圖 8 中,作者展示了這項結果在多種訓練資料規模下的穩健性。
對於循環模型,可以選擇在訓練時改變每次前向傳遞的循環次數。這往往會提高模型測試時對較難任務的泛化能力,這也被稱為漸進式損失計算(progressive loss computation)。這個損失函數是兩個前向傳遞的損失值的凸組合,一個使用字面上的循環數(1 × 16 模型為 16),另一個使用隨機的較小循環數。
接下來,作者探討了在保持有效深度固定的同時改變循環區塊大小的效果。他們將循環區塊中的層數減半,循環次數增加一倍,從區塊中有16 層、循環次數只有一次(16 × 1,即標準transformer)的模型,過渡到區塊中只有一層、循環次數有16 次(1 × 16)的模型。
透過圖 4 分析這些結果,作者發現在某些情況下,結合循環和 Abacus 嵌入可以進一步提高效能。具體來說,在OOD 問題上,有兩個循環的模型(8 × 2)產生的誤差是純非循環模型(16 × 1)的一半,而在100 + 的OOD 問題上,其準確率也有所提高。
最後,在附錄 A.7.3 中,作者改變了模型的有效深度,以分析參數數量對這項任務的影響,包括 Abacus、FIRE 和 NoPE 嵌入。雖然圖 4 中的實驗是對不同深度的公平比較,但純粹的標準 transformer 模型比相應的循環模型擁有更多的參數。在附錄的表 3 中,作者記錄了最接近百萬的參數量。
實驗
#研究者不僅對加法問題進行了探討,也對乘法和排序進行了研究。
整數乘法
圖5 展示了Abacus 嵌入模型在15 位數乘法的分佈內準確率超過了先前的工作,且不需要用零將每個操作數填入相同長度。特別地,研究強調,與僅使用 FIRE 的基線相比,將 Abacus 嵌入與 FIRE 結合也提高了分佈問題中最難的分佈準確率 (右下)。
陣列排序
表 1 展示了使用不同嵌入 ——FIRE、Abacus 及其組合 —— 訓練的標準 transformer(八層)的表現。結果顯示,組合嵌入法增強了模型的泛化能力。
如表2 所示,研究者觀察到在將Abacus+FIRE 嵌入組合與不同的模型架構(有效深度為8)配對時,結果表現出混合性。
Abacus 和相關嵌入
#圖6 展示了將Abacus 嵌入整合到更通用系統中的真正潛力,顯示出Abacus 嵌入與FIRE 結合可以解鎖遠超FIRE 嵌入解決問題的能力。
更多研究細節,請參考原文。
以上是單GPU訓練一天,Transformer在100位數字加法上就達能到99%準確率的詳細內容。更多資訊請關注PHP中文網其他相關文章!