Skip to content

SON-GOKU: Conflict-Free Scheduling for Multi-Task Learning via Graph Coloring

Conference: ECCV 2026
arXiv: 2509.16959
Code: https://anonymous.4open.science/r/SON-GOKU-Impl/
Area: Optimization
Keywords: Multi-Task Learning, Gradient Conflict, Graph Coloring, Task Scheduling, Negative Transfer

TL;DR

SON-GOKU dynamically partitions tasks into low-conflict groups by constructing a conflict graph based on real-time gradient conflict measurements in multi-task learning and applying a greedy graph coloring algorithm. By activating only one group of tasks per training step and periodically recoloring to adapt to gradient evolution, it consistently improves multiple MTL baselines across six datasets.

Background & Motivation

Multi-Task Learning (MTL) aims to improve generalization performance by sharing network parameters and learning multiple related tasks simultaneously. However, when the optimization directions of different tasks diverge, their gradients with respect to shared parameters conflict—a phenomenon widely known as negative transfer. This conflict intuitively manifests as gradient vectors pointing in opposing directions, causing their sum to cancel out. This hinders convergence, slows down training, or even reverses the optimization progress. Existing strategies fall into two main categories: gradient manipulation (e.g., PCGrad projects conflicting gradients onto orthogonal directions, CAGrad finds a balanced direction within the convex hull of gradients) to reduce conflicts by altering the geometric shape of shared gradients; and loss reweighting (e.g., GradNorm dynamically adjusts task weights, AdaTask learns unique learning rates for each task) to alleviate the dominance of certain tasks by modulating their influence. Both approaches activate all tasks at every training step, attempting to mitigate conflicts by performing "surgery" on the shared update.

Recently, an alternative paradigm has emerged: partitioning tasks into subgroups and updating only one group at a time. The intuition is straightforward—since conflicting gradients cancel out when combined, separating them allows compatible tasks to mutually reinforce each other. Early works such as TaskGrouping and Selective Task Group Update have demonstrated that grouped updates can yield more stable training than gradient manipulation. However, these methods suffer from three critical limitations: first, they rely on dense pairwise task affinity matrices, which scale quadratically with the number of tasks and are highly noisy; second, the groupings are typically determined at the beginning of training and rarely updated, whereas gradient relationships evolve continuously, rendering static groupings ineffective in later stages; third, the grouping strategies are often local heuristics that lack global compatibility guarantees and structured rotation mechanisms, failing to ensure fair scheduling for all tasks.

The core insight of SON-GOKU is: if each task is treated as a node and the degree of conflict between tasks acts as edge weights, the problem of "which tasks should not be updated together" naturally maps to a graph coloring problem—conflicting tasks cannot share the same color, while tasks with the same color can be updated together. From this perspective, this paper designs a complete pipeline from gradient measurement and coloring scheduling to periodic recoloring, backed by rigorous convergence theory justifying the grouping rationality and scheduling efficiency. Core Idea: Translate the task scheduling in multi-task learning into a dynamic graph coloring problem—construct a sparse conflict graph in real-time using EMA-smoothed gradient vectors, partition tasks into low-conflict color groups using a greedy graph coloring algorithm, cyclically activate one group per training step, and periodically recolor to adapt to evolving task relationships while leveraging classic graph theory results to guarantee each task is updated at least once every \(\Delta+1\) steps in the worst case (where \(\Delta\) is the maximum degree of conflict).

Method

Overall Architecture

