PaperlayerView exact revision on arXiv

Source-linked AI summary

OptiMer: Optimal Distribution Vector Merging Is Better than Data Mixing for Continual Pre-Training

Haiyue Song, Masao Utiyama

arXiv:2603.28858v2cs.CLcs.AIcs.LG

TL;DR

Continual pre-training needs effective data mixtures, but ratios are costly to tune and must traditionally be fixed before training. OPTIMER trains one CPT model per dataset, extracts distribution vectors, and optimizes their composition weights post-hoc with Bayesian optimization. On Gemma 3 27B across languages and domains, it consistently outperforms data-mixture and model-merging baselines with 15–35× lower search cost, while supporting objective-specific models without retraining.

  • Problem

    Mixture ratios are sensitive CPT hyperparameters that must be fixed before training, so poor choices can waste days or weeks of GPU time.

  • Method

    OPTIMER trains separate per-dataset CPT models, extracts distribution vectors, and uses Bayesian optimization to select post-hoc merge weights.

  • Results

    OPTIMER consistently outperforms data-mixture CPT and four model-merging methods across 16 benchmarks and three dataset combinations with 15–35× lower search cost.

  • Takeaways & Limitations

    A single distribution-vector pool can be re-optimized for different objectives, producing multiple target-tailored models without re-CPT.

  • Takeaways & Limitations

    At larger-scale CPT, the authors identify preventing excessive divergence from the base model as necessary and leave iterative OPTIMER for future work.

Abstract

from arXiv · show

Continual pre-training is widely used to adapt LLMs to target languages and domains, yet the mixture ratio of training data remains a sensitive hyperparameter that is expensive to tune: they must be fixed before training begins, and a suboptimal choice can waste weeks of compute. In this work, we propose OptiMer, which decouples ratio selection from training: we train one CPT model per dataset, extract each model's distribution vector, which represents the parameter shift induced by that dataset, and search for optimal composition weights post-hoc via Bayesian optimization. Experiments on Gemma 3 27B across languages (Japanese, Chinese) and domains (Math, Code) show that OptiMer consistently outperforms data mixture and model averaging baselines with 15-35 times lower search cost. Key findings reveal that 1) the optimized weights can be interpreted as data mixture ratios, and retraining with these ratios improves data mixture CPT, and 2) the same vector pool can be re-optimized for a given objective without any retraining, producing target-tailored models on demand. Our work establishes that data mixture ratio selection, traditionally a pre-training decision, can be reformulated as a post-hoc optimization over distribution vectors, offering a more flexible paradigm for continual pre-training.

1 Introduction

OPTIMER reframes continual-pretraining data-ratio selection as post-hoc optimization over independently trained dataset distribution vectors. Across language and domain adaptation, it improves over data-mixture and merging baselines while greatly reducing search cost.

  • Motivation: OPTIMER addresses sensitive mixture ratios that must otherwise be fixed before training, risking days or weeks of wasted GPU time when suboptimal.Prior ratio estimates from proxy models or small experiments cannot be corrected after training begins.
  • Approach: OPTIMER decouples data-ratio selection from training by optimizing merge weights over distribution vectors after separate per-dataset CPT.Each distribution vector represents the parameter shift from the base pretrained model, and Bayesian optimization selects the composition weights.
  • Implications: Optimized merge weights can be interpreted as data-mixture ratios, and a single vector pool can be re-optimized for different objectives without retraining.Retraining with optimized ratios improves DataMix, while direct OPTIMER merging remains strongest in the reported experiments.
  • Results: 15–35× lower search cost accompanies OPTIMER’s consistent improvement over data-mixture CPT and four model-merging methods across 16 benchmarks.The benchmarks span English, Japanese, Chinese, Math, and Code, across three dataset combinations.
  • Analysis: Distribution vectors are approximately orthogonal, with cosine values of 0.03–0.31, supporting composition with limited interference.The analysis also reports approximately linear CPT trajectories and allows negative optimized weights to remove cross-distribution interference.

2 Related Work

