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

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

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

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

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

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

項(xiàng)目地址：https://github.com/EleutherAI/math-lm

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

具體來(lái)說(shuō)，本文貢獻(xiàn)如下：

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

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

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

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

方法概覽

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

DATA: Proof-Pile-2

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

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

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

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

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

模型和訓(xùn)練

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

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

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

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

評(píng)估結(jié)果

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

訓(xùn)練語(yǔ)言模型時(shí)，一種常見的做法是根據(jù)混合權(quán)重對(duì)訓(xùn)練數(shù)據(jù)的高質(zhì)量子集進(jìn)行上采樣。研究者在幾個(gè)精心挑選的混合權(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é)和評(píng)估結(jié)果參閱原論文。