Velocitune: A Velocity-based Dynamic Domain Reweighting Method for Continual Pre-training
This addresses the challenge of efficiently training models on diverse corpora for practitioners in NLP, though it appears incremental as it builds on existing domain reweighting methods.
The paper tackles the problem of domain-adaptive continual pre-training for large language models by proposing Velocitune, a framework that dynamically adjusts data proportions based on learning velocity, resulting in performance gains in math and code reasoning tasks and command-line generation benchmarks.
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.