Research map / Learning theory and algorithms / Prediction and decision losses

GO4Align: Group Optimization for Multi-Task Alignment

NeurIPS 2024 · Learning theory and algorithms › Prediction and decision losses · approach: Adaptive testing and experiments · 2 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes GO4Align, a loss-oriented multi-task optimization method that dynamically clusters tasks into groups (K-means on risk-guided indicators) and assigns group-level loss weights to align learning progress across tasks at low training-time cost.

Problem

Task imbalance in multi-task optimization: some tasks are under-optimized, and gradient-based balancing methods are computationally expensive, while loss-based methods underperform.

Approach

Adaptive group risk minimization: tasks are dynamically clustered into K groups via K-means over risk-guided group indicators (scale-balanced risks smoothed by an exponential moving average of past risks), and intra-group tasks share a group weight in the loss, avoiding per-task gradient computation.

Key results

• On NYUv2 (3 tasks), GO4Align reaches Delta m% of -6.08 and average rank 2.11, versus -4.10 and 4.33 for FAMO; it is the only method that beats STL on every task there.
• On QM9 (11 tasks), GO4Align obtains the lowest Delta m% (52.7 vs 58.5 for FAMO) among compared methods.
• Per-epoch training time on NYUv2 is 87.78 s versus 86.21 s for LS and 201.30 s for NashMTL, so it stays near scalarization cost.

Measured effects reported in the paper

Surprising findings

• Replacing K-means with SDP-based or spectral clustering gave weaker results on NYUv2, with K-means performing best.
• On CityScapes (2 tasks) and CelebA the method does not achieve the lowest average performance drop, which the authors attribute to too few tasks limiting the grouping mechanism.

Limitations

• Number of groups K is a heuristic hyperparameter chosen by elbow method, requiring up to M runs to search.
• Gains are weaker on datasets with few tasks (CityScapes, 2 tasks) and on CelebA average drop.

Experimental setup

Datasets: NYUv2 CityScapes QM9 CelebA

Benchmarks: NYUv2 CityScapes QM9 CelebA

Metrics: Delta m% (average per-task relative performance drop vs STL) average rank training time per epoch

Baselines: MGDA PCGrad CAGrad IMTL-G GradDrop NashMTL LS SI DWA UW RLW FAMO

Builds on: FAMO Multi-Task Attention Network MGDA Nash-MTL

Closest papers by meaning

• Selective Task Group Updates for Multi-Task Optimization (ICLR 2025 · similarity 0.86)
• Towards Principled Task Grouping for Multi-Task Learning (NeurIPS 2026 · similarity 0.84)
• Continual Optimization with Symmetry Teleportation for Multi-Task Learning (NeurIPS 2025 · similarity 0.84)
• DMTG: One-Shot Differentiable Multi-Task Grouping (ICML 2024 · similarity 0.84)
• Vulnerability-Aware Alignment: Mitigating Uneven Forgetting in Harmful Fine-Tuning (ICML 2025 · similarity 0.83)
• The Geometry of Alignment Collapse: When Fine-Tuning Breaks Safety (NeurIPS 2026 · similarity 0.83)