Cross-Attention Gradient Transplantation (CAGT): Mitigating Gradient Conflict in Multi-Task Learning Through Cross-Blockwise Attention
Abstract
Multi-task learning tries to use shared structure across tasks that are related by optimizing a joint objective. In practice though, gradients from individual tasks losses can cause conflict. This causes destructive interference that slows convergence and reduces performance. Previous approaches like PCGrad and CAGrad address this by projecting away or reweighting conflicting gradients at the vector level, but they do not explicitly exchange or reuse gradient components across tasks. This paper proposes Cross-Attention Gradient Transplantation (CAGT), which is a method that finds conflicting gradient components and replaces them using deterministic cosine similarity based cross-attention over gradient subblocks from other tasks. Task gradients are partitioned into subblocks at each optimization step, and for each subblock showing conflict, attention weights are computed from cosine similarities with other tasks' subblocks. The resulting weighted combination is rescaled to keep the original magnitude, and it is interpolated with the original subgradient. This produces adjusted task gradients that reduce adversarial interactions while keeping constructive signals. Experimental results show that CAGT outperforms traditional approaches such as PCGrad and CAGrad, achieving roughly 10% lower loss values on some multi task datasets, such as ROT-MNIST and CelebA.
Fields of Study & Subjects
Domain Adaptation and Few-Shot Learning
Stochastic Gradient Optimization Techniques
Advanced Neural Network Applications
Publication details
- Publication date
- 5 February 2026
- Document type
- Preprint