AI 已經(jīng)在幫助數(shù)學(xué)家和科學(xué)家做研究了，比如著名數(shù)學(xué)家陶哲軒就曾多次分享自己借助 GPT 等 AI 工具研究探索的經(jīng)歷。AI 要在這些領(lǐng)域大戰(zhàn)拳腳，強大可靠的因果推理能力是必不可少的。

本文要介紹的這項研究發(fā)現(xiàn)：在小圖譜的因果傳遞性公理演示上訓(xùn)練的 Transformer 模型可以泛化用于大圖譜的傳遞性公理。

也就是說，如果讓 Transformer 學(xué)會執(zhí)行簡單的因果推理，就可能將其用于更為復(fù)雜的因果推理。該團隊提出的公理訓(xùn)練框架是一種基于被動數(shù)據(jù)來學(xué)習(xí)因果推理的新范式，只有演示足夠就能用于學(xué)習(xí)任意公理。

引言

因果推理（causal reasoning）可以定義成一組推理流程并且這組推理流程要符合專門針對因果性的預(yù)定義公理或規(guī)則。舉個例子，d-separation（有向分離）和 do-calculus 規(guī)則可被視為公理，而 collider set 或 backdoor set 的規(guī)范則可被看作是由公理推導(dǎo)出的規(guī)則。

通常來說，因果推理使用的數(shù)據(jù)對應(yīng)于一個系統(tǒng)中的變量。通過正則化、模型架構(gòu)或特定的變量選擇，可以歸納偏置的形式將公理或規(guī)則集成到機器學(xué)習(xí)模型中。

根據(jù)可用數(shù)據(jù)種類的差異（觀察數(shù)據(jù)、干預(yù)數(shù)據(jù)、反事實數(shù)據(jù)），Judea Pearl 提出的「因果階梯」定義了因果推理的可能類型。

由于公理是因果性的基石，因此我們不禁會想是否可以直接使用機器學(xué)習(xí)模型來學(xué)習(xí)公理。也就是說，如果學(xué)習(xí)公理的方式不是學(xué)習(xí)通過某個數(shù)據(jù)生成流程得到的數(shù)據(jù)，而是直接學(xué)習(xí)公理的符號演示（并由此學(xué)習(xí)因果推理），哪又會如何呢？

相較于使用特定的數(shù)據(jù)分布構(gòu)建的針對特定任務(wù)的因果模型，這樣的模型有一個優(yōu)勢：其可在多種不同的下游場景中實現(xiàn)因果推理。隨著語言模型具備了學(xué)習(xí)以自然語言表達的符號數(shù)據(jù)的能力，這個問題也就變得非常重要了。

事實上，近期已有一些研究通過創(chuàng)建以自然語言編碼因果推理問題的基準(zhǔn)，評估了大型語言模型（LLM）是否能夠執(zhí)行因果推理。

微軟、MIT 和印度理工學(xué)院海得拉巴分校（IIT Hyderabad）的研究團隊也朝這個方向邁出了重要一步：提出了一種通過公理訓(xùn)練（axiomatic training）學(xué)習(xí)因果推理的方法。

• 論文標(biāo)題：Teaching Transformers Causal Reasoning through Axiomatic Training
• 論文地址：https://arxiv.org/pdf/2407.07612

公理訓(xùn)練

他們假設(shè)，可將因果公理表示成以下符號元組 ?premise, hypothesis, result?。其中 hypothesis 是指假設(shè)，即因果陳述；premise 是前提，是指用于確定該陳述是否為「真」的任意相關(guān)信息；result 自然就是結(jié)果了。結(jié)果可以是簡單的「是」或「否」。

舉個例子，來自論文《Can large language models infer causation from correlation?》的 collider 公理可以表示成：，結(jié)論就為「是」。

基于這個模板，可通過修改變量名稱、變量數(shù)量和變量順序等來生成大量合成元組。

