HBO: Hierarchical Balancing Optimization for Fine-Tuning Large Language Models

TL;DR

This work tackles data imbalance and heterogeneity in fine-tuning large language models by introducing Hierarchical Balancing Optimization (HBO). HBO uses a bilevel optimization framework with a Global Actor and multiple Local Actors to adapt data sampling globally across datasets and locally within each dataset, guided by two reward signals derived from training dynamics. The global reward measures learning progress via the $L_2$ norm of gradients, while the local reward tracks relative improvement using perplexity ratios across difficulty-grouped examples. Across three backbones and nine tasks in multilingual and multitask settings, HBO consistently outperforms baselines, demonstrates robust global-local balancing, and reveals emergent curriculum-like sampling patterns. The approach offers a scalable, reward-guided method to leverage heterogeneous data mixtures for more effective LLM fine-tuning with practical implications for cross-domain and cross-language generalization.

Abstract

Fine-tuning large language models (LLMs) on a mixture of diverse datasets poses challenges due to data imbalance and heterogeneity. Existing methods often address these issues across datasets (globally) but overlook the imbalance and heterogeneity within individual datasets (locally), which limits their effectiveness. We introduce Hierarchical Balancing Optimization (HBO), a novel method that enables LLMs to autonomously adjust data allocation during fine-tuning both across datasets (globally) and within each individual dataset (locally). HBO employs a bilevel optimization strategy with two types of actors: a Global Actor, which balances data sampling across different subsets of the training mixture, and several Local Actors, which optimizes data usage within each subset based on difficulty levels. These actors are guided by reward functions derived from the LLM's training state, which measure learning progress and relative performance improvement. We evaluate HBO on three LLM backbones across nine diverse tasks in multilingual and multitask setups. Results show that HBO consistently outperforms existing baselines, achieving significant accuracy gains. Our in-depth analysis further demonstrates that both the global actor and local actors of HBO effectively adjust data usage during fine-tuning. HBO provides a comprehensive solution to the challenges of data imbalance and heterogeneity in LLM fine-tuning, enabling more effective training across diverse datasets.

Figures (4)

• Figure 1: The bilevel optimization framework of HBO. Global Actor and Local Actors jointly adjust data sampling probabilities both globally (across datasets) and locally (within datasets) to optimize the LLM parameters. Based on the LLM's training state, the global reward and the local reward are computed to guide the optimization of the global and local actors, respectively.
• Figure 2: The variation of sampling probabilities given by (a) global actor in the multilingual setup, (b) global actor in the multitask setup, (c) local actor in the English subset, and (d) local actor in the Math subset. The model backbone is Llama-3.1-8B.
• Figure 3: Improvements of HBO and MoS compared to the Prop. with Llama-3.1-8B backbone on the MMMLU testset.
• Figure 4: (a) The absolute performance gains of HBO compared to Prop. with different settings of updating frequency for global actor and local actor. (b) The relative runtime overhead introduced by HBO compared to Prop. with different settings of updating frequency for global actor and local actor.