Improving Multi-Task Generalization via Regularizing Spurious Correlation
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Multi-Task Learning (MTL) is a powerful learning paradigm to improve generalization performance via knowledge sharing.
However, existing studies ﬁnd that MTLcould sometimes hurt generalization, especially when two tasks are less correlated.
One possible reason that hurts generalization is spurious correlation, i.e., someknowledge is spurious and not causally related to task labels, but the model couldmistakenly utilize them and thus fail when such correlation changes.
In MTL setup,there exist several unique challenges of spurious correlation.
First, the risk ofhaving non-causal knowledge is higher, as the shared MTL model needs to encodeall knowledge from different tasks, and causal knowledge for one task could bepotentially spurious to the other.
Second, the confounder between task labelsbrings in a different type of spurious correlation to MTL.
Given such label-labelconfounders, we theoretically and empirically show that MTL is prone to takingnon-causal knowledge from other tasks.
To solve this problem, we propose MultiTask Causal Representation Learning (MT-CRL) framework.
MT-CRL aims torepresent multi-task knowledge via disentangled neural modules, and learn robusttask-to-module routing graph weights via MTL-speciﬁc invariant regularization.
Experiments show that MT-CRL could enhance MTL model’s performance by5.5 % on average over Multi-MNIST, MovieLens, Taskonomy, CityScape, and
NYUv2, and show it could indeed alleviate spurious correlation problem.
Introduction
Multi-Task Learning (MTL), a learning paradigm (Caruana, 1997; Zhang & Yang, 2018)
aiming totrain a single model for multiple tasks, is expected to beneﬁt the overall generalization performancethan single-task learning (Maurer et al., 2016; Tripuraneni et al., 2020) given the assumption thatthere exists some common knowledge to handle different tasks.
However, recent studies observedthat, when two tasks are less correlated, MTL could lead to even worse overall performance (Parisottoet al., 2016; Zhang et al., 2021).
A line of works (Yu et al., 2020;
Wang et al., 2021; Fifty et al.,2021) resort performance drop to optimization challenge because conﬂicting tasks might competefor model capacity.
However, both Standley et al.
and our analysis in Section 2.2 show that,even with an over-parameterized model that achieves low MTL training loss, the ﬁnal generalizationperformance could be worse than single-task learning.
This ﬁnding motivates us to think about thefollowing question: Are there any intrinsic problems in MTL that hurt generalization?
One widely studied issue that inﬂuences generalization is the spurious correlation problem (Geirhoset al., 2019, 2020), i.e., correlation that only existed in training datasets due to unobserved confounders (Lopez-Paz, 2016), but not causally correct.
For example, as Beery et al.
(2018) discussed,when we train an image classiﬁcation model to identify cows with a biased dataset where cows mostlyappear in pastures, the trained cow classiﬁcation model could exploit the features of background (e.g., pastures) to make prediction.
Thus, when we apply the classiﬁer to another dataset where cows alsoappear in other locations such as farms or rivers, it will fail to generalize (Nagarajan et al., 2021).
When it comes to MTL setting, there exist several unique challenges to handle spurious correlationproblem.
First, the risk of having non-causal features is higher .
Suppose each task has differentsets of causal features.
To train a single model for all these tasks, the shared representation shouldencode all required features.
Consequently, the causal features for one task could be potentiallyspurious to the other tasks, and such risk could be even higher with an increasing number of tasks.
Second, the confounder that leads to spurious correlation is different .
Instead of the standardconfounders between feature and label, the nature of MTL brings in a unique type of confoundersbetween task labels, e.g., correlation between tasks’ labels could change in different distributions.
Given such label-label confounders that are unique for MTL, we theoretically prove that MTL isprone to taking non-causal knowledge learned from other tasks.
We then conduct empirical analysisto validate the hypothesis.
In summary, we point out the unique challenges of spurious correlation in
MTL setup, and show that it indeed inﬂuences multi-task generalization.
Following the analysis, we try to solve the spurious correlation problem in MTL.
Among all theknowledge learned in the shared representation layer through end-to-end training, an ideal MTLframework should learn to leverage only the causal knowledge to solve each task by identifyingthe correct causal structure.
We propose a Multi-Task Causal Representation Learning (MT-CRL)framework, aiming to represent the multi-task knowledge via a set of disentangled neural modulesinstead of a single encoder, and learn the task-to-module causal relationship jointly.
We adopt decorrelation and sparsity regularization over popular Mixture-of-Expert (MoE) architecture (Shazeeret al., 2017).
The most critical and challenging step is to learn the causal graph in the MTLsetup, which requires distinguishing the genuine causal correlation from spurious ones for all tasks.
Motivated by the recent studies that invariance could lead to causality (Ahuja et al., 2020; Koyama& Yamaguchi, 2021), we propose to penalize the variance of gradients w.r.t.
causal graph weightsacross different distributions.
On a high level, this invariance regularization encourages the causalgraph to assign higher weights to the modules that are consistently useful.
In contrast, the modulesencoding spurious knowledge that cannot consistently achieve graph optimality are assigned lowerweights and be discarded by task predictors.
We evaluate our method on existing MTL benchmarks, including Multi-MNIST, MovieLens, Taskonomy, CityScape, and NYUv2.
For each dataset, to mimic distribution shifts, we adopt some attributeinformation given in the dataset, such as the released time of the movie or district of a building, tosplit train/valid/test datasets.
The results show that MT-CRL could consistently enhance the MTLmodel’s performance by 5.5 % on average, and outperform both the MTL optimization and robustmachine learning baselines.
We also conduct case studies to show that MT-CRL indeed alleviatespurious correlation problem in MTL setup.
The key contributions of this paper are as follows:
• We are the ﬁrst to analyze spurious correlation problem in MTL setup, and point out severalkey challenges unique to MTL with theoretical and empirical analysis.•
We propose MT-CRL with MTL-speciﬁc invariant regularizers to elleviate spurious correlation problem, and enhances the performance on several MTL benchmarks.
Analyzing Spurious Correlation in MTL
To systematically analyze the spurious correlation problem in MTL, we ﬁrst assume that data andtask labels are generated by ground-truth causal mechanisms described in Suter et al.
Wedenote X as the variable of observed data, and each data is associated with K latent generativefactors F = { F i } Ki =1 representing different semantics of the data (e.g., color, shape, background ofan image).
We assume that the data X is generated by disentangled causal mechanisms P ( X | F i ) ,such that P (cid:0) X
| F (cid:1) =
(cid:81) K i =1 P (cid:0) X
| F i (cid:1) .
As F represents high-level knowledge of the data, we could naturally deﬁne task label variable Y tfor task t as the cause of a subset of generative factors.
We denote F Ct as a subset of causal featurevariables within F that are causally related to each task variable Y t , and we could deﬁne F St = F \ F Ctas a subset of non-causal feature variables to task t , such that P ( Y t | F )
=
P ( Y t | F Ct ) .
In other words,changing the values of any non-causal factors in F St does not change the conditional distribution.
Learning is mainly caused by factor-label confounders C STLdist .
We could remove spurious factors F S from representation Z .
Learning could be caused by label-label confounders C MTLdist .
Factors for both tasks F Ca and
F Cb need to be encoded and potentially spurious.
Note that the discussion so far is based on the assumption that the ground-truth causal generativemodel is known.
In a real-world learning setting, however, we are only given a supervised dataset( X, Y ) without access to generative factors F .
To solve the task, a neural encoder Φ( · ) is required toextract representation Z from the data that encodes the information about the causal factors, on top ofwhich a task predictor f ( · ) could predict the label.
2.1 Spurious Correlation Problem
Based on the ground-truth generative model, an ideal predictor for each task should only utilizethe causal factors, and keep invariant to any intervention on non-causal factors.
However, in realworld problems, it is hard to achieve an invariant predictor due to the spurious correlation issuedue to unobserved confounders C dist (Lopez-Paz, 2016).
Formally, confounders are variables thatinﬂuence the two connected variables’ correlation, and such correlation could change under differentdistribution (different value of C dist ), thus the model exploiting such spurious correlation will fail togeneralize.
Below we summarize the differences of spurious correlation problems for single-task andmulti-task learning settings:
Single-Task Learning (STL).
As illustrated in Figure 1, the label-factor confounders for singletask learning C STLdist connects non-causal factors F ∈ F S and task label Y , bringing in spuriouscorrelation.
For example, temperature could confound crime and ice cream consumption.
Whenthe weather is hot, both crime rates and ice cream sales increase, but these two phenomena are notcausally related.
Based on the proof by Nagarajan et al.
(2021); Khani & Liang (2021), such spuriouscorrelation could lead the model to use non-causal factors, and thus hurt generalization performance.
Multi-Task Learning (MTL).
In the MTL setting, there exist several unique challenges to handlespurious correlation.
First, the risk of having non-causal features is higher.
As is illustrated in
Figure 2, the shared encoder Φ needs to encode all the factors causally related to each task in therepresentation Z .
Therefore, for each task, all non-overlapping factors from other tasks could bepotentially spurious.
Second, besides the standard label-factor confounders C STLdist for each singletask introduced above, we deﬁne label-label confounders C MTLdist connecting multiple tasks’ label{ Y } .
Such confounder is unique to MTL setting.
As an example, consider two binary classiﬁcation tasks, with Y a
and Y b as variables from {± 1 } fortask label.
The two labels’ correlation P ( Y a = Y b ) =
m C could change with different confounder
C MTLdist = C .
We assume the two tasks have non-overlapping factors F a and F b drawn from Gaussiandistribution.
We then show MTL model with both two factors as input will utilize non-causal factors:
Proposition 1 Given m C (cid:54) = 0 .
5 , the Bayes Optimal per-task classiﬁer has non-zero weights tonon-causal factor.
Given m C = 0 .
5 and limited training dataset, the trained per-task classiﬁer willassign non-zero weights to non-causal factor as noise.
Detailed proof is in Appendix A.
Therefore, in this linear classiﬁcation example, when we deploy thetrained model to a new distribution with changed label-label confounder C MTLdist , the model trained by
MTL that utilizes non-causal factors generalize relatively worse.
On the contrary, the model trainedby STL don’t need to encode all causal factors from two tasks.
Assuming there is no task-labelconfounder C STLdist in each task’s dataset, the trained model could remove non-causal factors fromrepresentation.
[Table]
2.2 Empirical Experiments
In the following, we conduct experiments to validate the claims.
As there is no existing MTLdatasets speciﬁcally designed to analyze spurious correlation problem, we construct synthetic MultiSEM (Rosenfeld et al., 2021) and Multi-MNIST (Harper & Konstan, 2016) datasets with known causalstructure to study whether the model trained by MTL indeed exploits more non-causal factors, andhow the spurious correlation inﬂuences multi-task generalization.
Dataset details are in Appendix D.1.
Spurious Score.
As we know the ground-truth causal structure for the two datasets, we couldquantify how much a model utilizes the non-causal factors.
Following the gradient saliency mapproposed by Simonyan et al.
(2014), we calculate the average absolute gradients w.r.t each factor as
Grad
[Equation]
, which measures how much a model leverage this factor tomake prediction.
We then deﬁne the spurious score ρ spur as the proportion of average gradients overnon-causal feature ρ spur =
[Equation]
[Table]
Empirical Results.
We train a shared-bottom model via Multi-task learning (MTL) and single-tasklearning (STL) over the two datasets and report both the training and test accuracy with spurious ratioρ spur in Table 1.
As illustrated, the test accuracies of MTL for both Multi-SEM and Multi-MNISTdatasets are both worse than STL.
The training accuracies of MTL are very similar to STL, meaningthat the performance drop is not due to the optimization difﬁculty that many previous works try toaddress.
The spurious ratio ρ spur of MTL is much higher than the STL, which means that it exploitsmore non-causal factors.
To give a more straightforward illustration, we plot the gradient saliencymap of the right-side digit classiﬁer for Multi-MNIST in Figure 3.
The model trained by MTLutilizes more left-side pixels, which are non-causal to the ﬁnal prediction.
We also show the resultsof Multi-SEM with more than 2 tasks in Appendix B.
These ﬁndings support our hypothesis that withspurious correlation caused by label-label confounder C MTLdist , models trained by MTL is more proneto leverage non-causal knowledge than STL, and thus inﬂuence generalization performance.
Method
Based on the previous analysis of the spurious correlation problem in MTL, we now introduce a
Multi-Task Causal Representation Learning (MT-CRL) framework with the goal that the per-taskpredictor only leverages the causal knowledge instead of spurious correlation.
The high-level ideaof the framework is to reconstruct the ground-truth causal mechanisms introduced in section 2through end-to-end representation learning.
To accomplish this goal, the framework aims to 1) modelmulti-task knowledge via a set of disentangled neural modules; 2) learn the task-to-module causalgraph that is optimal across different distributions.
We introduce the two crucial designs as follows.
3.1 Modelling via Disentangled Neural Modules
In order to alleviate spurious correlation, an ideal MTL model should learn the multi-task knowledgein the shared representation while identifying which part of the knowledge is causally related toeach task.
However, directly conducting causal discovery is impossible if all the knowledge isfused in a single shared encoder.
Thus, we seek to adopt a modularized architecture in which eachmodule encodes disentangled knowledge, and thus enable modeling causal relationship betweentask and modules.
We adopt Multi-gate Mixture-of-Experts (MMoE) (Ma et al., 2018), a variantof MoE (Shazeer et al., 2017) architecture tailored for MTL setting, as our underlying model.
Speciﬁcally, we have K different neural modules as shared encoders Φ = (cid:2)
Φ i ( · ) (cid:3) K i
=1 .
Given abatch of input data X = { x n } Bn =1 with batch size B , we extract k representations via differentneural modules, i.e., Z i =
Φ i ( X ) ∈ R B × d .
On top of the learned neural modules, we learn atask-to-module routing graph, aiming to estimate which module is causally related to each task.
Wemodel the bipartite adjacency (a.k.a. biadjacency) matrix A = sigmoid ( θ ) ∈ [0 , 1] T × K by applyingsigmoid over a learnable parameter θ to enforce the range constraint.
Note that original MMoE adoptssoftmax to get gate vector, which encourages only a small portion of modules being utilized for eachtask.
Our graph modelling allows multiple modules utilized for each task.
With the correct graphweights A as routing layer, we could utilize only the causally related modules and make predictionswith per-task predictor f t ( · ) as ˆ Y t ( X ) =
f t (cid:0) (cid:80) i
A t,i · Φ i ( X ) (cid:1) .
Disentangling Modules.
One of the main properties of the causal mechanisms we introducedin section 2 is disentanglement, such that each factor represents a different view of the data, andchanging the value of one factor does not inﬂuence the others.
If without explicit constraints, thelearned modules’ outputs could still be correlated and hinder the causal structure learning.
Therefore,we need to add regularization to disentangle these modules during training.
Most existing disentangled representation learning methods are under the generative modeling framework, e.g. VAE (Higgins et al., 2017) or GAN (Chen et al., 2016).
However, Locatello et al.
(2019)argues that without explicit supervision, it is hard for generative models to learn correct disentangledfactors.
We therefore only borrow the regularization methods utilized in existing generative disentangled representation works (Cheung et al., 2015; Cogswell et al., 2016) to directly penalize thecorrelation of learned modules.
Speciﬁcally, we regularize the in-batch Pearson correlation ρ ( Z i , Z j )between every pair of output dimensions from different representation matrices Z i and Z j , as:
[Equation]
i ,
[Equation]
j
[Equation]
) , where Cov
[Equation]
By minimizing the Frobenius norm of the correlation matrix ρ for every two different representationpairs, we could enforce the encoder Φ to extract disentangled representations.
[Equation]
Task-to-Module Graph Regularization.
Based on sparsity assumption of the causal mechanisms (Parascandolo et al., 2018; Bengio et al., 2020; Lachapelle et al., 2021), each task is causallyrelated to only a few modules.
To learn the graph structure, existing works (Zheng et al., 2018; Nget al., 2019; Lachapelle et al., 2020) propose to to ﬁt structural equation model (SEM) with sparsityregularization over the graph weights.
We adopt a similar sparse regularization with an entropybalancing term (Hainmueller, 2012) over the biadjacency matrix A weights of the task-to-modulerouting graph:
[Equation]
Note that the entropy term aims at keeping the causal weights for each module i summing over allthe tasks to be balanced.
This could help avoid degenerate solutions in which only a few modulesare utilized.
By combining the two regularizations in Eq.(2) and Eq.(3) with per-task supervised riskterm R t (cid:0) Φ , A t , f t (cid:1) = (cid:80) ( X , Y t ) ∈D
L t (cid:0) ˆ Y t ( X ) , Y t (cid:1) , we get the regularized loss as:
[Equation]
3.2 Causal Learning via Graph-Invariant Regularization
It is critical and challenging to learn the correct causal graph, which requires distinguishing the truecausal correlation from spurious ones.
Motivated by the recent studies of robust machine learningthat a predictor invariant to multiple distributions could learn causal correlation (Ahuja et al., 2020;
, 2021), we assume the true causal relationship to be optimal across different distributions.
To do so, we assume to have access to multiple slices of datasets collected from differentenvironments e ∈ E in which the confounder C MTLdist that controls task correlation might change.
For example, one natural choice is to consider train/valid dataset split (the setting we utilize inexperiment), or assume the training set is split into multiple slices with different attributes.
We desirethe task-to-module graph weights A and per-task predictor f t to be optimal across all environmentse ∈ E .
Formally, we aim to solve the following bi-level optimization problem:
[Equation]
where R et denotes the risk over data slice in environment e .
This optimization problem could beregarded as a multi-task version of IRM.
Based on Theorem 9 described in Ahuja et al.
(2020), byenforcing invariance over a sufﬁcient number of environments that exhibit distribution shifts (i.e.,changes of confounder C MTLdist ), per-task predictors should only utilize modules that are consistentlyhelpful to the task, and assign zero weights to modules that encode non-causal factors to the task, andthus alleviate spurious correlation and help out-of-distribution generalization.
Even if all data aresampled from the same distribution and there are no distribution shifts, invariance could also helpeliminate noisy correlation due to the limited training dataset and help in-distribution generalization.
Invariant Optimality of Task-to-Module Graph for MTL.
As discussed in IRM, the bi-leveledoptimization problem in Eq.(5) is highly intractable, especially with complex and non-linear Φ .
Toimplement a practical optimization objective, IRM proposes to softly regularize the gradient of thetask-predictor at different environments to enforce it to be optimal:
[Equation]
However, as is discussed in IRM paper, if the complexity of a task-predictor f is much larger thanthe number of environments, it could learn an over-ﬁtted solution that makes gradient zero but doesnot achieve invariance.
IRM adopts a ﬁxed all-one vector as predictor to reduce complexity.
Thisapproach is not applicable to MTL setup , as the optimal task-predictors f ∗ t for different task t couldbe very distinctive and complex, and we cannot use a ﬁxed uniform predictor for all tasks.
To strike a balance between invariance and complexity of multi-task predictors, we propose only toregularize the gradient of the task-to-module routing graph while assuming the complex predictorf t for each task is ﬁxed at each iteration.
We call this modiﬁcation as Graph-Invariant Risk
Minimization (G-IRM) , which is designed speciﬁcally to MTL setup:
[Equation] [Equation]
By adopting the similar gradient penalty term as adopted in IRM, we deﬁne L NormG - IRM as:
[Equation]
[Equation]
As we assume f t is ﬁxed for invariance regularization term L NormG - IRM , we only calculate gradientand optimize for Φ and A , but not updating f t .
This could avoid the over-parametrized predictor f tﬁnding a trivial solution to achieve zero gradients instead of learning the correct causal correlation.
Similar trick is utilized in (Ahmed et al., 2021).
Note that the gradient w.r.t each graph weight meanswhether a module could help reduce the risk for this task.
Therefore, by penalizing the invarianceregularization, the modules containing non-causal factors will be assigned zero weights.
In the experiments, we observe that at the early optimization stage, the model has non-zero gradientsfor all parameters, including the graph weights, thus directly regularizing the gradient norm mightinﬂuence the optimization.
Therefore, we propose a modiﬁed version of gradient regularization
L V arG - IRM that penalizes the variance of the task-to-module graph’s gradient on different environments:
[Equation] [Equation]
By minimizing L V arG - IRM , we force all the learned modules to have similar gradients across differentenvironments, and not overﬁt only to some of the environments.
It still allows some modules to have non-zero gradients as long as it’s the same across environments, and relies on loss term ˜ L toupdate these weights, while L NormG - IRM forces all gradient to be zero.
Therefore, L V arG - IRM is a looseregularization that not inﬂuences the overall optimization too much.
It shares similar intuition of
REx (Krueger et al., 2021) that penalizes risk variance, while L V arG - IRM penalize gradient variance.
We provide pseudo-code of MT-CRL framework in Appendix C.
Experiment
In this section, we evaluate whether MT-CRL could beneﬁt the performance of MTL models onexisting benchmark datasets, and study whether it could indeed alleviate spurious correlation.
Experimental Setup.
One key ingredient of our MT-CRL is to achieve the optimality of causalgraph over different distributions.
However, we might not access multiple environmental labels inmost real-world multi-task learning datasets.
Therefore, we adopt a more realistic setup, such that weonly assume to have a single validation set that contains unknown distribution shifts (i.e. change ofconfounder C MTLdist ) compared to the training dataset.
We thus could utilize training and valid setsas two environments to calculate invariance regularization, while we only utilize the training set tocalculate task loss to avoid the task predictor overﬁts.
Note that in this way, our method could getaccess to the label information in the validation set.
To avoid the possibility that the performanceimprovement is brought by additional label, for all the other baseline methods, we also add thevalidation data into the training set to calculate task loss and learn MTL model.
Dataset.
We choose ﬁve widely-used real-world MTL benchmark datasets, i.e., Multi-MNIST (Sun,2019), MovieLens (Harper & Konstan, 2016), Tasknomy (Zamir et al., 2018), NYUv2 (Silbermanet al., 2012) and CityScape (Cordts et al., 2016), and try to determine train/valid/test split such thatthere exist distribution shifts between these sets.
Dataset details are in Appendix D.2.
Note thatexcept NYUv2, our data split is the same as the default split settings of these datasets, which also tryto test model’s capacity to generalize across domains.
Baselines.
As MT-CRL is a regularization framework built upon modular MTL architecture (inthis paper we choose MMoE as instantiation, but it can be applied to other modular networks), wemainly compare with two gradient-based multi-task optimization baselines:
PCGrad (Yu et al.
,2020) and GradVac (Wang et al., 2021).
We also compare with two domain generalization baselines:
IRM (Ahuja et al., 2020) and DANN (Ganin et al., 2016).
For IRM we adopt different per-taskpredictors instead of all-one vector to adapt MTL setup, and calculate penalty via Eq.
Hyper-Parameter Selection.
For a fair comparison, all methods are based on the same MMoEarchitecture.
Our methods contain a lot of hyper-parameters, including some model speciﬁc onessuch as number of modules ( K ) and regularization speciﬁc ones.
To avoid the case that performanceimprovement is caused by extensive hyper-parameter tuning, we mainly search optimal model hyperparameter on Vanilla MTL setting, and use for all baselines.
For regularization speciﬁc parameters,we take Multi-MNIST, the simplest dataset among the testbed, to ﬁnd a optimal combination, and usefor all other datasets.
Detailed selection procedure and results are shown in Appendix H.
4.1 Experiment Results
As each task has a different evaluation metric and cannot be directly compared, we calculate therelative performance improvement of each method compared to vanilla MTL, and then averagethe relative improvement for all tasks of each dataset.
As summarized in Table 2, the averageimprovement of MT-CRL with L V arG - IRM is 5 .
5% , signiﬁcantly higher than all other baseline methods.
The most critical step of MT-CRL is to learn correct causal graph.
We therefore report MT-CRLwith different invariance regularization.
As is shown in the last block, L V ar G - IRM achieve better resultsfor most datasets than L NormG - IRM , while removing the invariance regularization could signiﬁcantlydrop the relative performance.
Compared to IRM which calculate gradient and update per-taskpredictors, MT-CRL uses disentangled modules and G-IRM to avoid overﬁtting to achieve invariance.
Results show that for datasets with large amount of tasks, e.g., Taskonomy and NYUv2, MT-CRLsigniﬁcantly outperform IRM, showing the modiﬁcation is more suitable for MTL setup.
[Table]
are shown in
[Table]
Ablation Studies.
We then study the effectiveness of the other two components in MT-CRL, i.e.,disentangled and graph regularization.
We mainly report the ablation studies on Multi-MNIST intable 3 as it’s relatively small so that we could quickly get the results of all combinations.
For disentangled regularization, after removing L decor , the performance drops from 0.915 to 0.882,which ﬁts our discussion that we cannot conduct causal learning over entangled modules.
We alsoexplore one classical generative disentangled representation method, i.e., β -VAE.
As shown in thetable, the results of using β -VAE are 0.896, lower than our utilized decorrelation regularization..
Wehypothesize that this is probably because not all generative factors are useful for downstream tasks.
Generative objectives might compete for the model capacity and in addition, the unused factors couldbe potentially spurious.
Another key component is graph regularization.
After removing both L sps and L bal , the performancedrops to 0 . 891 .
This show that even if invariance regularization could penalize non-causal modules, itwould be better to force their weights to be zero via sparsity regularization, and to be non-degeneratevia balance regularization.
We also conduct ablation studies to remove either L sps or L bal , and resultsshow both are important, and combining the two could help to achieve the best results.
Case Study.
To show that real-world MTL problem indeed have spurious correlation problem andour MT-CRL could alleciate it, we take MovieLens as an example to conduct case study.
Each taskis for different movie types, and bag-of-word of movie title is one of the features.
We calculatethe task-to-module gradients ∂
( f (Φ( x ))
[ y ]) ∂F of the vanilla MMoE model without MT-CRL.
We thenvisualize ‘train’ gradients, which shows how much each module is utilized to ﬁt the training set,and ‘valid-train’ gradients, which shows how generalizable each module is.
We ﬁnd that module 5is utilized for children movie, but harmful in valid set, indicating it is a spurious feature.
We thenuse Grad-CAM to show that top words of module 5 include strip and die , which is not relevant tochildren movies.
One possible reason is that some children movies contain the words club , whichis often co-occurred with strip and die in crime and war movies.
After adding our MT-CRL, themodule assigned to ‘children’ movie attends Pink , Parenthood , Alice and Jungle .
We also show the saliency map of MT-CRL for Multi-MNIST in Figure 8 in Appendix.
Both examples show MT-CRLcould indeed alleviate spurious correlation in real MTL problems.
Related Work
Multi-Task Generalization.
A deep neural model often requires a large number of training samplesto generalize well (Arora et al., 2019).
To alleviate the sample sparsity problem, MTL could leveragemore labeled data from multiple tasks (Zhang & Yang, 2018).
Most works studying multi-taskgeneralization are based on a core assumption that the tasks are correlated.
Earlier research directlydeﬁne the task relatedness with statistical assumption (Baxter, 2000; Ben-David & Borbely, 2008;
Lampinen & Ganguli, 2019).
With the increasing focus on deep learning models, recent researchdecompose ground-truth MTL models into a shared representation and different task-speciﬁc layersfrom a hypothesis family (Maurer et al., 2016).
With such decomposition, Tripuraneni et al.
(2020)and Du et al.
(2021) prove that a diverse set of tasks could help learn more generalizable representation.
Wu et al.
(2020) study how covariate shifts inﬂuence MTL generalization.
Despite these ﬁndings, thecore assumption of task relatedness might not be satisﬁed in many real-world applications (Parisottoet al., 2016; Zhang et al., 2021), in which tasks could even conﬂict with each other to compete modelcapacity, and the generalization performance of MTL could be worse than single-task training.
To solve the task conﬂict problem, a number of MTL model architectures have utilized modular (Misra et al., 2016;
Lu et al., 2017; Rosenbaum et al., 2018; Ma et al., 2018; Guo et al.,
2020)or attention-based (Liu et al., 2019; Maninis et al., 2019) design to enlarge model capacity whilepreserving information sharing.
Another line of research alleviate task conﬂict during optimization.
Some propose to balance the task weight via uncertainty estimation (Kendall et al., 2018), gradientnorm (Chen et al., 2018), convergence rate (Liu et al., 2019), or pareto optimality (Sener & Koltun,2018).
Others directly modulate task gradients via dropping part of the conﬂict gradient (Chen et al.
,2020) or project task’s gradient onto other tasks’ gradient surface (Yu et al., 2020; Wang et al., 2021).
Though these works successfully facilitate MTL model to converge easier, our analysis show thatwith spurious correlation, the MTL model with low training loss could still generalize bad.
Spurious Correlation Problem.
Due to the selection bias (Torralba & Efros, 2011; Gururanganet al., 2018) or unobserved confounding factors (Lopez-Paz, 2016), training datasets always containspurious correlations between non-causal features and task labels, with which trained models oftenleverage non-causal knowledge and may fail to generalize Out-Of-Distribution (OOD) when suchcorrelation changes (Nagarajan et al., 2021).
To solve the spurious correlation problem, somefairness research pre-deﬁne a set of non-causal features (e.g., gender and underrepresented identity)and then explicitly remove them from the learned representation (Zemel et al., 2013; Ganin et al.
,2016; Wang et al., 2019).
Another line of robust machine learning research does not assume toknow spurious features, but regularize the model to perform equally well under different distribution.
Distributionally Robust Optimization (DRO) optimizes worst-case risk (Sagawa et al., 2020).
Invariant
Causal Prediction (ICP) learns causal relations via invariance testing (Peters et al., 2016).
Invariant
Risk Minimization (IRM) forces the ﬁnal predictor to be optimal across different domains (Arjovskyet al., 2019).
Risk Extrapolation (REx) directly penalizes the variance of training risk in differentdomains (Krueger et al., 2021).
Another line of work aim at learning causal representation (Sch¨olkopfet al., 2021), i.e., high-level variables representing different aspect of knowledge from raw data input.
Most of these works try to recover disentangled causal generative mechanisms (Parascandolo et al.
,2018; Bengio et al., 2020; Liu et al., 2020; Mitrovic et al., 2021).
Conclusion
In this paper, we study spurious correlation problem in the Multi-Task Learning (MTL) setting.
Wetheoretically and experimentally shows that task correlation can introduce special type of spuriouscorrelation in MTL, and the model trained by MTL is more prone to leverage non-causal knowledgefrom other tasks than single-task learning.
To solve the problem, we propose Multi-Task Causal
Representation Learning (MT-CRL) which consists of: 1) a decorrelation regularizer to learn disentangled modules; 2) a graph regularizer to learn sparse and non-degenerate task-to-module graph; 3)
G-IRM invariant regularizer.
We move limitation statement to Appendix Sec I.
[-We use 
-]