summaries: The paper tackles the problem of spurious correlations caused by observed confounders and offers two potential solutions: implicit and explicit counterfactual maximum likelihood estimation. Specifically the paper deals with the setup of predicting the outcome Y of some action T on X and considers X to potentially be a confounder. The paper compares the performance of the proposed methods to on a natural language inference and an image captioning task and shows improvements in terms of human evaluation but little difference in terms of automatic evaluation.

summaries: This work proposes counterfactual maximum likelihood estimation (CMLE), a new training objective for suppressing the spurious effect of confounding variables. The authors derive two types of upperbound formulation for interventional log-likelihood (i.e. implicit CMLE and explicit CMLE). Using natural language inference (NLI) and image captioning as the intervention modeling tasks, the proposed method demonstrated superior performance in terms of human-perceived performance.

summaries: This paper applies methodology from CATE estimation to do maximum likelihood with respect to samples from the interventional distributions Y|do(T),X, which they call Counterfactual MLE (CMLE). They provide two algorithms, Implicit CMLE and Explicit CMLE. The implicit variant uses the CATE generalization bounds from Shalit's CATE generalization bounds paper to get bounds on the likelihood from the desired counterfactual sampling distribution. The Explicit variant writes bounds in terms of the distribution T|Y,X. The two methods are applied to natural language inference and image captioning tasks. The authors suggest that this methodology and the surrounding literature on causally-inspired prediction should play a large role in the standard approach of deep learning in order to avoid shortcut learning / relying on spurious correlations.

summaries: In this work, authors propose the Counterfactual Maximum Likelihood Objective CMLE, where the parameters of the statistical model are inferred by maximizing the likelihood of the data under the interventional distribution, as opposed to the observational data. For the assumed model X -> Y <- T and X -> T, where X is the covariate, T is the treatment and Y is the outcome, the objective is to learn the parameters \theta that predict Y|X=x,T=t using just the causal link T->Y and not the spurious features in X=x. Ideally, if one could intervene on T and set do(T=t), it is possible to collect data where the outcome is observed upon the intervention done on T. Then maximizing the following likelihood: (CMLE objective) \argmax_\theta \sum_t E_{X|t}  E_{Y|x,t} p_\theta(Y|X=x,do(T=t)) would give us the MLE estimate of the parameter \theta under the set of interventional distributions do(T=t), thereby eliminating any effect of the spurious features in X on the prediction of the outcome label Y. The paper proposes two upper bounds for the above CMLE objective: i) Implicit MLE: focuses on obtaining better representations of the covariates X, such that the distribution of these representations is similar (by Wasserstein distance) for every value of T; and ii) Explicit MLE: where counterfactual examples are generated during training. Expectations in the theoretical derivations of the above upper bounds, can be estimated by Monte Carlo evaluations using only the observational data. Results are provided on two downstream tasks where models often rely upon spurious correlations: i) Natural Language Inference (NLI); and ii) Image Captioning (IC).

summaries: This paper proposed to use counterfactual maximum likelihood estimation to learn deep learning models that are less susceptible to spurious correlation relationships. In particular, it proposes two general algorithms, Implicit CMLE and Explicit CMLE, for learning causal predictions of DL models under observational data. The proposed method shows improvement over the regular MLE method in two real data sets.

abstractText: Although deep learning models have driven state-of-the-art performance on a wide array of tasks, they are prone to spurious correlations that should not be learned as predictive clues. To mitigate this problem, we propose a causality-based training framework to reduce the spurious correlations caused by observed confounders. We give theoretical analysis on the underlying general Structural Causal Model (SCM) and propose to perform Maximum Likelihood Estimation (MLE) on the interventional distribution instead of the observational distribution, namely Counterfactual Maximum Likelihood Estimation (CMLE). As the interventional distribution, in general, is hidden from the observational data, we then derive two different upper bounds of the expected negative log-likelihood and propose two general algorithms, Implicit CMLE and Explicit CMLE, for causal predictions of deep learning models using observational data. We conduct experiments on both simulated data and two real-world tasks: Natural Language Inference (NLI) and Image Captioning. The results show that CMLE methods outperform the regular MLE method in terms of out-of-domain generalization performance and reducing spurious correlations, while maintaining comparable performance on the regular evaluations.1

authors:
J. Jo
Y. Bengio

title: Measuring the tendency of cnns to learn surface statistical regularities, 2017

authors:
S. Beery
G. Van Horn
P. Perona

title: Recognition in terra incognita

venue: In Proceedings of the European Conference on Computer Vision (ECCV)

authors:
A. Poliak
J. Naradowsky
A. Haldar

title: Hypothesis only baselines in natural language inference

venue: In Proceedings of the Seventh Joint Conference on Lexical and Computational Semantics,

authors:
A. Rohrbach
L.A. Hendricks
K. Burns

title: Object hallucination in image captioning

venue: In Empirical Methods in Natural Language Processing (EMNLP)

authors:
R. Geirhos
P. Rubisch
C. Michaelis

title: Imagenet-trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness, 2019

authors:
Y. Zhang
H. Tan
M. Bansal

title: Diagnosing the environment bias in vision-and-language navigation

venue: Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence,

authors:
D. Kaushik
Z.C. Lipton

title: How much reading does reading comprehension require? a critical investigation of popular benchmarks

venue: In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing,

authors:
J. Quiñonero-Candela
M. Sugiyama
N.D. Lawrence

title: Dataset shift in machine learning

authors:
C. Szegedy
W. Zaremba
I. Sutskever

title: Intriguing properties of neural networks

authors:
Y. Ovadia
E. Fertig
J. Ren

title: Can you trust your model’s uncertainty? evaluating predictive uncertainty under dataset shift

venue: NeurIPS

authors:
A. Filos
P. Tigkas
R. Mcallister

title: Can autonomous vehicles identify, recover from, and adapt to distribution shifts

venue: Proceedings of the 37th International Conference on Machine Learning,

authors:
J. Peters
P. Buhlmann
N. Meinshausen

title: Causal inference using invariant prediction: identification and confidence intervals

venue: arXiv: Methodology,

authors:
S. Chang
Y. Zhang
M. Yu

title: Invariant rationalization

venue: In International Conference on Machine Learning,

authors:
D. Kaushik
E. Hovy
Z.C. Lipton

title: Learning the difference that makes a difference with counterfactually augmented data

venue: International Conference on Learning Representations (ICLR),

authors:
K. Lu
P. Mardziel
F. Wu

title: Gender bias in neural natural language processing

venue: arXiv preprint arXiv:1807.11714,

authors:
R. Zmigrod
S.J. Mielke
H. Wallach

title: Counterfactual data augmentation for mitigating gender stereotypes in languages with rich morphology

venue: In Association for Computational Linguistics (ACL)

authors:
R.H. Maudslay
H. Gonen
R. Cotterell

title: It’s all in the name: Mitigating gender bias with name-based counterfactual data substitution

authors:
S. Pitis
E. Creager
A. Garg

title: Counterfactual data augmentation using locally factored dynamics

venue: Advances in Neural Information Processing Systems,

authors:
S. Greenland
J. Pearl
J.M. Robins

title: Causal diagrams for epidemiologic research

venue: Epidemiology, pages 37–48,

authors:
I. Shpitser
J. Pearl

title: Identification of conditional interventional distributions

venue: In Proceedings of the Twenty-Second Conference on Uncertainty in Artificial Intelligence,

authors:
U. Shalit
F.D. Johansson
D. Sontag

title: Estimating individual treatment effect: generalization bounds and algorithms

authors:
S.R. Bowman
G. Angeli
C. Potts

title: A large annotated corpus for learning natural language inference

venue: arXiv preprint arXiv:1508.05326,

authors:
A. Williams
N. Nangia
S. Bowman

title: A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pages 1112–1122

venue: Association for Computational Linguistics,

authors:
Y. Zhang
J. Baldridge
L. He

title: PAWS: Paraphrase Adversaries from Word Scrambling

venue: In Proc. of NAACL

authors:
W. Lan
W. Xu

title: Neural network models for paraphrase identification, semantic textual similarity, natural language inference, and question answering

venue: In Proceedings of COLING

authors:
A. Suhr
M. Lewis
J. Yeh

title: A corpus of natural language for visual reasoning. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pages 217–223

authors:
L. Gatys
A. Ecker
M. Bethge

title: A neural algorithm of artistic style

venue: Journal of Vision,

authors:
A. Madaan
A. Setlur
T. Parekh

title: Politeness transfer: A tag and generate approach

venue: In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics,

authors:
M. Mirza
S. Osindero

title: Conditional generative adversarial nets

venue: arXiv preprint arXiv:1411.1784,

authors:
Z. Hu
Z. Yang
X. Liang

title: Toward controlled generation of text

venue: In ICML,

authors:
M. Khalifa
H. Elsahar
M. Dymetman

title: A distributional approach to controlled text generation

venue: In International Conference on Learning Representations

authors:
L. Chen
Z. Jiang
J. Xiao

title: Human-like controllable image captioning with verb-specific semantic roles

venue: arXiv preprint arXiv:2103.12204,

authors:
C. Deng
N. Ding
M. Tan

title: Length-controllable image captioning

venue: arXiv preprint arXiv:2007.09580,

authors:
Q. You
H. Jin
Z. Wang

title: Image captioning with semantic attention

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
J. Peters
P. Bühlmann
N. Meinshausen

title: Causal inference by using invariant prediction: identification and confidence intervals

venue: Journal of the Royal Statistical Society: Series B (Statistical Methodology),

authors:
A. Ghassami
S. Salehkaleybar
N. Kiyavash

title: Learning causal structures using regression invariance

venue: In Advances in Neural Information Processing Systems (NeurIPS)

authors:
M. Srivastava
T. Hashimoto
P. Liang

title: Robustness to spurious correlations via human annotations

venue: International Conference on Machine Learning (ICML),

authors:
D. Kaushik
A. Setlur
E. Hovy

title: Explaining the efficacy of counterfactually augmented data

authors:
D. Teney
E. Abbasnedjad
A. v. d. Hengel

title: Learning what makes a difference from counterfactual examples and gradient supervision

venue: arXiv preprint arXiv:2004.09034,

authors:
Z. Liang
W. Jiang
H. Hu

title: Learning to contrast the counterfactual samples for robust visual question answering

venue: In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP),

authors:
N. Hassanpour
R. Greiner

title: Learning disentangled representations for counterfactual regression

authors:
R. Pryzant
D. Card
D. Jurafsky

title: Causal effects of linguistic properties

venue: arXiv preprint arXiv:2010.12919,