Prior work optimizes data mixtures before large-scale training or automates merge-weight search under computational constraints. OPTIMER instead targets post-hoc ratio adjustment using distribution-vector merging.

  • Continual Pre-training: CPT adapts pretrained LLMs to languages and domains, but mixture ratios strongly affect performance and motivate data-mixture optimization.The related work describes language and domain adaptation as established CPT applications.
  • Data Mixture Optimization: DoReMi, RegMix, and related predictive methods estimate optimal ratios through proxy models, many small models, or cross-scale transfer before training.These approaches still require fixing the ratio before the target training run.
  • Task Vectors and Model Merging: Task-vector research composes fine-tuning shifts through linear arithmetic, with later methods addressing sign conflicts, sparsity, and instruction-following transfer.The cited work motivates treating dataset-induced parameter shifts as composable vectors.
  • Automatic Merge Weight Search: Automated merge-weight methods use entropy minimization, evolutionary search, or output-divergence objectives, but generally remain limited to few or small models.High-dimensional search and population-based iterations create computational barriers.
  • Automatic Merge Weight Search: Grid search becomes exponentially more expensive as vector count increases, motivating OPTIMER’s Bayesian-optimization replacement.The related work contrasts grid-search scaling with more efficient search over merge weights.

3 Methodology

OPTIMER represents each dataset’s continual-pretraining effect as a distribution vector and searches post-hoc for merge weights that optimize a development-set score. It uses Bayesian optimization with TPE instead of impractical grid search.

  • Distribution Vectors: Distribution vectors extend task vectors to CPT by encoding the parameter change induced by each training distribution.
  • Multi-Vector Composition: A merged model combines the pretrained base parameters with an instruction-tuning vector and weighted distribution vectors.
  • Multi-Vector Composition: Uniform weighting is a baseline, but different distributions contribute unequally to the target objective, motivating optimized merge weights.
  • Problem Formulation: The weight search maximizes a development-set evaluation score over a black-box objective with no gradient.
  • Bayesian Optimization via TPE: Grid search requires O(G^(n+1)) evaluations, whereas TPE typically converges in O(10^n) trials and supports parallel execution.
  • Bayesian Optimization via TPE: TPE separates observed configurations into high- and low-scoring groups, then samples candidates by maximizing their density ratio.

8 Re-evaluate top-K trials on full Ddev

OPTIMER iteratively updates TPE from scored candidate merges, then re-evaluates the top-K configurations on the full development set to select the final model.

  • Each subsequent trial uses TPE to propose a candidate, constructs and scores its merged model, and updates the density models.
  • After T trials, the top-K configurations are re-evaluated on the full development set to obtain the final model.

4 Experimental Settings

Experiments train separate CPT models on 1B-token language and domain datasets, merge them with instruction-tuning and distribution vectors, and evaluate OPTIMER against data-mixing and averaging baselines.

  • Continual Pre-Training Configuration: The study constructs Japanese, Chinese, Math, and Code CPT datasets, each containing 1B tokens, from LLM-jp Corpus v4.
  • Merge Settings: Merges use DARE-Linear while excluding embedding, language-model head, and rotary layers to preserve the base model’s token representations.
  • OPTIMER Hyperparameters: OPTIMER runs 100 TPE trials with 20 random startup trials in parallel across 8 GPUs.
  • Evaluation Settings: Proxy-task scoring uses the first 100 samples per task during search, followed by top-three re-evaluation on the first 300 samples per task.
  • Baseline Settings: Baselines include equal-ratio DataMix, DataMixOptiMer models trained with OPTIMER-derived ratios, and equal-weight Task Arithmetic, TIES, and DARE merges.
  • Evaluation Settings: Evaluation uses 1-shot prompting across English, Japanese, Chinese, Math, and Code task groups, with Table 1 comparing OPTIMER and baselines across dataset combinations.

5 Results and Analysis

OPTIMER outperforms data-mixture and model-merging baselines while reducing ratio-search cost, and its distribution-vector analysis explains composition, training dynamics, and efficient optimization. The same vector pool can also be re-optimized for different objectives without retraining.

  • Main Results: OPTIMER achieves the highest average score across dataset combinations, outperforming DataMix by 2.1–6.7 points in each group.It also avoids the catastrophic Code-task failures observed with DARE-Linear and preserves stronger TruthfulQA scores of 51–55 versus 30–49 for other methods.
  • Efficiency: 15–35× lower search time makes OPTIMER substantially more efficient than DataMix for ratio optimization.A 100-trial search takes 8.6 hours with OPTIMER versus 128.9 hours for one DataMix run; training costs are identical.
  • Flexibility: OPTIMER re-optimizes the same distribution-vector pool for different objectives, producing target-tailored models without retraining.Objective-specific models generally achieve the highest scores on their target tasks, while the Japanese-optimized model achieves the best overall performance; the cross-lingual transfer effect remains for future work.
  • Distribution Vectors: Distribution vectors exhibit low pairwise similarity, with different CPT vectors at cosine 0.29–0.31, supporting approximately independent subspaces and linear composition.CPT and instruction-tuning vectors are nearly orthogonal at cosine ≈ 0.03; PCA visualizations and optimized-weight charts provide complementary views.
  • Continual Pre-Training Dynamics: CPT trajectories are approximately linear, making merge weight α_i analogous to effective training duration.Performance peaks early at small vector norm and then declines gradually, possibly as the model diverges from the base model.
  • Optimization Dynamics: TPE converges within 100 trials because high-scoring combinations occupy a narrow region, making grid search impractical.Extending weights to [−1, 1] improves Japanese and Chinese objectives, and negative English weights suggest removal of interference.
  • Japanese Optimization: On Gemma-SEA-LION-v4-27B, OPTIMER improves Japanese leaderboard performance from 66.34 to 74.40 (+8.1) while maintaining ASEAN-language performance.Overall average also rises from 54.37 to 70.19 (+15.8).

