Domain Mixture Design via Log-Likelihood Differences for Aligning Language Models with a Target Model

TL;DR

This paper introduces a novel method for aligning language models with a target model by designing domain mixtures based on log-likelihood differences, improving model similarity without direct distillation.

Contribution

The study presents a new approach to domain mixture design using log-likelihood differences to align models, offering an alternative to traditional knowledge distillation.

Findings

The method reduces KL divergence to the target model compared to uniform weighting.

It achieves meaningful model alignment and improves downstream task performance.

Knowledge distillation remains more effective when available.

Abstract

Instead of directly distilling a language model, this study addresses the problem of aligning a base model with a target model in distribution by designing the domain mixture of training data for pretraining or continued pretraining as a fixed training recipe. We propose a method for determining domain weights by viewing models as points in log-likelihood space and aligning the training update direction with the direction toward the target model. Experiments with NanoGPT show that the proposed method consistently reduces the KL divergence to the target model compared with uniform weighting over the Pile. Although knowledge distillation remains more effective when available, the proposed method still achieves meaningful alignment, and downstream task performance also tends to become closer to that of the target model.