MoL for LLMs: Dual-Loss Optimization to Enhance Domain Expertise While Preserving General Capabilities

TL;DR

MoL addresses the core problem of domain-specific CPT degrading general capabilities and the costly tuning of corpus mixtures. It introduces a dual-loss framework where domain data are trained with CE and general data with KL divergence relative to a base model, guided by a metadata cue; the recommended balance is a 1:1 domain/general corpus ratio. Empirical results on medical-domain CPT show MoL improves domain performance without sacrificing general reasoning, achieving substantial gains on Math-500 and AIME benchmarks. The approach offers a scalable, principled pathway for multi-domain LLM training with practical impact for real-world deployment.

Abstract

Although large language models (LLMs) perform well in general tasks, domain-specific applications suffer from hallucinations and accuracy limitations. Continual Pre-Training (CPT) approaches encounter two key issues: (1) domain-biased data degrades general language skills, and (2) improper corpus-mixture ratios limit effective adaptation. To address these, we propose a novel framework, Mixture of Losses (MoL), which decouples optimization objectives for domain-specific and general corpora. Specifically, cross-entropy (CE) loss is applied to domain-corpus to ensure knowledge acquisition, while Kullback-Leibler (KL) divergence aligns general-corpus training with the base model's foundational capabilities. This dual-loss architecture preserves universal skills while enhancing domain expertise, avoiding catastrophic forgetting. Empirically, we validate that a 1:1 domain-to-general corpus ratio optimally balances training and overfitting without the need for extensive tuning or resource-intensive experiments. Furthermore, our experiments demonstrate significant performance gains compared to traditional CPT approaches, which often suffer from degradation in general language capabilities; our model achieves 27.9% higher accuracy on the Math-500 benchmark in the non-think reasoning mode, and an impressive 83.3% improvement on the challenging AIME25 subset in the think mode, underscoring the effectiveness of our approach.

Figures (3)

• Figure 1: Schematic illustration of the MoL framework architecture. Unlike existed single-objective pre-training approaches, our MoL framework introduces an additional metadata input ("add-info") to distinguish between domain-specific and general corpora during training. This information determines the loss function selection: CE loss for domain corpora and KL divergence loss for general corpora (highlighted in red). The model's forward computation and backpropagation mechanisms retain the standard implementation pipeline of traditional LLMs.
• Figure 2: Training loss evolution across aggregated datasets and individual subsets, depicting CE loss dynamics for both domain-specific and general corpora. "train_MoL" represents the loss on the training set under the MOL framework. The domain-corpus include medical-o1-sft, MedicationQA, and MedQA, while Light-R1 is the general-corpus. The "_val" indicates the validation set, and "_CE" denotes CE loss. The validation set's CE loss is calculated every 10 steps, resulting in a smoother curve compared to the training set loss. The CE loss for general-corpus remains nearly constant throughout training, while the domain-corpus exhibits a steady decline in loss until reaching convergence at nearly 1.3 epochs (marked by the dashed Line).
• Figure 3: (A) Comparison of CE loss across various validation sets during training on Qwen3-8B. The main plot depicts the loss trends across different domain-corpus, including medical-o1-sft, MedQA and MedicationQA, while the inset offers a magnified view focusing on the performance of the medical-o1-sft under the two configurations. Specifically, D&G 1:1 refers to the setup described in Table \ref{['table1']}, while D&G 1:1 CE denotes the alternative configuration that replaces KL divergence with CE throughout the training. The "_val" indicates the validation set, and "_CE" denotes CE loss. (B) Comparison of gradient norm during training. The main plot illustrates the gradient norm dynamics of the D&G 1:1 and D&G 1:1 CE configurations across training epochs. While both configurations exhibit similar temporal evolution patterns, the D&G 1:1 consistently shows smaller magnitudes than D&G 1:1 CE until convergence, after which the two curves align closely. The inset provides a zoomed-in view centered at 1.3 training epochs.