🤖 machine learning

CORE-MTL: Rethinking Gradient Balancing via Causal Orthogonal Representations

The paper proposes CORE-MTL, a causally motivated, representation-centric framework that improves multi-task learning by factorizing shared representations into semantic and residual streams to disentangle task-relevant structures from spurious context, thereby achieving superior generalization and reduced gradient interference compared to existing optimization-centric methods.

Original authors: Chengfeng Wu, Tao Zou, Yanru Wu, Jingge Wang

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Chengfeng Wu, Tao Zou, Yanru Wu, Jingge Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a single student to do three different jobs at once: driving a car, painting a portrait, and solving a math problem.

In the world of Artificial Intelligence, this is called Multi-Task Learning (MTL). The goal is to have one "brain" (a shared model) that learns a common language to understand all three jobs.

The Problem: The "Noisy Classroom"

The paper argues that current methods for teaching this student are flawed. They focus on balancing the grades (gradients) of the three tasks. If the student gets a bad grade on math, the teacher adjusts the lesson plan to focus more on math. If the painting grade drops, they shift focus there.

But the paper says: "You can't fix a bad student just by changing the grading curve."

The real issue is what the student is actually learning.

The Good Stuff (Semantics): The student learns the shape of a car, the concept of a face, or the logic of math. This is useful and stable.
The Bad Stuff (Nuisance/Spurious Correlations): The student also accidentally learns shortcuts. For example, they might think "all cars are red" because every car in their textbook happened to be red. Or they might think "all faces are smiling" because the dataset only had happy people.

When the student encounters a black car or a sad face (a new environment), they fail miserably. They relied on the "red" or "smiling" shortcut, not the actual concept of a car or a face. Current methods try to juggle the tasks, but they don't stop the student from memorizing these useless shortcuts.

The Solution: CORE-MTL (The "Two-Stream" Brain)

The authors propose a new way to build the student's brain, called CORE-MTL. Instead of just juggling grades, they restructure the classroom into two separate streams:

The "Semantics" Stream (The Serious Student): This stream is strictly for learning the important, universal rules (the shape of the car, the logic of math).
The "Residual" Stream (The Note-Taker): This stream is a trash can for all the junk, noise, and specific details that don't matter (the color of the car, the background scenery, the lighting).

The Golden Rule: The student is only allowed to use the "Semantics" stream to answer the test questions. The "Residual" stream is forced to hold all the junk, but it is never allowed to influence the final answer.

How Do They Force This Separation?

You can't just tell the student "don't learn the junk." They need a way to prove they are separating the two. The paper uses two clever tricks:

1. The "Physics" Test (Hard Grounding)

For tasks like driving or scene understanding, the authors use the laws of physics.

Imagine the student is looking at a 3D object.
The Semantics stream must figure out the shape (geometry).
The Residual stream must figure out the lighting and color (shading).
They are then forced to rebuild the image using a physics formula: Shape + Light = Image.
If the student tries to put the "lighting" into the "shape" stream, the image they rebuild will look wrong (e.g., a shadow floating in mid-air). This forces the brain to keep the shape and light separate.

2. The "What If?" Test (Counterfactuals)

This is like a "mad libs" game for the brain.

The teacher takes a picture of a car in the rain (Residual = Rain) and swaps the rain for a sunny day (Residual = Sun).
The "Semantics" stream must still correctly identify the car, even though the background changed completely.
If the student fails this test, it means they were relying on the rain to identify the car. The system punishes them until they learn to ignore the rain and focus only on the car.

Why Is This Better?

The paper claims that by physically separating the "good stuff" from the "junk" inside the brain:

No More Gradient Fighting: You don't need complex math to balance the tasks because the tasks naturally stop interfering with each other. The "Semantics" stream is clean, so the math works itself out.
Better at New Things: Because the student isn't memorizing shortcuts (like "cars are red"), they can handle a black car, a rainy day, or a different city much better than before.
Scalable: As you add more tasks (driving, painting, math, cooking), the system doesn't get slower or more confused. It just keeps the "junk" in the trash can and the "good stuff" in the clean stream.

The Bottom Line

Current AI tries to fix multi-task learning by adjusting the volume of each task (turning the math up, turning the painting down).

CORE-MTL says: "No, let's build a soundproof room." Put the important music (semantics) in one room and the noise (nuisance) in another. Only let the music out. This way, no matter how loud the noise gets outside, the music remains clear and perfect.

The paper proves this works better on standard tests (like driving scenes and facial attributes) and, crucially, when the AI is tested on things it has never seen before (like a different weather condition or a different city).

Technical Summary: CORE-MTL

Problem Statement

Multi-task learning (MTL) aims to improve sample efficiency and robustness by training a shared encoder across multiple related tasks. However, MTL frequently suffers from negative transfer, where optimizing for one task degrades performance on others, and exhibits poor generalization under distribution shifts (Out-of-Distribution or OOD).

Existing approaches primarily operate at the optimization level, attempting to reconcile task conflicts through:

Gradient Balancing: Reweighting or projecting gradients (e.g., GradNorm, PCGrad, STCH).
Architecture Modification: Routing tasks to specific slices of backbone features (e.g., MTAN).

The authors argue that these methods treat feature learning as a "black box," ignoring the internal structure of the shared representation. When the shared representation entangles invariant semantic factors (task-relevant) with spurious residual factors (nuisance variations like background or lighting), gradient manipulation alone cannot resolve the resulting conflicts. This entanglement leads to reliance on "nuisance-driven shortcuts," causing brittle OOD generalization.

