「可解釋性」一直是深度學(xué)習(xí)難以繞過的難題，用戶無法理解模型的運(yùn)行機(jī)制，也就無法放心地把模型應(yīng)用到實(shí)際場景中。

最近來自蘇黎世聯(lián)邦理工學(xué)院和DeepMind的研究人員提出了一種全新的模型構(gòu)造工具Tracr，直接由人來根據(jù)「已知的機(jī)制」針對不同任務(wù)編寫代碼，然后由Tracr編譯成模型的權(quán)重，讓模型的解釋變得更輕松！

論文鏈接：?https://arxiv.org/pdf/2301.05062.pdf?

代碼鏈接：?https://github.com/deepmind/tracr?

Tracr的輸入為采用領(lǐng)域特定語言RASP編寫的代碼，輸出為標(biāo)準(zhǔn)的、僅包含解碼器的、類似GPT的Transformer結(jié)構(gòu)的權(quán)重。

在實(shí)驗(yàn)部分，研究人員使用Tracr創(chuàng)建了一系列g(shù)round truth Transformers，實(shí)現(xiàn)了包括計(jì)算token頻率、排序和Dyck-n括號檢查等程序。

如何解釋模型？

可解釋性是理解機(jī)器學(xué)習(xí)模型的重要手段，但由于模型的實(shí)際運(yùn)行方式仍然不清晰，所以目前的研究成果大多難以評估。

其中一個(gè)工作機(jī)理可解釋性（Mechanistic interpretability）嘗試對神經(jīng)網(wǎng)絡(luò)進(jìn)行逆向工程（reverse engineering）來對模型實(shí)現(xiàn)的算法給出一個(gè)機(jī)理上的解釋，并在一系列任務(wù)中取得進(jìn)展，包括用于圖像分類的卷積神經(jīng)網(wǎng)絡(luò)，Transformer語言模型等。

但該方法仍然存在相關(guān)工具少、對模型機(jī)制解釋不深、仍然需要研究人員進(jìn)行創(chuàng)造性的解釋等問題。

評估機(jī)理解釋的標(biāo)準(zhǔn)方法結(jié)合了許多臨時(shí)實(shí)驗(yàn)的證據(jù)，然而，由于這樣做的成本很高，許多方法只能在玩具模型中進(jìn)行評估，或在真實(shí)模型中的少數(shù)不重要的circuits上進(jìn)行評估。

Tracr的解決方案則是，通過將人類可讀的代碼「編譯」成神經(jīng)網(wǎng)絡(luò)的權(quán)重，直接解決缺乏基礎(chǔ)機(jī)理解釋的問題。

也就是說，Tracr實(shí)際上起到的就是類似編譯器的作用。

 Tracr中主要涉及三個(gè)組件：

1. RASP代碼

RASP即Restricted Access Sequence Processing Language，是2021年提出的一個(gè)用于表達(dá)Transformer計(jì)算的語言，可以作為計(jì)算模型用來描述Transformers，并配有相應(yīng)的解釋器運(yùn)行RASP代碼。

可以把RASP程序視為一個(gè)計(jì)算圖，當(dāng)根據(jù)給定的輸入token序列時(shí)，圖上的每個(gè)節(jié)點(diǎn)都會取一個(gè)特定值。

RASP語言中包括兩個(gè)基本的節(jié)點(diǎn)類型；序列操作（Sequence Operations， s-op），返回輸入值的token序列及indices序列；元素操作（Elementwise operations）；選擇-聚合操作等。

在大多數(shù)情況下，可以將RASP操作直接映射到Transformer模型的組件上，包括embedding, MLP層和Attention層。

2. 對RASP語言的修改

雖然RASP的操作可以直接映射到Transformers中，但仍然需要對RASP語言進(jìn)行修改以對模型權(quán)重進(jìn)行翻譯。

3. craft，Transformer的匯編語言

如果說RASP算是編譯的高級語言，那么craft就是匯編語言，它提供了比在純權(quán)重矩陣上操作更多的抽象性。

craft可以表示具有基維度（basis dimensions）的向量空間，及其相應(yīng)的操作，可以用基方向的標(biāo)簽來定義投影或其他線性操作。重要的是，craft抽象化了跟蹤權(quán)重矩陣中的填充的需要。

Tracr: Transformer編譯器

Tracr的代碼使用Python進(jìn)行編寫，并將RASP實(shí)現(xiàn)嵌入到Python中，從而可以在Python中直接編寫RASP程序，比較方便地對變量編碼（variable encoding）進(jìn)行標(biāo)注。

在Tracr中，RASP程序是一個(gè)數(shù)據(jù)結(jié)構(gòu)，通過向每個(gè)操作傳遞依賴關(guān)系來逐步構(gòu)建，同時(shí)對RASP程序做了一些基本的簡化。

Tracr將RASP程序翻譯成Transformer權(quán)重的過程主要包括六個(gè)步驟：

1. 構(gòu)建計(jì)算圖

追蹤整個(gè)程序以創(chuàng)建一個(gè)代表計(jì)算過程的有向圖。對于輸出的s-op，圖中包括表示tokens和indices的source節(jié)點(diǎn)，和代表輸出s-op的sink節(jié)點(diǎn)。

2. 推理s-op值

對于每個(gè)s-op，需要決定如何將其嵌入到剩余流中；為了使用類別編碼，需要知道一個(gè)s-op可以取哪些值。

