When Losses Align: Gradient-Based Composite Loss Weighting for Efficient Pretraining
For practitioners pretraining large models with composite losses, this method reduces the computational cost of hyperparameter tuning, which is otherwise expensive due to multiple training runs.
The paper proposes a gradient-based bilevel method for online learning of loss weights in composite pretraining objectives, reducing hyperparameter tuning overhead to ~30% above a single training run while matching or improving tuned baselines on event-sequence modeling and self-supervised computer vision.
Modern deep models are often pretrained on large-scale data with missing labels using composite objectives, where the relative weights of multiple loss terms act as hyperparameters. Tuning these weights with random search or Bayesian optimization is computationally expensive, as it requires many independent training runs. To address this, we propose a gradient-based bilevel method that learns pretraining loss weights online by aligning the composite pretraining gradient with a downstream objective. By exploiting the structure of the loss, the method avoids the multiple backward passes typically required by truncated backpropagation through the full model, reducing the overhead of hyperparameter tuning to approximately 30% above a single training run. We evaluate the approach on event-sequence modeling and self-supervised computer vision, where it matches or improves upon carefully tuned baselines while substantially reducing the cost of hyperparameter tuning compared to random or Bayesian search.