擴散語言模型（Diffusion Language Model, DLM）作為自迴歸模型的一種有前途的替代方案，近年來受到了廣泛關注。然而，在DLM的後訓練階段，現有方法大多沿用傳統的獎勵最大化目標，這導致了一個被稱為“軌跡鎖定”（trajectory locking）的關鍵失敗模式。具體而言，當模型透過獎勵驅動的更新進行最佳化時，機率質量會過度集中在少數幾條去噪路徑上，從而在重複取樣時顯著降低了對其他潛在正確解的覆蓋能力。這種模式崩潰問題嚴重限制了DLM在實際應用中的表現。
為了解決這一問題，來自研究團隊的Saba Ahmadi等人提出了一種新的後訓練方法——TraFL（Trajectory Flow baLancing，軌跡流平衡）。TraFL的核心思想是引入一個軌跡平衡目標，該目標將策略訓練至一個獎勵傾斜的目標分佈，同時透過錨定到一個凍結的參考模型來防止過度最佳化。這種方法巧妙地避免了單純追求獎勵最大化所導致的模式崩潰。為了使該技術適用於擴散語言模型，研究人員還設計了擴散相容的序列級替代目標以及一個學習得到的提示相關歸一化項，從而確保了方法的實用性和穩定性。
在數學推理和程式碼生成等標準基準測試中，TraFL展現出了顯著的效能優勢。實驗結果表明，TraFL是唯一一種在每種基準長度設定下都優於基礎模型的後訓練方法，並且其效能提升隨著取樣預算的增加而持續擴大。更令人印象深刻的是，這些改進能夠很好地泛化到留出評估集上：在Minerva Math基準中，TraFL的效能始終保持在基礎模型之上，而在LiveCodeBench的每個難度等級上，它都成為表現最強的方法。
這項研究為擴散語言模型的後訓練開闢了新的方向，透過平衡軌跡分佈而非盲目追求高獎勵，有效提升了模型的探索能力和最終效能。未來，這一方法有望被進一步擴充套件至更廣泛的應用場景，為語言模型的訓練和最佳化提供新的思路。