Resolving Token-Space Gradient Conflicts: Token Space Manipulation for Transformer-Based Multi-Task Learning [ICCV 2025]

May 24, 2026 · View on GitHub

Official PyTorch implementation of Resolving Token-Space Gradient Conflicts: Token Space Manipulation for Transformer-Based Multi-Task Learning [ICCV 2025].

Wooseong Jeong & Kuk-Jin Yoon, Korea Advanced Institute of Science and Technology (KAIST)

Multi-Task Learning (MTL) enables multiple tasks to be learned within a shared network, but differences in objectives across tasks can cause negative transfer, where the learning of one task degrades another task's performance. While pre-trained transformers significantly improve MTL performance, their fixed network capacity and rigid structure limit adaptability. Previous dynamic network architectures attempt to address this but are inefficient as they directly convert shared parameters into task-specific ones. We propose Dynamic Token Modulation and Expansion (DTME-MTL), a framework applicable to any transformer-based MTL architecture. DTME-MTL enhances adaptability and reduces overfitting by identifying gradient conflicts in token space and applying adaptive solutions based on conflict type. Unlike prior methods that mitigate negative transfer by duplicating network parameters, DTME-MTL operates entirely in token space, enabling efficient adaptation without excessive parameter growth. Extensive experiments demonstrate that DTME-MTL consistently improves multi-task performance with minimal computational overhead, offering a scalable and effective solution for enhancing transformer-based MTL models.

Method overview

DTME-MTL operates inside the self-attention of a shared transformer backbone:

Conflict detection (token space). For each attention module, the covariance of the input tokens is decomposed by SVD into a range subspace (high-energy directions) and a null subspace. Per-task gradients w.r.t. the tokens are projected into these subspaces, and inter-task gradient conflicts (negative agreement) are accumulated separately for each subspace.
Adaptive, automatic placement. Modules are ranked by the amount of conflict they exhibit, and DTME is applied where conflicts concentrate (not at a fixed, hand-picked layer):
- Token modulation for range-space conflicts — a lightweight per-task affine transform of the tokens.
- Token expansion for null-space conflicts — a few learnable per-task tokens appended to the sequence.
Training. The expanded structure is trained end-to-end. Because adaptation lives entirely in token space, the parameter overhead is small.

The implementation lives in models/transformers/vit.py (the Attention module: update_cov, svd, store_grad, cal_conflict, expand_network) and utils/train_utils.py (measure_conflict, which performs the ranking and automatic placement).

Installation

conda create -n dtme python=3.8 -y
conda activate dtme
pip install -r requirements.txt

Note on timm. The ViT backbone relies on the timm.models.helpers / timm.models.layers API of timm 0.5–0.6 (pinned to timm==0.6.13 in requirements.txt). Newer timm releases reorganized these modules.

Datasets

We use the PASCAL-Context and NYUD-v2 multi-task datasets in the layout prepared by InvPT. Download and unpack them so that you have:

<DB_ROOT>/
├── PASCALContext/
└── NYUDv2/

Point the code to your data root in one of two ways:

set the environment variable export DB_ROOT=/path/to/data, or
edit db_root in configs/mypath.py.

Usage

Training follows three stages. The example scripts in scripts/ run the full pipeline on ViT-Tiny:

bash scripts/run_nyud.sh      # NYUD-v2
bash scripts/run_pascal.sh    # PASCAL-Context

Equivalently, run the stages manually (single GPU shown; use --nproc_per_node=N for multi-GPU):

LAUNCH="python -m torch.distributed.launch --nproc_per_node=1 --master_port 29500"

# Stage 0 — multi-task baseline (conflicts are measured on a trained model)
CUDA_VISIBLE_DEVICES=0 $LAUNCH train_baseline.py \
    --config_exp ./configs/nyud/baseline_vitT.yml --run_mode train

# Stage 1 — measure token-space gradient conflicts and expand the network
#           (loads the baseline via `resume`, writes mod.txt / exp.txt to the results dir)
CUDA_VISIBLE_DEVICES=0 $LAUNCH measure_and_expand.py \
    --config_exp ./configs/nyud/dtme_vitT.yml --run_mode train

# Stage 2 — train the expanded DTME model
CUDA_VISIBLE_DEVICES=0 $LAUNCH train_dtme.py \
    --config_exp ./configs/nyud/dtme_vitT.yml --run_mode train

Inference / evaluation on the trained model:

CUDA_VISIBLE_DEVICES=0 $LAUNCH train_dtme.py \
    --config_exp ./configs/nyud/dtme_vitT.yml --run_mode infer

Outputs (logs, checkpoints, results) are written to ./output/<version_name>/.

Configs and backbones

Configs are provided for ViT-Tiny (vitT) and ViT-Large (vitL, with the InvPT decoder) on both datasets:

configs/
├── nyud/   { baseline_vitT, baseline_vitL, dtme_vitT, dtme_vitL }.yml
└── pascal/ { baseline_vitT, baseline_vitL, dtme_vitT, dtme_vitL }.yml

The key DTME hyper-parameters (in the dtme_* configs) are:

field	meaning
`R_svd`	SVD energy ratio that splits the token covariance into range / null subspaces
`R_r`	fraction of the most-conflicting modules (range space) that receive token modulation
`R_n`	fraction of the most-conflicting modules (null space) that receive token expansion
`num_token`	number of expansion tokens added per task in an expanded module

Boundary (edge) evaluation

Edge metrics during training use a fast surrogate. The paper's boundary scores are computed offline with the SEISM MATLAB toolbox via eval_edge.py. Set export SEISM_ROOT=/path/to/seism and run:

python eval_edge.py --database NYUD --save_dir ./output/<version_name>/results --exp_name my_run

@inproceedings{jeong2025dtme,
  title     = {Resolving Token-Space Gradient Conflicts: Token Space Manipulation for Transformer-Based Multi-Task Learning},
  author    = {Jeong, Wooseong and Yoon, Kuk-Jin},
  booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
  year      = {2025}
}

Contact

Wooseong Jeong: stk14570@kaist.ac.kr