Graph Coloring for Multi-Task Learning
This addresses performance degradation in multi-task learning due to conflicting gradients, offering a practical solution for AI applications requiring efficient multi-task optimization.
The paper tackles gradient interference in multi-task learning by introducing SON-GOKU, a scheduler that groups tasks using graph coloring to activate only compatible tasks per step, improving model performance without extra tuning. Empirical results on six datasets show it consistently outperforms baselines and state-of-the-art optimizers.
When different objectives conflict with each other in multi-task learning, gradients begin to interfere and slow convergence, thereby potentially reducing the final model's performance. To address this, we introduce SON-GOKU, a scheduler that computes gradient interference, constructs an interference graph, and then applies greedy graph-coloring to partition tasks into groups that align well with each other. At each training step, only one group (color class) of tasks are activated, and the grouping partition is constantly recomputed as task relationships evolve throughout training. By ensuring that each mini-batch contains only tasks that pull the model in the same direction, our method improves the effectiveness of any underlying multi-task learning optimizer without additional tuning. Since tasks within these groups will update in compatible directions, multi-task learning will improve model performance rather than impede it. Empirical results on six different datasets show that this interference-aware graph-coloring approach consistently outperforms baselines and state-of-the-art multi-task optimizers. We provide extensive theory showing why grouping and sequential updates improve multi-task learning, with guarantees on descent, convergence, and accurately identifying what tasks conflict or align.