SON-GOKU is a plug-and-play scheduler that does not modify the underlying multi-task optimizer, only determining which tasks to activate at each step. The entire workflow operates cyclically with a refresh period of \(R\) steps, reconstructing the schedule every \(R\) steps.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["K Tasks<br/>T₁ … T_K"] --> B["① Gradient Sampling<br/>Each Step: Calculate g_k(t)<br/>for each task on independent mini-batches"]
    B --> C["② EMA Smoothing<br/>g̃_k(t) = β·g̃_k(t-1)<br/>+ (1-β)·g_k(t)"]
    C --> D["③ Conflict Graph Construction<br/>Every R steps: ρ_ij = -cos(g̃_i, g̃_j)<br/>Connect edge if ρ_ij > τ"]
    D --> E["④ Greedy Coloring<br/>Welsh-Powell heuristic<br/>yields m color groups C₁...C_m"]
    E --> F["⑤ Cyclic Scheduling<br/>S_t = C_{(t mod m)+1}<br/>Activate one color group per step"]
    F --> G["Shared Parameter Update<br/>θ ← θ - η·Σ_{k∈S_t} g_k"]
    G -->|Recolor every R steps| C

Key Designs

1. EMA-Smoothed Gradient Conflict Estimation: Filtering Batch Noise via Historical Averages

Single-step minibatch gradients suffer from high directional noise—estimating task relationships based on cosine similarity of only one or two batches is highly unreliable. SON-GOKU maintains an Exponential Moving Average (EMA) gradient vector for each task:

\[g̃_k^{(t)} = β g̃_k^{(t-1)} + (1-β) g_k^{(t)}\]

Only during schedule refreshes (every \(R\) steps) pair-wise cosine similarities are calculated based on the EMA vectors to define the interference coefficient \(ρ_{ij} = -⟨g̃_i, g̃_j⟩/(‖g̃_i‖‖g̃_j‖)\), where positive values denote conflict. EMA makes the gradient direction estimation robust to step-level noise, while the storage overhead requires only two buffers (current + previous state). To further reduce GPU memory footprint, the practical implementation employs a low-dimensional sketch approximation (sketch width \(d_{sk} ≪ d\)) to compress the \(K \times d\) gradient matrix into \(K \times d_{sk}\). Consequently, the scheduler's memory overhead scales with the number of tasks \(K\) rather than the model dimension \(d\), maintaining feasibility on large backbone networks.

2. Conflict Graph Construction and Greedy Graph Coloring: Transforming Scheduling into Graph Theory

Given a tolerance threshold \(τ \in (0,1)\), an undirected conflict graph \(G_τ = (V, E_τ)\) is constructed, where vertices represent the \(K\) tasks, and an undirected edge connects \(i\) and \(j\) if \(ρ_{ij} > τ\), indicating they should not be updated simultaneously. Here, \(τ\) controls the sparsity of the graph—larger values of \(τ\) yield sparser graphs with fewer colors, activating more tasks per step; conversely, smaller values of \(τ\) yield denser graphs with more colors, minimizing within-step conflicts at the cost of longer update intervals for each task. Once the graph is built, the classical Welsh-Powell largest-degree-first greedy algorithm is applied: nodes are sorted in descending order of their degrees and sequentially assigned the smallest available color channel. This heuristic guarantees that the number of colors used does not exceed \(\Delta+1\) (where \(\Delta\) is the maximum degree of the graph), meaning the schedule cycle length is determined by the most conflicting task rather than the total number of tasks—allowing highly compact schedules when \(\Delta \ll K\).

3. Periodic Recoloring and Schedule Generation: Adapting to Gradient Evolution

Gradient relationships between tasks are not static; as model parameters update, the optimization directions of tasks continuously drift. Fixed groupings gradually degrade in effectiveness past the early stages of training. SON-GOKU reconstructs the conflict graph and recolors every \(R\) steps using the latest EMA gradients, generating a cyclic schedule of length \(m\) (the current number of colors): \(S_t = C_{(t \bmod m) + 1}\). This guarantees that all tasks are updated at least once within a single cycle. For single-element color groups produced by greedy coloring (i.e., a task heavily conflicting with others, forced to occupy a color alone), the scheduler replicates them into other training steps that do not share any conflicting edges to prevent their update frequencies from dropping too low.

4. Warm-up Annealing and Plug-and-Play Design: Smooth Startup and Compatibility

