Mix-CPT: A Domain Adaptation Framework via Decoupling Knowledge Learning and Format Alignment

TL;DR

The paper tackles domain adaptation for large language models by decoupling knowledge learning from format alignment. It introduces Mix-CPT, a two-stage framework: first, knowledge mixture continual pre-training (Mix-CPT) that blends domain data with general instruction and alignment data in a unified format, using Logit Swap Self-Distillation to mitigate catastrophic forgetting; second, an efficient format alignment stage that selects easy, perplexity-sorted instruction and preference samples to tune format without introducing new knowledge. Across encyclopedic, mathematical, and coding domains, and for both base and chat LLMs, Mix-CPT consistently improves domain performance while preserving or enhancing general capabilities, outperforming traditional CPT-followed-by-SFT/DPO baselines. The approach reduces data requirements and avoids excessive retraining, offering a practical and scalable pathway for domain adaptation of large language models. Key contributions include the unified data format $\mathcal{D}_{\text{MIX}}$, the LSSD objective, and the perplexity-based sample selection strategy that together enable efficient, robust domain adaptation with minimal loss of general knowledge.

Abstract

Adapting general large language models (LLMs) to specialized domains presents great challenges due to varied data distributions. This adaptation typically requires continual pre-training on massive domain-specific corpora to facilitate knowledge memorization, followed by training to apply this knowledge following human instructions and preferences. However, this method may result in inefficient knowledge memorization due to a lack of awareness of knowledge utilization and imposes substantial demands on LLMs to simultaneously learn knowledge utilization and format alignment with limited training samples. To facilitate the domain adaptation of LLM, we revise this process and propose a new domain adaptation framework including domain knowledge learning and general format alignment, called Mix-CPT. Specifically, we first conduct a knowledge mixture continual pre-training that concurrently focuses on knowledge memorization and utilization, allowing for mutual reinforcement. To avoid catastrophic forgetting during the continual pre-training process, we further incorporate a logit swap self-distillation constraint. Subsequently, leveraging the knowledge and capabilities acquired during continual pre-training, we efficiently perform instruction tuning and alignment with a few general training samples to achieve format alignment. Extensive experiments demonstrate that our proposed Mix-CPT framework can simultaneously improve the task-solving capabilities of LLMs on the target and general domains compared to the traditional adaptation methods.

Figures (6)

• Figure 1: Comparison of traditional domain adaptation approaches (top) and our proposed rescheduled domain adaptation paradigm (bottom). "[EOS]" is the special token representing the end of the document. "[ST]", "[UT]" and "[AT]" denote the system, user, assistant chat template, repesctively.
• Figure 2: The illustration of our proposed rescheduled domain adaptation paradigm, including first conducting knowledge mixture continual pre-training, then selecting top-$K$ easy training samples with the lowest perplexity for performing supervised fine-tuning and direct preference optimization.
• Figure 3: (Left) Pass@K on code and Average results w.r.t. Proportion of code data. (Right) Accuracy on math and Average results w.r.t. Quality score of math data.
• Figure 4: Accuracy on math and Average results w.r.t. Ratio of SFT and DPO data.
• Figure 5: Accuracy on math and Average results w.r.t. selection strategy.