OptiMer：最优分布向量合并优于持续预训练中的数据混合方法

通过向量合并后优化数据比例解决持续预训练中数据混合比例难调的问题

前置知识

持续预训练

持续预训练是指在预训练好的基础模型上使用特定领域或语言的数据继续训练模型以使模型适应新的任务或应用场景。与从头训练相比CPT能够显著降低训练成本同时保留了基础模型的通用能力。在CPT中训练数据通常来自多个分布如不同语言或领域需要合理混合这些数据以获得最佳性能。

本文提出的OptiMer方法专门针对CPT场景下的数据混合比例优化问题理解CPT是理解本文方法和实验设计的基础。

任务向量

任务向量定义为τ等于θft减去θbase即微调后模型参数与基础模型参数的差值。这个向量捕获了微调过程中参数的变化可以被视为模型在特定任务上学习的知识表示。任务向量可以进行算术运算通过线性组合来添加或删除任务能力。这一概念被扩展到CPT场景中形成了分布向量的概念。

本文提出的分布向量是任务向量在CPT场景下的扩展理解任务向量的数学定义和线性组合原理是理解OptiMer方法的核心。

贝叶斯优化和树结构Parzen估计器

贝叶斯优化是一种用于黑盒函数优化的方法适用于目标函数梯度不可获取或计算成本高昂的场景。TPE是贝叶斯优化的一种具体实现通过核密度估计对历史试验进行建模将试验分为好和差两个集合然后采样最大化好除以坏密度比的点。相比网格搜索TPE通常在十乘n次试验中收敛而网格搜索需要G的n次方次评估。

OptiMer使用TPE来搜索最优的向量合并权重理解TPE的工作原理和效率优势是理解OptiMer为何比传统数据混合方法快15到35倍的关键。

DARE-Linear合并

DARE即Drop And REscale是一种模型合并方法通过在合并前对模型参数进行稀疏化随机丢弃一部分参数和重新缩放来减少不同模型之间的干扰。DARE-Linear将DARE与线性任务算术结合即θmerge等于θbase加上Σαi乘以τi其中αi是合并权重。本文使用DARE-Linear作为基础的向量合并函数因为它在处理多个模型合并时表现优于简单的任务算术。

OptiMer基于DARE-Linear进行向量合并理解其原理有助于理解为什么OptiMer能够有效合并多个分布向量而不会产生严重干扰。

研究动机

持续预训练是将大型语言模型适配到特定语言和领域的常用方法但训练数据的混合比例是一个敏感且难以调优的超参数。现有的数据混合方法必须在训练开始前固定混合比例而次优的选择会导致模型性能显著下降。例如在Xie等人2023年的研究中不当的数据混合比例可以使模型性能下降多达10到20个百分点。更严重的是由于CPT训练通常需要几天到几周的GPU集群时间一个糟糕的比例选择可能浪费大量计算资源后才能被发现。虽然近期有一些方法如DoReMi和RegMix通过代理模型或小规模实验来估计混合比例但这些估计仍然需要在训练前固定无法在训练后进行校正。

本文的目标是本文的目标是提出一个新的框架将数据混合比例的选择与模型训练解耦使得比例可以在训练后通过优化来确定而无需重新训练。具体来说作者希望实现以下几点第一能够独立训练每个数据集的CPT模型提取其分布向量第二能够通过优化找到最优的向量合并权重且优化过程快速高效第三优化后的权重能够解释为数据混合比例并可用于指导传统的数据混合训练第四同一组分布向量能够针对不同目标进行重新优化产生多个定制化模型而无需重新训练。

与已有工作不同的是，本文的独特切入角度是将传统上作为预训练决策的数据混合比例选择问题重新表述为训练后对分布向量的优化问题。这与现有的数据混合优化方法有本质区别DoReMi等方法仍然需要固定比例后训练而OptiMer允许在训练后调整比例。现有的模型合并方法如Task Arithmetic、TIES、DARE主要关注任务特定的能力转移或少数模型的合并而OptiMer专注于多分布CPT场景下的通用性能提升。此外OptiMer的创新之处在于利用了CPT分布向量之间近似正交的特性使得线性组合成为可能并通过贝叶斯优化高效地找到最优权重。

核心方法

OptiMer的整体思路是先独立训练后优化合并。直觉上不同分布的数据在参数空间中修改模型的方向是不同的如果这些方向之间互不干扰即近似正交那么就可以通过线性组合这些方向来实现对多个分布的适配而无需在训练时固定数据比例。技术路线分为三个步骤第一对每个数据集Di独立进行CPT训练得到模型θCPTi第二提取分布向量τi等于θCPTi减去θpt和IT向量τit等于θit减去θpt第三使用贝叶斯优化TPE搜索最优的合并权重α*等于αit、α1到αn使得合并后的模型θmerge等于θpt加上αit乘以τit加上Σαi乘以τi在验证集上获得最优性能。这种方法的优势是将昂贵的比例搜索从训练过程转移到训练后的向量合并过程大大降低了搜索成本。