為了用 Transformer 學(xué)習(xí)因果公理，實現(xiàn)公理訓(xùn)練，該團隊采用了以下方法構(gòu)建數(shù)據(jù)集、損失函數(shù)和位置嵌入。

公理訓(xùn)練：數(shù)據(jù)集、損失函數(shù)和位置編制

訓(xùn)練數(shù)據(jù)

基于一個特定公理，可根據(jù)「前提」將「假設(shè)」映射成合適的標(biāo)簽（Yes 或 No）。要創(chuàng)建訓(xùn)練數(shù)據(jù)集，該團隊的做法是在特定的變量設(shè)置 X、Y、Z、A 下枚舉所有可能的元組 {(P, H, L)}_N，其中 P 是前提，H 是假設(shè)，L 是標(biāo)簽（Yes 或 No）。

給定一個基于某個因果圖譜的前提 P，如果可通過使用特定的公理（一次或多次）推導(dǎo)出假設(shè) P，那么標(biāo)簽 L 就為 Yes；否則為 No。

舉個例子，假設(shè)一個系統(tǒng)的底層真實因果圖譜具有鏈?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，而另一個假設(shè) X_3 → X_1 有標(biāo)簽 No。上述公理可被歸納式地多次用于生成更復(fù)雜的訓(xùn)練元組。

對于訓(xùn)練設(shè)置，使用傳遞性公理生成的 N 個公理實例構(gòu)建一個合成數(shù)據(jù)集 D。D 中的每個實例都構(gòu)建成了 (P_i, H_ij, L_ij) 的形式，，其中 n 是每第 i 個前提中的節(jié)點數(shù)量。P 是前提，即某種因果結(jié)構(gòu)的自然語言表達（如 X 導(dǎo)致 Y，Y 導(dǎo)致 Z）；之后是問題 H（如 X 導(dǎo)致 Y 嗎？）；L 為標(biāo)簽（Yes 或 No）。該形式能有效覆蓋給定因果圖譜中每條獨特鏈的所有成對節(jié)點。

損失函數(shù)

給定一個數(shù)據(jù)集，損失函數(shù)的定義基于每個元組的基本真值標(biāo)簽，表示為：分析表明，相比于下一 token 預(yù)測，使用該損失能得到很有希望的結(jié)果。

位置編碼

除了訓(xùn)練和損失函數(shù)，位置編碼的選擇也是另一個重要因素。位置編碼能提供 token 在序列中絕對和相對位置的關(guān)鍵信息。

著名論文《Attention is all you need》中提出了一種使用周期函數(shù)（正弦或余弦函數(shù)）來初始化這些編碼的絕對位置編碼策略。

絕對位置編碼能為任何序列長度的所有位置提供確定的值。但是，有研究表明絕對位置編碼難以應(yīng)對 Transformer 的長度泛化任務(wù)。在可學(xué)習(xí)的 APE 變體中，每個位置嵌入都是隨機初始化的，并使用該模型完成了訓(xùn)練。該方法難以應(yīng)對比訓(xùn)練時的序列更長的序列，因為新的位置嵌入依然未被訓(xùn)練和初始化。

有趣的是，近期有發(fā)現(xiàn)表明移除自回歸模型中的位置嵌入可以提升模型的長度泛化能力，而自回歸解碼期間的注意力機制足以編碼位置信息。該團隊使用了不同的位置編碼來理解其對因果任務(wù)中的泛化的影響，包括可學(xué)習(xí)位置編碼（LPE）、正弦位置編碼（SPE）、無位置編碼（NoPE）。

為了提升模型的泛化能力，該團隊也采用了數(shù)據(jù)擾動，其中包括長度、節(jié)點名稱、鏈順序和分支情況的擾動。

實驗

下面問題又來了：如果使用這些數(shù)據(jù)訓(xùn)練一個模型，那么該模型能否學(xué)會將該公理應(yīng)用于新場景？

