Transformer 模型是 AI 系統(tǒng)的基礎(chǔ)。已經(jīng)有了數(shù)不清的關(guān)于 "Transformer 如何工作" 的核心結(jié)構(gòu)圖表。

但是這些圖表沒有提供任何直觀的計(jì)算該模型的框架表示。當(dāng)研究者對(duì)于 Transformer 如何工作抱有興趣時(shí)，直觀的獲取他運(yùn)行的機(jī)制變得十分有用。

Thinking Like Transformers 這篇論文中提出了 transformer 類的計(jì)算框架，這個(gè)框架直接計(jì)算和模仿 Transformer 計(jì)算。使用 RASP 編程語言，使每個(gè)程序編譯成一個(gè)特殊的 Transformer。

在這篇博客中，我用 python 復(fù)現(xiàn)了 RASP 的變體 (RASPy)。該語言大致與原始版本相當(dāng)，但是多了一些我認(rèn)為很有趣的變化。通過這些語言，作者 Gail Weiss 的工作，提供了一套具有挑戰(zhàn)性的有趣且正確的方式可以幫助了解其工作原理。

在說起語言本身前，讓我們先看一個(gè)例子，看看用 Transformers 編碼是什么樣的。這是一些計(jì)算翻轉(zhuǎn)的代碼，即反向輸入序列。代碼本身用兩個(gè) Transformer 層應(yīng)用 attention 和數(shù)學(xué)計(jì)算到達(dá)這個(gè)結(jié)果。

本文內(nèi)容目錄

• 部分一：Transformers 作為代碼

• 部分二：用 Transformers 編寫程序

Transformers 作為代碼

我們的目標(biāo)是定義一套計(jì)算形式來最小化 Transformers 的表達(dá)。我們將通過類比，描述每個(gè)語言構(gòu)造及其在 Transformers 中的對(duì)應(yīng)。(正式語言規(guī)范請(qǐng)?jiān)诒疚牡撞坎榭凑撐娜逆溄?。

這個(gè)語言的核心單元是將一個(gè)序列轉(zhuǎn)換成相同長(zhǎng)度的另一個(gè)序列的序列操作。我后面將其稱之為 transforms。

輸入

在一個(gè) Transformer 中，基本層是一個(gè)模型的前饋輸入。這個(gè)輸入通常包含原始的 token 和位置信息。

在代碼中，tokens 的特征表示最簡(jiǎn)單的 transform，它返回經(jīng)過模型的 tokens，默認(rèn)輸入序列是 "hello":

tokens


如果我們想要改變 transform 里的輸入，我們使用輸入方法進(jìn)行傳值。

作為 Transformers，我們不能直接接受這些序列的位置。但是為了模擬位置嵌入，我們可以獲取位置的索引:

前饋網(wǎng)絡(luò)

經(jīng)過輸入層后，我們到達(dá)了前饋網(wǎng)絡(luò)層。在 Transformer 中，這一步可以對(duì)于序列的每一個(gè)元素獨(dú)立的應(yīng)用數(shù)學(xué)運(yùn)算。

在代碼中，我們通過在 transforms 上計(jì)算表示這一步。在每一個(gè)序列的元素中都會(huì)進(jìn)行獨(dú)立的數(shù)學(xué)運(yùn)算。

結(jié)果是一個(gè)新的 transform，一旦重構(gòu)新的輸入就會(huì)按照重構(gòu)方式計(jì)算:

該運(yùn)算可以組合多個(gè) Transforms，舉個(gè)例子，以上述的 token 和 indices 為例，這里可以類別 Transformer 可以跟蹤多個(gè)片段信息:

我們提供了一些輔助函數(shù)讓寫 transforms 變得更簡(jiǎn)單，舉例來說，where?提供了一個(gè)類似?if?功能的結(jié)構(gòu)。

map?使我們可以定義自己的操作，例如一個(gè)字符串以?int?轉(zhuǎn)換。（用戶應(yīng)謹(jǐn)慎使用可以使用的簡(jiǎn)單神經(jīng)網(wǎng)絡(luò)計(jì)算的操作）

函數(shù) (functions) 可以容易的描述這些 transforms 的級(jí)聯(lián)。舉例來說，下面是應(yīng)用了 where 和 atoi 和加 2 的操作

注意力篩選器

到開始應(yīng)用注意力機(jī)制事情就變得開始有趣起來了。這將允許序列間的不同元素進(jìn)行信息交換。

我們開始定義 key 和 query 的概念，Keys 和 Queries 可以直接從上面的 transforms 創(chuàng)建。舉個(gè)例子，如果我們想要定義一個(gè) key 我們稱作?key。

對(duì)于?query?也一樣

標(biāo)量可以作為?key?或?query?使用，他們會(huì)廣播到基礎(chǔ)序列的長(zhǎng)度。

我們創(chuàng)建了篩選器來應(yīng)用 key 和 query 之間的操作。這對(duì)應(yīng)于一個(gè)二進(jìn)制矩陣，指示每個(gè) query 要關(guān)注哪個(gè) key。與 Transformers 不同，這個(gè)注意力矩陣未加入權(quán)重。

一些例子：

• 選擇器的匹配位置偏移 1:

• key 早于 query 的選擇器:

• key 晚于 query 的選擇器:

  


選擇器可以通過布爾操作合并。比如，這個(gè)選擇器將 before 和 eq 做合并，我們通過在矩陣中包含一對(duì)鍵和值來顯示這一點(diǎn)。

