CORE-MTL: Rethinking Gradient Balancing via Causal Orthogonal Representations
Abstract
Multi-task learning (MTL) aims to construct a joint model for multiple tasks by sharing a common representation across domains. To achieve this goal, existing optimization-centric methods either balance task gradients or modify the shared architecture. However, as these approaches remain agnostic to the content of the shared representation, they fail to disentangle task-relevant structure from spurious context, leading to negative transfer and poor generalization. To overcome this limitation, we propose Causal Orthogonal Representations for Multi-Task Learning (CORE-MTL) , a representation-centric framework that structurally disentangles the shared representation into semantic and residual streams, concentrating task-relevant structure in the semantic stream while relegating nuisance variation to the residual stream. We instantiate this framework in the visual domain by leveraging physical priors for structured scenes and statistical constraints for attributes. Theoretically, our method enjoys a tighter out-of-distribution generalization bound than optimization-centric methods and reduces task gradient interference without explicit gradient projection or reweighting. Empirically, CORE-MTL consistently outperforms existing methods on visual multi-task benchmarks in both in-distribution and out-of-distribution settings.
Lay Summary
Modern AI vision systems often need to solve several tasks at once, such as recognizing objects, estimating depth, and understanding the shape of surfaces in a scene. Training one shared model for all these tasks can be efficient, but it can also cause problems: the model may rely on visual details such as lighting, texture, or background that are useful in the training data but unreliable in new environments. This can make the tasks interfere with each other and reduce performance when conditions change. We propose CORE-MTL, a method that encourages the model to organize its internal image understanding into two parts. One part captures stable scene information that is useful for the tasks, while the other absorbs changing visual details that should not drive the final predictions. The model is trained to make predictions mainly from the stable part, reducing its reliance on misleading shortcuts. Our experiments show that this approach improves performance across several visual tasks and makes the model more robust when tested on new domains or corrupted images. This could help build more reliable perception systems for applications such as autonomous driving, robotics, and assistive technologies.