Commute Your Domains: Trajectory Optimality Criterion for Multi-Domain Learning

TL;DR

The paper tackles how the order of multi-domain data exposure affects training outcomes by introducing a Lie-bracket framework over gradient vector fields. It derives a local optimality criterion and an explicit second-order expansion showing how noncommuting losses can bias training trajectories, often favoring single-domain focus unless a zero-projection condition holds. The authors validate their theory with a toy quadratic MDL and a bilingual LLM pre-training setup, demonstrating that the direction and magnitude of performance changes align with predictions from the Lie-bracket-based quantity $P(L_i-L_j,\sum_k w_k L_k;L)$. While providing a principled tool for analyzing data mixing, the work stops short of prescribing an explicit optimal schedule and acknowledges limitations related to optimizers and stochasticity, pointing to future work on online data mixing and broader domain settings.

Abstract

In multi-domain learning, a single model is trained on diverse data domains to leverage shared knowledge and improve generalization. The order in which the data from these domains is used for training can significantly affect the model's performance on each domain. However, this dependence is under-studied. In this paper, we investigate the influence of training order (or data mixing) in multi-domain learning using the concept of Lie bracket of gradient vector fields. By analyzing the infinitesimal effects of changing the training order, we identify regions in the parameter space where altering the order between two training domains can benefit the target loss. We validate the predictions of our theoretical framework on the influence of training order (or data mixing) both on a toy example and bilingual LLM pre-training.

Figures (4)

• Figure 1: Example of vector fields Lie bracket. (a), (b): the level curves and gradients for two functions $L_1, L_2$. Vectors in (c) correspond to Lie bracket of gradient vector fields. In the green area, the Lie bracket has positive dot products with both $\nabla L_1$ and $\nabla L_2$, and in purple area similar dot products are negative. It means (see Corollary \ref{['corollary']}) that in these areas some reordering of the default training would benefit both losses.
• Figure 2: The results of intervening into domain weight schedule for bilingual LLM training. The blue points constitute loss values (negative log likelihoods for English and Russian data) for a training trajectory with constant domain weight schedule $w(t)=(0.5,0.5)$, from a checkpoint at step 4000 (right, top) to step 28000 (bottom, left). The red points correspond to weight schedules from formula (\ref{['w12_schedule']}). The total amount of English and Russian tokens used for training is constant across each red curve: equal of 5000, 8000, 11000 and 14000 batches for each language for ●, ✚, ■ and ★, respectively; only the degree of intervention $\Delta w$ changes from $0.15$ for annotated points to $0.45$ for the most distant points.
• Figure 3: The results of intervening into domain weight schedule for bi-lingual LLM training with language imbalance. The blue points constitute loss values for a training trajectory with constant domain weight schedule $w(t)=(0.9,0.1)$, from a checkpoint at step 4000 (right, top) to step 25000 (bottom, left). The red points correspond to weight schedules from formula (\ref{['w12_schedule']}). The total amount of English and Russian tokens used for training is constant across each series of points: (9000, 1000), (14400, 1600), (19800, 2200), (24200, 2800) for ●, ✚, ■ and ★, respectively; only the degree of intervention $\Delta w$ changes from $0.03$ for annotated points to $0.09$ for the most distant points.
• Figure 4: Domain losses dynamics. Here X and Y axes correspond to values of $L_1$ and $L_2$ and curves are trajectories for same constant weight schedule but different initial points. All trajectories converge to a single point which corresponds to the global optimum of $\frac{L_1+L_2}{2}$, though we see that loss behavior is non-monotonous for some of the trajectories.

Theorems & Definitions (5)

• Remark
• Theorem 3.1
• Corollary 3.2
• proof : Proof of Theorem \ref{['thm_comm']}
• proof : Proof of Corollary \ref{['corollary']}