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 consistently improves various MTL baselines across six datasets by measuring real-time gradient conflicts between tasks, constructing a conflict graph, and applying a greedy graph coloring algorithm to dynamically partition tasks into low-conflict groups. Only one task group is activated per training step, with periodic recoloring to adapt to gradient evolution.
Background & Motivation¶
Multi-Task Learning (MTL) aims to leverage commonalities across related tasks to improve generalization performance by simultaneously learning them via shared network parameters. However, when the optimization directions of different tasks are inconsistent, their gradients on the shared parameters conflict, a phenomenon known as negative transfer. Intuitively, this conflict manifests as gradient vectors of two tasks pointing in opposite directions, causing their resultant force to cancel out. This not only fails to accelerate convergence but also slows down or even reverses the learning progress. Existing strategies primarily fall into two categories: gradient manipulation (e.g., PCGrad projects conflicting gradients onto orthogonal directions, and CAGrad finds a balanced direction within the convex hull of gradients) to reduce conflict by modifying the geometry of the shared gradients; and loss reweighting (e.g., GradNorm dynamically adjusts task weights, and AdaTask learns independent learning rates for each task) to alleviate the issue of dominant tasks suppressing localized tasks by regulating the influence of each task. These two types of methods activate all tasks at each training step and merely perform "surgery" on the shared updates to mitigate conflicts.
Recently, a different line of thought has emerged: partitioning tasks into subgroups and updating only one group at a time. The intuition is straightforward—since conflicting gradients cancel out when placed together, separating them allows compatible tasks to promote 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 key limitations: first, they rely on a dense, pairwise task affinity matrix, which incurs a computational overhead scaling quadratically with the number of tasks and is highly noisy; second, the grouping is determined at the beginning of training and is rarely updated, whereas the gradient relationships between tasks continuously evolve throughout training, causing static grouping to fail in the mid-to-late stages; third, grouping strategies are mostly local heuristics, lacking global compatibility guarantees and a formal rotation mechanism, making it difficult to guarantee fair scheduling for each task.
The core insight of SON-GOKU is: if each task is viewed as a node, the conflict degree between tasks naturally serves as edge weights, and "which tasks should not be updated together" becomes a classic graph coloring problem—conflicting tasks cannot share the same color, while tasks with the same color can be updated together. Based on this perspective, this paper designs a complete pipeline from gradient measurement to coloring scheduling and periodic recoloring, proving the rationality of grouping and the effectiveness of scheduling with rigorous convergence theory. Core Idea: Formulate task scheduling in multi-task learning as 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 via a greedy graph coloring algorithm, cyclically activate one color group per training step, and periodically recolor to adapt to the evolution of task relationships. Meanwhile, classical graph theory results are utilized to guarantee that each task is updated at least once within the worst-case \(\Delta+1\) steps (where \(\Delta\) is the maximum degree of conflict).
Method¶
Overall Architecture¶
SON-GOKU is a plug-and-play scheduler that does not alter the underlying multi-task optimizer and is only responsible for determining which tasks to activate at each step. The entire workflow operates cyclically with a refresh period of \(R\) steps, rebuilding 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: tasks compute g_k(t)<br/>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 if ρ_ij > τ"]
D --> E["④ Greedy Coloring<br/>Welsh-Powell heuristic yields<br/>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 with Historical Averages
Single-step minibatch gradient directions are extremely noisy—relying on cosine similarity from just one or two batches to judge task relationships is highly unreliable. SON-GOKU maintains an Exponential Moving Average (EMA) gradient vector for each task:
Only during schedule refreshing (every \(R\) steps) is the pairwise cosine similarity calculated based on the EMA vectors, defining the interference coefficient as \(ρ_{ij} = -⟨g̃_i, g̃_j⟩/(‖g̃_i‖‖g̃_j‖)\), where a positive value indicates conflict. The EMA makes gradient direction estimation robust to single-step noise, while the storage overhead requires only two buffers (current and previous steps). To further reduce memory usage, the actual implementation adopts a low-dimensional sketch approximation (sketch width \(d_{sk} ≪ d\)), compressing the \(K×d\) gradient matrix into \(K×d_sk\). Consequently, the scheduler's extra memory scales only with the number of tasks \(K\) rather than the model dimension \(d\), rendering it highly feasible for large backbone networks.
2. Conflict Graph Construction and Greedy Graph Coloring: Translating Scheduling into Graph Theory
Given a tolerance threshold \(τ \in (0,1)\), an undirected conflict graph \(G_τ = (V, E_τ)\) is constructed: nodes represent the \(K\) tasks; if \(ρ_{ij} > τ\), an edge is connected between \(i\) and \(j\), indicating they should not be updated simultaneously. \(τ\) controls the sparsity of the graph—a larger \(τ\) leads to a sparser graph, fewer colors, and more tasks activated per step; conversely, a smaller \(τ\) leads to a denser graph, more colors, and fewer step-wise conflicts, but at the cost of longer update intervals for each task. After graph construction, the classical Welsh-Powell largest-degree-first greedy algorithm is applied: nodes are sorted in descending order of their degrees, and each node is sequentially assigned the smallest available color ID. 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 scheduling cycle length is determined by the most conflicting task rather than the total number of tasks—when \(\Delta \ll K\), the scheduling is extremely compact.
3. Periodic Recoloring and Schedule Generation: Adapting to Gradient Evolution
Gradient relationships between tasks are not static—as model parameters undergo updates, the optimization directions of different tasks continuously drift. Fixed grouping gradually loses effectiveness in the mid-to-late stages of training. SON-GOKU reconstructs the conflict graph and performs recoloring every \(R\) steps using the latest EMA gradients, and subsequently generates a periodic schedule of length equal to the current number of colors \(m\): \(S_t = C_{(t \bmod m) + 1}\). This guarantees that all tasks are updated at least once within one cycle. For single-element color groups generated by greedy coloring (i.e., tasks that conflict with many others and are forced to occupy a color exclusively), the scheduler duplicates them into other training steps without conflicting edges, ensuring that their update frequencies do not fall too low.
4. Warm-up Annealing and Plug-and-Play: Smooth Initialization and Compatibility
At the beginning of training, gradient direction information is highly unstable, and introducing strict grouping too early can be detrimental. SON-GOKU sets \(τ\) to 1 (all-pass graph, no edges, all tasks updated simultaneously) during the first \(T_{warm}\) steps, and then logarithmically anneals it to the target threshold \(τ^*\). This allows the model to undergo an "all-task joint training" warm-up phase before smoothly transitioning to grouped scheduling. More importantly, SON-GOKU is a plug-and-play scheduler—it is independent of the underlying optimizer and can be stacked on top of any existing MTL methods (e.g., PCGrad + SON-GOKU, AdaTask + SON-GOKU). The scheduler first eliminates large-scale conflicts, after which the optimizer operates under cleaner intra-group conditions, achieving complementary effects.
Loss & Training¶
SON-GOKU does not modify 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 alters the activation set. Theoretical analysis provides three key guarantees: (1) when \(τ(|S_t|-1) < 1\), the direction of group updates remains descent and does not reverse to ascent due to conflicts; (2) under standard non-convex smoothness assumptions, SON-GOKU preserves the \(O(1/√T)\) convergence rate of SGD, introducing only an additional constant factor of \((1+τ)\); (3) within a refresh window, the expected descent of sequentially updating each color group is strictly no less than that of one-time mixed updates, with the advantage being more pronounced when inter-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 with improvements ranging from 10% to 20%. Compared to strong baselines such as FAMO, AdaTask, and Nash-MTL, SON-GOKU achieves optimal or comparable results across multiple metrics. When combined with AdaTask or PCGrad, performance is further enhanced—stacking with AdaTask yields the best performance on classification tasks (its task-specific learning rates smooth out sudden spikes in classification gradients), while stacking with PCGrad performs best on regression/dense prediction tasks (its projection operation removes remaining minor conflicts within groups).
Ablation Study¶
| Configuration | CIFAR-10 | F&B | NYUv2 Angle | Description |
|---|---|---|---|---|
| SON-GOKU (Full) | 65 | 69 | 19.8 | EMA + Dynamic Coloring + Threshold Graphing |
| Static One-Shot | 61 | 66 | 20.5 | Coloring only once at the beginning, then freezing groups |
| Single-Step (H=1) | 40 | 59 | 26.4 | No EMA, using only current batch gradient at each step |
| Signed-Only | 56 | 63 | 24.0 | Connecting edges based only on signs, ignoring magnitudes |
| kNN-Symmetric | 60 | 65 | 22.1 | Each task connects to its top-k most conflicting tasks |
Key Findings¶
- Dynamic recoloring is crucial: Static One-Shot generally falls behind by 3–5 percentage points across various metrics, proving that gradient relationships drift significantly during training.
- Historical average is indispensable: Single-Step performance on CIFAR-10 plunges from 65% to 40%, indicating that single-step noise is extremely high, and EMA smoothing is a necessary prerequisite for reliable grouping.
- Experiments on graph construction rules show that the simple global threshold method (\(ρ_{ij} > τ\)) and quantile method yield the best performance; Signed-Only and kNN-Symmetric perform worse due to ignoring conflict magnitudes or altering the graph structure.
- Regarding time overhead, SON-GOKU (with R=32) takes about 12 seconds per step when K=40, whereas PCGrad requires 1127 seconds and Nash-MTL takes 1014 seconds, demonstrating a massive efficiency advantage; extra memory scales only with K² rather than the model dimension d.
Highlights & Insights¶
- A highly natural graph coloring perspective: Equating "which tasks should not be updated together" to graph coloring directly invokes theoretical guarantees from classical graph theory. The \(\Delta+1\) coloring bound couples the scheduling cycle length to the worst-case conflict degree rather than the total number of tasks, rendering scheduling extremely efficient when most tasks do not conflict.
- Plug-and-play design philosophy: SON-GOKU augments rather than replaces existing optimizers. The scheduler first eliminates large-scale conflicts, enabling the optimizer to work under cleaner conditions—PCGrad only needs to handle remaining minor intra-group conflicts, and AdaTask is spared from adversarial gradient interference, resulting in synergistic gains.
- Dual noise reduction with EMA and sketch: EMA filters single-step noise, while sketching reduces memory from \(O(Kd)\) to \(O(K^2 + Kd_{sk})\), making the scheduler viable on large backbones.
- Comprehensive theoretical foundation: The paper simultaneously provides three proofs: descent guarantee (direction does not reverse when \(τ(|S|-1) < 1\)), convergence rate preservation (\(O(1/√T)\) + small constant factor), and exact group recovery guarantee (recovering the ideal conflict graph when conflict estimation error is below the tolerance). Such rigorous backing is rare in scheduling-based MTL literature.
Limitations & Future Work¶
- Coloring complexity scales quadratically with the number of tasks (\(O(K^2)\) + coloring \(O(K^2)\)). Although it only executes every \(R\) steps, it can still become a bottleneck when \(K\) is extremely large (hundreds of tasks). The paper discusses acceleration strategies, but experiments only cover up to 40 tasks.
- The initial value of \(τ\) and the annealing strategy rely on manual configuration, lacking an adaptive mechanism. Retuning may be necessary when the number of tasks changes or the distributions vary significantly.
- SON-GOKU itself only schedules shared parameter gradients. The updates of task-specific heads remain unrestricted and their gradients are not incorporated into the conflict graph—this warrants extension if task heads exhibit significant conflict.
- The experiments only cover vision and tabular data. The effectiveness in scenarios such as multi-task training in language models or multi-task reinforcement learning remains to be verified.
Related Work & Insights¶
- vs PCGrad / CAGrad: They perform projection/balancing on all task gradients at each step without altering the activation set; SON-GOKU filters tasks first before letting the optimizer process them. The two are complementary rather than mutally exclusive, achieving the best performance when combined.
- vs TaskGrouping / Selective Task Update: Existing grouping methods rely on dense affinity matrices or static grouping; SON-GOKU achieves dynamic sparse scheduling via EMA and graph coloring with rigorous convergence guarantees.
- vs AdaTask / GradNorm: These methods adjust numerical scales (weights/learning rates), whereas SON-GOKU eliminates large-scale conflicts structurally. Combining them yields better results than using either in isolation.
Rating¶
- Novelty: ⭐⭐⭐⭐ Introducing graph coloring to MTL scheduling is not completely unprecedented, but the comprehensive design and validation of the full pipeline (EMA + sketch + dynamic recoloring + theoretical analysis) offer significant novelty
- Experimental Thoroughness: ⭐⭐⭐⭐⭐ Highly thorough; includes evaluation across six datasets, over ten baselines, and systematic ablation of 6 graph construction rules under 2 density settings, along with comparisons of time and memory overhead
- Writing Quality: ⭐⭐⭐⭐ The main text is well-structured, and placing extensive theoretical proofs in the appendix preserves readability; diagrams effectively convey the core workflow
- Value: ⭐⭐⭐⭐ The plug-and-play scheduler can be directly integrated into existing MTL pipelines, validating the synergistic "scheduling + optimization" paradigm with strong potential for practical deployment