DoReMi：最佳化資料混合比例加速語言模型預訓練

語言模型的效能高度依賴預訓練資料的組成，但最佳的資料混合比例事先無從得知，且取決於下游用途。本文提出 DoReMi（Domain Reweighting with Minimax Optimization，極小化極大最佳化的領域重新加權），一種利用小型代理模型 (proxy model) 來判斷如何對各領域的資料進行加權以訓練大型模型的演算法。DoReMi 首先在給定的初始領域權重上訓練一個小型參考模型 (reference model)，然後使用分組式分散穩健最佳化 (Group DRO) 訓練一個小型代理模型，在所有領域上極小化超額損失 (excess loss) 的最壞情況。最後使用代理模型產生的領域權重來訓練一個大得多的完整模型。在實驗中，DoReMi 使用 280M 參數的代理模型來改善 8B 參數模型的訓練效率。在 The Pile 資料集上，DoReMi 在所有領域都改善了困惑度 (perplexity)，即使是那些被降低了權重的領域也是如此，且平均下游準確率相較基準提升了 6.5%，達到 2.6 倍的訓練加速效果。在 GLaM 資料集上，DoReMi 在不需要知道下游任務的情況下，就達到了以下游任務為依據微調出的領域權重的效能水準。