During early training stages, gradient direction information is highly unstable, and introducing strict grouping too early can be detrimental. SON-GOKU sets \(τ = 1\) during the first \(T_{warm}\) steps (representing a fully disconnected graph with no edges, where all tasks update simultaneously) and then logarithmically anneals it to the target threshold \(τ^*\). This allows the model to undergo a 'collaborative training' warm-up phase before smoothly transitioning to grouped scheduling. Crucially, SON-GOKU is a plug-and-play scheduler that does not tie to any underlying optimizer, meaning it can be overlaid on top of any existing MTL methods (e.g., PCGrad + SON-GOKU, AdaTask + SON-GOKU). The scheduler first eliminates large-scale conflicts, allowing the optimizer to operate under cleaner within-group conditions, rendering them highly complementary.

Loss & Training

SON-GOKU does not alter the loss function of any task. The training objective remains the standard weighted MTL objective \(F(θ, φ_1, ..., φ_K) = Σ w_k L_k(θ, φ_k)\). The scheduler only modifies the active set. Theoretical analysis establishes three key guarantees: (1) When \(τ(|S_t|-1) < 1\), the group update direction remains a descent direction, preventing conflicts from turning it into an ascent direction; (2) Under standard non-convex smoothness assumptions, SON-GOKU maintains the \(O(1/\sqrt{T})\) convergence rate of SGD, introducing only an additional constant factor of \((1+τ)\); (3) Within the refresh window, the expected descent of sequentially updating each color group is strictly no less than that of a combined mixed update, with the advantage becoming more pronounced when cross-group conflicts are negative.

Key Experimental Results

Main Results

Dataset Metric Uniform FAMO SON-GOKU ++AdaTask ++PCGrad
CIFAR-10 Acc. (%) ↑ 55 64 65 67 65
F&B Acc. (%) ↑ 63 70 69 71 70
AV-MNIST Acc. (%) ↑ 52 60 58 59 62
NYUv2 Angle Error ↓ 21.6 19.9 19.8 20.1 19.7
MM-IMDb Acc. (%) ↑ 52 60 58 59 62
HEALTH Acc. (%) ↑ 56 61 61 63 60

SON-GOKU consistently outperforms the Uniform baseline across all six datasets, yielding improvements of 10%-20%. Compared to strong baselines such as FAMO, AdaTask, and Nash-MTL, SON-GOKU achieves state-of-the-art or comparable results on multiple metrics. Performance is further boosted when integrated with AdaTask or PCGrad—combining with AdaTask achieves the best performance on classification tasks (where task-specific learning rates smooth out sudden spikes in classification gradients), while combining with PCGrad excels in regression and dense prediction tasks (where the projection operator removes remaining micro-conflicts within groups).

Ablation Study

Configuration CIFAR-10 F&B NYUv2 Angle Description
SON-GOKU (Full) 65 69 19.8 EMA + Dynamic Coloring + Threshold Mapping
Static One-Shot 61 66 20.5 Only colored once at start, frozen groupings thereafter
Single-Step (H=1) 40 59 26.4 No EMA, uses current batch gradient at each step
Signed-Only 56 63 24.0 Edges built solely on sign parity, ignoring magnitude
kNN-Symmetric 60 65 22.1 Each task connected to top-k most conflicting tasks

Key Findings

  • Dynamic recoloring is critical: Static One-Shot consistently lags by 3-5 percentage points across all metrics, proving that gradient relationships drift significantly during training.
  • Historical average is indispensable: Single-Step performance on CIFAR-10 plunges from 65% to 40%, demonstrating that single-step gradient noise is extremely high, making EMA smoothing a necessary prerequisite for reliable grouping.
  • Graph construction rule experiments show that the simple global threshold method (\(ρ_{ij} > τ\)) and the quantile method yield the best results; Signed-Only and kNN-Symmetric perform poorly due to ignoring conflict magnitudes or altering graph topology.
  • In terms of computational overhead, SON-GOKU (\(R=32\)) takes approximately 12 seconds per step when \(K=40\), whereas PCGrad requires 1127 seconds and Nash-MTL requires 1014 seconds, demonstrating a massive efficiency advantage; its auxiliary memory overhead scales with \(K^2\) rather than the model dimension \(d\).

