自 2017 年被提出以來，Transformer 已成為 AI 大模型的主流架構(gòu)，一直穩(wěn)站 C 位。

但所有研究者都不得不承認(rèn)的是，Transformer 在算數(shù)任務(wù)中表現(xiàn)非常糟糕，尤其是加法，這一缺陷在很大程度上源于 Transformer 無法跟蹤大范圍數(shù)字中每個數(shù)字的確切位置。

為了解決這個問題，來自馬里蘭大學(xué)、CMU 等機(jī)構(gòu)的研究者向這一問題發(fā)起了挑戰(zhàn)，他們通過在每個數(shù)字中添加一個嵌入來解決這個問題，該嵌入編碼數(shù)字相對于開頭的位置。該研究發(fā)現(xiàn)，只用一天時間在單個 GPU 上訓(xùn)練 20 位數(shù)字，就可以達(dá)到最新的性能水平，100 位數(shù)字加法問題高達(dá) 99% 的準(zhǔn)確率。

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

項(xiàng)目地址：https://github.com/mcleish7/arithmetic

標(biāo)題：Transformers Can Do Arithmetic with the Right Embeddings

具體而言，研究者建議對數(shù)據(jù)表示進(jìn)行一個簡單的修改，就能解決這個缺點(diǎn)。他們提出了 Abacus 嵌入用于編碼每個數(shù)字符號 token 范圍內(nèi)的位置。將 Abacus 嵌入與標(biāo)準(zhǔn)位置嵌入結(jié)合使用后，該研究觀察到 Transformer 在算數(shù)任務(wù)上的準(zhǔn)確率有顯著提高，以至于最多只訓(xùn)練了 20 位數(shù)操作數(shù)的模型可以泛化到 120 位數(shù)操作數(shù)的問題。這一數(shù)字代表了 6 倍的 SOTA 泛化因子，而以前的最先進(jìn)的泛化因子也只有 2.5 倍。據(jù)了解，這是迄今為止被證明的最長的學(xué)習(xí)加法序列。

此外，本文還研究了幾種其他方法來改善 transformer 在算術(shù)和泛化方面的性能，他們發(fā)現(xiàn)結(jié)合輸入注入（input injection），即在輸入層和每個解碼器層之間插入跳躍連接，可以在 Abacus 嵌入基線上減少 50% 的泛化誤差。本文還發(fā)現(xiàn)，與嵌入結(jié)合使用的 looped transformer 架構(gòu)可以在加法問題上實(shí)現(xiàn)幾乎完美的泛化。

本文的貢獻(xiàn)可以總結(jié)如下：

• 本文提出了一種新的位置嵌入，稱為 Abacus 嵌入，以更好地捕獲每個數(shù)字的重要性，從而實(shí)現(xiàn)近乎完美的分布內(nèi)泛化；
• 研究表明，當(dāng)將 Abacus 嵌入與輸入注入和 looped transformer 相結(jié)合時，性能會進(jìn)一步提高，分布外準(zhǔn)確率從 92.9% 提高到 99.1%，與單獨(dú)使用標(biāo)準(zhǔn)架構(gòu)的嵌入相比，誤差降低了 87%；
• 研究者將這些發(fā)現(xiàn)擴(kuò)展到更復(fù)雜的問題，包括乘法和排序，在這些領(lǐng)域也展現(xiàn)出了長度泛化。

實(shí)現(xiàn)加法的長度泛化

作者研究了一系列方法，旨在提高從頭開始訓(xùn)練的語言模型在算術(shù)能力上的表現(xiàn)。他們主要關(guān)注兩個假設(shè)：1）數(shù)字內(nèi)各個位數(shù)的位置信息正在丟失；2）循環(huán)可以提高 Transformer 架構(gòu)在多步算術(shù)推理問題上的推理能力。在詳細(xì)描述每項(xiàng)改進(jìn)之前，作者簡要討論了訓(xùn)練和評估設(shè)置。

實(shí)驗(yàn)設(shè)置

作者訓(xùn)練了僅包含解碼器的因果語言模型來解決加法問題。

他們考慮了兩種標(biāo)準(zhǔn) transformer 架構(gòu)。首先，他們使用一個標(biāo)準(zhǔn)的自回歸 transformer 模型，多個解碼器層以前饋方式堆疊。其次，他們通過輸入注入（input injection）增強(qiáng)了這一標(biāo)準(zhǔn) transformer 模型，即把嵌入的輸入添加到每個解碼器層的輸入中。作者在圖 20 中直觀地描述了這些架構(gòu)。

Abacus 嵌入幫助對齊數(shù)字

通過之前的研究和初步實(shí)驗(yàn)，作者發(fā)現(xiàn)，即使輸入的數(shù)字是先顯示最不重要的數(shù)字，訓(xùn)練數(shù)據(jù)是分層的、豐富的（幾百萬個例子），標(biāo)準(zhǔn) transformer 也很難學(xué)習(xí)多位數(shù)加法。他們還觀察到，人類在進(jìn)行長加法運(yùn)算時，會先將數(shù)位相同的數(shù)字排列成列。因此，作者的第一個假設(shè)是，對于 transformer 來說，每個數(shù)字的數(shù)位并不容易表示，而且這個子問題比實(shí)際加法本身帶來的障礙更大。

為了解決 transformer 在表示位置信息方面的局限性，作者設(shè)計(jì)了一種特殊的位置嵌入，它可以編碼每個數(shù)字相對于當(dāng)前數(shù)字起始位置的位置。作者將其稱之為 Abacus 嵌入。他們將相同的位置嵌入應(yīng)用于所有具有相同數(shù)位的數(shù)字，從而提供一個顯式的信號，供模型用于對齊數(shù)字，如圖 2 所示。

