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

arXiv: https://arxiv.org/abs/2605.07756 · HTML: https://mlsanity.com/paper/2605.07756/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 34 references that are in the map.
- First posted on arXiv 8 May 2026.

## Builds on (most cited first, 28 of 28)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [Bilevel Programming for Hyperparameter Optimization and Meta-Learning](https://mlsanity.com/paper/1806.04910/) (2018)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [With a Little Help from My Friends: Nearest-Neighbor Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2104.14548/) (2021)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Truncated Back-propagation for Bilevel Optimization](https://mlsanity.com/paper/1810.10667/) (2018)
- [Solo-learn: A Library of Self-supervised Methods for Visual Representation Learning](https://mlsanity.com/paper/2108.01775/) (2021)
- [Adapting Auxiliary Losses Using Gradient Similarity](https://mlsanity.com/paper/1812.02224/) (2018)
- [Tabular Transformers for Modeling Multivariate Time Series](https://mlsanity.com/paper/2011.01843/) (2020)
- [Scalable Gradient-Based Tuning of Continuous Regularization Hyperparameters](https://mlsanity.com/paper/1511.06727/) (2015)
- [Independent Component Alignment for Multi-Task Learning](https://mlsanity.com/paper/2305.19000/) (2023)
- [CoLES: Contrastive Learning for Event Sequences with Self-Supervision](https://mlsanity.com/paper/2002.08232/) (2020)
- [Jacobian Descent for Multi-Objective Optimization](https://mlsanity.com/paper/2406.16232/) (2024)
- [All4One: Symbiotic Neighbour Contrastive Learning via Self-Attention and Redundancy Reduction](https://mlsanity.com/paper/2303.09417/) (2023)
- [DrMAD: Distilling Reverse-Mode Automatic Differentiation for Optimizing Hyperparameters of Deep Neural Networks](https://mlsanity.com/paper/1601.00917/) (2016)
- [LDC-MTL: Balancing Multi-Task Learning through Scalable Loss Discrepancy Control](https://mlsanity.com/paper/2502.08585/) (2025)
- [Sample-Level Weighting for Multi-Task Learning with Auxiliary Tasks](https://mlsanity.com/paper/2306.04519/) (2023)
- [HT-Transformer: Event Sequences Classification by Accumulating Prefix Information with History Tokens](https://mlsanity.com/paper/2508.01474/) (2025)
- [Unprejudiced Training Auxiliary Tasks Makes Primary Better: A Multi-Task Learning Perspective](https://mlsanity.com/paper/2412.19547/) (2024)

## Abstract (arXiv)

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.