PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Alexey Rukhovich, Alexander Podolskiy, Irina Piontkovskaya

arXiv:2501.15556v1cs.LGcs.CL

TL;DR

Domain order and data mixing remain under-studied despite the diversity of training sources and the lack of a rigorous approach. The paper develops a criterion for assessing order changes and guiding domain-weight shifts; it reports validation on a toy example and bilingual LLM pre-training. For two domains, the criterion can favor training on one domain at a time rather than mixing when its stated condition holds.

  • Problem

    How to mix diverse domain data and when to use each domain remains an under-studied question without a rigorous approach.

    Source: §1 · p. 1

  • Method

    The paper develops a theoretical framework for predicting training-order effects and guiding domain-weight adjustments along a schedule.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    For two domains with P(L1, L2; L) ≠ 0, the criterion favors training entirely on one domain rather than mixing, and its predictions were validated on a toy example and bilingual LLM pre-training.

    Source: §1 · p. 2, §3.2 · p. 4

  • Takeaways & Limitations

    With total domain sizes fixed, the criterion indicates a direction for locally improving an existing schedule with respect to a target loss.

    Source: §1 · p. 2, §3.2 · p. 4

    The method does not produce an optimal weight schedule; it only indicates how to shift an existing suboptimal one.

    Source: §6.1 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames domain order as an under-studied training choice without a rigorous approach, then proposes a framework for predicting its effects and guiding weight adjustments. It motivates the question with varied domain data and validates predictions on a toy example and bilingual LLM pre-training.

1 Introduction

  • Domain mixing and timing remain without a rigorous approach, while practitioners generally rely on handcrafted solutions.The question includes when to use each domain, not only the proportions used during training.Source: §1 · p. 1
  • Prior multilingual evidence [6] reports better results from high-resource-first exposure followed by an equal high- and low-resource mix under dataset imbalance.The paper takes up domain order under a fixed training budget.Source: §1 · p. 1
  • Unlike multi-task work [28] [21] [40] [36], the paper does not interfere with domain gradients; unlike domain adaptation [19], it does not interpolate between domains.It focuses on interactions between domains to improve the training trajectory.Source: §1 · p. 1
  • With total data fixed per domain, the paper varies the interleaving order; using the same high-quality-data proportion in every batch may be suboptimal.A common practice is to increase the proportion of high-quality data later in training.Source: §1 · p. 2
  • The paper introduces a framework for predicting how training-order changes affect performance and for guiding domain-weight adjustments along a training trajectory.The framework concerns the direction in which to shift an existing schedule, rather than an explicit schedule-generation algorithm.Source: §1 · p. 2
  • The framework’s predictions are validated on a toy example and bilingual LLM pre-training.The paper studies bilingual pre-training as a setting where languages can differ substantially in available data.Source: §1 · pp. 1–2

2 Multi-Domain Learning

This section describes multi-domain training through domain losses and weights, including a weighted-loss formulation and an equivalent sampling view for forming mini-batches.

2 Multi-Domain Learning

  • The section distinguishes multi-objective optimization over the Pareto front from loss averaging, which uses a weighted sum of dataset losses.For the loss-averaging setup, training reweights domain gradients.Source: §2 · p. 2
  • At each step, domain weights can be interpreted as probabilities for sampling examples from the domains into a mini-batch.The weights are indexed by training step and constrained to sum to one.Source: §2 · p. 2
  • The paper uses the sampling-based formulation in its experiments, reporting that it leads to better performance in practice.This formulation expresses domain weights through batch composition.Source: §2 · p. 2

3 Vector Fields and Their Commutators

Domain-training order can alter the resulting trajectory when gradient flows do not commute. This section develops a Lie-bracket criterion for assessing whether local order changes improve a target loss.

3 Vector Fields and Their Commutators

  • Smooth domain losses define gradient flows; commuting flows make the training result depend only on the amount of data per domain, while noncommuting flows make order consequential.The analysis assumes an identity Euclidean metric for gradients in standard coordinates.Source: §3 · pp. 2–3