authors:
S. Li
Y. Fu

title: Matching on balanced nonlinear representations for treatment effects estimation

authors:
F. Johansson
U. Shalit
D. Sontag

title: Learning representations for counterfactual inference

venue: In International conference on machine learning,

authors:
A.M. Alaa
M. van der Schaar

title: Bayesian inference of individualized treatment effects using multi-task gaussian processes

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
J. Yoon
J. Jordon
M. Van Der Schaar

title: Ganite: Estimation of individualized treatment effects using generative adversarial nets

authors:
L. Yao
S. Li
Y. Li

title: Representation learning for treatment effect estimation from observational data

authors:
D. Rezende
S. Mohamed

title: Variational inference with normalizing flows

venue: Proceedings of the 32nd International Conference on Machine Learning,

authors:
D. Tran
K. Vafa
K. Agrawal

title: Discrete flows: Invertible generative models of discrete data

title: Integral probability metrics and their generating classes of functions

venue: Advances in Applied Probability,

title: Optimal transport: old and new, vol. 338

venue: Springer Science & Business Media,

authors:
B.K. Sriperumbudur
K. Fukumizu
A. Gretton

title: On the empirical estimation of integral probability metrics

venue: Electronic Journal of Statistics,

authors:
M. Cuturi
A. Doucet

title: Fast computation of wasserstein barycenters

authors:
E. Jang
S. Gu
B. Poole

title: Categorical reparameterization with gumbel-softmax

venue: arXiv preprint arXiv:1611.01144,

authors:
A. Vaswani
N. Shazeer
N. Parmar

title: Attention is all you need

venue: arXiv preprint arXiv:1706.03762,

authors:
J. Devlin
M.-W. Chang
K. Lee

title: Bert: Pre-training of deep bidirectional transformers for language understanding

venue: arXiv preprint arXiv:1810.04805,

authors:
H. Tan
M. Bansal

title: Lxmert: Learning cross-modality encoder representations from transformers

venue: In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP),

authors:
Y. Liu
M. Ott
N. Goyal

title: Roberta: A robustly optimized bert pretraining approach

venue: arXiv preprint arXiv:1907.11692,

authors:
A. Williams
N. Nangia
S.R. Bowman

venue: arXiv preprint arXiv:1704.05426,

authors:
Y. Nie
A. Williams
E. Dinan

title: Adversarial NLI: A new benchmark for natural language understanding

venue: In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics. Association for Computational Linguistics,

authors:
M. Lewis
Y. Liu
N. Goyal

title: BART: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension

authors:
T. Wolf
L. Debut
V. Sanh

title: Transformers: State-of-the-art natural language processing

venue: In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations,

authors:
X. Chen
H. Fang
T.-Y. Lin

title: Microsoft coco captions: Data collection and evaluation server

venue: arXiv preprint arXiv:1504.00325,

authors:
A. Karpathy
L. Fei-Fei

title: Deep visual-semantic alignments for generating image descriptions

authors:
R. Luo
B. Price
S. Cohen

title: Discriminability objective for training descriptive captions

venue: arXiv preprint arXiv:1803.04376,

authors:
K. Papineni
S. Roukos
T. Ward

title: Bleu: a method for automatic evaluation of machine translation

venue: In Proceedings of the 40th annual meeting of the Association for Computational Linguistics,

authors:
S. Banerjee
A. Lavie

title: Meteor: An automatic metric for mt evaluation with improved correlation with human judgments. In Proceedings of the acl workshop on intrinsic and extrinsic evaluation measures for machine translation and/or summarization, pages 65–72

authors:
R. Vedantam
C. Lawrence Zitnick
D. Parikh

title: Cider: Consensus-based image description evaluation

authors:
P. Anderson
B. Fernando
M. Johnson

title: Spice: Semantic propositional image caption evaluation

venue: In European conference on computer vision,

authors:
D.P. Kingma
M. Welling

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114,

title: Simple statistical gradient-following algorithms for connectionist reinforcement learning

venue: Machine learning,

heading: 1 Introduction

text:
Deep neural networks have been tremendously successful across a variety of tasks and domains in recent years. However, studies have shown that deep learning models trained with traditional supervised learning framework tend to learn spurious correlations as predictive clues [1–4]. For example, in computer vision, deep learning models can rely on surface-level textures [1, 5] or background environment [2, 6] instead of the object of interest in images. In natural language processing (NLP), question-answering models are insensitive to the choice of question [7] and natural language inference models are surprisingly accurate at predicting the logical relationship between a pair of sentences from just one of them [3]. In image captioning, a phenomenon called object hallucination is observed where models tend to include nonexistent objects in a caption based on their common association with other objects that are actually present in the image [4].
Another related problem in the current associative learning framework, like maximum likelihood estimation (MLE), is that neural networks can achieve almost perfect performance on one dataset but dramatically fail to generalize on another because of a naturally occurring or adversarially enforced distribution shift [8–11]. One explanation is that the model learns ‘fake’ features induced by spurious correlations instead of invariant features that would hold in different domains for the same task [12]; this kind of invariance can be explained by the underlying causal mechanism of the task [13].
1Our code is released at https://github.com/WANGXinyiLinda/CMLE.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
It is usually impractical to identify and disentangle all of the potential sources of spurious correlations. Various training algorithms have been proposed to reduce the influence of spurious correlations by learning the underlying causal mechanism. One approach for discovering the underlying causal mechanism is by invariant prediction based on the invariance of causal mechanisms across different environments [13, 12, 14]. However, constructing such diverse environments is usually unrealistic. Another approach is counterfactual data augmentation [15–19], which directly modifies the part of the input that causes the target variable, but usually involves expensive human efforts in generating the counterfactual examples. We propose Counterfactual Maximum Likelihood Estimation (CMLE), which tries to perform MLE on an interventional distribution instead of the observational distribution. Here, obervational distribution means the distribution that the train data is sampled from. Interventional distribution means the distribution with one or some of the variables intentionally intervened (set to some fixed value). More specifically, we aim to remove observable confounders and reduce spurious correlations by directly intervening on the label variable, without requiring human annotators or the curation of diverse environments.
In this paper, we focus on the setting of predicting the outcome Y of a certain action T on X , with the underlying causal model as X → Y ← T and an observed train dataset in the form of (X,Y, T ). Such a setting enables us to assume that there exist some observed confounders in X that influence both Y and T . This means there is potentially a false causal relation from X to T , indicated by a red arrow in Figure 1a. In our framework, instead of trying to identify all the confounders, we try to rule out the effect of the confounders by considering the interventional distribution that directly deleting the false causal edge from X to T with a resulting causal graph as shown in Figure 1d. This causal graph can be viewed as a Randomized Controlled Trial (RCT), which is the gold standard approach for estimating the treatment effect [20].
Our method is inspired by individual treatment effect (ITE) prediction [21, 22], which tries to predict the expected effect E[Y1 − Y0|X = x] (e.g. differnce in blood pressure) of a treatment T ∈ {0, 1} (e.g. drug/surgery) on a individual unit X = x (e.g. a patient). In this case, the observed confounder in X can be explained by selection bias [22], which means the treatment T applied to an individual X is dependent on X . For example, young patients are more likely to be treated by surgery, while elder patients are more likely to be treated by drugs.
There are two broad kinds of tasks to which our proposed CMLE framework can apply: relationship prediction tasks that predicts the relation T given two inputsX and Y (e.g. natural language inference [23, 24], paraphrase identification [25, 26], natural language for visual reasoning [27], etc.), and conditional generation tasks that generate Y given preconditionX and constraint T (e.g. style transfer [28, 29], controllable image/text generation [30–32], controllable image captioning [33, 34] etc.). For conditional generation tasks, the causal relation of X,Y, T fits in our setting. For relation prediction tasks, we first assume there is an underlying conditional generation process of Y for a given pair of X and T . This is a natural assumption as this is usually how the datasets for this type of task is generated. Then we augment the original dataset using the counterfactual examples generated by our method. In this paper, we try to reduce the spurious correlations contained in both kinds of tasks.
Since we generally cannot observe the interventional distribution, we derive two different upper bounds of the expected interventional negative log-likelihood using only the observational distribution: one is Implicit CMLE, as it does not involve any explicit generation of counterfactual examples; while the other is Explicit CMLE, as it explicitly generates counterfactual examples during the training process. We test our framework with deep neural networks on a simulated dataset and two real-world tasks with well-known spurious correlations: natural language inference [23] and image captioning [35]. Compared to regular MLE, we improve the out-of-domain accuracy of NLI on a hard adversarial dataset by 2.9% relatively and beat baselines on human preference evaluations by more than 10%, while maintaining comparable performance on automated evaluations. Our results show that our learning framework can better capture the underlying causal mechanism and reduce spurious correlations without degrading performance.

heading: 2 Related Work

text:
A growing body of work has investigated algorithmic improvements for machine learning model training by leveraging underlying causal mechanisms. One line of work tries to utilize the invariance of causal mechanisms across different environments. Invariant Causal Prediction (ICP) [36] aims to
X Y T
(a) Observed pO(X,Y, T )
X
Y1 Y2 Ym
T=1 T=2
T=m
T
… YT …
(b) Expand all possible outcomes
X
Y1 Y2 Ym
T=1 T=2 T=m
… ...
(c) Intervene on T
X
Y T
(d) Intervened pI(X,Y, T )
Figure 1: The Structural Causal Model (SCM) assumed by Counterfactual MLE, where (a) and (b) have the same distribution, (c) is the resulting causal graph of intervening on T . (d) is the intervened distribution by only removing the X → T edge from (a). Shaded nodes means being observed and red arrows denote the causal relation causing the spurious correlation.
learn a linear invariant causal predictor to identify the causal parents of given target variables with data obtained from different environments. Further refinements on this line of work include nonlinear and nonparametric ICP assessment methods [37] and Invariant Risk Minimization (IRM) where out-of-distribution generalization is enabled by learning a nonlinear invariant causal predictor shared across different environments [12]. However, these methods generally require a set of environments that are both sufficiently diverse to eliminate confounders while maintaining the underlying causal structure. Though this can be achieved by intervening in the non-causal variables in a synthetic or interactive setting, it is very difficult to obtain or even define such environments with complex high-dimensional observational data like text or images. Our proposed CMLE framework only requires a single observational dataset, making it much more practical.
Another line of work approaches this problem with data augmentation. Methods including programmatically editing text to address gender bias [16–18], employing crowd workers to augment training data to capture potential unmeasured variables [38], and manually editing minimal text spans to flip the predicted label [15] have all been proposed to produce counterfactually augmented data (CAD) for improving out-of-domain generalization. Conceptually, in contrast to the first approach that involves interventions on a potentially large number of non-causal variables, CAD directly intervenes on a target variable [39], which is more practical for a specific task. Refinements on this work attempt to better utilize the manually edited or synthetically generated counterfactual data by pairing the factual and counterfactual examples together [40, 41], and leverage local causal structures to generate counterfactual experience in training reinforcement learning algorithms [19]. While human annotators can generate high-quality counterfactual examples that are effective in reducing spurious correlations, they are expensive to obtain and do not easily scale. Meanwhile, our proposed CMLE framework is fully automated, requiring no human annotations.
Our CMLE framework is inspired by a technique for predicting the Individual Treatment Effect (ITE) from observational data [22]. This method only observes one possible outcome of an individual X = x with one binary treatment T = 0 or T = 1, and it assumes there are no hidden confounders but potentially observed ones. Assuming a similar underlying causal model as ours (m = 2) gives an error bound as the sum of the standard generalization error and the distance between the control and treated distributions in a representation space. Improvements on this work include disentangling the observable confounders in ITE using representation learning [42], and estimating the causal effects of linguistic properties [43]. While in this work we do not aim to separate the confounders from the causal factors in X , we try to reduce the spurious correlations caused by the possible confounders contained in X using our proposed CMLE training framework.

