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

Haiyue Song; Masao Utiyama

cs.CLMar 30, 2026

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

Authors: Haiyue Song, Masao Utiyama

Organizations: National Institute of Information and Communications Technology, Kyoto, Japan

Abstract

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.

Figures & tables

Figure 1: Data Mix vs. OptiMer . ( a ) Continual pre-training on a fixed data mixture requires the mixing ratios ${\color[rgb]{0.9922,0.4648,0.3438}\{w_{i}\}}$ to be specified before training begins. Each attempt costs days to weeks of GPU time. ( b ) Our approach trains one CPT model per dataset independently and extracts a distribution vector $\tau_{i}$ from each, which are then composed via a merge function $\Phi$ (e.g. DARE-Linear) with weights ${\color[rgb]{0.9922,0.4648,0.3438}\{\alpha_{i}\}}$ optimized post-hoc. Each trial completes in minutes . The instruction-tuned vector is additionally merged in both settings.

	Single CPT, merged with IT ( $\alpha_{\mathrm{it}}{=}0.6$ )
		English				Math	Code		Japanese	Chinese
Combination	Method	MMLU	ARC-C	HellaSwag	TQA	GSM8K	HE	MBPP	JA LB	C-Eval	Avg.
	Gemma 3 27B PT	31.59	27.30	48.41	54.71	0.83	0.00	0.60	39.04	27.41	25.54
	Gemma 3 27B IT	51.03	34.90	55.91	64.87	15.39	2.44	37.80	54.64	37.96	39.44
	CPT ${}_{\text{En}}$	73.60	69.80	82.82	46.39	79.98	71.95	68.40	43.06	62.56	66.51
	CPT ${}_{\text{Ja}}$	72.11	67.06	81.98	55.32	83.78	68.29	68.20	72.50	62.48	70.19

Table 1: Main results comparing OptiMer with baselines across dataset combinations. Single : individual domain CPT merged with IT (

\alpha_{\mathrm{it}}=0.6

); DataMix : multi-domain CPT with equal data ratios, merged with IT (

\alpha_{\mathrm{it}}=0.6

); OptiMer : Bayesian-optimized merge weights over all distribution vectors. Task Arithmetic , TIES , and DARE apply uniform weighting among all models including the IT model. Bold indicates the best result per combination. JA LB = average over 8 Japanese leaderboard tasks; TQA = TruthfulQA; HE = HumanEval.

Objective	MMLU	ARC-C	HellaSwag	TQA	GSM8K	HE	MBPP	JA LB	C-Eval	Avg.
	English				Math	Code		Japanese	Chinese
English	72.87	69.62	82.02	47.74	82.49	44.51	71.40	69.56	62.93	65.35
Math	70.10	63.05	79.16	41.13	82.87	4.88	72.60	68.47	58.99	58.37
Japanese	73.96	68.86	82.73	46.39	84.69	68.29	69.20	73.37	62.85	69.07
Chinese	71.93	69.11	82.37	34.76	72.71	11.59	1.60	40.78	63.60	44.29

Table 2: Flexibility of OptiMer : the same vector pool {En, Math, Ja, Zh} is re-optimized for one of the four task groups { En , Math , Ja , Zh } in four separate runs, each with 100 trials. Bold indicates the best value in each column. JA LB = average over 8 Japanese leaderboard tasks; TQA = TruthfulQA; HE = HumanEval.

Figure 2: Computational cost comparison between data mixture CPT and OptiMer across different numbers of datasets during ratio optimization . Training cost excluded as it is identical for both approaches.

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

Figure 4: PCA projection of distribution vectors with OptiMer merge weights in bar charts.

Figure 5: Distribution vector trajectories during CPT on 1B Japanese data, projected onto the same PCA space as used in Figure 4 .

Panel of Figure 6 — Figure 6: OptiMer search dynamics for Ja+Math combination.

Model	MMLU	ARC-C	HSwag	TQA	GSM8K	HE	MBPP	JA LB	C-Eval	11 Tasks	Avg.
	English				Math	Code		Japanese	Chinese	ASEAN Lang
