自生成錯誤訓練:擴散語言模型中的令牌編輯改進
本研究針對LLaDA2.1擴散語言模型中的令牌到令牌(T2T)編輯模組存在的訓練-推理不匹配問題,提出自生成T2T方法。該方法透過無梯度草稿傳遞、用預測令牌填充掩碼位置,並在第二遍中監督恢復,實現了在不改變推理引數的情況下提升準確性並減少編輯強度。在多個基準測試上,該方法有效緩解了最終數字轉錄錯誤和過度自我修正等問題。
擴散語言模型是一種新興的文本生成技術,它透過逐步去噪的方式從隨機噪聲中恢復出完整的文本序列。與自迴歸模型不同,擴散模型允許並行解碼,並在解碼過程中透過編輯模組對已生成的令牌進行修正。LLaDA2.1是這類模型中的代表,其令牌到令牌(Token-to-token, T2T)編輯模組專門用於在塊擴散解碼過程中修訂已提交的令牌。然而,該模組的訓練和推理之間存在一個關鍵的不匹配:訓練時,編輯器被訓練以應對隨機詞彙損壞——即隨機替換一些令牌為錯誤令牌;但在推理時,編輯器實際面對的是模型自身生成的錯誤,這些錯誤通常是流暢的、高置信度的草稿錯誤,而非隨機噪聲。這種差異導致編輯器的效能受到限制,尤其是在需要精確數字或簡短事實回答的任務中,容易出現最終數字轉錄錯誤或在正確答案前進行不必要的自我修正。
為了解決這一訓練-推理不匹配問題,Lin Yao等人在arXiv上發表的論文中提出了自生成錯誤訓練(Self-Generated Error Training)方法。該方法的核心在於讓編輯器在訓練過程中模擬推理階段的實際錯誤分佈。具體步驟如下:首先,模型進行一次無梯度的草稿生成,即普通的前向傳播得到初步輸出;然後,將這些輸出中的某些位置用模型預測的令牌進行填充,從而構造出類似推理階段的錯誤模式;最後,在第二次前向傳播中,模型被監督學習如何從這些自生成的錯誤中恢復到正確的結果。整個過程不涉及額外的梯度計算,而是透過對LLaDA2.1-mini模型進行簡短的LoRA(Low-Rank Adaptation)持續預訓練來實現。LoRA是一種高效的微調方法,透過更新一小部分引數來適應新的任務,從而避免了全模型微調的高計算成本。
在實驗部分,研究者在多個基準測試上對自生成T2T方法進行了評估,包括需要精確數字輸出和事實回答的任務。他們採用了官方的Q-Mode T2T流程,並保持所有推理引數不變,以確保對比的公平性。實驗結果表明,與基線相比,自生成T2T方法在整體準確率上有所提升,同時T2T編輯的強度——即模型對已生成令牌進行修改的頻率——有所降低。更重要的是,該方法有效緩解了兩種典型的失敗模式:一是模型在正確推理後卻將最終數字轉錄錯誤,例如在算術題中得出正確結果但輸出時數字出錯;二是在回答簡短事實問題前進行過度的自我修正,導致原本正確的答案被改寫錯誤。這些改進表明,透過使訓練資料更接近推理時的錯誤分佈,可以顯著提升編輯器的魯棒性和準確性。
這項研究對於擴散語言模型的發展具有重要意義。它揭示了訓練資料分佈對模型效能的深遠影響,並提供了一種簡單而有效的解決方案。未來,該方法有望被整合到更大的模型和更廣泛的任務中,進一步推動擴散模型在文本生成領域的應用。研究人員表示,相關程式碼和模型將在未來公開,以便社群復現和進一步發展。