Uncovering Intermediate Variables in Transformers using Circuit Probing
使用電路探測揭示Transformer中的中間變量
https://arxiv.org/pdf/2311.04354
摘要
神經(jīng)網(wǎng)絡(luò)模型在各種復(fù)雜任務(wù)上取得了高性能，但它們實(shí)現(xiàn)的算法卻難以解釋。通常需要假設(shè)網(wǎng)絡(luò)計(jì)算中涉及的中間變量，以便理解這些算法。例如，語言模型在生成句子時是否依賴特定的句法屬性？然而，現(xiàn)有的分析工具使得測試此類假設(shè)變得困難。我們提出了一種新的分析技術(shù)——電路探測（circuit probing），它能夠自動發(fā)現(xiàn)計(jì)算假設(shè)中間變量的低級電路。這使得可以通過對模型參數(shù)的針對性消融進(jìn)行因果分析。
我們將這種方法應(yīng)用于在簡單算術(shù)任務(wù)上訓(xùn)練的模型，證明了它在以下三個方面的有效性：（1）解讀模型學(xué)習(xí)到的算法，（2）揭示模型內(nèi)的模塊化結(jié)構(gòu)，以及（3）跟蹤訓(xùn)練過程中電路的發(fā)展。在這三項(xiàng)實(shí)驗(yàn)中，我們證明了電路探測結(jié)合并擴(kuò)展了現(xiàn)有方法的能力，為各種分析提供了一種統(tǒng)一的方法。最后，我們在現(xiàn)實(shí)世界的用例中展示了電路探測：揭示了GPT2-Small和Medium中負(fù)責(zé)主謂一致和反身代詞的電路。
1 引言
Transformer模型是現(xiàn)代機(jī)器學(xué)習(xí)的主要工具，推動了從自然語言處理（NLP）（Devlin等人，2018；Radford等人，2019；Brown等人，2020）到計(jì)算機(jī)視覺（Dosovitskiy等人，2020）再到強(qiáng)化學(xué)習(xí)（Chen等人，2021）等不同子領(lǐng)域的突破。盡管取得了成功，但人們對它們學(xué)習(xí)實(shí)現(xiàn)的算法知之甚少。這個核心問題激發(fā)了一系列分析和可解釋性研究，試圖“打開黑箱”（Rogers等人，2021；Elhage等人，2021；Belikov，2022）。盡管付出了相當(dāng)大的努力，這些模型仍然幾乎完全不透明。
解釋一個在復(fù)雜任務(wù)上取得成功的模型所固有的一個挑戰(zhàn)是，研究人員通常沒有他們試圖揭示的算法的完整圖景。然而，他們可能能夠提出構(gòu)成此類算法的高級因果變量。例如，人們可能會直覺地認(rèn)為，計(jì)算句子中主語名詞的句法數(shù)量可能對語言建模有用（Chomsky，1965；Linzen等人，2016）。這個變量必須在解決語言建模任務(wù)的算法中具有因果作用，因?yàn)樗捎谝恢滦砸?guī)則（即主動詞的句法數(shù)量必須與主語的句法數(shù)量匹配）而限制了句子的其他部分。然而，它也為影響下一個標(biāo)記預(yù)測的其他變量留下了無限的可能性。雖然我們專注于語言建模，但這種討論更廣泛地適用于神經(jīng)網(wǎng)絡(luò)應(yīng)用的任何復(fù)雜領(lǐng)域，從視覺（Dosovitskiy等人，2020）到天文學(xué)（?iprijanovi?等人，2020）到蛋白質(zhì)折疊（Jumper等人，2021）。
我們提出電路探測（circuit probing）以便于研究Transformers（Vaswani等人，2017）中的中間變量。電路探測在模型權(quán)重上引入了一個可訓(xùn)練的二進(jìn)制掩碼，該掩碼被優(yōu)化以揭示（如果存在的話）計(jì)算高級中間變量的電路。這種技術(shù)結(jié)合了標(biāo)準(zhǔn)探測和因果分析方法的最佳特點(diǎn)，使研究人員能夠（1）測試模型是否表示高級中間變量，（2）測試它們是否在模型行為中具有因果作用（而不僅僅是從模型表示中解碼出來），以及（3）揭示計(jì)算這些變量的特定模型權(quán)重子集。我們首先使用簡單的算術(shù)任務(wù)來展示電路探測相對于現(xiàn)有方法的優(yōu)勢，表明它比現(xiàn)有方法更忠實(shí)于底層模型（即，它只提供支持模型實(shí)際表示的因果變量的證據(jù)），并可靠地揭示了在模型行為中具有因果作用的電路。然后，我們使用電路探測來分析GPT2-Small和Medium（Radford等人，2019）上的兩個句法現(xiàn)象，揭示了負(fù)責(zé)主謂一致和反身代詞一致性的特定電路。
2 電路探測
在嘗試解釋執(zhí)行某項(xiàng)任務(wù)的模型時，通常需要假設(shè)模型正在計(jì)算某個中間變量。電路探測試圖（1）測量該中間變量是否被計(jì)算，以及（2）識別負(fù)責(zé)計(jì)算它的模型組件（即電路）。
與之前的研究（Conmy等人，2023；Cao等人，2021）類似，我們嘗試通過優(yōu)化凍結(jié)模型權(quán)重上的二進(jìn)制掩碼來揭示模型組件。最近的研究表明，神經(jīng)網(wǎng)絡(luò)通常在子網(wǎng)絡(luò)層面上表現(xiàn)出結(jié)構(gòu)（Csordás等人，2020；Lepori等人，2023b；Hod等人，2021）。鑒于此，我們嘗試在單個注意力和MLP（多層感知機(jī)）塊內(nèi)揭示電路。這些層對殘差流產(chǎn)生加性更新（Elhage等人，2021）。直觀上，計(jì)算中間變量的電路應(yīng)該產(chǎn)生根據(jù)該變量劃分的輸出。
例如，如果一個電路正在計(jì)算主語名詞的句法數(shù)量，那么該電路應(yīng)該產(chǎn)生屬于兩個等價(jià)類的輸出，對應(yīng)于單數(shù)主語和復(fù)數(shù)主語。因此，我們優(yōu)化一個二進(jìn)制掩碼，以便——如果變量在特定層內(nèi)由模型計(jì)算——該層的輸出根據(jù)該變量進(jìn)行聚類。對此的自然優(yōu)化目標(biāo)是軟最近鄰損失（見附錄D），這是一種對比損失，它最小化來自同一類別的輸出之間的（余弦）距離，并最大化來自不同類別的輸出之間的（余弦）距離（Salakhutdinov & Hinton，2007；Frosst等人，2019）。圖1中展示了電路探測。
在這項(xiàng)工作中，我們在MLP和注意力塊的神經(jīng)元層面（即構(gòu)成Transformer中MLP和注意力塊的線性變換矩陣的列）上進(jìn)行掩碼。我們使用連續(xù)稀疏化（Savarese等人，2020），這是一種在訓(xùn)練過程中將軟掩碼退火成離散掩碼的剪枝技術(shù)，來學(xué)習(xí)二進(jìn)制掩碼（見附錄C）。我們還使用l0正則化來鼓勵稀疏的二進(jìn)制掩碼。如果這個過程成功，它將在模型組件內(nèi)產(chǎn)生一個計(jì)算假設(shè)中間變量的稀疏電路。見算法1的偽代碼。
電路探測評估：我們通過兩種方式評估電路探測：（1）我們在發(fā)現(xiàn)的電路產(chǎn)生的輸出向量上訓(xùn)練一個1-最近鄰分類器，然后在保留的數(shù)據(jù)上測試這個分類器。如果電路探測成功找到了計(jì)算特定中間變量的電路，那么它的輸出向量將根據(jù)該變量的可能標(biāo)簽進(jìn)行劃分，我們期望這個分類器能夠取得高性能。我們在這里采用了一種相當(dāng)保守的策略，隨機(jī)采樣每個標(biāo)簽的1個輸出向量來訓(xùn)練最近鄰分類器。見算法2的偽代碼。（2）我們消融發(fā)現(xiàn)的電路（即反轉(zhuǎn)學(xué)習(xí)的二進(jìn)制掩碼）并分析模型行為的變化。
因果關(guān)系：這種消融相當(dāng)于提出了一個反事實(shí)問題：“如果模型在特定塊中不計(jì)算特定中間變量z，模型的輸出將如何變化？”，假設(shè)電路僅在計(jì)算z時具有因果作用。我們注意到，這個假設(shè)不太可能完全正確，因?yàn)樯窠?jīng)元通常是多義性的（Elhage等人，2022）。然而，（1）我們試圖通過使用L0正則化來鼓勵最大稀疏子網(wǎng)絡(luò)來減輕這一點(diǎn)，以及（2）先前的研究表明，稀疏電路可能出人意料地是單義性的（Hamblin等人，2022）。
3 基線比較
電路探測實(shí)現(xiàn)了兩個不同的目標(biāo)：（1）它允許人們測量電路根據(jù)中間變量對輸入進(jìn)行劃分的程度（即它提供了探測準(zhǔn)確性測量），（2）它允許通過消融電路來進(jìn)行因果分析。先前的方法只實(shí)現(xiàn)了這些目標(biāo)中的一個，而不是兩者。因此，我們對每種類型的分析使用不同的基線。在分析探測準(zhǔn)確性時，我們將電路探測與“普通探測”（以下簡稱“探測”）和對比探測進(jìn)行比較。探測涉及學(xué)習(xí)一個分類器，從模型激活中解碼關(guān)于假設(shè)中間變量的信息（Tenney等人，2019；Hewitt & Manning，2019；Ettinger，2020；Li等人，2022；Nanda等人，2023）。先前的研究表明，探測常常錯誤地描述了模型執(zhí)行的底層計(jì)算（Hewitt & Liang，2019；Zhang & Bowman，2018；Voita & Titov，2020）。我們通過實(shí)驗(yàn)證明，電路探測比線性或非線性探測更忠實(shí)于模型的計(jì)算（見第4.2節(jié)、4.3節(jié)）。
對比探測是電路探測的一種消融方法，我們使用類似的對比目標(biāo)訓(xùn)練線性探測。我們在附錄R中展示了結(jié)果，并發(fā)現(xiàn)這種方法大多失敗。
在進(jìn)行因果分析時，我們將與三種現(xiàn)有方法進(jìn)行比較：因果抽象分析、健忘探測和反事實(shí)嵌入。因果抽象分析干預(yù)Transformer層產(chǎn)生的激活向量，將中間變量定位到特定的向量子空間（Geiger等人，2021；2023；Wu等人，2023）。然而，它需要假設(shè)一個完整的因果圖——一個高級描述，說明輸入如何映射到預(yù)測，包括中間變量之間的所有交互。這對于我們希望應(yīng)用神經(jīng)網(wǎng)絡(luò)的大多數(shù)現(xiàn)實(shí)世界任務(wù)來說是不可能的（例如，我們將在第4.4節(jié)中對語言建模應(yīng)用電路探測）。我們通過實(shí)驗(yàn)證明了無界分布式對齊搜索（無界DAS；Wu等人，2023），一種最先進(jìn)的因果抽象分析技術(shù)，的效用，并表明電路探測得出了相同的結(jié)果（見第4.1節(jié)、4.2節(jié)、4.3節(jié)）。健忘探測試圖抹去關(guān)于中間變量的線性可解碼信息，然后觀察這種抹除對下游行為的影響（Elazar等人，2021）。如果這種行為發(fā)生了變化，那么就意味著原始網(wǎng)絡(luò)內(nèi)存在中間變量。盡管已經(jīng)采用了多種技術(shù)來抹去線性可解碼信息（Ravfogel等人，2020；2022；Shao等人，2023），我們在我們的健忘探測實(shí)驗(yàn)中使用了最先進(jìn)的LEACE方法（Belrose等人，2023）。我們通過實(shí)驗(yàn)證明，電路探測比健忘探測更忠實(shí)于模型的計(jì)算（見第4.2節(jié)和4.3節(jié)）。最后，反事實(shí)嵌入（Tucker等人，2021）被引入以使用特定探測進(jìn)行因果分析。我們發(fā)現(xiàn)，在我們的實(shí)驗(yàn)中，反事實(shí)嵌入相當(dāng)沒有信息量（見第4.1節(jié)和4.2節(jié)）。
4 實(shí)驗(yàn)
我們展示了四個不同的實(shí)驗(yàn)，以說明電路探測可以幫助解決的問題的廣泛性。實(shí)驗(yàn)1、2和3研究了在簡單算術(shù)任務(wù)上訓(xùn)練的玩具模型，其中完整的因果圖很容易構(gòu)建。這些實(shí)驗(yàn)既（1）細(xì)化或再現(xiàn)了先前工作的結(jié)果，（2）將電路探測與現(xiàn)有的分析方法進(jìn)行了比較。在評估電路探測解碼中間變量的能力時，我們將其與線性和非線性探測進(jìn)行了比較，并在進(jìn)行因果分析時，將其與健忘探測、因果抽象分析和反事實(shí)嵌入進(jìn)行了比較。實(shí)驗(yàn)4將電路探測應(yīng)用于語言模型，以證明該方法可以擴(kuò)展到更現(xiàn)實(shí)的模型和設(shè)置。
4.1 實(shí)驗(yàn)1：解讀神經(jīng)網(wǎng)絡(luò)算法
目標(biāo)：可解釋性研究的一個核心目標(biāo)是描述模型實(shí)現(xiàn)的算法（Olah，2022）。當(dāng)我們能夠在兩個假設(shè)的替代方案之間進(jìn)行裁決時，這一崇高目標(biāo)就變得更加可行。我們展示了所有探測方法和大多數(shù)因果分析方法在描述在簡單算術(shù)任務(wù)上訓(xùn)練的模型實(shí)現(xiàn)的算法時產(chǎn)生了一致的結(jié)果。
4.2 實(shí)驗(yàn)2：中間變量的模塊化
目標(biāo)：我們現(xiàn)在應(yīng)用電路探測來分析Transformer模型的內(nèi)部組織，這是最近幾項(xiàng)研究的主題（Lepori等人，2023b；Csordás等人，2020；Hod等人，2021；Mittal等人，2022）。我們展示了電路探測可以用來描述計(jì)算是否以模塊化和可重用的方式在Transformer內(nèi)實(shí)現(xiàn)，而其他方法未能揭示這種結(jié)構(gòu)。
因果分析：因果抽象分析與電路探測和線性探測一致（見附錄K），因果地涉及每個任務(wù)的自由和共享變量。另一方面，從非線性探測器生成的反事實(shí)嵌入再次未能產(chǎn)生任何變量具有因果性的證據(jù)（見附錄L）。我們得出結(jié)論，反事實(shí)嵌入在我們的訓(xùn)練非線性探測器中更多地表現(xiàn)為對抗性示例，而不是模型的有意義的反事實(shí)輸入。接下來，我們進(jìn)行因果分析以了解模型是否表現(xiàn)出模塊化。我們預(yù)期，消融計(jì)算任一任務(wù)共享變量的電路應(yīng)該破壞兩個任務(wù)的性能。另一方面，我們預(yù)期，消融任務(wù)1的自由變量應(yīng)該破壞任務(wù)1的性能，同時對任務(wù)2的性能影響較?。ǚ粗嗳唬?。在附錄M中，我們分析了電路探測返回的兩個自由變量電路的形態(tài)，發(fā)現(xiàn)電路僅在一個注意力塊內(nèi)的張量中有所不同。因此，我們只消融該張量內(nèi)的電路權(quán)重。從表1中，我們看到消融計(jì)算任一任務(wù)共享變量的電路破壞了兩個任務(wù)的性能。另一方面，消融任務(wù)1的自由變量破壞了任務(wù)1的性能，同時保持了任務(wù)2的一些性能。我們觀察到任務(wù)2的自由變量也有相同的趨勢。我們使用健忘探測進(jìn)行類似的分析，發(fā)現(xiàn)健忘探測沒有揭示這種內(nèi)部結(jié)構(gòu)。我們在表1中展示了在注意力塊之后的殘差流上進(jìn)行健忘探測的結(jié)果?？偟膩碚f，我們的結(jié)果表明，電路探測在描述模型如何構(gòu)建其計(jì)算方面優(yōu)于競爭技術(shù)。
4.3 實(shí)驗(yàn)3：電路探測作為進(jìn)度測量
目標(biāo)：電路探測使我們能夠在中間變量的層面上深入了解Transformer的訓(xùn)練動態(tài)。最近的研究表明，模型可能在過擬合訓(xùn)練數(shù)據(jù)后很長時間突然學(xué)會泛化（Power等人，2022）。盡管這種從過擬合到泛化的相當(dāng)不連續(xù)的轉(zhuǎn)換（通常稱為“頓悟”），Nanda等人（2022）揭示了計(jì)算可泛化算法的電路是在訓(xùn)練過程中連續(xù)形成的。雖然他們的工作需要逆向工程整個算法以獲得這種對電路形成的洞察，但我們使用電路探測（只需要我們假設(shè)一個高級中間變量）在略有不同的環(huán)境中重現(xiàn)了他們的發(fā)現(xiàn)。
4.4 實(shí)驗(yàn)4：語言模型中的電路探測
目標(biāo)：先前的實(shí)驗(yàn)集中在可以指定完整因果圖的玩具任務(wù)上。然而，我們開發(fā)可解釋性工具的興趣在于分析在實(shí)際中使用的模型，這些模型用于無法構(gòu)建因果圖的任務(wù)。在這里，我們使用電路探測來研究預(yù)訓(xùn)練的GPT2-Small和GPT2-Medium如何執(zhí)行語言建模。特別是，我們調(diào)查了兩種依賴于句法數(shù)的語言學(xué)依賴關(guān)系：主謂一致和反身代詞。
主謂一致性是指英語中句子的主語必須與句子的主要動詞在句法數(shù)上保持一致的現(xiàn)象。例如：“The keys are on the table”是符合語法的，而“The keys is on the table”則不符合語法。我們假設(shè)在預(yù)測句子的主要動詞時，會計(jì)算代表主語名詞句法數(shù)的中間變量。反身代詞確保反身代詞與其指代對象在數(shù)上保持一致。例如：“the consultants injured themselves”是符合語法的，而“the consultants injured herself”則不符合語法。我們假設(shè)在預(yù)測反身代詞時，會計(jì)算代表指代對象句法數(shù)的中間變量。
任務(wù)：我們使用Marvin & Linzen（2018）的模板生成句子前綴，這些前綴的延續(xù)可能是一個主要動詞（研究主謂一致性時）或一個反身代詞（研究反身代詞時）。見附錄S中的例子前綴。我們在包含一個干擾項(xiàng)的句子上運(yùn)行電路探測，這個干擾項(xiàng)是一個非主語或非指代名詞（例如“the keys to the cabinet are on the table”中的“cabinet”）。對于每種現(xiàn)象，我們在句子前綴的最后一個標(biāo)記上運(yùn)行電路探測，以揭示計(jì)算主語名詞或指代對象句法數(shù)的電路。然后我們消融發(fā)現(xiàn)的電路，并評估模型繼續(xù)保留句子前綴語法能力。具體來說，我們評估模型是否更有可能預(yù)測與主語/指代對象句法數(shù)一致或不一致的標(biāo)記。對于主謂一致性，我們檢查“is”和“are”的logits——如果當(dāng)主語是單數(shù)時（例如“The officer...”），“is”的logit高于“are”的logit，那么我們認(rèn)為模型在該句子前綴上成功了。對于反身代詞，我們運(yùn)行兩次分析，一次比較“herself”和“themselves”的logits，另一次比較“himself”和“themselves”的logits。我們在包含一個干擾名詞的IID句子和包含兩個干擾名詞的OOD數(shù)據(jù)集上評估模型（見附錄S中的例子）。我們假設(shè)相同的電路計(jì)算兩種句子結(jié)構(gòu)的相關(guān)語言依賴性。即使我們從這個分析中恢復(fù)了積極的結(jié)果，也有可能我們只是在破壞整個模型，而不是消融一個專門的電路。作為對照，我們從電路中的神經(jīng)元的補(bǔ)集集合中隨機(jī)采樣5個隨機(jī)子網(wǎng)絡(luò)，并重新運(yùn)行消融分析。隨機(jī)采樣的子網(wǎng)絡(luò)總是包含與我們的電路相同數(shù)量的神經(jīng)元。
探測：見附錄T.1和U中對電路探測準(zhǔn)確性的調(diào)查。通常，我們發(fā)現(xiàn)大多數(shù)注意力層可以計(jì)算正確的句法數(shù)，但MLP只有在GPT2-Small和Medium的中間層才開始表現(xiàn)良好。見附錄Y中所有GPT2設(shè)置的線性探測準(zhǔn)確性。
因果分析：對于兩種句法依賴性，我們預(yù)期消融發(fā)現(xiàn)的電路將使模型在區(qū)分主語/指代對象的句法數(shù)方面變得更差。我們預(yù)期消融隨機(jī)子網(wǎng)絡(luò)不應(yīng)該損害模型在任何數(shù)據(jù)集上的性能。我們在正文中展示了GPT2-Small的結(jié)果，并在附錄U中展示了GPT2-Medium的結(jié)果。對于這兩種現(xiàn)象，我們發(fā)現(xiàn)依賴性是在第6層的注意力塊中計(jì)算的。消融電路探測返回的電路顯著降低了IID和OOD數(shù)據(jù)集上兩種現(xiàn)象的性能，而消融隨機(jī)子網(wǎng)絡(luò)則不影響模型性能（見圖4）。其他塊沒有表現(xiàn)出這種特征模式（見附錄T.2）。這與先前的工作一致，表明句法依賴性在Transformer的中間層中表示（Tenney等人，2019；Vig & Belinkov，2019）。見附錄V中的電路重疊分析，附錄W中的定性結(jié)果，以及附錄X中的結(jié)果，證明模型中較早計(jì)算單個標(biāo)記的句法數(shù)。
5 討論
相關(guān)工作：電路探測與近期在機(jī)制解釋性方面的努力相關(guān)——這是一個蓬勃發(fā)展的領(lǐng)域，試圖逆向工程神經(jīng)網(wǎng)絡(luò)算法。通過大量的人工努力，研究人員已經(jīng)揭示了玩具模型（Olsson等人，2022；Nanda等人，2022；Chughtai等人，2023）以及更現(xiàn)實(shí)模型（Wang等人，2022；Hanna等人，2023；Merullo等人，2023）正在實(shí)現(xiàn)的算法。更廣泛地說，已經(jīng)有很多工作分析了語言模型的句法（Linzen & Baroni，2021；Goldberg，2019；Tenney等人，2018；McCoy等人，2018）和語義能力（Pavlick，2022；2023；Yu & Ettinger，2020；Hupkes等人，2020；Dziri等人，2023）。電路探測還與試圖將模型分解為功能性子網(wǎng)絡(luò)的工作相關(guān)（Csordás等人，2020；Hamblin等人，2022；Lepori等人，2023b；Zhang等人，2021；Panigrahi等人，2023；Hod等人，2021；Cao等人，2021）。電路探測的成功進(jìn)一步證明了子網(wǎng)絡(luò)是分析模型的有用視角。
結(jié)論：我們介紹了電路探測，這是一種揭示計(jì)算高級中間變量的低級電路的新方法。通過四個實(shí)驗(yàn)，我們已經(jīng)展示了如何深入了解模型正在實(shí)現(xiàn)的底層算法，這些算法在模型中的結(jié)構(gòu)，以及它們在整個訓(xùn)練過程中如何發(fā)展。電路探測結(jié)合并擴(kuò)展了現(xiàn)有方法的能力，并在多個設(shè)置中超越了它們。然而，目前尚不清楚在一個給定塊內(nèi)多個電路如何組合以創(chuàng)建對殘差流的一個加性更新，因此無法替換單個變量以執(zhí)行反事實(shí)干預(yù)。未來的工作可能會尋求理解電路如何相互組合以實(shí)現(xiàn)這一目的。
原文鏈接： https://arxiv.org/pdf/2311.04354
熱門跟貼