This repository contains code for Gradient Surgery for Multi-Task Learning in TensorFlow v1.0+ (PyTorch implementation forthcoming).
PCGrad is a form of gradient surgery that projects a task’s gradient onto the normal plane of the gradient of any other task that has a conflicting gradient, which achieves substantial gains in efficiency and performance on a range of supervised multi-task learning and multi-task reinforcement learning domains. Moreover, it is model-agnostic and can be combined with previously-proposed multitask architectures for enhanced performance.
Clone this repo and copy PCGrad_tf.py to your codebase.
optimizer=PCGrad(tf.train.AdamOptimizer()) # wrap your favorite optimizerlosses=# a list of per-task lossesassertlen(losses) ==num_taskstrain_op=optimizer.minimize(losses)Our experiments in the paper were based on the following repositories.
CIFAR-100-MTL: RoutingNetworks
NYUv2: MTAN
MultiMNIST: MultiObjectiveOptimization
MT10/MT50/goal-conditioned pushing in MetaWorld: softlearning with modifications (per-task temperature and per-task replay buffers). We will release modified multi-task softlearning code soon.
Please cite as:
@article{yu2020gradient,
title={Gradient surgery for multi-task learning},
author={Yu, Tianhe and Kumar, Saurabh and Gupta, Abhishek and Levine, Sergey and Hausman, Karol and Finn, Chelsea},
journal={arXiv preprint arXiv:2001.06782},
year={2020}
}