為了解答這個問題，該團隊使用這個因果無關(guān)型公理的符號演示從頭開始訓(xùn)練了一個 Transformer 模型。

為了評估其泛化性能，他們在簡單的大小為 3-6 個節(jié)點的因果無關(guān)公理鏈上進行了訓(xùn)練，然后測試了泛化性能的多個不同方面，包括長度泛化性能（大小 7-15 的鏈）、名稱泛化性能（更長的變量名）、順序泛化性能（帶有反向的邊或混洗節(jié)點的鏈）、結(jié)構(gòu)泛化性能（帶有分支的圖譜）。圖 1 給出了評估 Transformer 的結(jié)構(gòu)泛化的方式。

具體來說，他們基于 GPT-2 架構(gòu)訓(xùn)練了一個基于解碼器的有 6700 萬參數(shù)的模型。該模型有 12 個注意力層、8 個注意力頭和 512 嵌入維度。他們在每個訓(xùn)練數(shù)據(jù)集上從頭開始訓(xùn)練了該模型。為了理解位置嵌入的影響，他們還研究了三種位置嵌入設(shè)置：正弦位置編碼（SPE）、可學(xué)習(xí)位置編碼（LPE）和無位置編碼（NoPE）。

結(jié)果如表 1、圖 3 和圖 4 所示。

表 1 給出了在訓(xùn)練時未曾見過的更大因果鏈上評估時不同模型的準(zhǔn)確度?？梢钥吹剑履Ｐ?TS2 (NoPE) 的表現(xiàn)能與萬億參數(shù)規(guī)模的 GPT-4 相媲美。

圖 3 是在有更長節(jié)點名稱（長于訓(xùn)練集的）的因果序列上的泛化能力評估結(jié)果以及不同位置嵌入的影響。

圖 4 評估的是在更長的未見過的因果序列上的泛化能力。

他們發(fā)現(xiàn)，在簡單鏈上訓(xùn)練的模型可以泛化到在更大的鏈上多次應(yīng)用公理，但卻無法泛化到順序或結(jié)構(gòu)泛化等更復(fù)雜的場景。但是，如果在簡單鏈以及帶有隨機逆向邊的鏈組成的混合數(shù)據(jù)集上訓(xùn)練模型，則模型可以很好地泛化到各種評估場景。

通過擴展在 NLP 任務(wù)上的長度泛化研究結(jié)果，他們發(fā)現(xiàn)了位置嵌入在確保在長度和其它方面實現(xiàn)因果泛化的重要性。他們表現(xiàn)最佳的模型沒有位置編碼，但他們也發(fā)現(xiàn)正弦編碼在某些情況下也很好用。

這種公理訓(xùn)練方法還能泛化用于一個更困難的問題，如圖 5 所示。即以包含統(tǒng)計獨立性陳述的前提為基礎(chǔ)，任務(wù)目標(biāo)是根據(jù)因果關(guān)系分辨相關(guān)性。解決該任務(wù)需要多個公理的知識，包括 d-separation 和馬爾可夫性質(zhì)。

該團隊使用與上面一樣的方法生成了合成訓(xùn)練數(shù)據(jù)，然后訓(xùn)練了一個模型，結(jié)果發(fā)現(xiàn)在包含 3-4 個變量的任務(wù)演示上訓(xùn)練得到的 Transformer 能學(xué)會解決包含 5 個變量的圖譜任務(wù)。并且在該任務(wù)上，該模型的準(zhǔn)確度高于 GPT-4 和 Gemini Pro 等更大型的 LLM。

該團隊表示：「我們的研究提供了一種通過公理的符號演示教模型學(xué)習(xí)因果推理的新范式，我們稱之為公理訓(xùn)練（axiomatic training）。」該方法的數(shù)據(jù)生成和訓(xùn)練流程是普適的：只要一個公理能被表示成符號元組的格式，就可使用此方法學(xué)習(xí)它。