3.1 Toy example

  • For quadratic domain losses, noncommuting matrices yield order-dependent flow compositions, and even commuting matrices can produce different results when domain centers differ.The latter difference is the constant vector (1 − e^t1A1)(1 − e^t2A2)(b1 − b2).Source: §3.1 · p. 3

3.2 General case

  • The Lie-bracket criterion tests schedule optimality: when its pairwise value is nonzero, at least one of the two domain weights must be zero; violating the criterion signals a locally better ordering.It can be used to assess checkpointed schedules and suggest weight-schedule corrections, but it does not itself construct an optimal trajectory.Source: §3.2 · p. 4

4 Experiments

The experiments compare domain-order interventions with constant-weight training on a quadratic toy problem and bilingual LLM pre-training. In the LLM experiments, changing the schedule affects language losses even when token totals are fixed, while theoretical predictions are directionally correct at most checkpoints but noisy.

4.1 Quadratic optimization

  • The toy experiment compares reordered domain training trajectories against a constant-weight baseline using equal total training data, then compares predicted and actual excess losses.The setup varies starting steps and intervention lengths, but the supplied passages do not report the numerical comparison.Source: §4.1 · p. 5

4.2 Bilingual LLM pre-training

  • The bilingual study uses a small GPT-2 pretrained on English C4 and Russian mC4, following the standard Adam pipeline despite theory stated for gradient descent.The weight interventions start from a constant (0.5, 0.5) schedule and reorder periods with shifted language weights.Source: §4.2 · pp. 5–6
  • Schedule interventions affect bilingual LLM losses despite fixed token totals; later checkpoints form Pareto fronts, and intervention generally improves the language receiving more weight while worsening the other.Large interventions harm performance at early checkpoints.Source: §4.2 · p. 6
  • Estimating the theoretical proxy uses averaged domain gradients and Hessian-vector products because direct Hessian computation is intractable.The experiments average over 100 batches for each computation.Source: §4.2 · p. 6
  • For three of four bilingual checkpoints, the predicted direction of intervention-induced loss change is correct, though its magnitude is noisy.The authors cite Adam, gradient stochasticity, and the non-local 4,000-step intervention as possible sources of discrepancy.Source: §4.2 · p. 7
  • With imbalanced bilingual data, changing the schedule may benefit the low-resource language while causing a negligible decrease in the high-resource language.Source: §4.2 · p. 7

5 Related Work

Related work addresses multi-domain weighting through scalarized or adaptive methods, gradient manipulation, and influence-based analyses. This paper instead studies how changing domain order affects training.

5 Related Work

  • Adaptive multi-domain methods set weights using loss-based balancing or gradient-based recalculation, while scalarization uses fixed weights.Gradient-based approaches require multiple backward operations per step.Source: §5 · p. 7
  • Multi-task methods often manipulate gradients to avoid conflicting directions, whereas this work changes domain order without interfering with gradients.The related work also notes domain adaptation and curriculum learning as neighboring directions.Source: §5 · p. 7
  • Influence-function approaches use an inverse Hessian, unlike this paper’s visually similar formula, which uses a Hessian.The Lie-bracket analysis also requires Hessian-vector products.Source: §5 · pp. 7–8

6 Conclusion

The paper proposes a theoretical framework for analyzing and predicting how training order affects multi-domain learning, with demonstrations in toy examples and multilingual LLM pre-training. Its current limitations include the lack of an optimal-schedule algorithm, optimizer differences, and stochasticity.

6 Conclusion

  • The conclusion presents a framework for predicting training-order effects and reports demonstrations on toy examples and multilingual LLM pre-training.The LLM results illustrate variation associated with dataset order alone.Source: §6 · p. 8

6.1 Limitations

  • The method indicates which direction to shift an existing suboptimal schedule but does not produce an explicit optimal weight schedule.Source: §6.1 · p. 8
  • The method does not account for differences between optimizers, although optimizer choice can affect training dynamics and loss landscapes.Optimizer modifications for multi-domain or multi-task learning would also change the commutator formulae.Source: §6.1 · p. 8
  • The analysis does not properly handle training stochasticity, whose order-dependent effects may alter both drift and diffusion.Formal derivation and empirical validation of the diffusion dependence are left for future work.Source: §6.1 · p. 8

