Transformer 已成為各種機器學習任務(wù)的熱門選擇，并且取得了很好的效果，那它還能怎么用？腦洞大開的研究者竟然想用它來設(shè)計可編程計算機！

這篇論文的作者來自普林斯頓大學和威斯康星大學，標題為《Looped Transformers as Programmable Computers》，旨在探索如何用 Transformer 來實現(xiàn)通用計算機。

具體來說，作者提出了一個將 transformer 網(wǎng)絡(luò)用作通用計算機的框架，方法是使用特定權(quán)重對它們進行編程并將它們置于循環(huán)（loop）中。在這個框架中，輸入序列充當穿孔卡片，由指令和用于數(shù)據(jù)讀 / 寫的內(nèi)存組成。

作者證明了恒定數(shù)量的編碼器層可以模擬基本計算塊。使用這些構(gòu)建塊，他們模擬了一個小型指令集計算機。這使得他們能夠?qū)⒌惴ㄓ成涞娇梢杂裳h(huán)的（looped）13 層 transformer 執(zhí)行的程序。他們展示了這個 transformer 如何在其輸入的指導(dǎo)下模擬一個基本的計算器、一個基本的線性代數(shù)庫和使用反向傳播的 in-context learning 算法。這項工作突出了注意力機制的多功能性，并證明即使是淺層 transformer 也可以執(zhí)行成熟的通用程序。

論文概述

Transformer (TF) 已成為各種機器學習任務(wù)的熱門選擇，它在自然語言處理和計算機視覺等領(lǐng)域很多問題上取得了最先進的成果。Transformer 成功的一個關(guān)鍵原因是能夠通過注意力機制捕捉 higher-order 關(guān)系和 long-range 依賴關(guān)系。這使得 TF 可以對上下文信息進行建模，并使其在機器翻譯和語言建模等任務(wù)中更加有效，在這些任務(wù)中 Transformer 的表現(xiàn)一直優(yōu)于其它方法。

GPT-3（175B 參數(shù)）和 PaLM（540B 參數(shù)）等具有上千億參數(shù)的語言模型在許多自然語言處理任務(wù)上都取得了最先進的性能。有趣的是，這些大型語言模型 (LLM) 中的一些還可以執(zhí)行 in-context learning (ICL) ，根據(jù)簡短 prompt 和一些示例即時適應(yīng)和執(zhí)行特定任務(wù)。LLM 的 ICL 能力是在沒有可以訓(xùn)練的情況下就具備，并允許這些大型模型有效地執(zhí)行新任務(wù)而無需更新權(quán)重。

 令人驚訝的是，通過 ICL，LLM 可以執(zhí)行算法任務(wù)和推理，并且 [Nye et al. [2021], Wei et al. [2022c], Lewkowycz et al. [2022], Wei et al. [2022b], Zhou et al. [2022]] 等人的工作已經(jīng)證明其可行性。[Zhou et al. [2022] ] 等的工作證明，當給出多位加法算法和一些加法示例的 prompt 時，LLM 可以成功地對未知用例進行加法運算。這些結(jié)果表明，LLM 可以根據(jù)算法原理并在推理時對給定輸入執(zhí)行 pre-instructed 的命令，就好像將自然語言解釋為代碼一樣。

 有證據(jù)表明，Transformer 可以模擬具有足夠深度的圖靈機或 Attention 層之間的遞歸鏈接 [Pérez et al. [2021], Pérez et al. [2019], Wei et al. [2022a]]。這證明了 Transformer 網(wǎng)絡(luò)精確遵循輸入指定的算法指令的潛力。然而，這些構(gòu)造相對籠統(tǒng)，無法深入理解如何創(chuàng)建能夠執(zhí)行特定算法任務(wù)的 Transformer。

然而更專業(yè)的設(shè)計可以讓 TF 執(zhí)行更高級的程序。如 [Weiss et al. [2021]] 設(shè)計了一種計算模型和一種編程語言，可將簡單的選擇和聚合命令映射到索引輸入 token。這種語言可用于創(chuàng)建多種有趣的算法，例如 token 計數(shù)、排序、創(chuàng)建直方圖和識別 Dyck-k 語言。然后可以將用受限訪問序列處理語言 (RASP) 編寫的程序映射到 Transformer 網(wǎng)絡(luò)中，Transformer 網(wǎng)絡(luò)的大小通常會隨著程序的大小而變化。

 另一項研究展示了選擇 Transformer 模型權(quán)重的方法，以用作動態(tài)學習線性回歸模型的優(yōu)化算法，在給定訓(xùn)練數(shù)據(jù)作為輸入時在推理時執(zhí)行隱式訓(xùn)練。這些方法通常需要與學習算法的迭代次數(shù)成比例的層數(shù)，并且僅限于一個損失函數(shù)和模型集合。

