大模型家族來了一個專門解決數(shù)學(xué)問題的「新成員」——Llemma。

如今，在各種文本混合數(shù)據(jù)上訓(xùn)練出來的語言模型會顯示出非常通用的語言理解和生成能力，可以作為基礎(chǔ)模型適應(yīng)各種應(yīng)用。開放式對話或指令跟蹤等應(yīng)用要求在整個自然文本分布中實現(xiàn)均衡的性能，因此更傾向于通用模型。

不過如果想要在某一領(lǐng)域（如醫(yī)學(xué)、金融或科學(xué)）內(nèi)最大限度地提高性能，那么特定領(lǐng)域的語言模型可能會以給定的計算成本提供更優(yōu)越的能力，或以更低的計算成本提供給定的能力水平。

普林斯頓大學(xué)、 EleutherAI 等的研究者為解決數(shù)學(xué)問題訓(xùn)練了一個特定領(lǐng)域的語言模型。他們認(rèn)為：首先，解決數(shù)學(xué)問題需要與大量的專業(yè)先驗知識進(jìn)行模式匹配，因此是進(jìn)行領(lǐng)域適應(yīng)性訓(xùn)練的理想環(huán)境；其次，數(shù)學(xué)推理本身就是 AI 的核心任務(wù)；最后，能夠進(jìn)行強(qiáng)數(shù)學(xué)推理的語言模型是許多研究課題的上游，如獎勵建模、推理強(qiáng)化學(xué)習(xí)和算法推理。

因此，他們提出一種方法，通過對 Proof-Pile-2 進(jìn)行持續(xù)的預(yù)訓(xùn)練，使語言模型適應(yīng)數(shù)學(xué)。Proof-Pile-2 是數(shù)學(xué)相關(guān)文本和代碼的混合數(shù)據(jù)。將這一方法應(yīng)用于 Code Llama，可以得到 LLEMMA：7B 和 34B 的基礎(chǔ)語言模型，其數(shù)學(xué)能力得到了大幅提高。

論文地址：https://arxiv.org/pdf/2310.10631.pdf

項目地址：https://github.com/EleutherAI/math-lm

LLEMMA 7B 的 4-shot Math 性能遠(yuǎn)超谷歌 Minerva 8B，LLEMMA 34B 在參數(shù)少近一半的情況下性能逼近 Minerva 62B。

具體來說，本文貢獻(xiàn)如下：

• 1. 訓(xùn)練并發(fā)布了 LLEMMA 模型：專門用于數(shù)學(xué)的 7B 和 34B 語言模型。LLEMMA 模型是在 MATH 上公開發(fā)布的基礎(chǔ)模型的最新水平。

• 2. 發(fā)布了代數(shù)堆棧（AlgebraicStack），這是一個包含 11B 專門與數(shù)學(xué)相關(guān)的代碼 token 的數(shù)據(jù)集。

• 3. 證明了 LLEMMA 能夠使用計算工具來解決數(shù)學(xué)問題，即 Python 解釋器和形式定理證明器。

• 4. 與之前的數(shù)學(xué)語言模型（如 Minerva）不同，LLEMMA 模型是開放式的。研究者開放了訓(xùn)練數(shù)據(jù)和代碼。這使得 LLEMMA 成為未來數(shù)學(xué)推理研究的一個平臺。

方法概覽

LLEMMA 是專門用于數(shù)學(xué)的 70B 和34B 語言模型。它由 Proof-Pile-2 上繼續(xù)對代碼 Llama 進(jìn)行預(yù)訓(xùn)練得到的。

DATA: Proof-Pile-2

研究者創(chuàng)建了 Proof-Pile-2，這是一個 55B token 的科學(xué)論文、包含數(shù)學(xué)的網(wǎng)絡(luò)數(shù)據(jù)和數(shù)學(xué)代碼的混合物。除了 Lean proofsteps 子集之外，Proof-Pile-2 的知識截止日期為 2023 年 4 月。

數(shù)值模擬、計算機(jī)代數(shù)系統(tǒng)和形式定理證明器等計算工具對數(shù)學(xué)家的重要性與日俱增。因此，研究者創(chuàng)建了代數(shù)堆棧（AlgebraicStack），這是一個包含 17 種語言源代碼的 11B token 數(shù)據(jù)集，涵蓋數(shù)值數(shù)學(xué)、符號數(shù)學(xué)和形式數(shù)學(xué)。該數(shù)據(jù)集由來自 Stack、GitHub 公共資源庫和形式證明步驟數(shù)據(jù)的過濾代碼組成。表9顯示了AlgebraicStack 中各語言的 token 數(shù)量。

AlgebraicStack 中各語言的 token 數(shù)。

研究者了使用 OpenWebMath，這是一個由高質(zhì)量網(wǎng)頁組成的 15B token 數(shù)據(jù)集，其中過濾了數(shù)學(xué)內(nèi)容。OpenWebMath 根據(jù)數(shù)學(xué)相關(guān)關(guān)鍵詞和基于分類器的數(shù)學(xué)評分過濾 CommonCrawl 網(wǎng)頁，保留數(shù)學(xué)格式（如 LATEX、AsciiMath），并包含額外的質(zhì)量過濾器（如 plexity、domain、length）和近似重復(fù)。

除此之外，研究者還使用了 RedPajama 的 ArXiv 子集，它是 LLaMA 訓(xùn)練數(shù)據(jù)集的開放再現(xiàn)。ArXiv 子集包含 29B 個詞塊。訓(xùn)練混合數(shù)據(jù)由少量一般領(lǐng)域數(shù)據(jù)組成，起到了正則化的作用。由于 LLaMA 2 的預(yù)訓(xùn)練數(shù)據(jù)集尚未公開，研究者使用 Pile 作為替代訓(xùn)練數(shù)據(jù)集。