heading: 3 Method

text: In this section, we first introduce our problem setting and the goal of Counterfactual Maximum Likelihood Estimation (CMLE). As we cannot observe the counterfactual outcomes in observational data, we propose two different ways of estimating the counterfactual likelihood by deriving two different upper bounds: Implicit CMLE, which is faster in training, and Explicit CMLE, which is more flexible. We also briefly introduce the deep learning architecture we use in our experiments.

heading: 3.1 Problem Setting

text:
We consider the scenario with one input variable X ∈ X that contain all the confounders, one discrete treatment variable T ∈ {1, 2, ...,m} where m ≥ 2, and one outcome variable Y ∈ Y ,
as shown in Figure 1a. Note that X and Y can be very complicated high dimensional space like text or image and all three variables are observed in the dataset. Suppose the data we observed is D = {(x1, t1, y1), (x2, t2, y2), ..., (xn, tn, yn)} and only xi’s are sampled i.i.d. In this paper, we adopt the causal notations from [44]. Below we give a formal definition of Structural Causal Model (SCM) [44] that we use to model our framework: Definition 1. A Structural Causal Model (SCM) is a pair 〈M,p(u)〉, where M is a triple 〈U, V, F 〉 and p(u) is a probability function defined over the domain of U . U = {U1, U2, ..., Un} is a set of background variables (exogenous) that are determined by factors outside the model. V = {V1, V2, ..., Vn} is a set of variables (endogenous) and F = {f1, f2, ..., fn} is a set of functions s.t. Vi = fi(Pa(Vi), Ui), i = 1, 2, ..., n, where Pa(Vi) ∈ V \Vi are the parents of Vi. We say that “Vj is a direct cause of Vi" if Vj ∈ Pa(Vi) and illustrate this relation as Vj → Vi in a causal graph.
In Figure 1, we omit the exogenous and only show the causal relations between endogenous variables. The shaded nodes indicate that the corresponding variables are observed. Intervention on a specific variable is formalized as fixing the value of this variable to a specific value. This is known as the do-operation, changing the causal graph by removing all edges incident to the intervened variable. We denote the outcome of Y corresponding to do(T = t) as Yt. We expand all the possible outcomes of Y in Figure 1b while we actually can only observe one outcome for each data point in the observational dataset. We also make the following standard assumption [22, 45–49]: Assumption 1. (Strong Ignorability Condition) (Y1, Y2, ..., Ym) ⊥ T |X and 0 < p(T |X) < 1.
Under this condition, the counterfactual outcomes are identifiable since p(Yt = y|X = x) = p(Y = y|X = x, T = t). i.e. The counterfactual likelihood can be written as observable conditional probabilities estimated from data.
Our goal is to eliminate the false causal relation from X to T represented by the red arrow in Figure 1a. To achieve this goal, we define our objective over the interventional distribution pI(X,Y, T ) as shown in Figure 1d, instead of over the observed data distribution pO(X,Y, T ) as shown in Figure 1a. To construct the interventional distribution, we only delete the causal edge between X and T and keep all the other causal relationships unchanged. In this distribution, T is uniformly sampled from {1, 2, ...,m}. Note that pI(Y |X,T ) = pI(YT |X) = pO(YT |X) = pO(Y |X,T ). As p(YT |X) is the same for both observational distribution and interventional distribution, we omit the superscript for simplicity. Then for predicting Y , the objective of counterfactual maximum likelihood estimation (CMLE) would be:
EpI(X,Y,T ) [ − log pIθ(Y |X,T ) ] = EX [ 1 m m∑ i=1 EYi|X [− log pθ(Yi|X)] ]
(1)
Similarly, for predicting T , we have:
EpI(X,Y,T ) [ − log pIθ(T |X,Y ) ] = EX [ 1 m m∑ i=1 EYi|X [− log p I θ(T = i|X,YT )] ] (2)
Here θ is the learnable parameters. In the remaining part of paper, the probabilities without superscripts are default to be observational probabilities. In the following sections, we first focus on deriving two different upper bounds of the CMLE objective for predicting Y : Implicit CMLE and Explicit CMLE. In Implicit CMLE, we add a regularizer to balance the distribution of X given T in the representation space. For Explicit CMLE, we directly generate counterfactual examples during training. To predict T , we train a classifier using normal MLE on a augmented dataset with the generated counterfactual examples to estimate the interventional distribution pI(X,Y, T ). We define a counterfactual example as follows: Definition 2. A counterfactual example is an example (x, t, y) with y sampled from p(Yt|X = x) s.t. (x, t, y) /∈ D and there is an example (x, t′, y′) ∈ D where t′ 6= t.
We can rewrite our objective function in Equation 1 as Ex [∑m t=1 Lθ(x, t) ] , by defining the expected loss function of a specific input and treatment as: Definition 3. We denote the loss function for predicting Y as Lθ(x, t, y) = − log pθ(Yt = y|X = x). Then the expected loss function given X = x and T = t is Lθ(x, t) = EYt|X=x[Lθ(x, t, Yt)].
Because of the Strong Ignorability Assumption, we have Lθ(x, t) > 0. Then we define an expected factual loss T=tF and expected counterfactual loss T=t CF when observing a specific treatment T = t:
Definition 4. The expected factual loss and expected counterfactual loss given T = t are T=tF = EX|T=t[Lθ(X, t)] and T=tCF = EX|T 6=t[Lθ(X, t)], respectively.
Here T=tF measures how well the model performs on the observational distribution while T=t CF measures how well the model would perform on an alternative world where a different treatment is chosen. Then we can rewrite our CMLE objective as:
Ex [ m∑ t=1 Lθ(x, t) ] = m∑ t=1 [p(T = t) T=tF + p(T 6= t) T=tCF ] (3)

heading: 3.2 Implicit CMLE

text:
We consider learning a representation function of X , Φ : X → Rd, where Rd is the representation space. Here we make the following assumption about Φ: Assumption 2. Φ : X → Rd is a invertible function. We denote Ψ : Rd → X as the inverse of Φ. For continuous X , Ψ also needs to be differentiable to calculate the Jacobian matrix.
Note that this assumption can be fulfilled for neural networks as long as all the activation functions are invertible. We denote the distribution induced by Φ by pΦ. We can easily transform p(X|T ) to pΦ(Φ(X)|T ) by a standard change of variable [50, 51]: Definition 5. For i ∈ {1, 2, ...,m} and r ∈ Rd, we denote pT=iΦ (r) := pΦ(r|T = i). For discrete X , we simply have pT=iΦ (r) = p(X = Ψ(r)|T = i), while for continuous X , we have pT=iΦ (r) = p(X = Ψ(r)|T = i) det ∣∣∂Ψ ∂r
∣∣. We adopt a similar idea as [22] and use the Integral Probability Metric (IPM) to derive an upper bound of the CMLE objective. IPM is a class of metrics between probability distributions with the following definition [52]: Definition 6. For two probability density functions p and q defined over S, we have
IPMG(p, q) := sup g∈G ∣∣∣∣∫ S g(s)(p(s)− q(s))ds ∣∣∣∣ Where G is a family of functions g : S → R.
Note that here (and in the rest of our paper) we abuse the notation of integration. If X is discrete, the integration would become a sum over it. Then we can derive an upper bound for our CMLE objective by Assumption 2, Definition 5 and Equation 3: (See Appendix A.1 for full proof.) Theorem 1. For a function familyG and a constantBΦ > 0 s.t. 1BΦLθ(Ψ(r), t) ∈ G for any r ∈ R d and t ∈ {1, 2, ...,m}, the CMLE objective is bounded by:
Ex [ m∑ t=1 Lθ(x, t) ] ≤ m∑ t=1 [ T=tF + p(T 6= t)BΦIPMG(pT=tΦ , p T 6=t Φ )]
In this upper bound, we are only left with observational/factual quantities that can be estimated from the observational data. We choose G to be 1-Lipschitz functions, which implies IPMG to be the Wasserstein distance [53, 54], denoted by WASS(·, ·). Then the empirical CMLE objective using an unbiased estimator of this upper bound can be written as:
arg min θ
1
n n∑ i=1 n uti Lθ(xi, ti, yi) + α m∑ j=1 (1− uj n )WASS({Φ(xi)}i:ti=j , {Φ(xi)}i:ti 6=j) (4)
Where ui = ∑n j=1 1tj=i and 1 is the indicator function. α is a hyperparameter accounting for BΦ and we abuse the notation of WASS(·, ·) to denote an empirical version of Wasserstein distance. The resulting empirical training objective consists of a weighted version of the normal factual loss for empirical risk minimization and a regularization term to draw pT=iΦ and p T 6=i Φ closer.
n uti accounts
for different number of examples with each label and 1 − ujn account for p(T 6= j) (see Theorem 1). The empirical Wasserstein distance is calculated among examples with different labels in the representation space. In practice, to perform back propagation, we directly compute the gradient of WASS(pT=iΦ , p T 6=i Φ ) over a single mini-batch, using the algorithm proposed by [55]. (See Appendix B.1 for details.)

heading: 3.3 Explicit CMLE

