JADORE: Online Data Mixing Through the Lens of Optimization Geometry

Training data mixtures greatly impact the generalization performance of language models. Existing gradient-based mixing methods rely on raw gradients, ignoring the optimizer dynamics. We introduce *JAcobian DOmain REweighting* (JADORE), an online data mixing framework that estimates the local effect of mixtures on optimization geometry. It defines the domain scores based on the derivative of the optimizer mapping from raw gradients to actual updates. We derive the action of the optimizer's Jacobian for multiple update geometries and show its efficient computation. To ensure scalability, we develop a single backward pass implementation for both training and scoring. Extensive experiments demonstrate accuracy improvements across diverse benchmarks in both pretraining and fine-tuning on multiple datasets. These gains are consistent across optimizers and model scales.