Home

Awesome

[CVPR 2023] Independent Component Alignment for Multi-Task Learning


Preview animation

Dmitriy Senushkin, Nikolay Patakin, Arsneii Kuznetsov, Anton Konushin

Paper | Supplementary | Arxiv | Video

We introduce an Aligned-MTL - gradient optimization method for training multi-task neural networks. We address multi-task optimization issues by examining the stability of a gradient system as measured by its condition number. Using this criterion as a design principle, we present an Aligned-MTL method that eliminates instability in the training process by aligning the principal components of a gradient system.

This is an official Pytorch implementation of the paper

Installation

Besides torch installation, please ensure that the following libraries are installed:

pip install scipy cvxpy matplotlib seaborn tqdm

Data prepration

  1. For NYU (3 task) benchmark please download data from this link (~9GB, provided by MTAN repository). Unzip and specify --data-path path to a folder containing train and val subfolders.

  2. For Cityscapes (2 task) benchmark: link. Unzip and specify --data-path path to a folder containing train and val subfolders.

  3. For Cityscapes (3 task) benchmark please refer to the original Cityscapes benchmark website and download the following files:

    • disparity_trainvaltest.zip (3.8 GB)
    • gtFine_trainvaltest.zip (250 MB)
    • leftImg8bit_trainvaltest.zip (11.6 GB)

    Unzip them into the same folder (as a result, should contain leftImg8bit, disparity and gtFine)

MTL methods

We provide implementations of several multi-task optimization methods. Gradient optimization method can be used during training by setting --balancer parameter to one of the following values:

Parameter valueMethodConference
lsUniform weighting (aka linear scalarization)-
uncertaintyUncertainty weghtingCVPR 2018
gradnormGradient normalizationICML 2018
mgda, mgdaubMGDA and MGDA-UBNeurIPS 2018
dwaDynamic Weight AverageCVPR 2019
pcgradProjecting Conflicting GradientsNeurIPS 2020
graddropGradient Sign DropoutNeurIPS 2020
imtlImpartial Multi-Task LearningICLR 2021
gradvacGradient VaccineICLR 2021
cagradConflict-Averse Gradient descentNeurIPS 2021
nashNash-MTLICML 2022
rlwRandom Loss Weighting with normal distributionTMLR 2022
amtl, amtlubOurs, Aligned-MTL and Aligned-MTL-UBCVPR 2023

Benchmarks

Our repository provides four benchmarks:

  1. cityscapes_pspnet Cityscapes 3-task. PSPNet model. Semantic segmentation + Instance segmentation + Depth estimation
  2. cityscapes_mtan Cityscapes 2-task. MTAN model. Semantic segmentation + Depth estimation
  3. nyuv2_pspnet NYUv2 3-task. PSPNet model. Semantic segmentation + Depth estimation + Surface Normal Estimation
  4. nyuv2_mtan NYUv2 3-task. MTAN model. Semantic segmentation + Depth estimation + Surface Normal Estimation

To run the training and evaluation code use (change parameter values according to your needs):

python train.py --benchmark cityscapes_pspnet --balancer amtl --data-path /path/to/cityscapes

To compare methods on a synthetic two-task benchmark use:

python optimize_toy.py --balancer amtl --scale 0.5

Scale parameter controls the balance between problems, i.e. L0 = scale * L1 + (1-scale) * L2.

Links

Our synthetic benchmark and visualisation code is based on the NashMTL repository.

Citation

If you find our code repository or paper useful, please cite us:

@InProceedings{Senushkin_2023_CVPR,
    author    = {Senushkin, Dmitry and Patakin, Nikolay and Kuznetsov, Arseny and Konushin, Anton},
    title     = {Independent Component Alignment for Multi-Task Learning},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2023},
    pages     = {20083-20093}
}