公理訓(xùn)練讓LLM學(xué)會(huì)因果推理:6700萬(wàn)參數(shù)模型比肩萬(wàn)億參數(shù)級(jí)GPT-4
AI 已經(jīng)在幫助數(shù)學(xué)家和科學(xué)家做研究了,比如著名數(shù)學(xué)家陶哲軒就曾多次分享自己借助 GPT 等 AI 工具研究探索的經(jīng)歷。AI 要在這些領(lǐng)域大戰(zhàn)拳腳,強(qiáng)大可靠的因果推理能力是必不可少的。
本文要介紹的這項(xiàng)研究發(fā)現(xiàn):在小圖譜的因果傳遞性公理演示上訓(xùn)練的 Transformer 模型可以泛化用于大圖譜的傳遞性公理。
也就是說(shuō),如果讓 Transformer 學(xué)會(huì)執(zhí)行簡(jiǎn)單的因果推理,就可能將其用于更為復(fù)雜的因果推理。該團(tuán)隊(duì)提出的公理訓(xùn)練框架是一種基于被動(dòng)數(shù)據(jù)來(lái)學(xué)習(xí)因果推理的新范式,只有演示足夠就能用于學(xué)習(xí)任意公理。
引言
因果推理(causal reasoning)可以定義成一組推理流程并且這組推理流程要符合專門針對(duì)因果性的預(yù)定義公理或規(guī)則。舉個(gè)例子,d-separation(有向分離)和 do-calculus 規(guī)則可被視為公理,而 collider set 或 backdoor set 的規(guī)范則可被看作是由公理推導(dǎo)出的規(guī)則。
通常來(lái)說(shuō),因果推理使用的數(shù)據(jù)對(duì)應(yīng)于一個(gè)系統(tǒng)中的變量。通過(guò)正則化、模型架構(gòu)或特定的變量選擇,可以歸納偏置的形式將公理或規(guī)則集成到機(jī)器學(xué)習(xí)模型中。
根據(jù)可用數(shù)據(jù)種類的差異(觀察數(shù)據(jù)、干預(yù)數(shù)據(jù)、反事實(shí)數(shù)據(jù)),Judea Pearl 提出的「因果階梯」定義了因果推理的可能類型。
由于公理是因果性的基石,因此我們不禁會(huì)想是否可以直接使用機(jī)器學(xué)習(xí)模型來(lái)學(xué)習(xí)公理。也就是說(shuō),如果學(xué)習(xí)公理的方式不是學(xué)習(xí)通過(guò)某個(gè)數(shù)據(jù)生成流程得到的數(shù)據(jù),而是直接學(xué)習(xí)公理的符號(hào)演示(并由此學(xué)習(xí)因果推理),哪又會(huì)如何呢?
相較于使用特定的數(shù)據(jù)分布構(gòu)建的針對(duì)特定任務(wù)的因果模型,這樣的模型有一個(gè)優(yōu)勢(shì):其可在多種不同的下游場(chǎng)景中實(shí)現(xiàn)因果推理。隨著語(yǔ)言模型具備了學(xué)習(xí)以自然語(yǔ)言表達(dá)的符號(hào)數(shù)據(jù)的能力,這個(gè)問(wèn)題也就變得非常重要了。
事實(shí)上,近期已有一些研究通過(guò)創(chuàng)建以自然語(yǔ)言編碼因果推理問(wèn)題的基準(zhǔn),評(píng)估了大型語(yǔ)言模型(LLM)是否能夠執(zhí)行因果推理。
微軟、MIT 和印度理工學(xué)院海得拉巴分校(IIT Hyderabad)的研究團(tuán)隊(duì)也朝這個(gè)方向邁出了重要一步:提出了一種通過(guò)公理訓(xùn)練(axiomatic training)學(xué)習(xí)因果推理的方法。
- 論文標(biāo)題:Teaching Transformers Causal Reasoning through Axiomatic Training
- 論文地址:https://arxiv.org/pdf/2407.07612
公理訓(xùn)練
他們假設(shè),可將因果公理表示成以下符號(hào)元組 ?premise, hypothesis, result?。其中 hypothesis 是指假設(shè),即因果陳述;premise 是前提,是指用于確定該陳述是否為「真」的任意相關(guān)信息;result 自然就是結(jié)果了。結(jié)果可以是簡(jiǎn)單的「是」或「否」。
舉個(gè)例子,來(lái)自論文《Can large language models infer causation from correlation?》的 collider 公理可以表示成:,結(jié)論就為「是」。
基于這個(gè)模板,可通過(guò)修改變量名稱、變量數(shù)量和變量順序等來(lái)生成大量合成元組。
為了用 Transformer 學(xué)習(xí)因果公理,實(shí)現(xiàn)公理訓(xùn)練,該團(tuán)隊(duì)采用了以下方法構(gòu)建數(shù)據(jù)集、損失函數(shù)和位置嵌入。
公理訓(xùn)練:數(shù)據(jù)集、損失函數(shù)和位置編制
訓(xùn)練數(shù)據(jù)
基于一個(gè)特定公理,可根據(jù)「前提」將「假設(shè)」映射成合適的標(biāo)簽(Yes 或 No)。要?jiǎng)?chuàng)建訓(xùn)練數(shù)據(jù)集,該團(tuán)隊(duì)的做法是在特定的變量設(shè)置 X、Y、Z、A 下枚舉所有可能的元組 {(P, H, L)}_N,其中 P 是前提,H 是假設(shè),L 是標(biāo)簽(Yes 或 No)。
給定一個(gè)基于某個(gè)因果圖譜的前提 P,如果可通過(guò)使用特定的公理(一次或多次)推導(dǎo)出假設(shè) P,那么標(biāo)簽 L 就為 Yes;否則為 No。
舉個(gè)例子,假設(shè)一個(gè)系統(tǒng)的底層真實(shí)因果圖譜具有鏈?zhǔn)降耐負(fù)浣Y(jié)構(gòu):X_1 → X_2 → X_3 →???→ X_n。那么,可能的前提是 X_1 → X_2 ∧ X_2 → X_3,那么假設(shè) X_1 → X_3 有標(biāo)簽 Yes,而另一個(gè)假設(shè) X_3 → X_1 有標(biāo)簽 No。上述公理可被歸納式地多次用于生成更復(fù)雜的訓(xùn)練元組。
對(duì)于訓(xùn)練設(shè)置,使用傳遞性公理生成的 N 個(gè)公理實(shí)例構(gòu)建一個(gè)合成數(shù)據(jù)集 D。D 中的每個(gè)實(shí)例都構(gòu)建成了 (P_i, H_ij, L_ij) 的形式,,其中 n 是每第 i 個(gè)前提中的節(jié)點(diǎn)數(shù)量。P 是前提,即某種因果結(jié)構(gòu)的自然語(yǔ)言表達(dá)(如 X 導(dǎo)致 Y,Y 導(dǎo)致 Z);之后是問(wèn)題 H(如 X 導(dǎo)致 Y 嗎?);L 為標(biāo)簽(Yes 或 No)。該形式能有效覆蓋給定因果圖譜中每條獨(dú)特鏈的所有成對(duì)節(jié)點(diǎn)。
損失函數(shù)
給定一個(gè)數(shù)據(jù)集,損失函數(shù)的定義基于每個(gè)元組的基本真值標(biāo)簽,表示為:分析表明,相比于下一 token 預(yù)測(cè),使用該損失能得到很有希望的結(jié)果。
位置編碼
除了訓(xùn)練和損失函數(shù),位置編碼的選擇也是另一個(gè)重要因素。位置編碼能提供 token 在序列中絕對(duì)和相對(duì)位置的關(guān)鍵信息。
著名論文《Attention is all you need》中提出了一種使用周期函數(shù)(正弦或余弦函數(shù))來(lái)初始化這些編碼的絕對(duì)位置編碼策略。
絕對(duì)位置編碼能為任何序列長(zhǎng)度的所有位置提供確定的值。但是,有研究表明絕對(duì)位置編碼難以應(yīng)對(duì) Transformer 的長(zhǎng)度泛化任務(wù)。在可學(xué)習(xí)的 APE 變體中,每個(gè)位置嵌入都是隨機(jī)初始化的,并使用該模型完成了訓(xùn)練。該方法難以應(yīng)對(duì)比訓(xùn)練時(shí)的序列更長(zhǎng)的序列,因?yàn)樾碌奈恢们度胍廊晃幢挥?xùn)練和初始化。
有趣的是,近期有發(fā)現(xiàn)表明移除自回歸模型中的位置嵌入可以提升模型的長(zhǎng)度泛化能力,而自回歸解碼期間的注意力機(jī)制足以編碼位置信息。該團(tuán)隊(duì)使用了不同的位置編碼來(lái)理解其對(duì)因果任務(wù)中的泛化的影響,包括可學(xué)習(xí)位置編碼(LPE)、正弦位置編碼(SPE)、無(wú)位置編碼(NoPE)。
為了提升模型的泛化能力,該團(tuán)隊(duì)也采用了數(shù)據(jù)擾動(dòng),其中包括長(zhǎng)度、節(jié)點(diǎn)名稱、鏈順序和分支情況的擾動(dòng)。
實(shí)驗(yàn)
下面問(wèn)題又來(lái)了:如果使用這些數(shù)據(jù)訓(xùn)練一個(gè)模型,那么該模型能否學(xué)會(huì)將該公理應(yīng)用于新場(chǎng)景?
為了解答這個(gè)問(wèn)題,該團(tuán)隊(duì)使用這個(gè)因果無(wú)關(guān)型公理的符號(hào)演示從頭開(kāi)始訓(xùn)練了一個(gè) Transformer 模型。
為了評(píng)估其泛化性能,他們?cè)诤?jiǎn)單的大小為 3-6 個(gè)節(jié)點(diǎn)的因果無(wú)關(guān)公理鏈上進(jìn)行了訓(xùn)練,然后測(cè)試了泛化性能的多個(gè)不同方面,包括長(zhǎng)度泛化性能(大小 7-15 的鏈)、名稱泛化性能(更長(zhǎng)的變量名)、順序泛化性能(帶有反向的邊或混洗節(jié)點(diǎn)的鏈)、結(jié)構(gòu)泛化性能(帶有分支的圖譜)。圖 1 給出了評(píng)估 Transformer 的結(jié)構(gòu)泛化的方式。
具體來(lái)說(shuō),他們基于 GPT-2 架構(gòu)訓(xùn)練了一個(gè)基于解碼器的有 6700 萬(wàn)參數(shù)的模型。該模型有 12 個(gè)注意力層、8 個(gè)注意力頭和 512 嵌入維度。他們?cè)诿總€(gè)訓(xùn)練數(shù)據(jù)集上從頭開(kāi)始訓(xùn)練了該模型。為了理解位置嵌入的影響,他們還研究了三種位置嵌入設(shè)置:正弦位置編碼(SPE)、可學(xué)習(xí)位置編碼(LPE)和無(wú)位置編碼(NoPE)。
結(jié)果如表 1、圖 3 和圖 4 所示。
表 1 給出了在訓(xùn)練時(shí)未曾見(jiàn)過(guò)的更大因果鏈上評(píng)估時(shí)不同模型的準(zhǔn)確度??梢钥吹剑履P?TS2 (NoPE) 的表現(xiàn)能與萬(wàn)億參數(shù)規(guī)模的 GPT-4 相媲美。
圖 3 是在有更長(zhǎng)節(jié)點(diǎn)名稱(長(zhǎng)于訓(xùn)練集的)的因果序列上的泛化能力評(píng)估結(jié)果以及不同位置嵌入的影響。
圖 4 評(píng)估的是在更長(zhǎng)的未見(jiàn)過(guò)的因果序列上的泛化能力。
他們發(fā)現(xiàn),在簡(jiǎn)單鏈上訓(xùn)練的模型可以泛化到在更大的鏈上多次應(yīng)用公理,但卻無(wú)法泛化到順序或結(jié)構(gòu)泛化等更復(fù)雜的場(chǎng)景。但是,如果在簡(jiǎn)單鏈以及帶有隨機(jī)逆向邊的鏈組成的混合數(shù)據(jù)集上訓(xùn)練模型,則模型可以很好地泛化到各種評(píng)估場(chǎng)景。
通過(guò)擴(kuò)展在 NLP 任務(wù)上的長(zhǎng)度泛化研究結(jié)果,他們發(fā)現(xiàn)了位置嵌入在確保在長(zhǎng)度和其它方面實(shí)現(xiàn)因果泛化的重要性。他們表現(xiàn)最佳的模型沒(méi)有位置編碼,但他們也發(fā)現(xiàn)正弦編碼在某些情況下也很好用。
這種公理訓(xùn)練方法還能泛化用于一個(gè)更困難的問(wèn)題,如圖 5 所示。即以包含統(tǒng)計(jì)獨(dú)立性陳述的前提為基礎(chǔ),任務(wù)目標(biāo)是根據(jù)因果關(guān)系分辨相關(guān)性。解決該任務(wù)需要多個(gè)公理的知識(shí),包括 d-separation 和馬爾可夫性質(zhì)。
該團(tuán)隊(duì)使用與上面一樣的方法生成了合成訓(xùn)練數(shù)據(jù),然后訓(xùn)練了一個(gè)模型,結(jié)果發(fā)現(xiàn)在包含 3-4 個(gè)變量的任務(wù)演示上訓(xùn)練得到的 Transformer 能學(xué)會(huì)解決包含 5 個(gè)變量的圖譜任務(wù)。并且在該任務(wù)上,該模型的準(zhǔn)確度高于 GPT-4 和 Gemini Pro 等更大型的 LLM。
該團(tuán)隊(duì)表示:「我們的研究提供了一種通過(guò)公理的符號(hào)演示教模型學(xué)習(xí)因果推理的新范式,我們稱之為公理訓(xùn)練(axiomatic training)。」該方法的數(shù)據(jù)生成和訓(xùn)練流程是普適的:只要一個(gè)公理能被表示成符號(hào)元組的格式,就可使用此方法學(xué)習(xí)它。