Methodology: CORE-MTL

The authors propose Causal Orthogonal Representations for Multi-Task Learning (CORE-MTL), a representation-centric framework that shifts the focus from gradient surgery to causally motivated semantic-residual factorization.

1. Structural Decomposition

CORE-MTL decomposes the shared encoder output $\Phi_\theta(x)$ into two distinct streams:

Semantic Stream ( $\hat{Z}_s$ ): Captures invariant, task-relevant structure (e.g., geometry, object identity).
Residual Stream ( $\hat{Z}_r$ ): Absorbs task-irrelevant nuisance variation (e.g., texture, illumination, background).

Crucial Constraint: Task predictors ( $f_{\phi_t}$ ) are restricted to depend solely on the semantic stream $\hat{Z}_s$ . This structural isolation prevents the model from exploiting nuisance signals as shortcuts during learning.

2. Training Objectives

To enforce this decomposition and ensure the streams capture their intended roles, CORE-MTL employs a unified objective function comprising four components:

Task Supervision ( $L_{task}$ ): Standard loss for each task based on $\hat{Z}_s$ .
Structural Independence ( $L_{CKA}$ ): Uses Centered Kernel Alignment (CKA) to minimize statistical dependence between $\hat{Z}_s$ and $\hat{Z}_r$ , encouraging them to capture mutually exclusive information.
Counterfactual Feature Augmentation ( $L_{CFA}$ ): Synthesizes counterfactual instances by recombining the semantic stream of an image with a random residual vector from another image ( $\tilde{x} = D(\hat{Z}_s, \tilde{Z}_r)$ ). The model is trained to maintain consistent task predictions on these synthesized images, enforcing invariance of $\hat{Z}_s$ to nuisance changes.
Grounding via Reconstruction ( $L_{rec}$ ): An auxiliary decoder reconstructs the input from $(\hat{Z}_s, \hat{Z}_r)$ $(\hat{Z}_{s}, \hat{Z}_{r})$ . This provides a "grounding" signal to prevent arbitrary factorization. The authors implement two variants:
- Hard Grounding: Uses explicit physical priors (e.g., inverse rendering with albedo and lighting models) for structured scenes.
- Soft Grounding: Uses a generic convolutional decoder for tasks where physical priors are unavailable, relying on architectural bottlenecks and statistical constraints.

3. Theoretical Guarantees

The paper provides theoretical analysis under a Linear-Gaussian Structural Causal Model (SCM):

OOD Lower Bound: It proves that for any optimization-centric method operating on entangled representations, there exists an unavoidable OOD error floor proportional to the degree of entanglement ( $\sin^2 \psi$ ).
Disentangled Bound: Conversely, CORE-MTL's disentangled architecture guarantees a tighter OOD generalization bound that scales linearly with the "leakage coefficient" (sensitivity of the semantic stream to residual factors).
Gradient Orthogonality: The framework theoretically induces gradient orthogonality between task and residual losses as a structural by-product of disentanglement, resolving conflicts at the representation level rather than via post-hoc correction.

Key Contributions

Causal Perspective on MTL: The authors formalize the limitations of optimization-centric methods under a latent SCM, proving that manipulating gradients on entangled representations cannot eliminate OOD errors.
CORE-MTL Framework: A novel representation-centric framework that enforces semantic-residual factorization via physical or generic grounding, resolving negative transfer at the source.
Empirical Superiority: Extensive experiments show consistent improvements over ten representative baselines (including GradNorm, PCGrad, STCH, MTAN) on both in-distribution (NYUv2, Cityscapes) and out-of-distribution (GTA5→Cityscapes, Cityscapes-C, Colored-Cityscapes) benchmarks.
Scalability: The method maintains near-constant training costs as the number of tasks increases, avoiding the prohibitive linear scaling associated with gradient surgery methods like PCGrad.

Experimental Results

In-Distribution (ID): On NYUv2 and Cityscapes, CORE-MTL achieves the best overall trade-off, obtaining top performance in segmentation, depth estimation, and surface normal prediction.
Out-of-Distribution (OOD):
- Sim-to-Real (GTA5→Cityscapes): CORE-MTL achieves the best target-domain metrics and the smallest performance gaps compared to source performance.
- Corruption Robustness (Cityscapes-C): Outperforms baselines under various corruptions (contrast, blur, fog, noise), particularly in severe visibility reduction scenarios.
- Shortcut Stress Test (Colored-Cityscapes): In a benchmark where class colors are permuted between source and target, CORE-MTL remains robust, demonstrating it does not rely on spurious appearance correlations.
Scalability (CelebA): As the number of tasks increases (up to 40 attributes), CORE-MTL maintains high accuracy while gradient surgery methods (PCGrad) suffer from exploding training times.

Significance and Claims

The paper claims that entanglement in shared representations, rather than the gradient balancing scheme itself, is a primary source of task interference and OOD failure. By shifting the paradigm from "post-hoc gradient surgery" to "intrinsic representation disentanglement," CORE-MTL:

Resolves task conflicts structurally, allowing tasks to synergize naturally without rigid gradient correction.
Provides a tighter theoretical bound on OOD generalization.
Offers a scalable solution for multi-task learning that remains effective as the number of tasks grows.

The authors position this work as a step toward structure-aware representation design for robust multi-task learning, moving beyond purely optimization-centric coordination. They acknowledge limitations, noting their theoretical guarantees rely on stylized linear-Gaussian assumptions and that the current instantiation is primarily for dense scene understanding tasks.