When Losses Align: Gradient-Based Composite Loss Weighting for Efficient Pretraining

TL;DR

This paper introduces a gradient-based bilevel optimization method for efficiently tuning composite loss weights during pretraining, reducing hyperparameter tuning costs while maintaining or improving performance.

Contribution

It presents a novel online loss weighting technique that aligns pretraining gradients with downstream objectives, avoiding extensive hyperparameter search methods.

Findings

Method reduces hyperparameter tuning overhead to ~30% of a single training run.

Achieves comparable or better results than carefully tuned baselines.

Effective in event-sequence modeling and self-supervised vision tasks.

Abstract

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.