A Proofs

The proof expands the two training flows and compares their compositions to expose the order-dependent commutator. A schedule perturbation then yields a negative target-loss difference for sufficiently small perturbations, identifying the initial trajectory as suboptimal.

A Proofs

  • Expanding the two flow compositions gives their leading-order difference as a commutator term proportional to the product of the two step lengths.Source: A Proofs · pp. 10–11
  • The proof of Corollary 3.2 begins by deriving part (a) directly from the preceding formula.The supplied passage also states the contradiction setup for part (b).Source: A Proofs · p. 11
  • For a valid perturbed schedule, the target-loss difference is negative for small perturbations, so the initial trajectory is suboptimal.Source: A Proofs · p. 11

B Bi-lingual LLM pre-training with language imbalance

With 90% English and 10% Russian tokens, changing the training order can improve low-resource-language performance with negligible high-resource-language loss. The section also notes that the benefit of order changes for equal-sized domains remains unclear.

B Bi-lingual LLM pre-training with language imbalance

  • For equal-sized domains, the order change shifts the point along red curves, but whether it can lower L = L1+L2 remains unclear.This motivates the separate imbalanced-domain experiment.Source: B Bi-lingual LLM pre-training with… · p. 11
  • The experiment intervenes on a 90% English and 10% Russian training mix, testing weight changes of 0.03, 0.06, and 0.09.The total English and Russian tokens remain fixed within each Figure 3 series.Source: B Bi-lingual LLM pre-training with… · pp. 11–12, Figure 3 · p. 12
  • Figure 3 contrasts constant weights w(t) = (0.9, 0.1) with schedules from formula (9), across four fixed English/Russian token-count pairs.The four pairs are (9000, 1000), (14400, 1600), (19800, 2200), and (24200, 2800).Source: Figure 3 · p. 12
  • Table 3 reports excess losses and predicted excess losses for bilingual LLM pre-training with language imbalance.The supplied table caption gives the table’s scope, not its numerical values.Source: Table 3 · p. 12
  • Changing training order can significantly improve low-resource-language performance with negligible decrease in high-resource-language performance; the authors recommend increasing the low-resource share only at pre-training’s end.The experiment uses a 90% English and 10% Russian token mix.Source: B Bi-lingual LLM pre-training with… · pp. 11–12

D Toy example details

The toy example samples matrices and vectors to construct domain losses, then compares sequential domain updates with a combined-loss trajectory. Its trajectories converge to the combined-loss optimum, although some loss behavior is non-monotonic.

D Toy example details

  • The toy setup uses 100-dimensional matrices with power-law eigenvalues, conjugated by random orthogonal matrices, and Gaussian vectors b1 and b2.The vectors are sampled from a standard Gaussian with unit standard deviation in each direction.Source: D Toy example details · pp. 12–13
  • Table 4 reports predicted infinitesimal excess loss P12;i for the domain losses.The caption identifies the quantity but does not provide its values.Source: Table 4 · p. 13
  • All Figure 4 trajectories under the same constant weight schedule converge to the global optimum of L1+L2, despite different initial points.Some trajectories show non-monotonic loss behavior.Source: Figure 4 · p. 13, D Toy example details · p. 13
  • The experiment compares sequential domain-specific flows with the combined-loss flow and measures excess loss against the combined-loss trajectory.It reports the median and 90–10 percentile range and compares these values with the prediction.Source: D Toy example details · p. 13

E LLM pre-training details

The LLM pre-training setup follows a classical GPT-2 configuration using C4 and mC4 data. The authors also describe how they estimate gradients and Hessian-vector products.

E LLM pre-training details

  • Pre-training follows a classical GPT-2 setup with C4 and mC4, an mBERT tokenizer, global batch size 512, and sequence length 1024.The setup uses Megatron-Deepspeed, Adam with β1 = 0.9 and β2 = 0.999, and learning rate 1.5 × 10−4.Source: E LLM pre-training details · p. 13
  • The authors calculate Hessian-vector products with the reverse-on-reverse method and average gradients and Hessian-vector products over 600 random samples per domain.Source: E LLM pre-training details · p. 13
Loading 2501.15556v1…