OptiMer的核心创新点在于发现了CPT分布向量之间的近似正交性使得线性组合成为可能。具体来说不同分布如日语、数学、代码训练得到的CPT向量之间的余弦相似度仅为0.03到0.31而CPT向量与IT向量之间的余弦相似度更是低至0.03左右。这表明每个分布修改了模型参数空间中一个相对独立的子空间因此可以通过线性组合来叠加多个分布的效果而不会产生严重干扰。这与传统的数据混合方法形成鲜明对比数据混合在训练时就固定了比例如果比例不佳就会导致性能损失而OptiMer在训练后通过优化权重来调整不同分布的贡献度相当于动态调整了有效训练时长。此外OptiMer还支持负权重可以主动减去某些分布的干扰效应。

方法步骤详情

OptiMer的方法步骤包含以下五个主要步骤。第一准备阶段获取预训练基础模型θpt如gemma-3-27b-pt和指令调优模型θit如gemma-3-27b-it准备n个数据集Di每个数据集包含1B tokens。第二独立CPT训练对每个数据集Di从θpt开始进行1个epoch即2000步的持续预训练得到CPT模型θCPTi。训练使用AdamW优化器β1等于0.90、β2等于0.95峰值学习率4乘以10的负5次方余弦衰减到1%批次大小128梯度裁剪1.0。第三向量提取计算每个数据集的分布向量τi等于θCPTi减去θpt和IT向量τit等于θit减去θpt这些向量编码了对应分布引起的参数变化。第四贝叶斯优化搜索使用TPE算法在0.3到1.0即αit和0.0到1.0即αi范围内搜索最优权重α*。初始化N0等于20次随机试验然后进行T等于100次TPE试验每次试验构建合并模型θmerge并在代理任务的前100个样本上评估分数。第五最终选择从100次试验中选择top-K等于3个配置在完整验证集的前300个样本上重新评估选择分数最高的配置作为最终模型θ∗merge。

技术新颖性

OptiMer的技术新颖性体现在以下几个方面。首先它将任务向量的概念从微调扩展到了CPT场景提出了分布向量的新概念用于捕获特定数据分布引起的参数变化。其次它系统性地分析了CPT分布向量的性质包括近似正交性余弦相似度0.03到0.31训练轨迹的近似线性特征以及向量范数与性能的关系为向量合并提供了理论基础。第三它将贝叶斯优化TPE引入到模型合并权重搜索中相比传统的网格搜索DEM方法搜索效率显著提高从O的G次方降低到O的10n次方。第四OptiMer展示了优化后的权重可以解释为数据混合比例并且同一组向量池可以针对不同目标重新优化这为CPT提供了一个更加灵活的范式。最后OptiMer在实验中验证了支持负权重的有效性这相当于能够主动减去某些分布的干扰这是传统数据混合方法无法实现的。

**Figure 3**: Pairwise cosine similarity and norm of distribution vectors

**Figure 5**: Distribution vector trajectories during CPT

**Figure 6**: OptiMer search dynamics for Ja+Math combination

实验结果

OptiMer在Gemma 3 27B上的实验结果显示在所有三个数据集组合Ja加Math、Ja加Code、Ja加Zh加Math中OptiMer都取得了最佳的平均性能。具体来说在Ja加Math组合上OptiMer的平均分为69.98比DataMix基线67.86提高了2.12分。在Ja加Code组合上OptiMer的平均分为69.68比DataMix基线67.23提高了2.45分。在Ja加Zh加Math组合上OptiMer的平均分为70.37比DataMix基线63.71提高了6.66分。值得注意的是随着数据集数量增加OptiMer的优势更加明显从2-way提升的2.12到2.45分到3-way提升的6.66分。OptiMer还在TruthfulQA任务上表现优异得分保持在51到55之间而其他基线方法包括DataMix和各种模型合并方法的TruthfulQA分数都显著下降到30到49表明OptiMer更好地保留了基础模型的校准能力。在效率方面OptiMer的搜索成本比DataMix低15到35倍100次试验的OPTIMER搜索需要8.6小时而单次DataMix运行需要128.9小时。

**Table 1**: Main results comparing OPTIMER with baselines across dataset combinations

**Table 3**: Experiment results based on Gemma-SEA-LION-v4-27B model

**Table 4**: Optimized merge weights and proxy task scores

**Table 5**: Effect of the IT vector weight αit on downstream performance

**Table 6**: Effect of effective batch size on downstream performance

**Table 7**: Case study on HumanEval (sort_third problem)

**Table 8**: Case study on HumanEval (prime_length problem)

**Table 9**: Case study on JCommonsenseQA

**Table 10**: Case study on TruthfulQA (generative)

**Figure 2**: Computational cost comparison between data mixture CPT and OPTIMER

**Figure 4**: PCA projection of distribution vectors with OPTIMER merge weights in bar charts