因?yàn)橛?jì)算是確定的，基于有限的輸入詞匯量和上下文大小，所有節(jié)點(diǎn)都有一組有限的輸出值。

所以第二步的主要操作就是對圖進(jìn)行遍歷，并為每個(gè)節(jié)點(diǎn)標(biāo)注出其可能的輸出；標(biāo)注使用簡單的啟發(fā)式方法，確保找到的是s-op值集合的超集。

3. 獨(dú)立地翻譯s-ops

獨(dú)立考慮計(jì)算圖中的每個(gè)節(jié)點(diǎn)，并將其轉(zhuǎn)化為一個(gè)craft組件；元素操作翻譯為MLP塊，選擇-聚合操作翻譯為注意力塊。

使用人工設(shè)計(jì)的MLP和注意力模塊庫來近似數(shù)字和分類輸入和輸出的任意函數(shù)；將具有分類輸入和輸出的MLPs作為查找表使；帶有數(shù)字輸入和輸出的MLP使用基于通用函數(shù)近似定理的明確結(jié)構(gòu)。

對于注意層，把選擇器翻譯成??_{????}運(yùn)算符，把相應(yīng)的聚合運(yùn)算翻譯成??_{????}運(yùn)算符。

目前只支持對分類輸入的注意力。

4. 把組件分配到層上

為了構(gòu)建一個(gè)Transformer模型，需要將計(jì)算圖中的所有craft組件分配給各層。

理想情況下，目標(biāo)是找到最小的模型來進(jìn)行所需的計(jì)算，但一般可以將其表述為一個(gè)具有幾個(gè)約束條件的組合優(yōu)化問題：Transformer結(jié)構(gòu)具有交替的注意力和MLP層，而且所有相互依賴的計(jì)算都需要有正確的順序。

出于范圍上的考慮，可以用啟發(fā)式方法解決這個(gè)問題。

首先，計(jì)算出從輸入到一個(gè)給定節(jié)點(diǎn)的最長路徑，其路徑長度是可以分配給該節(jié)點(diǎn)的層數(shù)的一個(gè)上限；然后應(yīng)用額外的啟發(fā)式方法，將層與可以并行計(jì)算的塊結(jié)合起來。

這種方法可以返回一個(gè)正確但有時(shí)是次優(yōu)的層分配。

5. 構(gòu)造一個(gè)craft模型

直接對模型組件的輸入和輸出空間進(jìn)行求和作為構(gòu)建的殘差流空間（residual stream space ）。

換句話說，將每個(gè)s-op嵌入到它自己的正交子空間中，這個(gè)子空間在整個(gè)網(wǎng)絡(luò)中只保留給它使用。

然后按照層分配決定的順序遍歷計(jì)算圖，并將組件堆疊起來，得到一個(gè)用craft表示的完整Transformer。

6. 組裝Transformer權(quán)重

最后，將模型的craft表示轉(zhuǎn)換為具體的模型權(quán)重。

首先把并行的MLP層合并為一個(gè)層，再把并行的注意力頭合并為一個(gè)層。其中在注意層中，把??_{????}和??_{????}矩陣分成????、????、????、????重量矩陣。

然后調(diào)整所有權(quán)重的形狀，并將其連接到Transformer架構(gòu)中就可以推斷出模型配置（深度、層寬、殘差流大小等），以適應(yīng)所創(chuàng)建的元素。

只需重新實(shí)現(xiàn)步驟6，就可以直接將Tracr擴(kuò)展到支持任何其他Transformer的實(shí)現(xiàn)。

在可解釋性研究中的應(yīng)用

Tracr可以加速受控實(shí)驗(yàn)的驗(yàn)證過程，以測試關(guān)于Transformer的計(jì)算結(jié)構(gòu)的特定假設(shè)；通過這種方式，它也相當(dāng)于成了可解釋性研究的實(shí)驗(yàn)平臺。

研究人員對token計(jì)數(shù)、排序等例子編寫了RASP程序。

可解釋性工具的測試案例

編譯的模型可以很自然地作為測試「解釋忠實(shí)性」的基礎(chǔ)，并提供了一種偽造可解釋性技術(shù)給出的解釋的方法。

發(fā)展到最后，這些模型可以被用來建立可解釋性工具的測試案例庫，反過來又可以實(shí)現(xiàn)定量的評價(jià)指標(biāo)。

替換模型組件

評估對模型如何工作的理解的另一種方法是用手工編碼的組件替換模型的部分內(nèi)容。

例如有研究人員通過用他們自己理想化的實(shí)現(xiàn)來替換模型的組件以測試他們對Transformer如何實(shí)現(xiàn)模塊化加法的理解，結(jié)果發(fā)現(xiàn)該思路可以提高下游任務(wù)的性能，相當(dāng)于有力地證明了所提出的解釋是正確的。

雖然Tracr是將一個(gè)算法編譯成一個(gè)完整的Transformer模型，但也可以通過修改代碼調(diào)整為只編譯訓(xùn)練模型中的一部分，從而使得更容易評估對大型模型的理解。

理解模型現(xiàn)象和開發(fā)新技術(shù)

除了評估之外，編譯的模型可以作為研究circuits級別現(xiàn)象和開發(fā)解釋Transformer模型新方法的試驗(yàn)臺。