text:
While Implicit CMLE tries to reduce the spurious correlation by learning a better representation of X , we consider directly learning a better approximation of the log likelihood log pθ(Yt|X) by explicitly generating counterfactual examples at training. Here we assume that there is a neural network with parameter φ that approximates the interventional distribution pIθ(T |Y,X) defined as follow: Definition 7. pφ(T |Y,X) is an approximate of pIθ(T |Y,X) parameterized by φ s.t. KL(pIθ(T |X,Y )||pφ(T |X,Y )) < B, where B > 0 is a real constant. We also define an associated loss function Jφ(t, x, y) = − log pφ(T = t|Y = y,X = x).
Here KL(p||q) denotes the KL divergence between two distributions p and q. Our intuition is that it is directly feasible to learn the normal MLE objective EX,Y,T [log pφ(T |Y,X)] from the observation data, while is infeasible to directly learn the CMLE objective Ex [∑m t=1 Lθ(x, t) ] . Also, as T is a categorical variable, its distribution can be simplified as a Multinomial distribution. We then derive an upper bound of the CMLE objective using Bayes’ rule and Jensen’s Inequality: (See Appendix A.2 for full proof.) Theorem 2. By Assumption 1, there exist 0 < δ1 < δ2 < 1 s.t. δ1 < p(T |X) < δ2, then we have
Ex [ m∑ t=1 Lθ(x, t) ] ≤ Ex,t [ ηEpθ(Yt|X=x)[Lθ(x, t, y)] + m∑ i6=t Epθ(Yi|X=x)[Jφ(i, x, y)] ] + µ
Where η = (1 + 1p(t) (1− 1 m )), µ = (mδ2 −mδ1 −m+ 1) logm.
This theorem indicates that we can transfer the counterfactual likelihood of predicting Y into the counterfactual likelihood of T predicted by another model. Note that the resulting bound would be tighter when pφ(T |Y,X) is a better approximate of pIθ(T |Y,X) as we use the fact KL(pIθ(T |Y,X)||pφ(T |Y,X)) > 0 in the proof (see Appendix A.2). In practice, we can either pretrain a pφ(T |Y,X) on the observational data when pOθ (T |Y,X) and pIθ(T |Y,X) are not too far away, or we can simultaneously train a pφ(T |Y,X) with pθ(Yi|X) to take advantages of the generated counterfactual examples.
Since it is expensive to take m− 1 samples of yi from pθ(Yi|X) for each training example at train time, we introduce another random variable Kt for each t ∈ {1, 2, ...,m}, such that is Kt uniformly distributed on {1, 2, ...,m}\{t}. Then the Explicit CMLE objective can be rewritten as:
arg min θ
Ex,t [ Epθ(Yt|X=x) [ Lθ(x, t, Yt)] + m− 1 η EKt [Epθ(YKt |X=x)[Jφ(Kt, x, YKt)] ]]
Here we ignore the constant term µ in the upper bound and multiply the remaining terms by 1η . We can then write the Explicit CMLE objective empirically as:
arg min θ
1
n n∑ i=1 [Lθ(xi, ti, yi) + α(ti)Jφ(kti , xi, yθ(xi, kti))] (5)
Where α(ti) is a hyperparameter corresponding to m−1η , kti is randomly sampled from p(Kti) and yθ(xi, kti) is sampled from pθ(Ykti |X = x). We only sample Kt once and sample the corresponding YKt once for computational efficiency. In our implementation, to backpropagate through pθ(Ykti |X = x), we adopt the Gumbel-Softmax approach [56] to deal with the discrete text data, which creates a differentiable sample to replace the non-differentiable discrete variable. For alternative approaches for discrete and continuous Y , see Appendix B.2.

heading: 3.4 CMLE Implementation

text: While CMLE is a general causal training framework that can be applied to almost any deep learning architecture, in this paper, we primarily consider the Transformer [57] architecture, which excels at modeling sequential data and has been successfully applied to various domains, especially with pretraining on a large amount of data [58, 59]. Here we assume for any x ∈ X , we can write x as x = {x(0), x(1), ..., x(w)} and for any y ∈ Y , we have y = {y(0), y(1), ..., y(k)}. We input T = t by prepend a special token/vector corresponding to t as a prefix to the decoder input y. The model architectures of Implicit CMLE and Explicit CMLE are shown in Figure 2.

heading: 4 Experiments

text: We consider one synthetic dataset and two real-world tasks, Natural Language Inference and Image Captioning, to evaluate the effectiveness of our proposed CMLE causal training framework. Both real-world tasks have prominent datasets containing well-known spurious correlations; we aim to reduce these correlations’ influence on each target task. In the following sections, we will introduce our experiments in details and compare the performance of our proposed CMLE against vanilla MLE.

heading: 4.1 Simulated Experiments

text:
We consider a synthetic dataset generated by the following procedure inspired by [45]:
Given the followigng functions g1(x) = x−0.5, g2(x) = (x−0.5)2 +2, g3(x) = x2−1/3, g4(x) = −2 sin(2x), g5(x) = e−x−e−1−1, g6(x) = e−x, g7(x) = x2, g8(x) = x, g9(x) = 1x>0, g10(x) = cos(x), g11(x, t) = log(t+ x 2), g12(x, t) = e t+x, and g13(x, t) = sin(t+ x):
1. Sample xi from N (0, 1), i ∈ {1, 2, ..., 100}. 2. Let s = g1(x1) + g2(x2) + g3(x3) + g4(x4) + g5(x5). Let t = 1 if s ≤ 4; t = 2 if
4 < s ≤ 5; t = 3 if s > 5. 3. Let µ1 = g6(x1) + g7(x2) + g8(x3) + g11(x6, t) + g12(x7, t), µ2 = g9(x4) + g10(x5) + g11(x8) + g13(x9, t). Then sample y from N ((µ1, µ2), 1).
4. Repeat steps 1-3 for N times.
Here X is a real vector of length 100, T takes value in {1, 2, 3} and Y is a real vector of length 2. The task is to predict Y given X and T . Then we generate 10000 train data, and 5000 validation data and 5000 in distribution (observational) test data by setting N=10000 and 5000 respectively. We also generate corresponding ground truth counterfactual test data for testing how well our models can estimate the counterfactuals. OOD1, OOD2, OOD3 are three out-of-distribution test sets that generates T using different mechanisms as follows:
• OOD1: t is uniformly sampled from {1, 2, 3}. • OOD2: Let s = g1(x6) + g2(x7) + g3(x8) + g4(x9) + g5(x10). Let t = 1 if s ≤ 4; t = 2
if 4 < s ≤ 5; t = 3 if s > 5. • OOD3: Let s = g6(x1) + g7(x2) + g8(x3) + g9(x4) + g10(x5). Let t = 1 if s ≤ 4; t = 2
if 4 < s ≤ 5; t = 3 if s > 5. We use multi-layer perceptrons for all of our models and use the mean-squared-error (MSE) loss at training. This means we assume Y to follow a Gaussian distribution given X and T . In Table 1, we report MSE of predicting Y on different test sets.
Here Explicit CMLE* means using a pφ(T |X,Y ) that is trained on the interventional distribution instead of the observational distribution. The numbers are averaged over ten runs and the standard deviation is also reported. We take α = 0.01 for Implicit CMLE and α = 0.1 for Explicit CMLE.
We can see that all CMLE methods perform better than MLE on the counterfactual test set and the out-of-distribution test sets, while Explicit CMLE performs better than Implicit CMLE. And the performance of Explicit CMLE would be better when pφ(T |X,Y ) is closer to the ground truth interventional distribution.

heading: 4.2 Natural Language Inference (NLI)

text:
Natural language inference (NLI) [23] is a fundamental task in natural language understanding, which aims to infer the logical relation between a pair of sentences, premise and hypothesis. Unfortunately, NLI datasets are known to present significant spurious correlations. For example, [3] demonstrated that a machine learning model can non-trivially outperform the majority baseline using only the hypothesis without the premise. One known spurious correlation comes from the background knowledge contained in the premises. i.e. The prediction is based on the background knowledge instead of the logical relation between the premise and the hypothesis. For example, the hypothesis “TV viewing and consumption both begin at about the same time.” would usually be considered to be true from common knowledge. While a fine-tuned RoBERTa [60] also predict this example as entailment, it contradicts its premise “Regular TV viewing typically begins between 2 and 3 years of age, consuming about 10.”. We view the premise sentence as X , the hypothesis sentence as Y , and the label (entailment, neutral, or contradiction) as T . The NLI annotation procedure [23, 61] naturally fits in our causal framework in Figure 1, as the crowd workers are instructed to write a hypothesis that is entailed/neutral/contradicted to a given premise.
We consider three large-scale NLI datasets: SNLI [23], MNLI [24] and ANLI [62]. SNLI and MNLI are considered standard benchmarks, with each containing 550,152/10,000/10,000 examples and 392,702/20,000/20,000 examples for train/dev/test split respectively. ANLI was constructed as a hard adversarial dataset to eliminate spurious correlations and balance the data distribution, thus it serves as a great benchmark to evaluate a model’s generalizability. ANLI contains three subsets of increasing difficulty—A1, A2, and A3—with about about 392,702/20,000/20,000 examples, 45,460/1,000/1,000 examples and 100,459/1,200/1,200 examples each for train/dev/test split respectively.
For generating counterfactual hypotheses we fine-tune the pre-trained BART large model [63] provided by HuggingFace [64] with 12 encoder and decoder layers. We fine-tune BART on the combined dataset of SNLI and MNLI using both Implicit and Explicit CMLE. We choose α = 0.003 for Implicit CMLE and α = 0.1 for Explicit CMLE. We predict pφ(T |X,Y ) for Explicit CMLE with a conventionally fine-tuned RoBERTa large model with 24 layers. We call this model NLI-RoBERTa.
We evaluate how consistent our generated hypotheses Y are to their corresponding labels T using both NLI-RoBERTa and human evaluation2 in Table 2 (right). In both cases, we perform this assessment by comparing the model- or human-predicted t′ against the ground truth t. The accuracy of the NLIRoBERTa is shown in the first row of Table 2. Both of our CMLE methods significantly outperform the MLE baseline by a large margin on the labeling consistency of the generated hypothesis, which indicates that our CMLE framework can generate Yt that is truly influenced by T = t.
To construct a dataset that mimics sampling from the interventional distribution, we generate two counterfactual examples for each example in the original train dataset of SNLI and MNLI. We filter out the inconsistent examples using NLI-RoBERTa. We generate three augmented datasets using standard MLE, Implicit CMLE, and Explicit CMLE and fine-tune a copy of RoBERTa on each one. The results for these models are presented in rows MLE (w/ Aug), Implicit CMLE and Explicit CMLE of Table 2 respectively. The baseline MLE is NLI-RoBERTa itself. We test the accuracy of all the fine-tuned RoBERTas simultaneously on the regular test sets (MNLI, SNLI) and the out-of-domain adversarial test sets (ANLI). The results in Table 2 show that, while maintaining performance on
2All of our human evaluations are conducted via Amazon Mechanical Turk. See Appendix C.3 for details.
regular test sets, our proposed CMLE methods significantly outperform both MLE baselines on the adversarial test set. This indicates that CMLE can improve the out-of-domain generalization performance of large deep learning models without hurting the in-domain performance.