模型和訓(xùn)練

每個模型都是從 Code Llama 初始化而來，該模型又初始化自 Llama 2，使用僅解碼器（deconder only）的 transformer 結(jié)構(gòu)，在 500B 的代碼 token 上訓(xùn)練而成。研究者使用標(biāo)準(zhǔn)自回歸語言建模目標(biāo)，在 Proof-Pile-2 上繼續(xù)訓(xùn)練 Code Llama 模型。這里，LLEMMA 7B 模型有 200B token，LLEMMA 34B 模型有 50B token。

研究者使用 GPT-NeoX 庫在 256 個 A100 40GB GPU 上，以 bfloat16 混合精度來訓(xùn)練以上兩個模型。他們?yōu)?LLEMMA-7B 使用了世界大小為 2 的張量并行，為 34B 使用了世界大小為 8 的張量并行，以及跨數(shù)據(jù)并行副本的 ZeRO Stage 1 分片優(yōu)化器狀態(tài)。此外還使用 Flash Attention 2 來提高吞吐量并進(jìn)一步降低內(nèi)存需求。

LLEMMA 7B 經(jīng)過了 42000 步的訓(xùn)練，全局 batch 大小為 400 萬個 token，上下文長度為 4096 個 token。這相當(dāng)于 23000 個 A100 時。學(xué)習(xí)率在 500 步后預(yù)熱到了 1?10^?4，然后在 48000 步后將余弦衰減到最大學(xué)習(xí)率的 1/30。

LLEMMA 34B 經(jīng)過了 12000 步的訓(xùn)練，全局 batch 大小同樣為 400 萬個 token，上下文長度為 4096。這相當(dāng)于 47000 個 A100 時。學(xué)習(xí)率在 500 步后預(yù)熱到了 5?10^?5，然后衰減到峰值學(xué)習(xí)率的 1/30。

評估結(jié)果

在實驗部分，研究者旨在評估 LLEMMA 是否可以作為數(shù)學(xué)文本的基礎(chǔ)模型。他們利用少樣本評估來比較 LLEMMA 模型，并主要關(guān)注沒有在數(shù)學(xué)任務(wù)監(jiān)督樣本上進(jìn)行微調(diào)的 SOTA 模型。

研究者首先使用思維鏈推理和多數(shù)投票（majority voting）方法來評估 LLEMMA 求解數(shù)學(xué)題的能力，評估基準(zhǔn)包括了 MATH 和 GSM8k。然后探索使用少樣本工具和定理證明。最后研究了內(nèi)存和數(shù)據(jù)混合的影響。

使用思維鏈（CoT）求解數(shù)學(xué)題

這些任務(wù)包括為 LATEX 或自然語言表示的問題生成獨立的文本答案，而無需使用外部工具。研究者使用到的評估基準(zhǔn)有 MATH、GSM8k、 OCWCourses、SAT 和 MMLU-STEM。

結(jié)果如下表 1 所示，LLEMMA 在 Proof-Pile-2 語料庫上的持續(xù)預(yù)訓(xùn)練在 5 個數(shù)學(xué)基準(zhǔn)上均提升了少樣本性能，其中 LLEMMA 34B 在 GSM8k 上比 Code Llama 提高了 20 個百分點，在 MATH 上比 Code Llama 提高了 13 個百分點。同時 LLEMMA 7B 優(yōu)于專有的 Minerva 模型。

因此，研究者得到結(jié)論，在 Proof-Pile-2 上進(jìn)行持續(xù)預(yù)訓(xùn)練有助于提升預(yù)訓(xùn)練模型求解數(shù)學(xué)題的能力。

使用工具求解數(shù)學(xué)題

這些任務(wù)包括使用計算工具來解題。研究者使用到的評估基準(zhǔn)有 MATH+Python 和 GSM8k+Python。

結(jié)果如下表 3 所示，LLEMMA 在這兩項任務(wù)上均優(yōu)于 Code Llama。同時使用工具后在 MATH 和 GSM8k 上的性能也優(yōu)于沒有工具的情況。

形式數(shù)學(xué)

Proof-Pile-2 的 AlgebraicStack 數(shù)據(jù)集擁有 15 億 token 的形式數(shù)學(xué)數(shù)據(jù)，包括提取自 Lean 和 Isabelle 的形式化證明。雖然對形式數(shù)學(xué)的全面研究超出了本文的探討范圍，但研究者在以下兩個任務(wù)上評估了 LLEMMA 的少樣本性能。

非形式到形式證明任務(wù)，即在給定形式命題、非形式 LATEX 命題和非形式 LATEX 證明的情況下，生成一個形式證明；

形式到形式證明任務(wù)，即通過生成一系列證明步驟（或策略）來證明一個形式命題。

結(jié)果如下表 4 所示，LLEMMA 在 Proof-Pile-2 上的持續(xù)預(yù)訓(xùn)練在兩個形式定理證明任務(wù)上提升了少樣本性能。

數(shù)據(jù)混合的影響

訓(xùn)練語言模型時，一種常見的做法是根據(jù)混合權(quán)重對訓(xùn)練數(shù)據(jù)的高質(zhì)量子集進(jìn)行上采樣。研究者在幾個精心挑選的混合權(quán)重上進(jìn)行了短期訓(xùn)練，以此選擇混合權(quán)重。接著選擇了在一組高質(zhì)量 held-out 文本（這里使用了 MATH 訓(xùn)練集）上能夠最小化困惑度的混合權(quán)重。

下表 5 顯示了使用 arXiv、web 和代碼等不同數(shù)據(jù)混合訓(xùn)練后，模型的 MATH 訓(xùn)練集困惑度。

更多技術(shù)細(xì)節(jié)和評估結(jié)果參閱原論文。