對 Transformer 模型進行編程以模擬圖靈機的抽象計算的能力、RASP 等語言的專用命令以及 ICL 的特定算法，突出了 Transformer 網(wǎng)絡(luò)作為多功能可編程計算機的潛力。

本文作者的研究旨在研究探索這一充滿希望的前景，揭示 Attention 機制如何能夠模擬受指令集架構(gòu)啟發(fā)的通用計算機。

將 Transformer 作為可編程計算機

在本文中，作者展示了 Transformer 網(wǎng)絡(luò)可以通過使用特定權(quán)重對它們進行硬編碼并將它們置于一個循環(huán)中來模擬復(fù)雜的算法和程序。作者通過對 Attention 進行逆向工程來模擬基本計算塊來做到這一點，例如對輸入序列的編輯操作、非線性函數(shù)、函數(shù)調(diào)用、程序計數(shù)器和條件分支。作者的論文證明了使用單個循環(huán)或遞歸將 Transformer 的輸出序列連接回其輸入的重要性，從而避免對深度模型的需要。

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

作者通過設(shè)計一個 Transformer 來實現(xiàn)這一點，該 Transformer 可以執(zhí)行以單個指令的通用版本編寫的程序，稱為 SUBLEQ (A,B,C)，即如果小于或等于零，則相減并分支。SUBLEQ 是一種單指令語言，定義了單指令集計算機（one-instruction set computer，簡寫為 OISC）。SUBLEQ 由 3 個內(nèi)存地址操作數(shù)組成，執(zhí)行時用存儲器地址 B 的值減去存儲器地址 A 的值，結(jié)果存入 B。如果 B 的結(jié)果小于等于 0，則跳轉(zhuǎn)到地址 C，否則繼續(xù)執(zhí)行下一條指令。但這條指令能夠定義通用計算機。

作者構(gòu)造了實現(xiàn)類似 SUBLEQ 程序的顯式 Transformer，作者稱之為 FLEQ 的更靈活的單指令，其形式為

其中 f_m 可以從一組函數(shù)（矩陣乘法 / 非線性函數(shù) / 多項式等）中選擇，可以將其硬編碼到網(wǎng)絡(luò)中?？梢詧?zhí)行 FLEQ 程序的循環(huán) Transformer 的深度不依賴于程序的深度或代碼的行數(shù)，而是取決于實現(xiàn)單個 FLEQ 指令所需的深度，這是不變的。這是通過在輸入序列上循環(huán)運行 Transformer 來實現(xiàn)的，類似于 CPU 的運行方式。

作者使用這個框架，展示了在推理時模擬各種函數(shù)的能力，包括一個基本的計算器、一個基本的線性代數(shù)庫（矩陣轉(zhuǎn)置、乘法、求逆、冪迭代）和在隱式完全連接網(wǎng)絡(luò)上實現(xiàn)反向傳播的 ICL。輸入序列或 prompt 可當成穿孔卡片，包括 Transformer 需要執(zhí)行的指令，同時為存儲和處理程序中使用的變量提供空間。用于執(zhí)行這些程序的 Transformer 網(wǎng)絡(luò)的深度都小于或等于 13，并且提供了所有這些模型的權(quán)重矩陣。下面的定理總結(jié)了作者的主要發(fā)現(xiàn)：

定理 1：存在一個少于 13 層的循環(huán) Transformer，它可以模擬通用計算機（文章第 5 節(jié)）、基本計算器（文章第 7 節(jié)）、數(shù)值線性代數(shù)方法，如近似矩陣逆和冪迭代（文章第 8 節(jié)），以及基于神經(jīng)網(wǎng)絡(luò)的 ICL 算法（如 SGD）（文章第 9 節(jié)）。

圖 1：循環(huán) Transformer 架構(gòu)示意圖，其中輸入序列存儲命令，從中讀取 / 寫入數(shù)據(jù)的內(nèi)存以及存儲中間結(jié)果的暫存器。輸入由網(wǎng)絡(luò)處理，輸出用作新輸入，允許網(wǎng)絡(luò)迭代更新隱式狀態(tài)并執(zhí)行復(fù)雜計算。

