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

(相關資料圖)

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

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

作者證明了恒定數量的編碼器層可以模擬基本計算塊。使用這些構建塊，他們模擬了一個小型指令集計算機。這使得他們能夠將迭代算法映射到可以由循環(huán)的（looped）13 層 transformer 執(zhí)行的程序。他們展示了這個 transformer 如何在其輸入的指導下模擬一個基本的計算器、一個基本的線性代數庫和使用反向傳播的 in-context learning 算法。這項工作突出了注意力機制的多功能性，并證明即使是淺層 transformer 也可以執(zhí)行成熟的通用程序。

論文概述

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

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

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

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

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

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

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

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

將 Transformer 作為可編程計算機

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

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

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

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

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

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

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

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

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

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

構建面向通用計算的 Transformer 模塊

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

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

位置編碼、程序計數器和數據指針

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

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

讀 / 寫：將數據、指令復制到暫存器或從暫存器取出

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

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

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

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

條件分支?

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

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

模擬通用單指令集計算機

SUBLEQ Transformer

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

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

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

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

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

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

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

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

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

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

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

圖 6: 執(zhí)行 FLEQ 命令的輸入 X 的結構

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

關鍵詞： 程序計數器 通用計算機 語言模型 線性代數 指定位置