heading: 4.3 Image Captioning

text:
Image captioning is a fundamental task in studying multimodal interaction. However, image captioning models also suffer severely from spurious correlations [65] that lead to hallucinations. This type of spurious correlation is mainly caused by the typical co-occurrences of objects in images. For example, as toilets are frequently present in bathroom pictures, a captioning model trained by MLE is very likely to hallucinate a ‘toilet’ whenever it sees a ‘bathroom’, even if no toilet is present.
To reduce these hallucination artifacts, we apply our CMLE framework by regarding the image as X , the caption as Y , and whether the caption is hallucinated as T . In our setting, the nonhallucinated captions are ground truth captions written by a human, and the hallucinated captions are generated by image captioning models. The dataset we use is the MSCOCO 2014 dataset [66] which contains 123,287 images of common objects with 5 human-annotated captions per image. We use the Karpathy [67] split with 113,287/5,000/5,000 images in the train/validation/test set respectively. Our implementation is based on [68].
The Transformer architecture we use directly follows [57], with 6 layers in each of the encoder and decoder layers. To construct a dataset with both hallucinated and non-hallucinated captions, we synthetically generate captions by masking out one object word for each caption and use BERT (large) [58] to fill in an alternative word that aligns with the language model. We estimate the caption hallucination by the CHAIR score [65], which parses the object words in a caption and compares them to the ground truth objects annotated in MSCOCO images and captions. In total, we synthesize 648,550 hallucinated captions for 93,004 images in the training set.
Then we train the transformer model on this augmented dataset of MSCOCO captions with synthetic hallucinated captions using regular MLE and our proposed CMLE methods. We choose α = 0.0002 for Implicit CMLE and α = 0.0001 for Explicit CMLE. For Explicit CMLE, we fine-tune a pretrained LXMERT [59] on this augmented dataset as pφ(T |X,Y ). Here we consider two MLE baselines: The MLE baseline is the Transformer model regularly trained on the original MSCOCO dataset, while the MLE (w/ Aug) baseline is trained on the augmented dataset.
We conduct user studies on the generated captions as shown in Table 3. In the user study, we let the user choose the best caption among our two baselines and two proposed CMLE methods from the following aspects: Faithfulness, Expressiveness and Informativeness. Faithfulness means how faithful is the generated caption to the corresponding image, which partially measures the extend of hallucinations. Expressiveness means if the generated caption is coherent, grammatically, and semantically correct. Informativeness means if the caption includes most of the information in the image instead of giving a general description. We also let the user choose the caption that best describes the image in their opinion, which is referred to as Preference in Table 3. We can see that our proposed CMLE methods outperform the two MLE baselines with a large margin on all the evaluation dimensions which indicates that our methods can not only reduce the hallucination contents in the generated captions but also increase the quality of generated captions in general.
The automatic evaluation results are shown in Table 4, where CHAIRs is the percentage of hallucinated sentences and CHAIRi is the percentage of hallucinated instances (objects). We can see that our CMLE methods outperform both MLE baselines on CHAIR scores, which indicates that our methods
can effectively reduce hallucination in image captions. At the same time, the regular scores (BLEU@4, METEOR, CIDEr, SPICE) that measure the similarities between the generated captions and ground truth captions, are comparable between CMLE and the MLE baseline. We note that the MLE (w/ Aug) baseline on augmented data is significantly worse than the MLE baseline on both the automatic scores and human evaluations. We conjecture that this is because the model trained by vanilla MLE is not good at finding the causal relation between the hallucination label T and the caption Y and thus would be confused by the mixture of hallucinated captions and non-hallucinated captions.

heading: 4.4 Discussions and Limitations

text: From both NLI and image captioning experiments, we observe that Explicit MLE is significantly better than Implicit MLE in terms of human evaluations of the generated text, and comparable on other automatic evaluations. We conjecture that this is because that Implicit CMLE only incorporates counterfactual constraints into the representation of X while Explicit CMLE directly optimizes pθ(Yt|X) when generating the counterfactual examples. We also note that the improvement on automatic metrics is not that large for both real-world tasks. For NLI, our data augmentation cannot perfectly simulate sampling from the interventional distribution as we only augment two counterfactual examples for each factual example. For Image captioning, as it is difficult to simulate the true distribution of hallucinated captions, the models we trained on the augmented data can be affected by these newly introduced artifacts. On the other hand, the automatic metrics like BLEU only measure the similarity of generated sentences with the observed ground truth which cannot demonstrate the full strength of our methods.

heading: 5 Conclusion

text: Our proposed CMLE framework can be applied to a wide range of learning scenarios with paired inputs and corresponding labels or involving conditional generation. While currently we only consider the spurious correlations caused by the observable confounders under the Strong Ignorability assumption, it is often the case that the confounders are not observed. For example, annotators may have a biased pattern of annotation. As we do not record the information of the annotators, this confounder is considered unobserved. Expanding the current CMLE framework by including the unobserved confounders into the causal model would be an interesting and important future direction.

heading: Acknowledgments

text: This work was supported in part by the National Science Foundation Graduate Research Fellowship under Grant No. 1650114. This material is based on work that is partially funded by an unrestricted gift from Google.

title: Counterfactual Maximum Likelihood Estimation for Training Deep Networks


summaries: Authors propose a multitask learning approach to reduce gradient conflict by minimizing harm to the worst-performing task given any gradient update. Proofs are provided for convergence and pareto optimality of the resultant weight configuration. Experiments are run on a variety of settings, including computer vision and multitask reinforcement learning, with good results throughout.