Highlights & Insights

  • Natural Graph Coloring Perspective: Equating the problem of 'which tasks should not be updated together' with graph coloring directly inherits theoretical guarantees from classical graph theory. The \(\Delta+1\) color bound ties the cycle length to the maximum task conflict rather than the total task count, making the schedule highly efficient when most tasks do not conflict.
  • Plug-and-Play Design Philosophy: SON-GOKU enhances rather than replaces existing optimizers. By eliminating large-scale conflicts beforehand, it allows optimizers to operate under cleaner conditions—PCGrad only needs to handle remaining minor intra-group conflicts, and AdaTask is shielded from adversarial gradients, yielding synergistic gains.
  • Dual Denoising with EMA + Sketch: EMA filters out single-step gradient noise, while the sketch technique reduces memory overhead from \(O(Kd)\) to \(O(K^2 + Kd_{sk})\), keeping the scheduler highly viable on large-scale backbones.
  • Solid Theoretical Foundation: Standard descent guarantees (direction does not reverse when \(τ(|S|-1) < 1\)), convergence rate maintenance (\(O(1/\sqrt{T})\) + a small constant factor), and exact grouping recovery bounds (recovering the ideal conflict graph when estimation error is below tolerance) are simultaneously proven, which is rare for scheduling-oriented MTL literature.

Limitations & Future Work

  • The computational complexity of coloring scales quadratically with the number of tasks (\(O(K^2)\) construction + \(O(K^2)\) coloring). Although executed only every \(R\) steps, this can still become a bottleneck when \(K\) is extremely large (e.g., hundreds of tasks). While acceleration strategies are discussed in the paper, experiments only cover up to 40 tasks.
  • The initial value and annealing policy of \(τ\) rely on manual tuning without an adaptive mechanism. Parameter retraining might be required when task numbers or distributions vary widely.
  • SON-GOKU only schedules shared parameter gradients, leaving task-specific heads unconstrained and omitted from the conflict graph. If task-specific heads also exhibit significant conflicts, this framework may need further expansion.
  • The empirical validation is limited to vision and tabular datasets; its effectiveness in scenarios such as multi-task language modeling or multi-task reinforcement learning remains to be verified.
  • vs PCGrad / CAGrad: They project or balance all task gradients at each step without altering the active set; SON-GOKU filters tasks before letting the optimizer operate. The two approaches are complementary rather than mutually exclusive, performing best when combined.
  • vs TaskGrouping / Selective Task Update: Existing grouping methods rely on dense affinity matrices or static clusterings; SON-GOKU implements dynamic sparse scheduling via EMA + graph coloring, backed by rigorous convergence guarantees.
  • vs AdaTask / GradNorm: These methods modulate numeric scales (weights/learning rates), while SON-GOKU eliminates large-scale conflicts at a structural level. Their combination yields better results than using either in isolation.

Rating

  • Novelty: ⭐⭐⭐⭐ While introducing graph coloring to MTL scheduling is not entirely unique, the comprehensive design and validation of the entire pipeline (EMA + sketch + dynamic recoloring + theoretical analysis) offers significant novelty.
  • Experimental Thoroughness: ⭐⭐⭐⭐⭐ Extremely thorough, spanning six datasets, over ten baselines, systematic ablations across 6 graph construction rules × 2 density settings, and explicit runtime/memory overhead comparisons.
  • Writing Quality: ⭐⭐⭐⭐ The main text is well-structured, with tedious theoretical proofs relegated to the appendix to maintain readability, and diagrams effectively conveying the core workflow.
  • Value: ⭐⭐⭐⭐ The plug-and-play scheduler can be directly imported into existing MTL pipelines. It validates the synergy of 'scheduling + optimization,' exhibiting high potential for practical deployment.