查看结构化数据

任务	指标	本文	基线	提升
跨语言加数学	平均分（16个任务）	69.98	67.86 (DataMix)	+2.12
跨语言加代码	平均分（16个任务）	69.68	67.23 (DataMix)	+2.45
多语言加数学	平均分（16个任务）	70.37	63.71 (DataMix)	+6.66
TruthfulQA任务	分数（越高越好）	51到55	30到49 (其他方法)	+6到25
搜索效率	搜索时间（小时）	8.6	128.9 (DataMix)	15倍更快

局限与改进

作者承认的局限性包括第一OptiMer目前验证了在1B tokens规模的CPT上的有效性对于更大规模的CPT需要防止CPT模型偏离基础模型太远作者建议未来可以探索迭代训练合并方法。第二实验只在Gemma 3 27B和Gemma-SEA-LION-v4-27B两个模型上验证没有在其他架构如Llama-3、Qwen-3上验证需要进一步研究通用性。第三DataMix基线只使用了均匀混合比例虽然这是常见实践但与DoReMi和RegMix等先进的比例优化方法进行直接比较是未来的重要方向。第四实验采用统一的1-shot提示设置绝对分数可能与使用特定few-shot配置的公共排行榜有所不同。作者自己的观察包括OptiMer需要为每个数据集训练独立的CPT模型这虽然增加了存储成本但训练成本与传统DataMix方法相同都是训练nB tokens。此外OptiMer的成功依赖于分布向量之间的近似正交性如果不同分布之间存在强相关性合并效果可能会下降。

独立分析的弱点

OptiMer的主要弱点包括第一存储开销较大需要存储n个CPT模型的分布向量对于大模型如27B参数这可能需要数百GB的存储空间。改进方向是探索向量压缩或稀疏化技术或者只存储部分层的向量。第二搜索空间随向量数量线性增长当向量数量很多时如超过10个TPE搜索可能需要更多试验才能收敛。改进方向是引入分层搜索或先验知识约束搜索空间。第三当前方法假设分布向量之间近似正交但如果某些分布高度相关如都是编程语言合并可能会产生干扰。改进方向是在搜索时引入正则化项来惩罚高相关性组合或者使用更复杂的合并函数如基于注意力机制的动态合并。第四OptiMer主要优化了验证集上的分数但可能无法保证所有任务上的均衡性能。改进方向是引入多目标优化同时考虑多个任务组的分数。第五当前方法不支持在线学习即在训练过程中动态调整权重。改进方向是探索与训练合并再训练迭代方法结合的在线OptiMer。

未来方向

作者提出的未来工作方向包括第一将OptiMer扩展到中期训练场景在训练过程中动态决定大量领域特定数据集的混合比例。第二探索迭代OptiMer在CPT训练过程中多次进行向量提取和合并防止模型偏离基础模型太远。第三在更多模型架构上验证OptiMer的通用性包括Llama-3、Qwen-3等。第四与DoReMi、RegMix等先进的数据混合优化方法进行直接比较验证OptiMer在这些方法之上的优势。第五探索更复杂的合并函数超越线性组合如基于注意力机制或学习的合并方法。基于论文成果可延伸的方向包括第一将OptiMer应用于联邦学习场景每个客户端训练本地模型后通过OptiMer合并。第二探索OptiMer在模型编辑中的应用通过正负权重组合来添加或删除特定知识。第三研究OptiMer与其他模型合并方法如TIES、DARE的结合进一步提高合并质量。第四开发基于OptiMer的自动化模型库用户可以根据目标自动组合不同的能力向量。

复现评估

论文在复现性方面表现良好。第一作者承诺代码和模型将在GitHub上开源即https://github.com/shyyhs/optimer这提供了完整的实现细节。第二实验使用公开可用的基础模型和数据集包括Gemma 3 27B、LLM-jp Corpus v4以及标准的评估基准MMLU、ARC-Challenge、HellaSwag、GSM8K、HumanEval等。第三实验配置详细描述包括训练超参数学习率4乘以10的负5次方、批次大小128、2000步训练等优化器设置AdamWβ1等于0.90β2等于0.95和合并设置DARE-Linear排除嵌入和位置层。第四评估使用标准的lm-evaluation-harness框架和vLLM后端确保了评估的可重复性。然而复现OptiMer仍然面临一些挑战第一计算资源要求高训练5个1B tokens的CPT模型需要大量的GPU时间论文使用8个NVIDIA H200 GPU。第二搜索过程需要100次TPE试验虽然比DataMix快很多但仍需要约8.6小时。第三模型存储需求大需要存储多个27B参数的模型或其差分向量。总体而言对于有足够计算资源的实验室复现OptiMer是可行的但对于小团队来说仍有一定门槛。

OptiMer：最优分布向量合并优于持续预训练中的数据混合方法 OptiMer: Optimal Distribution Vector Merging Is Better than Data Mixing for Continual Pre-Training