summaries: Multi-task learning (MTL) is an efficient way of learning multiple tasks efficiently by training all tasks together, amortizing their cost by sharing some parameters across tasks. In order to perform this, a weighted sum of the task losses is usually optimized, which can lead to poor final results due to sub-optimal optimization, an effect produced by the competition between tasks for the shared resources. Specifically, when computing the gradient updates for the shared resources, gradients of different tasks can conflict with each other, cancelling each other out (an effect known as *conflicting gradients*).  This work introduces CAGrad, an algorithm which minimizes the average loss function (the network's loss), while maximizing the worst local improvement across all tasks, provably converging to a local minimum of the average loss. This algorithm results in a generalization of the usual gradient descent, as well as the MGDA algorithm. Empirical results show the effectiveness of the method.

summaries: This work proposes a method to reduce conflicting gradients during training of multi-task learning paradigms via gradient modification. Rather than simply take an average of the per-task gradients, the authors propose to learn a model-wide gradient direction $d$ to update the shared parameters by solving a dual objective to find per-task loss weights. Simply, the proposed method (CAGrad) finds a direction which maximizes the worst local improvement among all tasks, and then applies this gradient update to the shared parameters.

summaries: This work presents a first-order multi-task learning method which uses, in place of the average task gradient, a vector which maximises the minimum local single-task improvement in a neighborhood of the average task gradient.  In multi-task learning, the goal is to find the optimal point achieving low losses across all tasks. However, gradient descent methods commonly used for single tasks in computer vision and reinforcement learning achieve suboptimal performance when applied on the average task loss due to conflicting gradients (single tasks gradients pointing in opposite directions). Recently, several methods have proposed modification to the update direction that alleviate this problem. However, these methods are either heuristics or can converge to any point in the pareto optimal set of the task losses (e.g. MGDA [5]), which might not be a minimum of the average task loss.   The proposed method, called Conflict-Averse Gradient descent (CAGrad), introduces an additional hyperparameter $c \geq 0$ which allows to interpolate between gradient descent ($c = 0$) and MGDA ($c \to \infty$). Furthermore, CAGrad provably converges to a stationary point of the average task loss with the same rate as gradient descent (when $c < 1$) or to a pareto stationary point (when $c \geq 0$). At each step of the method, the update direction is obtained by solving  an optimization problem in $K$ dimensions, where $K$ is the number of tasks, but has an overhead comparable to state of the art multi-task learning methods of the same kind (see last plot of Figure 5). Experiments on several computer vision and reinforcement learning multi-task benchmarks, where the method is used to optimize deep network parameters, show that the method achieves state of the art performance in all the settings.

summaries: The paper proposes a new solution to multi-task learning  that should improve in cases with gradients conflict of a task loss and the average loss gradient. The solution optimizes for an update vector that maximizes the worst task loss relative improvement. Experiments show the superiority of the proposed solution over similar prior methods.

abstractText: The goal of multi-task learning is to enable more efficient learning than single task learning by sharing model structures for a diverse set of tasks. A standard multi-task learning objective is to minimize the average loss across all tasks. While straightforward, using this objective often results in much worse final performance for each task than learning them independently. A major challenge in optimizing a multi-task model is the conflicting gradients, where gradients of different task objectives are not well aligned so that following the average gradient direction can be detrimental to specific tasks’ performance. Previous work has proposed several heuristics to manipulate the task gradients for mitigating this problem. But most of them lack convergence guarantee and/or could converge to any Pareto-stationary point. In this paper, we introduce Conflict-Averse Gradient descent (CAGrad) which minimizes the average loss function, while leveraging the worst local improvement of individual tasks to regularize the algorithm trajectory. CAGrad balances the objectives automatically and still provably converges to a minimum over the average loss. It includes the regular gradient descent (GD) and the multiple gradient descent algorithm (MGDA) in the multi-objective optimization (MOO) literature as special cases. On a series of challenging multi-task supervised learning and reinforcement learning tasks, CAGrad achieves improved performance over prior state-of-the-art multi-objective gradient manipulation methods. Code is available at https://github.com/Cranial-XIX/CAGrad.

authors:
Vijay Badrinarayanan
Alex Kendall
Roberto Cipolla

title: Segnet: A deep convolutional encoder-decoder architecture for image segmentation

venue: IEEE transactions on pattern analysis and machine intelligence,

title: Multitask learning

authors:
Zhao Chen
Vijay Badrinarayanan
Chen-Yu Lee
Andrew Rabinovich

title: Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks

venue: In International Conference on Machine Learning,

authors:
Zhao Chen
Jiquan Ngiam
Yanping Huang
Thang Luong
Henrik Kretzschmar
Yuning Chai
Dragomir Anguelov

title: Just pick a sign: Optimizing deep multitask models with gradient sign dropout

title: The mnist database of handwritten digit images for machine learning research

venue: IEEE Signal Processing Magazine,

title: Multiple-gradient descent algorithm (mgda) for multiobjective optimization

venue: Comptes Rendus Mathematique,

authors:
Steven Diamond
Stephen Boyd

title: CVXPY: A Python-embedded modeling language for convex optimization

venue: Journal of Machine Learning Research,

authors:
Dibya Ghosh
Avi Singh
Aravind Rajeswaran
Vikash Kumar
Sergey Levine

title: Divide-andconquer reinforcement learning

venue: arXiv preprint arXiv:1711.09874,

authors:
Michelle Guo
Albert Haque
De-An Huang
Serena Yeung
Li Fei-Fei

title: Dynamic task prioritization for multitask learning

venue: In Proceedings of the European Conference on Computer Vision (ECCV),

authors:
Tuomas Haarnoja
Aurick Zhou
Pieter Abbeel
Sergey Levine

title: Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor

authors:
Kazuma Hashimoto
Caiming Xiong
Yoshimasa Tsuruoka
Richard Socher

title: A joint manytask model: Growing a neural network for multiple nlp tasks

venue: arXiv preprint arXiv:1611.01587,

authors:
Geoffrey Hinton
Oriol Vinyals
Jeff Dean

title: Distilling the knowledge in a neural network

venue: arXiv preprint arXiv:1503.02531,

authors:
Harold M Hochman
James D Rodgers

title: Pareto optimal redistribution

venue: The American economic review,

authors:
Adrián Javaloy
Isabel Valera

title: Rotograd: Dynamic gradient homogenization for multi-task learning

venue: arXiv preprint arXiv:2103.02631,

authors:
Alex Kendall
Yarin Gal
Roberto Cipolla

title: Multi-task learning using uncertainty to weigh losses for scene geometry and semantics

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

authors:
Yann LeCun
Léon Bottou
Yoshua Bengio
Patrick Haffner

title: Gradient-based learning applied to document recognition

venue: Proceedings of the IEEE,

authors:
Liyang Liu
Yi Li
Zhanghui Kuang
Jing-Hao Xue
Yimin Chen
Wenming Yang
Qingmin Liao
Wayne Zhang

title: Towards impartial multi-task learning

venue: In International Conference on Learning Representations,

authors:
Shikun Liu
Edward Johns
Andrew J Davison

title: End-to-end multi-task learning with attention

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Xingchao Liu
Xing Han
Na Zhang
Qiang Liu

title: Certified monotonic neural networks

venue: arXiv preprint arXiv:2011.10219,

authors:
Debabrata Mahapatra
Vaibhav Rajan

title: Multi-task learning with user preferences: Gradient descent with controlled ascent in pareto optimization

authors:
Kevis-Kokitsi Maninis
Ilija Radosavovic
Iasonas Kokkinos

title: Attentive single-tasking of multiple tasks

authors:
Ishan Misra
Abhinav Shrivastava
Abhinav Gupta
Martial Hebert

title: Cross-stitch networks for multi-task learning

authors:
Emilio Parisotto
Jimmy Lei Ba
Ruslan Salakhutdinov

title: Actor-mimic: Deep multitask and transfer reinforcement learning

venue: arXiv preprint arXiv:1511.06342,

authors:
Clemens Rosenbaum
Tim Klinger
Matthew Riemer

title: Routing networks: Adaptive selection of non-linear functions for multi-task learning

venue: arXiv preprint arXiv:1711.01239,

title: An overview of multi-task learning in deep neural networks

venue: arXiv preprint arXiv:1706.05098,

authors:
Andrei A Rusu
Sergio Gomez Colmenarejo
Caglar Gulcehre
Guillaume Desjardins
James Kirkpatrick
Razvan Pascanu
Volodymyr Mnih
Koray Kavukcuoglu
Raia Hadsell

title: Policy distillation

venue: arXiv preprint arXiv:1511.06295,

authors:
Ozan Sener
Vladlen Koltun

title: Multi-task learning as multi-objective optimization

venue: arXiv preprint arXiv:1810.04650,

authors:
Baifeng Shi
Judy Hoffman
Kate Saenko
Trevor Darrell
Huijuan Xu

title: Auxiliary task reweighting for minimum-data learning

venue: Advances in Neural Information Processing Systems,

authors:
Shagun Sodhani
Amy Zhang

title: Mtrl - multi task rl algorithms

venue: Github,

authors:
Shagun Sodhani
Amy Zhang
Joelle Pineau

title: Multi-task reinforcement learning with context-based representations

venue: arXiv preprint arXiv:2102.06177,

title: Inadmissibility of the usual estimator for the mean of a multivariate normal distribution

authors:
Kevin Swersky
Jasper Snoek
Ryan Prescott Adams

title: Multi-task bayesian optimization

authors:
Antti Tarvainen
Harri Valpola

title: Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
Yee Whye Teh
Victor Bapst
Wojciech Marian Czarnecki
John Quan
James Kirkpatrick
Raia Hadsell
Nicolas Heess
Razvan Pascanu

title: Distral: Robust multitask reinforcement learning

venue: arXiv preprint arXiv:1707.04175,

authors:
Simon Vandenhende
Stamatios Georgoulis
Wouter Van Gansbeke
Marc Proesmans
Dengxin Dai
Luc Van Gool

title: Multi-task learning for dense prediction tasks: A survey

venue: IEEE Transactions on Pattern Analysis and Machine Intelligence,

authors:
Han Xiao
Kashif Rasul
Roland Vollgraf

title: Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms

venue: arXiv preprint arXiv:1708.07747,

authors:
Ruihan Yang
Huazhe Xu
Yi Wu
Xiaolong Wang

title: Multi-task reinforcement learning with soft modularization

venue: arXiv preprint arXiv:2003.13661,

authors:
Tianhe Yu
Saurabh Kumar
Abhishek Gupta
Sergey Levine
Karol Hausman
Chelsea Finn

title: Gradient surgery for multi-task learning

venue: arXiv preprint arXiv:2001.06782,

authors:
Tianhe Yu
Deirdre Quillen
Zhanpeng He
Ryan Julian
Karol Hausman
Chelsea Finn
Sergey Levine

title: Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning

venue: In Conference on Robot Learning,

authors:
Amir R Zamir
Alexander Sax
William Shen
Leonidas J Guibas
Jitendra Malik
Silvio Savarese

title: Taskonomy: Disentangling task transfer learning

heading: 1 Introduction

text:
Multi-task learning (MTL) refers to learning a single model that can tackle multiple different tasks [11, 28, 44, 38]. By sharing parameters across tasks, MTL methods learn more efficiently with an overall smaller model size compared to learning with separate models [38, 40, 25]. Moreover, it has been shown that MTL could in principle improve the quality of the learned representation and therefore benefit individual tasks [35, 43, 34]. For example, an early MTL result by [2] demonstrated that training a neural network to recognize doors could be improved by simultaneously training it to recognize doorknobs.
However, learning multiple tasks simultaneously can be a challenging optimization problem because it involves multiple objectives [38]. The most popular MTL objective in practice is the average loss over all tasks. Even when this average loss is exactly the true objective (as opposed to only caring about a single task as in the door/doorknob example), directly optimizing the average loss could lead to undesirable performance, e.g. the optimizer struggles to make progress so the learning performance significantly deteriorates. A known cause of this phenomenon is the conflicting gradients [41]: gradients from different tasks 1) may have varying scales with the largest gradient dominating the update, and 2) may point in different directions so that directly optimizing the average loss can be quite detrimental to a specific task’s performance.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
To address this problem, previous work either adaptively re-weights the objectives of each task based on heuristics [3, 15] or seeks a better update vector [30, 41] by manipulating the task gradients. However, existing work often lacks convergence guarantees or only provably converges to any point on the Pareto set of the objectives. This means the final convergence point of these methods may largely depend on the initial model parameters. As a result, it is possible that these methods over-optimize one objective while overlooking the others (See Fig. 1).
Motivated by the limitation of current methods, we introduce Conflict-Averse Gradient descent (CAGrad), which reduces the conflict among gradients and still provably converges to a minimum of the average loss. The idea of CAGrad is simple: it looks for an update vector that maximizes the worst local improvement of any objective in a neighborhood of the average gradient. In this way, CAGrad automatically balances different objectives and smoothly converges to an optimal point of the average loss. Specifically, we show that vanilla gradient descent (GD) and the multiple gradient descent algorithm (MGDA) are special cases of CAGrad (See Sec. 3.1). We demonstrate that CAGrad can improve over prior state-of-the-art gradient manipulation methods on a series of challenging multi-task supervised, semi-supervised, and reinforcement learning problems.

heading: 2 Background

text: In this section, we first introduce the problem setup of multi-task learning (MTL). Then we analyze the optimization challenge of MTL and discuss the limitation of prior gradient manipulation methods.

heading: 2.1 Multi-task Learning and its Challenge

text:
In multi-task learning (MTL), we are given K ≥ 2 different tasks, each of which is associated with a loss function Li(θ) for a shared set of parameters θ. The goal is to find an optimal θ ∈ Rm that achieves low losses across all tasks. In practice, a standard objective for MTL is minimizing the average loss over all tasks:
θ∗ = argmin θ∈Rm
{ L0(θ) ≜ 1
K K∑ i=1 Li(θ)
} . (1)
Unfortunately, directly optimizing (1) using gradient descent may significantly compromise the optimization of individual losses in practice. A major source of this phenomenon is known as the conflicting gradients [41].
Optimization Challenge: Conflicting Gradients Denote by gi = ∇Li(θ) the gradient of task i, and g0 = ∇L0(θ) = 1K ∑K i gi the averaged gradient. With learning rate α ∈ R+, θ ← θ − αg0 is the steepest descent update that appears to be the most natural update to follow when optimizing (1). However, g0 may conflict with individual gradients, i.e. ∃ i, ⟨gi, g0⟩ < 0. When this conflict is large, following g0 will decrease the performance on task i. As observed by [41] and illustrated in Fig. 1, when θ is near a steep “valley", where a specific task’s gradient dominates the update, manipulating the direction and magnitude of g0 often leads to better optimization.

heading: 2.2 Prior Attempts and Convergence Issues

