Velocitune: A Velocity-based Dynamic Domain Reweighting Method for Continual Pre-training

TL;DR

Velocitune introduces a dynamic domain reweighting framework for continual pre-training, improving language model performance by adaptively balancing domain learning velocities based on a scaling law.

Contribution

It proposes Velocitune, a novel method that dynamically adjusts domain data proportions during pre-training using learning velocity assessments and a scaling law, addressing domain-adaptive continual pre-training challenges.

Findings

Improves performance on math and code reasoning tasks.

Enhances command-line generation benchmarks.

Effective due to target loss prediction and data ordering.

Abstract

It is well-known that a diverse corpus is critical for training large language models, which are typically constructed from a mixture of various domains. In general, previous efforts resort to sampling training data from different domains with static proportions, as well as adjusting data proportions during training. However, few methods have addressed the complexities of domain-adaptive continual pre-training. To fill this gap, we propose Velocitune, a novel framework dynamically assesses learning velocity and adjusts data proportions accordingly, favoring slower-learning domains while shunning faster-learning ones, which is guided by a scaling law to indicate the desired learning goal for each domain with less associated cost. To evaluate the effectiveness of Velocitune, we conduct experiments in a reasoning-focused dataset with CodeLlama, as well as in a corpus specialised for system command generation with Llama3 and Mistral. Velocitune achieves performance gains in both math and code reasoning tasks and command-line generation benchmarks. Further analysis reveals that key factors driving Velocitune's effectiveness include target loss prediction and data ordering.