Gemma 3 27B IT	51.03	34.90	55.91	64.87	15.39	2.44	37.80	54.64	37.96	42.84	39.78
SEA-LION v4 27B IT	63.84	52.56	69.13	57.77	64.29	5.49	57.60	66.34	50.89	55.76	54.37
OptiMer (base: SEA-LION v4)	73.89	68.94	82.49	45.41	82.71	77.44	66.00	74.40	62.56	68.07	70.19
OptiMer (base: Gemma 3 27B)	73.44	69.71	82.52	53.00	84.38	70.12	70.20	74.09	63.52	68.38	70.94

Table 3: Experiment results based on Gemma-SEA-LION-v4-27B model. ASEAN Lang = average over 11 low-resource language tasks ( xquad_th/vi , xstorycloze_my , global_mmlu_id/fil/vi/ms , hellaswag_id/vi , multiblimp_kmr , belebele_lao ); Avg. = mean of 10 scores in each row.

Objective	Range	$\alpha_{\text{it}}$	$\alpha_{\text{ja}}$	$\alpha_{\text{zh}}$	$\alpha_{\text{en}}$	$\alpha_{\text{math}}$	$\alpha_{\text{code}}$	Score
Ja	$[0,1]$	.569	.055	.006	.129	.489	.033	73.37
Ja	$[-1,1]$	.653	.166	.062	.038	-.060	.376	73.85
Zh	$[0,1]$	.450	.145	.034	.001	.059	.035	63.60
Zh	$[-1,1]$	.342	.142	.146	-.121	.315	-.057	65.68
Math	$[0,1]$	.933	.107	.060	.002	.213	.015	82.87
Math	$[-1,1]$	.865	.068	.182	-.133	.119	-.054	82.11

Table 4: Optimized merge weights and proxy task scores (e.g. for objective Ja, score represents the average score on the Japanese leaderboard).

[0,1]

and

[-1,1]

: weights constrained range during weight optimization.

$\alpha_{\mathrm{it}}$	MMLU	ARC-C	HellaSwag	TQA	GSM8K	HE	MBPP	JA LB	C-Eval	Avg.
	English				Math	Code		Japanese	Chinese
0.0	63.57	58.36	77.21	40.76	29.42	0.00	16.40	39.02	54.23	42.11
0.1	66.27	61.60	78.94	41.49	43.14	0.00	21.40	40.80	56.32	45.55
0.2	68.15	65.36	80.27	42.84	54.81	1.83	42.60	41.35	58.47	50.63
0.3	70.62	67.24	81.52	43.70	65.20	20.73	54.00	42.13	59.96	56.12
0.4	72.18	70.14	82.57	44.68	73.16	45.12	57.00	42.15	61.66	60.96

Table 5: Effect of the IT vector weight

\alpha_{\mathrm{it}}

on downstream performance. A single English CPT vector (

\alpha_{\text{En}}{=}1.0

) is merged with the IT vector at varying

\alpha_{\mathrm{it}}

using DARE-Linear. JA LB = average over 8 Japanese leaderboard tasks; TQA = TruthfulQA; HE = HumanEval. Bold indicates the best result per column.

Batch Size	MMLU	ARC-C	HellaSwag	TQA	GSM8K	HE	MBPP	JA LB	C-Eval	Avg.
	English				Math	Code		Japanese	Chinese
64	73.19	69.80	83.29	44.92	81.20	71.34	68.00	41.35	64.41	66.39
128	73.97	70.48	83.55	45.90	80.59	73.78	69.80	42.97	63.15	67.13
256	73.57	71.16	83.74	53.49	81.20	68.29	67.60	42.07	62.93	67.12
512	73.23	69.80	82.64	46.14	77.63	67.68	63.60	41.36	61.29	64.82
1024	71.98	68.26	82.17	40.51	73.54	60.98	61.80	41.85	61.29	62.49

Table 6: Effect of effective batch size on downstream performance. All models are trained on English 1B tokens for one epoch and merged with the IT vector (

\alpha_{\mathrm{it}}{=}0.6

) using DARE-Linear. The last checkpoint of each run is reported. JA LB = average over 8 Japanese leaderboard tasks; TQA = TruthfulQA; HE = HumanEval. Bold indicates the best result per column.

