加一個 scratchpad,就能讓 Transformer 更進一步?

大規模預訓練語言模型有多火?光看 Transformer 橫掃自然語言處理、計算機視覺等多個領域的勁頭就可見一斑了。但這類模型並不是全能的。它們往往在能夠「一次」(in one pass)完成的任務中表現優異,例如生成逼真文字或合成計算機程式。但是,這類模型在處理多步計算時卻犯了難,尤其是需要精確推理和無界計算(unbounded computation)的任務,如整數加法或程式執行。

為什麼大型語言模型難以處理演算法推理任務呢?近日 MTI 和谷歌釋出的一項研究認為,部分原因在於 Transformer 架構應用於這類任務的方式存在一定侷限性:模型以一次性的方式執行任務(in one forward pass)。給定固定的層數和計算時間,

模型無法適應在輸出結果前,將待解決問題的所需計算時間與其難度相適應。

那麼,如何解決這個問題呢?

之前的一些研究探討了針對不同子任務動態地選擇計算時間的神經架構,而 MIT 和谷歌的這項研究選擇反其道而行之:不改變現有的 Transformer 架構和大規模少次學習語言模型及其訓練步驟,而是

修改任務設計

該研究認為大規模預訓練語言模型也能夠執行復雜的多步計算,即使是在少次學習機制下,其解決方法就是

令模型「逐步」執行運算,並展示中間計算的結果。

研究人員使 Transformer 將中間計算步放入「便箋儲存器」(scratchpad),以此方式訓練 Transformer 來執行多步計算。

在從長加法到執行任意程式等不同複雜度的任務中進行實驗後,研究人員發現便箋儲存器極大地提升了語言模型執行多步計算的能力。

具體而言,這項研究:

引入了「便箋儲存器」(scratchpad)的概念,利用 scratchpad 在不改變 Transformer 底層架構的情況下,幫助模型更好地執行復雜離散計算;

證明 scratchpad 能夠幫助 Transformer 在微調機制下執行長加法,尤其是改進了對於較大問題例項的分佈外泛化情況;

發現 scratchpad 能夠幫助 Transformer 執行更高級別的任務,如多項式計算,且在少次和微調機制下均可行;

訓練 Transformer 逐行發出完整程式跟蹤,並標註區域性變數,這一做法能夠極大提升模型預測基於特定輸入執行給定計算機程式的結果。

方法

該研究的主要思路是:要想解決給定演算法任務,只需

將演算法的中間步編碼為文字,並訓練模型將其輸出到快取區(即 scratchpad)中。

以長加法為例。要想教會模型計算 29 + 57,只需將長加法的步驟顯示地寫出來,如圖 2 所示。

加一個 scratchpad,就能讓 Transformer 更進一步?

圖 2:利用 scratchpad 幫助語言模型做加法計算。

執行可編碼任務也可以採用類似的方式,不同之處在於需要在輸入、scratchpad 和期望輸出前新增原始碼 π_i。示例參見下圖 1:

加一個 scratchpad,就能讓 Transformer 更進一步?

圖 1:利用 scratchpad 技術預測程式碼執行,及其與直接執行的對比結果。

再來看多項式計算。研究表明,不管是在少次學習還是微調機制下,利用 scratchpad 技術執行計算都要比直接計算效果更好。示例和結果參見下圖。

加一個 scratchpad,就能讓 Transformer 更進一步?

scratchpad 的好處

新增 scratchpad 有何好處呢?該研究認為:

模型能夠適應性地調整計算時間,即模型可以根據任務複雜度按所需時間來處理資訊;

模型可以將計算的中間狀態儲存在便箋儲存器中,透過關注語境可以進行查詢;

強制模型透過從生成模型中取樣來輸出具體的中間狀態,能夠減少小誤差的傳播和複合,因為狀態可被量化為 token 嵌入;

檢視模型的 scratchpad 輸出有助於識別常見誤差,並透過修改 scratchpad 格式進行糾正。

總結

這項研究透過長加法、多項式計算和 Python 程式碼執行等多項實驗表明,令大規模預訓練語言模型從簡單的 scratchpad 中進行讀取和寫入,能夠提升其處理計算任務的效能。這類模型可能是將大型語言模型的知識壓縮能力與推理能力相結合的第一步,接下來或許就是構建能寫程式碼也能理解程式碼的模型了。這項研究對需要處理自然語言和程式語義推理的應用而言是有用的,如程式合成、神經指引程式分析,以及互動式程式設計助理。

論文連結:https://arxiv。org/pdf/2112。00114v1。pdf