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

Problem Formulation

Modern deep learning systems frequently employ a two-stage paradigm: pretraining on large-scale, often unlabeled data, followed by adaptation to downstream tasks. In many pretraining pipelines, the objective function is not a single loss but a weighted combination of multiple loss terms (e.g., next-token prediction, contrastive losses, or regression objectives). The relative weights of these terms are critical hyperparameters that significantly influence the quality of learned representations and downstream performance.

Currently, selecting these weights relies on manual tuning, random search, or Bayesian optimization. These methods are computationally prohibitive for large-scale pretraining because they require numerous independent training runs. Furthermore, naive defaults, such as uniform weighting, often fail to account for the varying contributions and conflicts between different loss terms. The paper addresses the need for an efficient method to tune linear loss weights in composite pretraining objectives to maximize downstream task performance without incurring the cost of repeated full training runs.

Methodology: Gradient-aligned Pretraining (GraP)

The authors propose Gradient-aligned Pretraining (GraP), a gradient-based bilevel optimization method that learns loss weights online. The core idea is to align the composite pretraining gradient with the gradient of a downstream objective.

1. Bilevel Optimization Formulation

The problem is cast as a bilevel optimization task where the outer loop optimizes the loss weights $w$ to minimize the downstream loss $L_{down}$, and the inner loop optimizes model parameters $\theta$ to minimize the composite pretraining loss $L_{pre}(w)$.
$$\min_{w} L_{down}(\theta^*(w)) \quad \text{subject to} \quad \theta^*(w) = \arg\min_{\theta} \sum_{k=1}^K w_k L_k(\theta)$$

2. Gradient-Based Update Rule

Using a truncated backpropagation approach (single-step SGD), the authors derive an update rule for the weights. The gradient of the downstream loss with respect to the weights is proportional to the inner product between the downstream gradient and the individual pretraining gradients:
$$\frac{\partial L_{down}}{\partial w_k} \propto -\eta \, g_{down}^\top g_k$$
where $g_{down}$ is the downstream gradient and $g_k$ is the gradient of the $k$-th pretraining loss. The goal is to maximize the alignment between the composite gradient $\sum w_k g_k$ and the downstream gradient.

3. Embedding-Space Approximation

A naive implementation would require computing $K$ full backward passes through the model to obtain $g_k$. To reduce computational cost, GraP computes gradients with respect to the shared backbone embedding $z$ rather than the full parameter space.
$$\tilde{g}_k = \nabla_z L_k(h_k(z))$$
Since the losses are applied via separate heads on top of the same embedding, computing these gradients is significantly cheaper than full backward passes. The optimization objective becomes maximizing the alignment in this shared representation space:
$$\arg\max_{w} \frac{w^\top \tilde{G} \tilde{g}_{down}}{\|w^\top \tilde{G}\|}$$

4. Normalization Strategy

The paper argues against normalizing the weights themselves (e.g., $\sum w_k = 1$) or their $L_2$ norm, as these can distort the optimal direction of the composite gradient. Instead, GraP imposes a constraint on the $L_2$-norm of the composite gradient itself ($\|\sum w_k g_k\| = 1$). This ensures that the optimization focuses on the direction of the update vector rather than its magnitude, recovering the optimal alignment direction under the local objective.

5. Algorithm Efficiency

The resulting algorithm requires only one forward pass and one backward pass through the backbone per iteration. The gradients computed for the heads are reused for both updating the loss weights and updating the backbone parameters. This reduces the overhead of hyperparameter tuning to approximately 30% above a single training run, compared to the orders-of-magnitude cost of Bayesian optimization.

Key Contributions

1. Problem Formulation: The authors formulate pretraining loss-weight selection as a bilevel optimization problem explicitly targeting downstream performance.
2. Normalization Insight: They analyze normalization strategies and demonstrate that constraining the norm of the update vector (composite gradient) is superior to constraining the weights themselves, preserving the optimal alignment direction.
3. Efficient Algorithm: They propose GraP, an algorithm with $O(1)$ complexity regarding the number of forward-backward steps relative to the number of hyperparameters. It leverages gradient alignment in the shared representation space to avoid multiple backward passes.
4. Empirical Efficiency: The method achieves performance comparable to or better than standard hyperparameter optimization (e.g., Bayesian search) while being nearly 50$\times$ more efficient, enabling tuning in settings where traditional methods are computationally infeasible.
5. Redundancy Identification: The approach demonstrates the ability to identify and down-weight redundant loss terms (e.g., assigning near-zero weight to specific contrastive losses in vision tasks) effectively acting as an automatic loss-selection mechanism.

Experimental Results

The method was evaluated on two domains: event-sequence modeling and self-supervised computer vision.

Event-Sequence Modeling

• Datasets: Churn, AgePred, AlfaBattle, MIMIC-III, and Taobao.
• Setup: Pretraining involved a composite objective of next-token prediction (NTP) across multiple fields and a contrastive loss.
• Performance: GraP matched or improved upon carefully tuned baselines (including GradNorm, DWA, MGDA, and PCGrad) and Optuna-tuned weights. It consistently achieved the best results on the Churn and MIMIC-III datasets across RNN, Transformer, and HT-Transformer backbones.
• Efficiency: Compared to Bayesian optimization (Optuna), which requires dozens of independent runs, GraP achieved similar performance with a single training run and a moderate overhead (average 44% slowdown in throughput).

Self-Supervised Computer Vision

• Setup: Evaluated on CIFAR-10, CIFAR-100, and ImageNet-100 using the All4One framework (combining Barlow Twins, BYOL, and NNCLR objectives).
• Performance: In this setting, where objectives are often well-balanced, GraP achieved performance comparable to the default All4One parameters and other multi-task learning methods.
• Loss Selection: Notably, GraP assigned near-zero weight to the Attn-NNCLR loss. Removing this loss manually yielded similar performance, validating the method's ability to identify low-impact or redundant objectives.

Ablation Studies

• Label Leakage: Experiments varying the fraction of labeled data used for weight tuning showed robustness, suggesting labels provide a weak guidance signal rather than direct supervision for the backbone.
• Normalization: The proposed composite-gradient normalization outperformed standard weight-sum normalization.
• Auxiliary Training: Pretraining with GraP followed by supervised fine-tuning outperformed joint auxiliary training, likely due to reduced overfitting on limited labeled data.

Significance and Claims

The paper claims that gradient-based loss-weight optimization is a practical and efficient alternative to classical multi-task learning and search-based hyperparameter optimization. Its primary significance lies in making hyperparameter tuning feasible for large-scale pretraining scenarios with multiple loss terms, where traditional methods are computationally prohibitive.

The authors emphasize that GraP does not require multiple training runs, making it suitable for settings with limited computational budgets. While the method assumes separate heads for each loss, it effectively handles the trade-offs between conflicting objectives by aligning the pretraining update direction with the downstream goal. The work suggests that in complex pretraining pipelines, automatic, online alignment of loss weights can yield representations that transfer as well as, or better than, those obtained through expensive manual tuning or black-box search.