作者的研究強調(diào)了 Attention 機制的靈活性和單循環(huán)的重要性，這使得設(shè)計可以模擬復(fù)雜迭代算法和執(zhí)行通用程序的模型成為可能。并進一步證明了 Transformer 模型高效執(zhí)行復(fù)雜數(shù)學和算法任務(wù)的能力。可以想象，現(xiàn)代 Transformer (如 GPT-3)，在執(zhí)行各種任務(wù)時使用類似的內(nèi)部子程序。在某種程度上，當給出上下文示例和說明時，可以啟發(fā)這些模型特定技巧或算法的能力，類似于函數(shù)調(diào)用。然而，應(yīng)該謹慎對待這個假設(shè)，因為作者設(shè)計結(jié)構(gòu)的方式與現(xiàn)實世界語言模型的訓(xùn)練方式?jīng)]有任何相似之處。

作者希望他們的研究能夠鼓勵進一步研究 Attention 機制的潛力，以及語言模型執(zhí)行算法指令的能力。作者提出的設(shè)計可以幫助確定執(zhí)行特定算法任務(wù)所需的最小 Transformer 網(wǎng)絡(luò)規(guī)模。此外，作者還希望其發(fā)現(xiàn)有助于啟發(fā)開發(fā)方法，從而通過利用更小的、逆向工程的 Transformer 網(wǎng)絡(luò)來完成特定的算法任務(wù)，從而增強訓(xùn)練語言模型的能力。

構(gòu)建面向通用計算的 Transformer 模塊

 要使用 Transformer 網(wǎng)絡(luò)構(gòu)建通用計算框架，需要專門的計算塊。將這些塊組裝來創(chuàng)建所需的最終功能。下面重點介紹 Transformer 層可以執(zhí)行的各種操作。這些操作將為創(chuàng)建更復(fù)雜的例程和算法提供基礎(chǔ)。這些操作的目的是彼此可互操作，利用 Attention 執(zhí)行各種任務(wù)的能力，例如生成近似置換矩陣和通過 sigmoid 函數(shù)逼近一般函數(shù)。

圖 2: 用作實現(xiàn)小型指令集計算機構(gòu)建塊的三個 Transformer 塊的示意圖。這些塊處理輸入序列中的編輯操作（例如從一個塊移動或復(fù)制到另一個塊），跟蹤程序計數(shù)器，并在滿足指定條件時執(zhí)行程序計數(shù)器跳轉(zhuǎn)。

位置編碼、程序計數(shù)器和數(shù)據(jù)指針

Transformer 通常需要執(zhí)行迭代算法或執(zhí)行一系列命令。為實現(xiàn)這一點，作者使用一個循環(huán)訪問命令的程序計數(shù)器。計數(shù)器包含存儲下一個命令的位置的編碼。此外，命令可能具有指向命令需要讀取和寫入的數(shù)據(jù)位置的數(shù)據(jù)指針。程序計數(shù)器和數(shù)據(jù)指針都使用與前一段中討論的相同的位置編碼。

作者的位置編碼方案也可用于指向特定數(shù)據(jù)位置以進行讀取或?qū)懭耄@將在下一節(jié)論述。這是通過使用相同的二進制向量作為程序計數(shù)器和數(shù)據(jù)指針的位置編碼來實現(xiàn)的。此外，這種指向特定數(shù)據(jù)位置的技術(shù)使 Transformer 能夠在執(zhí)行算法或構(gòu)建以實現(xiàn)的命令序列期間有效地讀取 / 寫入數(shù)據(jù)。

讀 / 寫：將數(shù)據(jù)、指令復(fù)制到暫存器或從暫存器取出

圖 3: 讀取操作的示意圖。箭頭顯示從分配給暫存器命令的輸入部分復(fù)制的命令塊。一條指令是一組指針。位置編碼和計數(shù)器用于跟蹤什么內(nèi)容被復(fù)制到哪里。

下面的引理說明，程序計數(shù)器指向的命令或當前命令中指定位置的數(shù)據(jù)可以復(fù)制到暫存器以供進一步計算。程序計數(shù)器的位置通常位于暫存器內(nèi)容的正下方，但可以任意更改。在整個計算過程中將其保持在特定位置有助于保持結(jié)構(gòu)的良好組織。

下一個引理解釋了存儲在暫存器中的向量 v 可以復(fù)制到存儲器中的指定位置，如暫存器本身指定的那樣。這允許將數(shù)據(jù)從暫存器傳輸?shù)絻?nèi)存中的特定位置以供進一步使用或存儲。

