Graph Coloring for Multi-Task Learning

When different objectives conflict with each other in multi-task learning, gradients begin to interfere and slow convergence, therebypotentially reducing the final model’s performance. To address this, weintroduce SON-GOKU, a scheduler that computes gradient interference,constructs an interference graph, and then applies greedy graph-coloringto partition tasks into groups that align well with each other. At eachtraining step, only one group (color class) of tasks are activated, and thegrouping partition is constantly recomputed as task relationships evolvethroughout training. By ensuring that each mini-batch contains onlytasks that pull the model in the same direction, our method improvesthe effectiveness of any underlying multi-task learning optimizer withoutadditional tuning. Since tasks within these groups will update in compati-ble directions, multi-task learning will improve model performance ratherthan impede it. Empirical results on six different datasets show thatthis interference-aware graph-coloring approach consistently outperformsbaselines and state-of-the-art multi-task optimizers. We provide extensivetheory showing why grouping and sequential updates improve multi-tasklearning, with guarantees on descent, convergence, and the ability toaccurately identify what tasks conflict or align.