text: Several methods have been proposed to manipulate the task gradients to form a new update vector and have shown improved performance on MTL. Sener et al. apply the multiple-gradient descent algorithm (MGDA) [6] for MTL, which directly optimizes towards the Pareto set [30]. Chen et al. dynamically re-weight each Li using a pre-defined heuristic [3]. More recently, PCGrad identifies conflicting gradients as the motivation behind manipulating the gradients and projects each task gradient to the normal plane of others to reduce the conflict [41]. While all these methods have shown success at improving the learning performance of MTL, they manipulate the gradient without respecting the original objective (1). Therefore, these methods could in principle converge to any point in the Pareto set (See Fig. 1 and Sec. 3.2). We provide the detailed algorithms of MGDA and PCGrad in Appendix A.1 and A.2, and a visualization of the update vector by each method in Fig. 2.

heading: 3 Method

text: We introduce our main algorithm, Conflict-Averse Gradient descent in Sec. 3.1, and then show theoretical analysis in Sec. 3.2.

heading: 3.1 Conflict-Averse Gradient Descent

text:
Assume we update θ by θ′ ← θ − αd, where α is a step size and d an update vector. We want to choose d to decrease not only the average loss L0, but also every individual loss. To do so, we consider the minimum decrease rate across the losses,
R(θ, d) = max i∈[K]
{ 1
α (Li(θ − αd)− Li(θ))
} ≈ − min
i∈[K] ⟨gi, d⟩, (2)
where we use the first-order Taylor approximation assuming α is small. If R(θ, d) < 0, it means that all losses are decreased with the update given a sufficiently small α. Therefore, R(θ, d) can be regarded as a measurement of conflict among objectives.
With the above measurement, our algorithm finds an update vector that minimizes such conflict to mitigate the optimization challenge while still converging to an optimum of the main objective L0(θ). To this end, we introduce Conflict-Averse Gradient descent (CAGrad), which on each optimization step determines the update d by solving the following optimization problem:
max d∈Rm min i∈[K] ⟨gi, d⟩ s.t. ∥d− g0∥ ≤ c ∥g0∥ , (3)
Here, c ∈ [0, 1) is a pre-specified hyper-parameter that controls the convergence rate (See Sec. 3.2). The optimization problem (3) looks for the best update vector within a local ball centered at the averaged gradient g0, which also minimizes the conflict in losses measured by (2). Since we focus on MTL and choose the average loss as the main objective, g0 is the average gradient. However, CAGrad also applies when g0 is the gradient of some other user-specified objective. We leave exploring this possibility as a future direction.
Dual Objective The optimization problem (3) involves decision variable d that has the same dimension as the number of parameters in θ, which could be millions for a deep neural network. It is not practical to directly solve for d on every optimization step. However, the dual problem of Eq. (3), as we will derive in the following, only involves solving for a decision variable w ∈ RK , which can be efficiently found using standard optimization libraries [7]. Specifically, first note that mini⟨gi, d⟩ = minw∈W⟨ ∑ i wigi, d⟩, where w = (w1, . . . , wK) ∈ RK and W denotes the
probability simplex, i.e. W = {w : ∑ i wi = 1 and wi ≥ 0}. Denote gw = ∑
i wigi and ϕ = c2 ∥g0∥2. The Lagrangian of the objective in Eq. (3) is
max d∈Rm min λ≥0,w∈W
g⊤wd− λ(∥g0 − d∥ 2 − ϕ)/2.
Since the objective for d is concave with linear constraints, by switching the min and max, we reach the dual form without changing the solution by Slater’s condition:
min λ≥0,w∈W max d∈Rm
g⊤wd− λ ∥g0 − d∥ 2 /2 + λϕ/2.
Algorithm 1 Conflict-averse Gradient Descent (CAGrad) for Multi-task Learning Input: Initial model parameter vector θ0, differentiable loss functions {Li}Ki=1, a constant c ∈ [0, 1) and learning rate α ∈ R+. repeat
At the t-th optimization step, define g0 = 1K ∑K i=1∇Li(θt−1) and ϕ = c2 ∥g0∥ 2. Solve
min w∈W
F (w) := g⊤wg0 + √ ϕ ∥gw∥ , where gw = 1
K K∑ i=1 wi∇Li(θt−1).
Update θt = θt−1 − α ( g0 + ϕ1/2
∥gw∥gw
) .
until convergence
We end up with the following optimization problem w.r.t. w after several steps of calculus,
w∗ = argmin w∈W
g⊤wg0 + √ ϕ ∥gw∥ ,
where the optimal λ∗ = ∥gw∗∥ /ϕ1/2 and the optimal update d∗ = g0 + gw∗/λ∗. The detailed derivation is provided in Appendix A.3 and the entire CAGrad algorithm is summarized in Alg. 1. The dimension of w equals to the number of objectives K, which usually ranges from 2 to tens and is much smaller than the number of parameters in a neural network. Therefore, in practice, we solve the dual objective to perform the update of CAGrad.
Remark In Alg. 1, when c = 0, CAGrad recovers the typical gradient descent with d = g0. On the other hand, when c→∞, then minimizing F (w) is equivalent to minw ∥gw∥. This coincides with the multiple gradient descent algorithm (MGDA) [6], which uses the minimum norm vector in the convex hull of the individual gradients as the update direction (see Fig. 2; second column). MGDA is a gradient-based multi-objective optimization designed to converge to an arbitrary point on the Pareto set, that is, it leaves all the points on the Pareto set as fixed points (and hence can not control which specific point it will converge to). It is different from our method which targets to minimize L0 while using gradient conflict to regularize the optimization trajectory. As we will analyze in the following section, to guarantee that CAGrad converges to an optimum of L0(θ), we have to ensure 0 ≤ c < 1.

heading: 3.2 Convergence Analysis

text:
In this section we first formally introduce the related Pareto concepts and then analyze CAGrad’s convergence property. Particularly, in Alg. 1, when c < 1, CAGrad is guaranteed to converge to a minimum point of the average loss L0.
Pareto Concepts Unlike single task learning where any two parameter vectors θ1 and θ2 can be ordered in the sense that either L(θ1) ≤ L(θ2) or L(θ1) ≥ L(θ2) holds, MTL could have two parameter vectors where one performs better for task i and the other performs better for task j ̸= i. To this end, we need the notion of Pareto-optimality [13].
Definition 3.1 (Pareto optimal and stationary points). Let L(θ) = {Li(θ) : i ∈ [K]} be a set of differentiable loss functions from Rm to R. For two points θ, θ′ ∈ Rm, we say that θ is Pareto dominated by θ′, denoted by L(θ′) ≺ L(θ), if Li(θ′) ≤ Li(θ) for all i ∈ [K] and L(θ′) ̸= L(θ). A point θ ∈ Rm is said to be Pareto-optimal if there exists no θ′ ∈ Rm such that L(θ′) ≺ L(θ). The set of all Pareto-optimal points is called the Pareto set. A point θ is called Pareto-stationary if we have minw∈W ∥gw(θ)∥ = 0, where gw(θ) = ∑K i=1 wi∇Li(θ), andW is the probability simplex on [K].
Similar to the case of single-objective differentiable optimization, a local Pareto optimal point θ must be Pareto stationary (see e.g., [6]).
Theorem 3.2 (Convergence of CAGrad). Assume the individual loss functions L0, L1, . . . , LK are differentiable on Rm and their gradients ∇Li(θ) are all H-Lipschitz, i.e. ∥∇Li(x)−∇Li(y)∥ ≤ H ∥x− y∥ for i = 0, 1, . . . ,K where 0 ≤ H ≤ ∞. Assume L∗0 = infθ∈Rm L0(θ) > −∞. With a fixed step size α satisfying 0 < α ≤ 1/H , we have for the CAGrad in Alg. 1:
1) For any c ≥ 1, all the fixed points of CAGrad are Pareto-stationary points of (L0, L1, . . . , LK). 2) In particular, if we take 0 ≤ c < 1, then CAGrad satisfies
T∑ t=0 ∥∇L0(θt)∥2 ≤ 2(L0(θ0)− L∗0) α(1− c2) .
This means that the algorithm converges to a stationary point of∇L0 if we take 0 ≤ c < 1. The proof is in Appendix A.3. As we discuss earlier, unlike our method, MGDA is designed to converge to an arbitrary point on the Pareto set, without explicit control of which point it will converges to. Another algorithm with similar property is PCGrad [41], which is a gradient-based algorithm that mitigates the conflicting gradients problem by removing the conflicting components of each gradient with respect to the other gradients before averaging them to form the final update; see Fig. 2, third column for the illustration. Similar to MGDA, as shown in [41], PCGrad also converges to an arbitrary Pareto point without explicit control of which point it will arrive at.

heading: 3.3 Practical Speedup

