Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Zheheng Luo, Xin Zhang, Xiao Liu, Haoling Li, Yeyun Gong, Chen Qi, Peng Cheng

TL;DR

Velocitune introduces a velocity-based domain reweighting framework for continual pre-training that dynamically allocates training data across multiple domains. By estimating per-domain learning velocity with initial and target losses derived from a scaling-law target, and updating domain weights via an exponential(Group DRO-like) rule, it balances learning progress and improves downstream performance. Across CodeLlama-7B reasoning data and Llama3/Mistral SystemStack benchmarks, Velocitune yields consistent gains in math, code, and system-command tasks, with ablations showing target-loss prediction and data ordering as key drivers. The approach highlights that dynamic data ordering plus velocity-aware reweighting can surpass static mixtures, though it incurs higher training costs and requires further evaluation in pre-training-from-scratch and SFT settings.

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.

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

TL;DR

Velocitune introduces a velocity-based domain reweighting framework for continual pre-training that dynamically allocates training data across multiple domains. By estimating per-domain learning velocity with initial and target losses derived from a scaling-law target, and updating domain weights via an exponential(Group DRO-like) rule, it balances learning progress and improves downstream performance. Across CodeLlama-7B reasoning data and Llama3/Mistral SystemStack benchmarks, Velocitune yields consistent gains in math, code, and system-command tasks, with ablations showing target-loss prediction and data ordering as key drivers. The approach highlights that dynamic data ordering plus velocity-aware reweighting can surpass static mixtures, though it incurs higher training costs and requires further evaluation in pre-training-from-scratch and SFT settings.

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.

Paper Structure

This paper contains 28 sections, 4 equations, 5 figures, 10 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Domain-adaptive continual pre-training with Velocitune
  3. The Velocitune algorithm
  4. Setup
  5. Target Estimation
  6. Velocity-Guided Training
  7. Experiment
  8. Experimental setup
  9. Training Corpus
  10. Training Setup
  11. Evaluation
  12. Reasoning training results
  13. Velocitune learns math and maintains coding ability.
  14. Velocitune dynamically regularises the learning velocity.
  15. SystemStack training results
  16. ...and 13 more sections

Figures (5)

  • Figure 1: The overall pipeline of Velocitune. Initially, a proxy model is trained using the original domain weights on a subset of the data. Following this, the initial loss is collected by evaluating the base model, while the target loss is determined by extrapolating the evaluation loss of the proxy model. In the second phase, we calculate the learning velocity by rescaling the learning progress between the initial and target losses. This learning velocity is then used to update the domain weights effectively.
  • Figure 2: Domain weights dynamic of Velocitune in training CodeLlama 7B on Reasoning. The vertical axis represents domain weights, while the horizontal axis denotes training steps.
  • Figure 3: Domain weights dynamic of Velocitune and DBL in training Llama-3 8B on SystemStack. The vertical axis represents domain weights, while the horizontal axis denotes training steps.
  • Figure 4: Domain weights dynamic of Velocitune without target loss in training CodeLlama 7B on Math&Code.
  • Figure : Velocitune