6 Conclusion and Future Work

OPTIMER reformulates CPT data-mixture ratio selection as post-hoc optimization over distribution vectors. It outperforms data-mixture baselines with lower search time and can produce multiple objective-tailored models from one vector pool without retraining.

  • Conclusion: OPTIMER decouples data-mixture ratio selection from CPT by optimizing weighted distribution-vector composition after training.Separate CPT models are trained per dataset, then their vectors are merged using Bayesian optimization.
  • Conclusion: OPTIMER consistently outperforms data-mixture baselines while requiring 15–35× lower search time.The same vector pool can be re-optimized for multiple objectives without retraining.
  • Future Work: Future work includes extending OPTIMER to mid-training with many domain-specific datasets.The paper identifies mixture-ratio selection in that setting as a central challenge.

Limitations

The experiments have several scope and evaluation boundaries, including limited architecture coverage, uniform DataMix ratios, and a controlled 1-shot prompting setup. Larger-scale CPT may require an iterative variant that the paper leaves for future work.

  • For larger-scale CPT, the paper states that models must be prevented from diverging too far from the base model and proposes iterative OPTIMER as future work.
  • OPTIMER was evaluated on Gemma 3 27B and Gemma-SEA-LION-v4-27B, so generalization to architectures such as Llama-3 and Qwen-3 remains unverified.
  • The DataMix baseline used only uniform mixing ratios, leaving direct 27B-scale comparisons with DoReMi and RegMix for future work.
  • All benchmarks used 1-shot prompting, so absolute scores should not be compared directly with results using task-specific few-shot configurations.

C More Results on Cosine Similarity

Cosine-similarity analyses support composing CPT distribution vectors with limited interference and show that OPTIMER merges preserve alignment with the instruction-tuned model. Search visualizations further indicate that high-performing configurations occupy a narrow IT-weight region.

  • Cosine similarity: CPT vectors have low cosine similarity with the IT vector and with one another, whereas OPTIMER vectors are highly similar to the IT model.This supports composition with limited conflict while preserving IT capability.
  • Layer-wise analysis: IT–CPT cosine similarity is nearly uniform across layers, while CPT–CPT similarity rises in early, middle, and final layers.The paper attributes early and final similarity to frozen position and embedding parameters and notes especially high similarity around middle layer 30.
  • Search dynamics: High-performing search configurations form a narrow layer of IT weights, indicating that the IT ratio is especially sensitive in the explored settings.

E Transfer Experiment on Gemma-SEA-LION-v4-27B

OPTIMER transfers to the Gemma-SEA-LION-v4-27B family by composing five language and domain vectors with the IT alignment vector under Bayesian optimization. The resulting configurations outperform their IT baselines and preserve IT capability, while case studies illustrate improvements over baseline behaviors.

  • Transfer setup: OPTIMER composes Japanese, English, Chinese, Math, and Code distribution vectors with the IT alignment vector on Gemma-SEA-LION-v4-27B.The larger five-vector search uses 500 trials under the same Bayesian-optimization procedure.
  • Transfer results: Both OPTIMER configurations outperform their respective IT baselines, confirming transfer to a different base-model family.
  • Representation analysis: SEA-LION IT and Gemma IT occupy similar PCA positions, while SEA-LION OPTIMER remains similar along the first component, suggesting preserved IT capability.
  • Search dynamics: The five-way search exhibits a sensitive optimal IT ratio: values below it are suboptimal, values above it are catastrophic, and random ratios are also catastrophic.
  • Case studies: Case studies cover HumanEval, JCommonsenseQA, and TruthfulQA, showing baseline failures involving code generation, shared Japanese commonsense errors, and misconceptions.OPTIMER avoids the shared JCommonsenseQA errors and is the only method described as avoiding the TruthfulQA misconception.
Loading 2603.28858v2…