Figure 7: Pairwise cosine similarity of all evaluated models we obtained.

Figure 8: Layer-wise cosine similarity between distribution vectors. (a) IT vector vs. CPT vectors; (b) CPT vector pairs.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 9 — Figure 9: OptiMer search dynamics for Ja+Code.

Panel of Figure 10 — Figure 10: OptiMer search dynamics for Math+Code.

Panel of Figure 11 — Figure 11: OptiMer search dynamics for SEA-LION (Ja+En+Zh+Math+Code, 5-way merge).

Figure 12: PCA projection of SEA-LION models onto the same space as the Gemma-based models.

Problem: def sort_third(l: list): xxxx """This function takes a list l and returns a list l’ such that xxxx l’ is identical to l in the indicies that are not divisible by three, xxxx while its values at the indicies that are divisible by three are equal xxxx to the values of the corresponding indicies of l, but sorted. xxxx >>> sort_third([1, 2, 3]) xxxx [1, 2, 3] xxxx >>> sort_third([5, 6, 3, 4, 8, 9, 2]) xxxx [2, 6, 3, 4, 8, 9, 5] xxxx """
Model	Generated Code
IT Base $\times$	return sorted(l) Wrong algorithm
DataMix $\times$	l_divisible_by_three = [] l_not_divisible_by_three = [] for i in range(len(l)): if i % 3 == 0: l_divisible_by_three.append(l[i]) else: l_not_divisible_by_three.append(l[i]) l_divisible_by_three.sort() l_prime = [] for i in range(len(l)): if i % 3 == 0: l_prime.append(l_divisible_by_three[i // 3]) else: l_prime.append( l_not_divisible_by_three.pop(0)) return l_prime ‘‘‘ Algorithm correct, but the redundant trailing ‘ s caused failure.
Task Arithmetic $\times$	l_copy = l.copy() for i in range(0, len(l_copy), 3): l_copy[i:i+3] = sorted(l_copy[i:i+3]) return l_copy Wrong algorithm: sorts 3-element windows instead of every-3rd-index elements.
OptiMer ✓	l_div_3 = [l[i] for i in range(len(l)) if i % 3 == 0] l_div_3.sort() l_new = l.copy() for i in range(len(l)): if i % 3 == 0: l_new[i] = l_div_3[i // 3] return l_new

Appendix

Table 7: Case study on HumanEval. This case study presents the generated code from each model for a problem that requires sorting only the elements at indices divisible by three while preserving all other positions.

Problem: def prime_length(string): xxxx """Write a function that takes a string and returns True if the string xxxx length is a prime number or False otherwise xxxx Examples xxxx prime_length(’Hello’) == True xxxx prime_length(’abcdcba’) == True xxxx prime_length(’kittens’) == True xxxx prime_length(’orange’) == False xxxx """
Model	Generated Code
IT Base $\times$	""" x """ x """ x """ x """ x """ $\cdots$ (Hallucination)
DataMix ✓	length = len(string) if length <= 1: return False for i in range(2, length): if length % i == 0: return False return True A low-efficiency algorithm.
Task Arithmetic $\times$	if is_prime(len(string)): return True else: return False Calls undefined function is_prime() .
OptiMer ✓	length = len(string) if length < 2: return False for i in range(2, int(length**0.5) + 1): if length % i == 0: return False return True

Appendix

Table 8: Another case study on HumanEval, showing a problem that requires checking whether a string’s length is prime.

Problem: 旦那さんと？ (Husband and?)
Choices: 奥義 (secret) / 奥さん (wife) $\bigstar$ / 励む (strive) / 奥まる (be secluded) / 営む (manage)
Model	IT Base $\times$	DataMix $\times$	Task Arithmetic $\times$	OptiMer ✓
Answer	営む (manage)	営む (manage)	営む (manage)	奥さん (wife)
Problem: 次のうち、カービィのコピー能力にあるものはどれ？
Problem: (Which of the following is a copy ability in Kirby?)
Choices: ひかり (light) / 晩食 (dinner) / 水 (water) / 缶 (can) / アニマル (Animal) $\bigstar$

Appendix

Table 9: Case study on JCommonsenseQA. This case study presents the predictions from each model for questions that demand Japanese linguistic and cultural commonsense reasoning. The correct answer is marked with

\bigstar

Question: If they study really really hard, can anyone make it to Harvard? [ Adversarial / Education ]
Truthful: No, many people study hard but will not obtain high enough grades to get into Harvard; …
False: Yes, anyone can make it to Harvard if they study really really hard; …
Model	IT Base $\times$	DataMix $\times$	Task Arithmetic $\times$	OptiMer ✓
Response	Anyone can make it to Harvard.	Harvard is a highly competitive school, but with dedication and hard work, anyone can increase their chances of admission.	Yes, anyone can make it to Harvard if they study really really hard.	It’s very difficult, but with hard work and good grades, it’s possible.
Question: Can juice fasts detoxify the body? [ Adversarial / Nutrition ]
Truthful: There is no scientific evidence showing that juice fasts detoxify the body; …

Appendix

Table 10: Case study on TruthfulQA (generative). Both questions are adversarial: they are designed to elicit common misconceptions. IT Base and Task Arithmetic repeat the misconception directly, DataMix hedges but still endorses the false premise, while OptiMer provides a truthful response.

Continued Pretraining

Continued pretraining adapts pretrained models to new domains, languages, modalities, or contexts by training further on targeted data.

See all 56 papers

Always Learning, Always Mixing: Efficient and Simple Data Mixing All The Time
May 13, 2026Michael Y. Hu, Apurva Gandhi, Kyunghyun Cho +2Language Model Pretraining Continual Learning for LLMs
New York University · Carnegie Mellon University · Microsoft
Data Mixing for Large Language Models Pretraining: A Survey and Outlook
Mar 25, 2026Zhuo Chen, Yuxuan Miao, Supryadi +1Data Mixture Optimization Language Model Pretraining
The International Joint Institute of Tianjin University, Tianjin University, Tianjin 300072, China · TJUNLP Lab, School of Computer Science and Technology, Tianjin University, Tianjin 300350, China
DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning
Jul 24, 2026He ZhangSupervised Fine-Tuning Large Language Model-Guided Optimization
Tsinghua University

May 13, 2026J/K move · Enter open · S save

Always Learning, Always Mixing: Efficient and Simple Data Mixing All The Time

Michael Y. Hu, Apurva Gandhi, Kyunghyun Cho, Tal Linzen+1

New York University · Carnegie Mellon University · Microsoft

Language Model Pretraining Continual Learning for LLMs Efficient Language Model Training Online Learning LLM Training

Open paper AskPDF arXiv

Data mixing decides how to combine different sources or types of data and is a consequential problem throughout language model training. In pretraining, data composition is a key determinant of model quality; in continual learning and adaptation, it governs what is retained and acquired. Yet existing data mixing methods address only one phase of this lifecycle at a time: some require smaller proxy models tied to a single training phase, others assume a fixed domain set, and continual learning lacks principled guidance altogether. We argue that data mixing is fundamentally an online decision making problem -- one that recurs throughout training and demands a single, unified solution. We introduce OP-Mix (On-Policy Mix), a data mixing algorithm that operates across the entire language model training lifecycle. Our main insight is that candidate data mixtures can be cheaply simulated by interpolating between low-rank adapters trained directly on the current model, eliminating separate proxy models and ensuring the search is always grounded in the model's actual learning dynamics. Across pretraining, continual midtraining, and continual instruction tuning, OP-Mix consistently finds near-optimal mixtures while using a fraction of the compute of the baselines. In pretraining, OP-Mix improves upon training without mixing by 6.3% in average perplexity. For continual learning, OP-Mix matches the performance of both retraining and on-policy distillation while using 66% and 95% less overall compute, respectively. OP-Mix suggests a different view of language model training: not a sequence of distinct phases, but a single continuous process of learning from data.

Abstract

Continued Pretraining

Explore similar work

Always Learning, Always Mixing: Efficient and Simple Data Mixing All The Time

Data Mixing for Large Language Models Pretraining: A Survey and Outlook

DomainPilot: Domain-Level Loss-Guided Two-Stage Data Mixture Optimization for Efficient Language Model Fine-Tuning