text:
A typical drawback of methods that manipulate gradients is the computation overhead. For computing the optimal update vector, a method usually requires K back-propagations to find all individual gradients gi, in addition to the time required for optimization. This can be prohibitive for the scenario with many tasks. To this end, we propose to only sample a subset of tasks S ⊆ [K], compute their corresponding gradients {gi | i ∈ S} and the averaged gradient g0. Then we optimize d in:
max d∈Rm min
( ⟨ Kg0 − ∑ i∈S gi
K − |S| , d⟩, min i∈S ⟨gi, d⟩
) s.t. ∥d− g0∥ ≤ c ∥g0∥ (4)
Remark Note that the convergence guarantee in Thm. 3.2 still holds for Eq. 4 as the constraint does not change (See Appendix A.3). The time complexity is O((|S|N + T ), where N denotes the time for one pass of back-propagation and T denotes the optimization time. For few-task learning (K < 10), usually T ≪ N . When S = [K], we recover the full CAGrad algorithm.

heading: 4 Related Work

text:
Multi-task Learning Due to its benefit with regards to data and computational efficiency, multi-task learning (MTL) has broad applications in vision, language, and robotics [11, 28, 22, 44, 38]. A number of MTL-friendly architectures have been proposed using task-specific modules [25, 11], attentionbased mechanisms [21] or activating different paths along the deep networks to tackle MTL [27, 40]. Apart from designing new architectures, another branch of methods focus on decomposing a large problem into smaller local problems that could be quickly learned by smaller models [29, 26, 37, 8]. Then a unified policy is learned from the smaller models using knowledge distillation [12].
MTL Optimization In this work, we focus on the optimization challenge of MTL [38]. Gradient manipulation methods are designed specifically to balance the learning of each task. The simplest form of gradient manipulation is to re-weight the task losses based on specific criteria, e.g., uncertainty [15], gradient norm [3], or difficulty [9]. These methods are mostly heuristics and their performance can be unstable. Recently, two methods [30, 41] that manipulate the gradients to find a better local update vector have become popular. Sener et al [30] view MTL as a multi-objective optimization problem, and use multiple gradient descent algorithm for optimization. PCGrad [41] identifies a major optimization challenge for MTL, the conflicting gradients, and proposes to project each task gradient to the normal plane of other task gradients before combining them together to form the final update vector. Though yielding good empirical performance, both methods can only guarantee convergence to a Pareto-stationary point, but not knowing where it exactly converges to. More recently, GradDrop [4] randomly drops out task gradients based on how much they conflict. IMTLG [20] seeks an update vector that has equal projections on each task gradient. RotoGrad [14] separately scales and rotates task gradients to mitigate optimization conflict.
Our method, CAGrad, also manipulates the gradient to find a better optimization trajectory. Like other MTL optimization techniques, CAGrad is model-agnostic. However, unlike prior methods, CAGrad converges to the optimal point in theory and achieves better empirical performance on both toy multi-objective optimization tasks and real-world applications.

heading: 5 Experiment

text:
We conduct experiments to answer the following questions:
Question (1) Do CAGrad, MGDA and PCGrad behave consistently with their theoretical properties in practice? (yes)
Question (2) Does CAGrad recover GD and MGDA when varying the constant c? (yes)
Question (3) How does CAGrad perform in both performance and computational efficiency compared to prior state-of-the-art methods, on challenging multi-task learning problems under the supervised, semi-supervised and reinforcement learning settings? (CAGrad improves over prior state-of-the-art methods under all settings)

heading: 5.1 Convergence and Ablation over c

text:
To answer questions (1) and (2), we create a toy optimization example to evaluate the convergence of CAGrad compared to MGDA and PCGrad. On the same toy example, we ablate over the constant c and show that CAGrad recovers GD and MGDA with proper c values. Next, to test CAGrad on more complicated neural models, we perform the same set of experiments on the Multi-Fashion+MNIST benchmark [19] with a shrinked LeNet architecture [18] (in which each layer has a reduced number of neurons compared to the original LeNet). Please refer to Appendix B for more details.
For the toy optimization example, we modify the toy example used by Yu et al. [41] and consider θ = (θ1, θ2) ∈ R2 with the following individual loss functions:
L1(θ) = c1(θ)f1(θ) + c2(θ)g1(θ) and L2(θ) = c1(θ)f2(θ) + c2(θ)g2(θ), where f1(θ) = log ( max(|0.5(−θ1 − 7)− tanh (−θ2)|, 0.000005) ) + 6,
f2(θ) = log ( max(|0.5(−θ1 + 3)− tanh (−θ2) + 2|, 0.000005) ) + 6,
g1(θ) = ( (−θ1 + 7)2 + 0.1 ∗ (−θ2 − 8)2 ) /10− 20,
g2(θ) = ( (−θ1 − 7)2 + 0.1 ∗ (−θ2 − 8)2) / 10− 20,
c1(θ) = max(tanh (0.5 ∗ θ2), 0) and c2(θ) = max(tanh (−0.5 ∗ θ2), 0).
The average loss L0 and individual losses L1 and L2 are shown in Fig. 1. We then pick 5 initial parameter vectors θinit ∈ {(−8.5, 7.5), (−8.5, 5), (0, 0), (9, 9), (10,−8)} and plot the corresponding optimization trajectories with different methods in Fig. 3. As shown in Fig. 3, GD gets stuck in 2 out of the 5 runs while other methods all converge to the Pareto set. MGDA and PCGrad converge to different Pareto-stationary points depending on θinit. CAGrad with c = 0 recovers GD and CAGrad with c = 10 approximates MGDA well (in theory it requires c→∞ to exactly recover MGDA).
Next, we apply the same set of experiments on the multi-task classification benchmark MultiFashion+MNIST [19]. This benchmark consists of images that are generated by overlaying an image from FashionMNIST dataset [39] on top of another image from MNIST dataset [5]. The two images are positioned on the top-left and bottom-right separately. We consider a shrinked LeNet as our model, and train it with Adam [16] optimizer with a 0.001 learning rate for 50 epochs using a batch size of 256. Due to the highly non-convex nature of the neural network, we are not able to visualize the entire Pareto set. But we provide the final training losses of different methods over three independent runs in Fig. 4. As shown, CAGrad achieves the lowest average loss with c = 0.2. In addition, PCGrad and MGDA focus on optimizing task 1 and task 2 separately. Lastly, CAGrad with c = 0 and c = 10 roughly recovers the final performance of GD and MGDA. By increasing c, the model performance shifts from more GD-like to more MGDA-like, though due to the non-convex nature of neural networks, CAGrad with 0 ≤ c < 1 does not necessarily converge to the exact same point.

heading: 5.2 Multi-task Supervised Learning

text:
To answer question (3) in the supervised learning setting, we follow the experiment setup from Yu et al. [41] and consider the NYU-v2 and CityScapes vision datasets. NYU-v2 contains 3 tasks: 13- class semantic segmentation, depth estimation, and surface normal prediction. CityScapes similarly contains 2 tasks: 7-class semantic segmentation and depth estimation. Here, we follow [41] and combine CAGrad with a state-of-the-art MTL method MTAN [21], which applies attention mechanism on top of the SegNet architecture [1]. We compare CAGrad with PCGrad, vanilla MTAN and CrossStitch [25], which is another MTL method that modifies the network architecture. MTAN originally experiments with equal loss weighting and two other dynamic loss weighting heuristics [15, 3]. For a fair comparison, all methods are applied under the equal weighting scheme and we use the same training setup from [3]. We search c ∈ {0.1, 0.2, . . . 0.9} with the best average training loss for CAGrad on both datasets (0.4 for NYU-v2 and 0.2 for Cityscapes). We perform a two-tailed, Student’s t-test under equal sample sizes, unequal variance setup and mark the results that are significant with an ∗. Following Maninis et al.[24], we also compute the average per-task performance drop of method m with respect to the single-tasking baseline b: ∆m = 1K ∑K i=1(−1)li(Mm,i −Mb,i)/Mb,i where li = 1 if a higher value is better for a criterion Mi on task i and 0 otherwise. The single-tasking baseline (independent) refers to training individual tasks with a vanilla SegNet. Results are shown in Tab. 1 and Tab. 2.
Given the single task performance, CAGrad performs better on the task that is overlooked by other methods (Surface Normal in NYU-v2 and Depth in CityScapes) and matches other methods’
performance on the rest of the tasks. We also provide the final test losses and the per-epoch training time of each method in Fig. 5 in Appendix B.2.

heading: 5.3 Multi-task Reinforcement Learning

text: To answer question (3) in the reinforcement learning (RL) setting, we apply CAGrad on the MT10 and MT50 benchmarks from the Meta-World environment [42]. In particular, MT10 and MT50 contains 10 and 50 robot manipulation tasks. Following [33], we use Soft Actor-Critic (SAC) [10] as the underlying RL training algorithm. We compare against Multi-task SAC (SAC with a shared model), Multi-headed SAC (SAC with a shared backbone and task-specific head), Multi-task SAC + Task Encoder (SAC with a shared model and the input includes a task embedding) [42] and PCGrad [41]. We also compare with Soft Modularization [40] that routes different modules in a shared model to form different policies. Lastly, we also include a recent method (CARE) that considers language metadata and uses a mixture of expert encoder for MTL. We follow the same experiment setup from [33]. The results are shown in Tab. 3. CAGrad outperforms all baselines except for CARE which benefits from extra information from the metadata. We also apply the practical speedup in Sec. 3.3 and sub-sample 4 and 8 tasks for MT10 and MT50 (CAGrad-Fast). CAGrad-fast achieves comparable performance against the state-of-the-art method while achieving a 2x (MT10) and 5x (MT50) speedup over PCGrad. We provide a visualization of tasks from MT10 and MT50, and the comparison of computational efficiency in Appendix B.3.

heading: 5.4 Semi-supervised Learning with Auxiliary Tasks

text:
Training with auxiliary tasks to improve the performance of a main task is another popular application of MTL. Here, we take semi-supervised learning as an instance. We combine different optimization
algorithms with Auxiliary Task Reweighting for Minimum-data Learning (ARML) [31], a state-ofthe-art semi-supervised learning algorithm. The loss function is composed of the main task and two auxiliary tasks:
L0 = LCE(θ;Dl) + w1L 1 aux(θ;Du) + w2L 2 aux(θ;Du), (5)
where LCE is the main cross-entropy classification loss on the labeled dataset Dl, and L1aux, L 2 aux are auxiliary unsupervised learning losses on the unlabeled dataset Du. We use the same w1 and w2 from ARML, and use the CIFAR10 dataset [17], which contains 50,000 training images and 10,000 test images. 10% of the training images is held out as the validation set. We test PCGrad, MGDA and CAGrad with 500, 1000 and 2000 labeled images. The rest of the training set is used for auxiliary tasks. For all the methods, we use the same labeled dataset, the same learning rate and train them for 200 epochs with the Adam [16] optimizer. Please refer to Appendix B.4 for more experimental details. Results are shown in Tab. 4. With all the different number of labels, CAGrad yields the best averaged test accuracy. We observed that MGDA performs much worse than the ARML baseline, because it significantly overlooks the main classification task. We also compare different gradient manipulation methods on the same task with GradNorm [3], which dynamically adjusts w1 and w2 during training. The results and conclusions are similar to those for ARML.

heading: 6 Conclusion

text: In this work, we introduce the Conflict-Averse Gradient descent (CAGrad) algorithm that explicitly optimizes the minimum decrease rate of any specific task’s loss while still provably converging to the optimum of the average loss. CAGrad generalizes the gradient descent and multiple gradient descent algorithm, and demonstrates improved performance across several challenging multi-task learning problems compared to the state-of-the-art methods. While we focus mainly on optimizing the average loss, an interesting future direction is to look at main objectives other than the average loss under the multi-task setting.

heading: Acknowledgements

text: The research was conducted in the statistical learning and AI group (SLAI) and the Learning Agents Research Group (LARG) in computer science at UT Austin. SLAI research is supported in part by CAREER-1846421, SenSE-2037267, EAGER-2041327, and Office of Navy Research, and NSF AI Institute for Foundations of Machine Learning (IFML). LARG research is supported in part by NSF (CPS-1739964, IIS-1724157, FAIN-2019844), ONR (N00014-18-2243), ARO (W911NF-19-2-0333), DARPA, Lockheed Martin, GM, Bosch, and UT Austin’s Good Systems grand challenge. Peter Stone serves as the Executive Director of Sony AI America and receives financial compensation for this work. The terms of this arrangement have been reviewed and approved by the University of Texas at Austin in accordance with its policy on objectivity in research. Xingchao Liu is supported in part by a funding from BP.

title: Conflict-Averse Gradient Descent for Multi-task Learning