使用注意力機(jī)制

給一個(gè)注意力選擇器，我們可以提供一個(gè)序列值做聚合操作。我們通過累加那些選擇器選過的真值做聚合。

(請(qǐng)注意：在原始論文中，他們使用一個(gè)平均聚合操作并且展示了一個(gè)巧妙的結(jié)構(gòu)，其中平均聚合能夠代表總和計(jì)算。RASPy 默認(rèn)情況下使用累加來使其簡(jiǎn)單化并避免碎片化。實(shí)際上，這意味著 raspy 可能低估了所需要的層數(shù)。基于平均值的模型可能需要這個(gè)層數(shù)的兩倍)

注意聚合操作使我們能夠計(jì)算直方圖之類的功能。

(key(tokens)?==?query(tokens)).value(1)


視覺上我們遵循圖表結(jié)構(gòu)，Query 在左邊，Key 在上邊，Value 在下面，輸出在右邊

一些注意力機(jī)制操作甚至不需要用到輸入 token 。舉例來說，去計(jì)算序列長(zhǎng)度，我們創(chuàng)建一個(gè) " select all " 的注意力篩選器并且給他賦值。

這里有更多復(fù)雜的例子，下面將一步一步展示。(這有點(diǎn)像做采訪一樣)

我們想要計(jì)算一個(gè)序列的相鄰值的和，首先我們向前截?cái)?

然后我們向后截?cái)?

兩者相交:

最終聚合:

這里有個(gè)可以計(jì)算累計(jì)求和的例子，我們這里引入一個(gè)給 transform 命名的能力來幫助你調(diào)試。

層

這個(gè)語言支持編譯更加復(fù)雜的 transforms。他同時(shí)通過跟蹤每一個(gè)運(yùn)算操作計(jì)算層。

這里有個(gè) 2 層 transform 的例子，第一個(gè)對(duì)應(yīng)于計(jì)算長(zhǎng)度，第二個(gè)對(duì)應(yīng)于累積總和。

用 transformers 進(jìn)行編程

使用這個(gè)函數(shù)庫，我們可以編寫完成一個(gè)復(fù)雜任務(wù)，Gail Weiss 給過我一個(gè)極其挑戰(zhàn)的問題來打破這個(gè)步驟，我們可以加載一個(gè)添加任意長(zhǎng)度數(shù)字的 Transformer 嗎？

例如:?給一個(gè)字符串 "19492+23919", 我們可以加載正確的輸出嗎？

如果你想自己嘗試，我們提供了一個(gè)版本你可以自己試試:
https://colab.research.google.com/github/srush/raspy/blob/main/Blog.ipynb

挑戰(zhàn)一 : 選擇一個(gè)給定的索引

加載一個(gè)在索引 i 處全元素都有值的序列

挑戰(zhàn)二 ：轉(zhuǎn)換

通過 i 位置將所有 token 移動(dòng)到右側(cè)。

挑戰(zhàn)三 ：最小化

計(jì)算序列的最小值。(這一步開始變得困難，我們版本用了 2 層注意力機(jī)制)

挑戰(zhàn)四：第一索引

計(jì)算有 token q 的第一索引 (2 層)

挑戰(zhàn)五 ：右對(duì)齊

右對(duì)齊一個(gè)填充序列。例："ralign().inputs('xyz___') ='—xyz'" (2 層)

挑戰(zhàn)六：分離

把一個(gè)序列在 token "v" 處分離成兩部分然后右對(duì)齊 (2 層):

挑戰(zhàn)七：滑動(dòng)

將特殊 token "<" 替換為最接近的 "<" value (2 層):

挑戰(zhàn)八：增加

你要執(zhí)行兩個(gè)數(shù)字的添加。這是步驟。

1. 分成兩部分。轉(zhuǎn)制成整形。加入

“683+345” => [0, 0, 0, 9, 12, 8]

1. 計(jì)算攜帶條款。三種可能性：1 個(gè)攜帶，0 不攜帶，< 也許有攜帶。

[0, 0, 0, 9, 12, 8] => “00<100”

1. 滑動(dòng)進(jìn)位系數(shù)

“00<100” => 001100"

1. 完成加法

這些都是 1 行代碼。完整的系統(tǒng)是 6 個(gè)注意力機(jī)制。（盡管 Gail 說，如果你足夠細(xì)心則可以在 5 個(gè)中完成?。?。

完美搞定！

參考資料 & 文內(nèi)鏈接:

• 如果你對(duì)這個(gè)主題感興趣想了解更多，請(qǐng)查看論文：Thinking Like Transformers
  https://arxiv.org/pdf/2106.06981.pdf

• 以及了解更多 RASP 語言
  https://github.com/tech-srl/RASP

• 如果你對(duì)「形式語言和神經(jīng)網(wǎng)絡(luò)」(FLaNN) 感興趣或者有認(rèn)識(shí)感興趣的人，歡迎邀請(qǐng)他們加入我們的線上社區(qū)！
  https://flann.super.site

• 本篇博文，包含庫、Notebook 和博文的內(nèi)容在
  https://github.com/srush/raspy

• 本博客文章由 Sasha Rush 和 Gail Weiss 共同編寫:Sasha Rush 博客地址
  https://rush-nlp.com/Gail Weiss 博客地址
  https://gailweiss.github.io/

英文原文: Thinking Like Transformers:?
https://srush.github.io/raspy/

中文譯者:?innovation64 (李洋)