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

2026-05-08

AI Summary

This work addresses the challenge of costly hyperparameter tuning for multitask loss weights in large-scale pretraining by proposing a gradient-based bilevel optimization method that learns optimal weights online. The approach aligns the composite pretraining gradient with the downstream target objective and innovatively leverages loss structure to avoid multiple backward passes, achieving weight adaptation at only ~130% the computational cost of a single training run. By integrating gradient alignment, bilevel optimization, and truncated backpropagation, the method matches or exceeds the performance of carefully tuned baselines on event sequence modeling and self-supervised vision tasks while substantially reducing tuning overhead.

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.

Problem

• composite loss
• loss weighting
• pretraining
• hyperparameter tuning
• missing labels

Innovation

• gradient-based bilevel optimization
• composite loss weighting
• efficient pretraining
• gradient alignment
• hyperparameter tuning