經(jīng)逆向工程，Transformer「翻譯」成數(shù)學(xué)框架 | 25位學(xué)者撰文
Transformer 是 Google 團隊在 2017 年 6 月提出的 NLP 經(jīng)典之作，由 Ashish Vaswani 等人在論文《 Attention Is All You Need 》中提出。自 Transformer 出現(xiàn)以來，便在 NLP、CV、語音、生物、化學(xué)等領(lǐng)域引起了諸多進展。
Transformer 在現(xiàn)實世界中的應(yīng)用越來越廣泛，例如 GPT-3 、LaMDA 、Codex 等都是基于 Transformer 架構(gòu)構(gòu)建的。然而，隨著基于 Transformer 模型的擴展，其開放性和高容量為意想不到的甚至有害的行為創(chuàng)造了越來越大的空間。即使在大型模型訓(xùn)練完成數(shù)年后，創(chuàng)建者和用戶也會經(jīng)常發(fā)現(xiàn)以前從來沒見過的模型問題。
解決這些問題的一個途徑是機械的可解釋性（mechanistic interpretability），即對 transformers 計算過程進行逆向工程，這有點類似于程序員如何嘗試將復(fù)雜的二進制文件逆向工程為人類可讀的源代碼。
如果逆向工程可行，那么我們就會有更系統(tǒng)的方法來解釋當(dāng)前模型的安全問題、識別問題，甚至可能預(yù)見未來尚未構(gòu)建的模型安全問題。這有點類似于將 Transformer 的黑箱操作進行逆向，讓這一過程變得清晰可見。之前有研究者開發(fā)了 Distill Circuits thread 項目，曾嘗試對視覺模型進行逆向工程，但到目前為止還沒有可比的 transformer 或語言模型進行逆向工程研究。
在本文中，由 25 位研究者參與撰寫的論文，嘗試采用最原始的步驟逆向 transformer。該論文由 Chris Olah 起草，Chris Olah 任職于 Anthropic 人工智能安全和研究公司，主要從事逆向工程神經(jīng)網(wǎng)絡(luò)研究。之后 Neel Nanda 對論文初稿進行了重大修改，Nanda 目前是 DeepMind 的一名研究工程實習(xí)生。Nelson Elhage 對論文進行了詳細(xì)的編輯以提高論文章節(jié)清晰度，Nelson Elhage 曾任職于 Stripe 科技公司。
左：Neel Nanda；右：Christopher Olah
考慮到語言模型的復(fù)雜性高和規(guī)模大等特點，該研究發(fā)現(xiàn)，從最簡單的模型開始逆向 transformer 最有效果。該研究旨在發(fā)現(xiàn)簡單算法模式、主題（motifs）或是框架，然后將其應(yīng)用于更復(fù)雜、更大的模型。具體來說，他們的研究范圍僅包括只有注意力塊的兩層或更少層的 transformer 模型。這與 GPT-3 這樣的 transformer 模型形成鮮明的對比，GPT-3 層數(shù)多達(dá) 96 層。
論文地址：https://transformer-circuits.pub/2021/framework/index.html#acknowledgments
該研究發(fā)現(xiàn)，通過以一種新的但數(shù)學(xué)上等效的方式概念化 transformer 操作，我們能夠理解這些小模型并深入了解它們的內(nèi)部運作方式。值得注意的是，研究發(fā)現(xiàn)特定的注意頭，本文稱之為歸納頭（induction heads），可以在這些小模型中解釋上下文學(xué)習(xí)，而且這些注意力頭只在至少有兩個注意層的模型中發(fā)展。此外，該研究還介紹了這些注意力頭對特定數(shù)據(jù)進行操作的一些示例。
各章節(jié)內(nèi)容概覽
為了探索逆向工程 transformers 面臨哪些挑戰(zhàn)，研究者對幾個 attention-only 的 toy 模型進行了逆向功能。
首先是零層 transformers 模型的二元統(tǒng)計。研究者發(fā)現(xiàn)，二元表可以直接通過權(quán)重訪問。
在討論更復(fù)雜的模型之前，考慮零層（zero layer）transformer 很有用。這類模型接受一個 token，嵌入，再取消嵌入，以生成預(yù)測下一個 token 的 logits
由于這類模型無法從其他 tokens 傳輸信息，因此只能從當(dāng)前 token 預(yù)測下一個 token。這意味著，W_UW_E 的最優(yōu)行為是近似二元對數(shù)似然。
零層 attention-only transformers 模型。
其次，單層 attention-only transformers 是二元和 skip 三元模型的集合。同零層 transformers 一樣，二元和 skip 三元表可以直接通過權(quán)重訪問，無需運行模型。這些 skip 三元模型的表達(dá)能力驚人，包括實現(xiàn)一種非常簡單的上下文內(nèi)學(xué)習(xí)。
對于單層 attention-only transformers 模型，有哪些路徑擴展（path expansion）技巧呢？研究者提供了一些。
如下圖所示，單層 attention-only transformers 由一個 token 嵌入組成，后接一個注意力層（單獨應(yīng)用注意力頭），最后是解除嵌入：
使用之前得到的張量標(biāo)記（tensor notation）和注意力頭的替代表征，研究者可以將 transformer 表征為三個項的乘積，具體如下圖所示：
研究者采用的核心技巧是簡單地擴展乘積，即將乘積（每個項對應(yīng)一個層）轉(zhuǎn)換為一個和，其中每個項對應(yīng)一個端到端路徑。他們表示，每個端到端路徑項都易于理解，可以獨立地進行推理，并能夠疊加組合創(chuàng)建模型行為。
最后，兩層 attention-only transformers 模型可以使用注意力頭組合實現(xiàn)復(fù)雜得多的算法。這些組合算法也可以直接通過權(quán)重檢測出來。需要注意的是，兩層模型適應(yīng)注意力頭組合創(chuàng)建「歸納頭」（induction heads），這是一種非常通用的上下文內(nèi)學(xué)習(xí)算法。
具體地，當(dāng)注意力頭有以下三種組合選擇：
- Q - 組合：W_Q 在一個受前面頭影響的子空間中讀取；
- K - 組合：W_K 在一個受前面頭影響的子空間中讀??；
- V - 組合：W_V 在一個受前面頭影響的子空間中讀取。
研究者表示，Q - 和 K - 組合與 V - 組合截然不同。前兩者都對注意力模式產(chǎn)生影響，允許注意力頭表達(dá)復(fù)雜得多的模式。而 V - 組合對一個注意力頭專注于某個給定位置時所要傳輸?shù)男畔a(chǎn)生影響。結(jié)果是，V - 組合頭變現(xiàn)得更像一個單一單元，并可以考慮用來創(chuàng)建額外的「虛擬注意力頭」。
對于 transformer 有一個最基礎(chǔ)的問題，即「如何計算 logits」？與單層模型使用的方法一樣，研究者寫出了一個乘積，其中每個項在模型中都是一個層，并擴展以創(chuàng)建一個和，其中每個項在模型中都是一個端到端路徑。
其中，直接路徑項和單個頭項與單層模型中的相同。最后的「虛擬注意力頭」項對應(yīng)于 V - 組合。虛擬注意力頭在概念上非常有趣，但在實踐中，研究者發(fā)現(xiàn)它們往往無法在小規(guī)模的兩層模型中發(fā)揮重大作用。
此外，這些項中的每一個都對應(yīng)于模型可以實現(xiàn)更復(fù)雜注意力模式的一種方式。在理論上，很難對它們進行推理。但當(dāng)討論到歸納頭時，會很快在具體實例中用到它們。