圖 4: 寫操作的示意圖。箭頭顯示數(shù)據(jù)塊正在從暫存器復(fù)制到分配給內(nèi)存的輸入部分中的指定位置。位置編碼用于跟蹤目標位置并確保數(shù)據(jù)寫入正確的內(nèi)存位置。

條件分支?

在這一部分，作者實現(xiàn)一個條件分支指令，該指令評估條件并在條件為真時將程序計數(shù)器設(shè)置到指定位置，或者在條件為假時將程序計數(shù)器遞增 1。 

命令的形式如下：如果 mem [a]≤0，則轉(zhuǎn)到 i，其中 mem [a] 是輸入序列的內(nèi)存部分中某個位置的值。該命令有兩部分：判斷不等式和修改程序計數(shù)器。

模擬通用單指令集計算機

SUBLEQ Transformer

Mavaddat 和 Parhami 早在 1988 年已經(jīng)證明存在一條指令，任何計算機程序都可以轉(zhuǎn)換為由這條指令的實例化組成的程序。這種指令的一個變體是 SUBLEQ，能訪問不同的寄存器或內(nèi)存位置。

SUBLEQ 的工作方式很簡單。它訪問內(nèi)存中的兩個寄存器，獲取它們內(nèi)容的差異并將其存儲回其中一個寄存器，然后如果結(jié)果為負，則它跳轉(zhuǎn)到不同的預(yù)定義代碼行，或者繼續(xù)執(zhí)行當前行的下一條指令。為執(zhí)行 SUBLEQ 程序而構(gòu)建的計算機稱為單指令集計算機，并且是通用計算機，即，如果可以訪問無限內(nèi)存，它就是圖靈完備的。

下面描述了一個循環(huán) Transformer 的構(gòu)造，它可以執(zhí)行用特定指令集編寫的程序。Transformer 跟蹤代碼行、內(nèi)存位置和程序計數(shù)器，將輸入的內(nèi)存部分用作內(nèi)存寄存器，將命令部分用作代碼行 / 指令。暫存器用于記錄每條指令涉及的加法和指針，讀、寫、條件分支操作等。

圖 5: 實現(xiàn) OISC 指令塊的圖形表示。前兩個塊將數(shù)據(jù) / 命令傳輸?shù)綍捍嫫?，第二個和第三個執(zhí)行減法并存儲結(jié)果，而最后一個執(zhí)行完成指令的 if goto 命令。

FLEQ：更靈活的基于 Attention 的計算機

在這一節(jié)，作者對 FLEQ 進行介紹，它是對 SUBLEQ 的推廣，它定義了一種更靈活的精簡指令集計算機。這一隱含的附加指令集基于 SUBLEQ 的更高級版本，允許在同一 Transformer 網(wǎng)絡(luò)中實現(xiàn)多種功能。作者使用術(shù)語 FLEQ 來指代指令、語言和它定義的基于 Attention 的計算機。

FLEQ 的設(shè)計允許通過生成比簡單減法更通用的函數(shù)來實現(xiàn)復(fù)雜的算法，如矩陣乘法、平方根計算、激活函數(shù)等。

 基于 Attention 的計算機執(zhí)行周期。在循環(huán) Transformer 的每次迭代中，根據(jù)程序計數(shù)器從輸入中的指令集中提取一條指令。然后指令被復(fù)制到暫存器。根據(jù)要實現(xiàn)的功能，使用不同的功能塊位置在局部記錄該功能的結(jié)果。一旦計算出結(jié)果，它就會被復(fù)制回指令提供的指定內(nèi)存位置。

執(zhí)行周期類似于上一節(jié)中的單指令集計算機 (OISC)，主要區(qū)別在于，對于每條指令，可以從預(yù)先選擇的函數(shù)列表中進行選擇，這些函數(shù)以任意數(shù)組的形式輸入，如矩陣、向量和標量。

輸入序列的格式。如圖 6 所示，循環(huán) Transformer 的輸入 X，它可以執(zhí)行一系列 FLEQ 指令構(gòu)成的程序（X 由暫存器、存儲器、指令等三部分構(gòu)成）。

基于 Transformer 的功能塊的格式。每個功能塊都位于輸入 X 的左下部分，如圖 6 所示。

圖 6: 執(zhí)行 FLEQ 命令的輸入 X 的結(jié)構(gòu)

更多技術(shù)細節(jié)和示例請參閱原論文。