Abacus 嵌入解決加法問題

對于標(biāo)準(zhǔn) transformer 架構(gòu)，Abacus 嵌入可將泛化性能提高到 100 位及以上。在圖 3（左）中，作者強(qiáng)調(diào)了 Abacus 嵌入與標(biāo)準(zhǔn) transformer 架構(gòu)和嵌入相比，在進(jìn)行加法運(yùn)算時所具有的比較優(yōu)勢，取三種模型在所有情況下的平均準(zhǔn)確度。

圖 1 還顯示了使用 FIRE 和 Abacus 訓(xùn)練的標(biāo)準(zhǔn) transformer 模型的準(zhǔn)確度結(jié)果，這些模型經(jīng)過了域內(nèi) (ID) 和域外 (OOD) 測試。

Transformer 中的循環(huán)提高了性能

在解決位置嵌入問題后，接下來作者探討了循環(huán)架構(gòu)能否進(jìn)一步提高 transformer 執(zhí)行多位數(shù)加法的能力。他們使用「循環(huán)塊（recurrent block）」一詞來指一組具有不同權(quán)重的解碼器層，而「循環(huán)（recurrence）」則指循環(huán)塊的重復(fù)次數(shù)。作者使用有效深度（effective depth）一詞來指 transformer 中使用的層數(shù)，無論其權(quán)重是否唯一。除非另有說明，否則他們使用的是最大循環(huán)架構(gòu)，即只循環(huán)一個唯一層來達(dá)到有效深度。他們還采用了輸入注入、 殘差連接的方式，將輸入的副本傳播到網(wǎng)絡(luò)中的每一層。

循環(huán)的優(yōu)勢

在圖 3（右）中，作者比較了使用 FIRE 和 NoPE 嵌入對操作數(shù)多達(dá) 40 位的加法進(jìn)行訓(xùn)練的所有架構(gòu)變體。盡管參數(shù)數(shù)量僅相當(dāng)于其他模型的 1/10，但可以看到，looped transformer（循環(huán)的、有輸入注入和漸進(jìn)損失）在使用任何一種位置嵌入時都取得了最佳的分布外性能。在圖 8 中，作者展示了這一結(jié)果在多種訓(xùn)練數(shù)據(jù)規(guī)模下的穩(wěn)健性。

對于循環(huán)模型，可以選擇在訓(xùn)練時改變每次前向傳遞的循環(huán)次數(shù)。這往往會提高模型測試時對較難任務(wù)的泛化能力，這也被稱為漸進(jìn)損失計(jì)算（progressive loss computation）。這個損失函數(shù)是兩個前向傳遞的損失值的凸組合，一個使用字面上的循環(huán)數(shù)（1 × 16 模型為 16），另一個使用隨機(jī)的較小循環(huán)數(shù)。

接下來，作者探討了在保持有效深度固定的同時改變循環(huán)塊大小的效果。他們將循環(huán)塊中的層數(shù)減半，循環(huán)次數(shù)增加一倍，從塊中有 16 層、循環(huán)次數(shù)只有一次（16 × 1，即標(biāo)準(zhǔn) transformer）的模型，過渡到塊中只有一層、循環(huán)次數(shù)有 16 次（1 × 16）的模型。

通過圖 4 分析這些結(jié)果，作者發(fā)現(xiàn)在某些情況下，結(jié)合循環(huán)和 Abacus 嵌入可以進(jìn)一步提高性能。具體來說，在 OOD 問題上，有兩個循環(huán)的模型（8 × 2）產(chǎn)生的誤差是純非循環(huán)模型（16 × 1）的一半，而在 100 + 的 OOD 問題上，其準(zhǔn)確率也有所提高。

最后，在附錄 A.7.3 中，作者改變了模型的有效深度，以分析參數(shù)數(shù)量對這項(xiàng)任務(wù)的影響，包括 Abacus、FIRE 和 NoPE 嵌入。雖然圖 4 中的實(shí)驗(yàn)是對不同深度的公平比較，但純粹的標(biāo)準(zhǔn) transformer 模型比相應(yīng)的循環(huán)模型擁有更多的參數(shù)。在附錄的表 3 中，作者記錄了最接近百萬的參數(shù)量。

實(shí)驗(yàn)

研究者不僅對加法問題進(jìn)行了探討，還對乘法和排序進(jìn)行了研究。

整數(shù)乘法

圖 5 展示了 Abacus 嵌入模型在 15 位數(shù)乘法的分布內(nèi)準(zhǔn)確率超過了之前的工作，且不需要用零將每個操作數(shù)填充到相同長度。特別地，該研究強(qiáng)調(diào)，與僅使用 FIRE 的基線相比，將 Abacus 嵌入與 FIRE 相結(jié)合也提高了分布問題中最難的分布準(zhǔn)確率 (右下)。

數(shù)組排序

表 1 展示了使用不同嵌入 ——FIRE、Abacus 及其組合 —— 訓(xùn)練的標(biāo)準(zhǔn) transformer（八層）的性能。結(jié)果顯示，組合嵌入方法增強(qiáng)了模型的泛化能力。

如表 2 所示，研究者觀察到在將 Abacus+FIRE 嵌入組合與不同的模型架構(gòu)（有效深度為 8）配對時，結(jié)果表現(xiàn)出混合性。

Abacus 和相關(guān)嵌入

圖 6 展示了將 Abacus 嵌入整合到更通用系統(tǒng)中的真正潛力，顯示出 Abacus 嵌入與 FIRE 結(jié)合可以解鎖遠(yuǎn)超 FIRE 嵌入解決問題的能力。 

更多研究細(xì)節(jié)，請參考原論文。

本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/zHx_pMk6sHmIm-AJCS_sRA??