A UXILIARY T ASK U PDATE D ECOMPOSITION : T HE G OOD , T HE B AD AND T HE N EUTRAL
Anonymous authors
Paper under double-blind review
A BSTRACT
While deep learning has been very beneﬁcial in data-rich settings, tasks with smaller training set often resort to pre-training or multitask learning to leverage data from other tasks.
In this case, careful consideration is needed to select tasks and model parameterizations such that updates from the auxiliary tasks actually help the primary task.
We seek to alleviate this burden by formulating a modelagnostic framework that performs ﬁne-grained manipulation of the auxiliary task gradients.
We propose to decompose auxiliary updates into directions which help, damage or leave the primary task loss unchanged.
This allows weighting the update directions differently depending on their impact on the problem of interest.
We present a novel and efﬁcient algorithm for that purpose and show its advantage in practice.
Our method leverages efﬁcient automatic differentiation procedures and randomized singular value decomposition for scalability.
We show that our framework is generic and encompasses some prior work as particular cases.
Our approach consistently outperforms strong and widely used baselines when leveraging out-of-distribution data for Text and Image classiﬁcation tasks.
I NTRODUCTION
Multitask learning (Caruana, 1997) and pretraining (Devlin et al., 2018; Caron et al., 2019) have transformed machine learning by allowing downstream tasks with small training sets to beneﬁt from statistical regularities from data-rich related tasks (Collobert & Weston, 2008; Zhang et al., 2014;
Liu et al., 2019; Kornblith et al., 2019).
Despite these advances, leveraging the mixing of tasks is still an art left to the practitioner.
When one is interested in a primary task, it is unclear how to select helpful auxiliary tasks, an appropriate parameter sharing architecture and a good way to ﬁlter out auxiliary data which might be detrimental to the primary tasks.
Without careful choices, pretraining might hurt end-task performance (Gururangan et al., 2020) or have limited impact (Raghu et al., 2019).
Prior work has examined these problems and proposed solutions, either to choose auxiliary tasks depending on their impact on the primary task (Du et al., 2018; Lin et al., 2019) or to equalize the impact of updates across tasks (Sener & Koltun, 2018;
Chen et al., 2018; Hessel et al., 2019).
Recently, several approaches (Sinha et al., 2018; Suteu & Guo, 2019; Yu et al., 2020) have been proposed that attempt to minimize interference between the updates across tasks.
Our work builds on this direction, but unlike these previous approaches, we do not consider a symmetric view of multi-task learning in the sense that our goal is not to train a model performing well on all tasks.
Instead, we focus on improving generalization for a single task, the primary task, and the other tasks, the auxiliary tasks are considered only through their impact on the problem of interest.
For that purpose, we introduce a framework which decomposes the gradient updates from the auxiliary tasks according to their impact on the primary task.
We analyze the auxiliary task gradients in the subspace spanned by the primary task per-example gradients.
This allows us to decompose auxiliary gradients into into three components : components that help, interfere or have no impact on the primary task according to the Taylor expansion of the expected primary loss.
This decomposition allows us to re-weight each component differently prior to the update.
Our framework enables us to treat each auxiliary update differently depending on its impact on the task of interest and it encompasses prior methods such as classical multitask learning (Caruana, 1997) or more novel gradient surgery techniques (Yu et al., 2020).
To achieve a tractable approach, we introduce an efﬁcient, robust algorithm (ATTITTUD, Auxiliary Task Training with Inﬂuence from Target Task Update Direction) to estimate the subspace spanned by the primary task gradients in an online manner and decompose the auxiliary updates appropriately.
As a result, we can integrate our approach with the stochastic training of large neural networks in various contexts.
The contribution of our work is four-fold.
To our knowledge, this paper proposes the ﬁrst approach to adapt auxiliary gradients using a decomposition built from the span of the primary task Jacobian.
In order to scale this approach to deep neural nets, we contribute a tractable and efﬁcient algorithm called ATTITTUD that leverages insights from randomized linear algebra and automatic differentiation such as the R-operator (Pearlmutter, 1994).
As our third contribution, we show that the ﬁne-grained manipulation of the auxiliary task gradients under ATTITTUD, represents a uniﬁed framework that encompasses several previous approaches to asymmetrical task learning as special cases.
Finally, we demonstrate the efﬁcacy of our approach in both data-rich and data-starved primary tasks, over both images and textual data.
R ELATED W ORK
Methods to leverage data outside of the task of interest are popular in machine learning since the inception of multitask learning (Caruana, 1997; Ruder, 2017; Vandenhende et al., 2020).
These methods address multiple task simultaneously and have been successful in various application domains (Collobert & Weston, 2008; Zhang et al., 2014; Misra et al., 2016).
The optimization problem induced by multitask learning is difﬁcult and solutions have been proposed for the various difﬁculties, including dealing with task gradients of different magnitude (Sener & Koltun, 2018;
Chen et al., 2018; Hessel et al., 2019), or interfering with each others (Sinha et al., 2018; Suteu & Guo, 2019;
Yu et al., 2020).
The speciﬁc problem of interference has been studied extensively in the context of continual learning.
Continual learning visits task in sequence and update interference is particularly problematic as it yields newer tasks to damage previously mastered tasks.
In particular, a family of methods to project the gradient of the new tasks to be orthogonal to the gradient of the previous tasks has been proposed (Lopez-Paz & Ranzato, 2017; Chaudhry et al., 2018; Farajtabar et al., 2019).
Different from many previous approaches, we are not interested in addressing multiple tasks per se.
In our setting, only the primary task matters and the other auxiliary task have the sole role of improving generalization on the primary task.
This is the setting considered by Du et al.
Lin et al.
(2019), who favor auxiliary tasks whose gradient directions are helpful to the primary task.
Unlike these works that use coarse properties like the cosine similarity between averaged gradients, our approach allows ﬁne-grained gradient manipulation within a subspace.
Also, in our case, we do not distinguish between the different auxiliary tasks.
Instead, we aim at correcting every auxiliary gradient in the same manner to improve the loss on the primary task.
This type of gradient correction is related to Yu et al.
(2020), which considers projecting multi-task gradients such that the directions of disagreement are removed.
This method is actually a special case of our framework.
Our work also shares some similarities with data selection and domain adaptation approaches.
In this case, the training data comes from a single task but its distribution is different from the validation/test distribution (Moore & Lewis, 2010; Axelrod et al., 2011; Ngiam et al., 2018).
This classical problem has recently been addressed by sampling training points whose gradient aligns well with the expected validation gradient (Wang et al., 2020b;a).
Instead of sampling individual points based on an estimated distribution of how helpful they will be to the primary task, our work avoids the use (and inherent challenges) of this reinforcement learning approach by operating on batch gradients of groups of points.
Our primary task/auxiliary task setting is also related to the pre-training then ﬁne-tuning paradigm in which the auxiliary tasks are visited ﬁrst (pre-training) to give an initialization for training on the primary task (ﬁne-tuning).
These methods have been very successful in settings where primary task data are rare.
In particular, it is common to ﬁrst rely on an unsupervised task over very large datasets prior to ﬁne tuning over a supervised task (Devlin et al., 2018; Liu et al., 2019; Kornblith et al., 2019;
Yang et al., 2019; Song et al., 2019; Caron et al., 2018).
A UXILIARY T ASK U PDATE D ECOMPOSITION
This section introduces a new method to improve generalization on a primary task T ∗ using training data from auxiliary tasks T
= { T 1 , . .
, T n } , where θ ∈ R D denote the parameters shared by all tasks.
Our approach leverages gradient updates from the auxiliary tasks, but unlike the traditional approach, we decompose these gradients to maximize their usefulness to T ∗ .
Precisely, we decompose the auxiliary task gradients into directions which decrease a ﬁrst-order approximation of the primary task loss, increase it or have no effect.
This decomposition allows weighting these three directions differently when learning from the auxiliary tasks.
In order to decompose the auxiliary gradient, we must collect more ﬁne-grained statistics about the primary task.
At each training step, we collect the gradient of the loss with respect to θ for individual examples from the primary task, {∇ θ L prim i , [Equation]
The span of these vectors,
[Equation]
deﬁnes a subspace in which any linear combination of primary task gradients lies, including the gradient of the expected primary task loss, i.e. g prim =
E ( ∇ θ L prim i ) ∈ S .
If we deﬁne the orthogonal complement of S as S ⊥ , any vector v ∈ S ⊥ , is therefore orthogonal to g prim , i.e. v · g prim = 0 .
This means that adding such a vector to the parameters has no impact on the expected primary task loss, according the order-1 Taylor expansion of L prim , i.e.
[Equation]
We propose to project auxiliary task gradients onto S and S ⊥ .
This allow to distinguish between the directions of the auxiliary task updates which impact the primary task loss and those which do not.
If we denote the averaged auxiliary task gradient as g aux = E ( ∇ θ L aux i ) , we can decompose this gradient as g aux =
g ± aux + g ⊥ aux .
where g ± aux ∈ S is the portion of the gradient that lies in the span of the primary task example gradients and g ⊥ aux ∈ S ⊥ is the portion that lies outside of it.
Since g ⊥ aux ∈ S ⊥ , it is orthogonal to the average primary task gradient and parameter updates along the direction of g ⊥ aux are expected to have limited impact on the primary task loss.
On the other hand, updates along the direction of g ± aux can potentially improve or damage the averaged primary task loss.
This component deserves a more careful treatment.
For that purpose, we introduce { u i , i = 1 , . . .
, k } an orthonormal basis of S .
In this basis, we can measure if the components of g ± aux agree or disagree with g prim .
We say that the two gradients agree along u i iif sign( g ± aux · u i )
=
sign( g prim · u i ) .
This means that we can decompose g ± aux =
g + aux + g − aux where g + aux refers to the projection of g ± aux onto the basis vectors where g ± aux and g prim agree.
By this decomposition, g + aux helps the primary task, g + aux · g prim > 0 , while g − aux interfere with the primary task, g − aux · g prim < 0 .
Guided by the primary task, we can therefore decompose the auxiliary task gradient as
[Equation]
which is described on Fig 1.
Our approach proposes to re-weight differently the components of g aux , i.e.
[Equation]
where η aux = (cid:0) η ⊥ , η + , η − ) are hyper-parameters adjusting the auxiliary gradient according to the impact on the main task.
If we also wish to include the primary task gradient in descent, as with multitasking, we can introduce η prim as a scalar control variable to control its weighting.
A consequence of introducing η aux is that speciﬁc conﬁgurations lead us to gradient updates that are guaranteed to do no harm to both tasks.
This is captured by Theorem 1 below.
Theorem 1.
Let L aux ( θ t ) and L prim ( θ t ) represent the full batch losses of the auxiliary tasks and primary task respectively at step t. We assume the gradients of L aux and L prim are Lipschitz continuous with constant L > 0 .
Following the update rule : θ t {++
+} = θ t − α · ˜ g aux , where α ≤ 1 L is the learning rate, we are guaranteed :
[Equation]
If η − = 0 and η ⊥ , η + ≥
Proof.
See Appendix A
This theorem focuses on a single update and guarantees progress on both auxiliary and primary tasks.
However, our asymmetric scenario is not interested in improving the auxiliary tasks per se and is amenable to more aggressive settings.
Ideally we want gradient updates during pre-training with T to not only do-no-harm to T ∗ when applied downstream but also to be along descent directions that are maximally beneﬁcial to T ∗ .
We can consider η − < 0 as in Fig 1.
Reversing the direction of g − aux by setting η − < 0 preserves the descent guarantee on L prim ( θ t +1 ) but no longer ensures descent on L aux ( θ t +1 ) .
There are other interesting settings for our control parameters.
One can recover the original gradient g aux with η ⊥ =
η − = η + = 1 . 0 .
One can choose to drop gradients orthogonal to the primary task gradient span with η ⊥ = 0 . 0 , or ignore those which conﬂict with the main task by setting η − = 0 . 0 .
Relationships to other approaches
Our framework is generic and encompasses other approaches as a particular case.
One can train solely on the primary task by selecting η aux = (cid:0) 0 . 0 , 0 . 0 , 0 . 0) and η prim = 1 . 0 .
Classical multitasking corresponds to η aux = (cid:0) 1 . 0 , 1 . 0 , 1 . 0) and η prim > 0 . 0 , while classical pre-training corresponds to performing a ﬁrst phase with η aux = (cid:0) 1 . 0 , 1 . 0 , 1 . 0) and η prim = 0 . 0 .
Interestingly, our formulation introduces novel variants of pre-training, for instance, one can consider pre-training with only auxiliary gradients helpful to the primary task, η aux = (cid:0) 0 . 0 , 1 . 0 , 0 . 0) and η prim = 0 . 0 , followed by ﬁne-tuning with η aux = (cid:0) 0 . 0 , 0 . 0 , 0 . 0) and η prim = 1 . 0 .
Our approach also instantiates PCGrad (Yu et al., 2020) as a particular case.
This method was introduced to address the issue of conﬂicting gradients in multitask settings.
PCGrad orthogonalizes the gradients of each task and removes conﬂicting gradients.
To recover PCGrad under our approach, note that it is equivalent to a speciﬁc choice of our decomposition in the 1-D subspace spanned by the g prim .
PCGrad then removes components of g aux that conﬂict with g prim which is equivalent to η aux = (cid:0) α aux , α aux , 0 . 0 (cid:1) and η aux = α prim .
I MPLEMENTATION
Equation 2 requires selecting a basis for the span of primary task gradients.
Multiple choices are possible to deﬁne the basis { u i } , to represent the span at each optimization time-step.
This choice is important since the components of g ± aux are labeled positive or negative depending on how they agree with the projection of the averaged primary task gradient onto the same basis.
A natural choice is to select the basis as the singular vectors of the matrix of primary task per-example gradients
J ∗ ∈ R m × D , also know as the Jacobian.
To improve efﬁciency and prevent over-ﬁtting on a few examples, we consider the span deﬁned by the k largest principal vectors of J ∗ .
Using the principal vectors as directions of descent instead of the mean induces a more robust algorithm since the minibatch average gradient is susceptible to outliers and skew from replicated data-points.
To the best of our knowledge, we are the ﬁrst to propose using the singular vectors of J ∗ as directions of descent.
We leave the theoretical implications of this algorithm to future work but note that its variance reduction properties may induce generalization beneﬁts (Namkoong & Duchi, 2017).
We also consider alternative choices of bases as baselines, including the canonical parameter basis.
This choice will examine the sign of every parameter update to verify whether it agrees with g prim .
Whilst Theorem 1 holds irrespective of the choice of basis, its proof reveals that the amount of progress made on each loss depends on the choice of basis.
Speciﬁcally, the reduction in
L prim ( θ t +1 ) , L aux ( θ t +1 ) after a gradient step along ˜ g aux is proportional to the fraction of the norms of g prim and g aux captured by the subspace spanned by our choice of basis.
We evaluate this fraction for different choice of basis in our experiments (see Appendix C).
We are interested in applying our approach to the training of large neural networks and must consider a scalable algorithmic solution.
As stochastic optimization is prevalent in this setting, we construct subspace S from a mini-batch of primary task data.
Similarly, the expected gradients g prim and g aux are deﬁned over a mini-batch.
Instead of computing the singular value decomposition (SVD) of {∇ θ L prim i , ∀ i } exactly, we rely on a randomized approximation (Halko et al., 2011; Rokhlin et al., 2010; Nakatsukasa, 2017).
This method does not require instantiating the vectors {∇ θ L prim i , ∀ i } and only needs a low dimensional projection onto a random subspace.
This is advantageous for high dimensional cases, i.e. when the number of model parameters is large.
In our case, this method also allows us to beneﬁt from memory-efﬁcient computation of Jacobian Vector product using the
R-operator (Pearlmutter, 1994) offered by automatic differentiation packages (Baydin et al., 2015) like Pytorch (Paszke et al., 2017).
This means that we can compute SVD with a limited computational and memory burden, albeit without sacriﬁcing approximation accuracy (Nakatsukasa, 2017).
Additionally, we do not recompute the basis at every optimization step but at every n steps, which is efﬁcient when training with small updates, e.g. when small learning rates and gradient clipping are used (Pascanu et al., 2013).
We study the impact of these choices in practice in Section 6.
Putting it all together results in the
ATTITTUD algorithm, Auxiliary Task Training with Inﬂuence from Target Task Update Direction, shown as Algorithm 1.
The sub-procedure randomized lowrank approx is detailed in Appendix B as Algorithm
Algorithm
Require : g aux , J ∗ : Auxiliary task average gradient, primary task Jacobian
Require : η aux = (cid:0) η ⊥ , η + , η − (cid:1) : Auxiliary task control parameters
Require : k : Size of subspace g prim = 1 m (cid:80) mi [-=
-]=1 J ∗ i, :
V ← randomized lowrank approx ( J ∗ , k ) p prim , p aux =
V t (cid:0) g prim (cid:1) T , V t (cid:0) g aux (cid:1) T // ◦ is the hadamard product operator p + aux , p
[Equation]
= (cid:18) [Equation] (cid:2)
[Equation]
p aux , (cid:18) 1 (cid:2) p prim ◦
[Equation]
aux < 0 (cid:3)(cid:19) ◦ p aux // Calculate the decomposition components g + aux , g − aux =
(cid:0) p + aux (cid:1) T V , (cid:0) p − aux (cid:1) T V //
Calculate the out of span component g ⊥ aux =
g aux − (cid:0) g + aux + g − aux (cid:1) ˜ g aux = (cid:0) η ⊥ · g ⊥ aux (cid:1) + (cid:0) η + · g + aux (cid:1) + (cid:0) η − · g − aux (cid:1)
Return : ˜ g
aux : Auxiliary task surrogate gradient
E XPERIMENTAL S ETUP
We compare ATTITTUD with previous methods on a variety of tasks and domains.
We rely on both text and image classiﬁcation tasks to conduct our analysis.
We also present ablation experiments to explain the impact of hyper-parameter selection.
Text Classiﬁcation.
We apply our method on binary sentiment classiﬁcation.
We consider the
Amazon Helpfulness (McAuley et al., 2015) and Imdb Movie Review (Maas et al., 2011) tasks.
The
Amazon Helpfulness task splits text reviews into 115k/5k/25k documents for train-validation-test split whilst the Imdb Review dataset has a 20k/5k/25k split.
The Imdb Review task also has 50k unlabeled reviews as extra data which we utilize.
For our models we build on top of Gururangan et al.
(2020)’s work where they introduce TaskAdaptive Pre-training (TAPT).
TAPT further pre-trains a generic model, Roberta (Liu et al., 2019), by performing Masked Language Modelling, MLM, (Devlin et al., 2018) on the task speciﬁc data (ignoring the labels) before doing supervised learning with the same data.
We replicate Gururangan et al.
(2020)’s experimental setup and re-use their hyper-parameters for our experiments.
We use the
TAPT task as our auxiliary task.
We extend TAPT to use our method by modifying the TAPT gradient with guidance from the supervised-learning task gradients.
As baselines, we compare against
TAPT and cross-TAPT: where we swap the masked language modelling pre-training data for the two tasks.
Cross-TAPT is a setting where one uses out-of-distribution data for pre-training.
Image Classiﬁcation.
We apply our method to both high-resource and limited-data image classiﬁcation tasks.
We use the Cifar100 dataset (Krizhevsky et al., 2009) to explore the high-resource setting.
We follow Rosenbaum et al.
and treat each of the 20 super-classes / coarse labels of Cifar100 as a separate task.
In our asymmetrical task setting, each of the 20 tasks is treated as a primary task, whilst the remaining 95 classes are grouped into a single auxiliary task.
Thus, for each coarse label, we have an auxiliary 95-way classiﬁcation task and a 5-way primary classiﬁcation task.
Moving forward, we refer to this setting as MultiCifar100.
We use a down-sampled version of Cifar10 (Krizhevsky et al., 2009) as a low-resource setting.
Speciﬁcally, we rely on Cat-vs-Dog for the primary task and use the remaining 8 classes for the auxiliary task.
Our auxiliary task is therefore an 8-way classiﬁcation task where each class has 5,000 examples.
We restrict the Cat and Dog classes to only 50 training examples from each class.
We use the low-resource setting to compare against other methods and for our ablation study.
For these vision experiments, we use a WideResNet-22 architecture (Zagoruyko & Komodakis, 2016) with a depth of k = 4 .
We compare our method to 4 different baselines : no pre-training, vanilla pre-training, multitasking and PCGrad (Yu et al., 2020).
Our architecture is more standard and allows gradient descent optimization unlike the routing network of Rosenbaum et al.
(2017) and (Yu et al., 2020), which requires reinforcement learning for training.
Medical Imaging Transfer.
We apply our method to cross-domain transfer for low-resource medical image classiﬁcation.
Speciﬁcally, we use 5k training examples from the ChexPert Dataset (Irvin et al., 2019) as our primary task and seek to identify 5 different thoracic pathologies: atelectasis, cardiomegaly, consolidation, edema and pleural effusion.
This setup has been used in several crossdomain pretraining studies (Raghu et al., 2019; Jaiswal et al., 2019).
Note that since we do not have access to the test set for this task, we use the validation set (231 images) as a proxy test set, and sample 100 images from the training data as a new validation set.
We rely on generic photographs (Imagenet) as an auxiliary task (Deng et al., 2009).
We use Tiny Imagenet Dataset (Le & Yang, 2015), a subset of Imagenet which consists of 500 examples each from 200 classes, instead of training on full Imagenet.
All approaches are applied to the Resnet18 model (He et al., 2016) trained with Adam (Kingma & Ba, 2014).
For all our experiments, we select the auxiliary task control parameters η aux within { (1 . 0 , 1 . 0 , − 1 . 0) , (1 . 0 , 1 . 0 , 0 . 0) , (1 . 0 , 0 . 0 , − 1 . 0) , (1 . 0 , 0 . 0 , 0 . 0) } .
For settings where we compare against multi-tasking, we select η prim within a small subset of the settings that worked best with multitasking baseline experiments.
These choices limit the overhead of hyper-parameter search but still allow us to show the empirical advantage of our method.
In all our experiments, we provide all methods with similar hyper-parameter search budgets, e.g. for Cifar10-Cat-vs-Dog, we ran a grid search with 16 conﬁgurations for regular pretraining, 16 conﬁgurations for PCGrad andconﬁgurations for ATTITUD.
More experimental details are available in Appendix C
[Table]
R ESULTS AND D ISCUSSION
Text Classiﬁcation.
Table 1 shows the results for text classiﬁcation.
When the same data is used both for the auxiliary task of MLM and the primary classiﬁcation task, TAPT and ATTITTUD both bring a similar improvement over Roberta (Imdb, Amazon columns).
When different data is used for the auxiliary task and the primary task (Imdb + Amazon MLM, Amazon + Imdb MLM columns),
TAPT does not perform as well as ATTITTUD.
This highlights the advantage of ATTITTUD when the auxiliary task data distribution differ from the primary task distribution.
Image Classiﬁcation.
Our results are presented in Table 2.
Both for MultiCifar100 (high resource setting) and Cifar10-Cat-vs-Dog (low resource setting), ATTITUD shows a strong improvement over baselines.
In general, we ﬁnd that primary-task aware pre-training
Multitasking, PCGrad,
Ours) is better than vanilla pre-training which also performs better than having no pre-training at all.
For MultiCifar100, we ﬁnd that using η aux = (1 . 0 , 1 . 0 , − 1 . 0) ,
η prim = 0 . 1 worked best for 11 out of the 20 Cifar100 super-classes tasks.
Note that η aux = (1 . 0 , 1 . 0 , − 1 . 0) is an aggressive but novel conﬁguration we introduce.
Multitask learning and PCGrad produce better models on 6 and 3 tasks respectively.
In the low-resource Cat-vs-Dog, setting ATTITUD produces a bigger boost in performance compared to baselines, with the best performing conﬁguration being η aux = (1 . 0 , 0 . 0 , 0 . 0) , η prim = 0 .
We posit that this conﬁguration is successful because removal of the in-span components makes overﬁtting less likely.
Applying the out-of-span components means the model learns features that do not harm the loss of the current mini-batch but could be useful later. Note that our best performing conﬁgurations are all novel and never an instantiation of PCGrad.
[Table] [Table]
Medical Imaging Transfer.
Table 3 shows our results on the ChexPert multi-label classiﬁcation task.
Per-pathology breakdowns are in Appendix C. Doing no pre-training at all performs worst.
Our method outperforms using a pre-trained Resnet18 model over Imagenet.
We apply the endtask-aware ATTITUD over 100k ImageNet images after the initial pretraining and we reach 83.3% AUC, an improvement over 81.4%.
Ablation Study.
Our approach relies on the top-k singular vectors from randomized svd to deﬁne the basis to identify the positive and negative component of the auxiliary task gradient, see Section 4.
This method is more accurate than several alternatives; see Table 4.
Namely, we compare our choice [Table] to random , the basis spanned by k randomly chosen orthogonal vectors in R D , unit avg grad , the basis spanned by the average primary task gradient, and canonical , the per-parameter basis.
This ablation was performed under a more limited tuning budget (we cross-validated on conﬁgurations (1 , 1 , 0) and (1 , 1 , − 1) only) than the full Cat-vs-Dog experiments from Table
[Figure] We also examine the number of samples to estimate the principal directions of the per-example primary task gradient.
Larger sample sizes involve more computation but have limited beneﬁt on average accuracy.
Large sample sizes however reduce variance, as shown in Figure 2 (left).
This is as expected since using more samples gives a higher ﬁdelity estimate of the top-k singular vectors.
Another parameter of our algorithm is the size of our subspace, k .
In general, we observe that in low-resource settings, it is better to operate on the auxiliary task gradient in a smaller dimensional subspace.
The opposite holds for high-resource settings.
This can be seen in Figure 2 (right).
Whilst using a larger dimensional subspace captures a richer description of the J ∗ , it also creates the risk of over-ﬁtting especially in a limited data setting.
This trade-off therefore has to be validated on a per-task basis.
C ONCLUSIONS
In this work, we propose a new approach to training a model with additional help from an auxiliary task.
Our method decomposes the gradients of the auxiliary task according to three directions, with positive, negative and neutral impact on the primary task.
This decomposition allows a ﬂexible re-weighting of the auxiliary task components and give rise to a family of training strategies, which encompasses novel and existing approaches.
We leverage insights from randomized linear algebra and automatic differentiation to scale the approach to large deep networks.
Experiments in multitasking, pretraining and domain transfer over vision and text classiﬁcation task demonstrate the empirical beneﬁt of our framework. A P ROOF OF T HEOREM Theorem.
Let L aux ( θ t ) and L prim ( θ t ) represent the full batch losses of the auxiliary tasks and primary task respectively at step t. We assume the gradients of L aux and L prim are Lipschitz continuous with constant L > 0 .
{+Following the update rule : θ t +
+}
=
θ t − α · ˜ g aux , where α ≤ 1 L is the learning rate, we are guaranteed :
[Equation]
If η − = 0 and η ⊥ , η + ≥
Proof.
Let V t ∈ R k × D be the orthonormal matrix whose rows span the per-example primary task gradients J ∗ at timestep t .
The projections of the average primary task gradient g prim
=
1 m (cid:80) mi =1 J ∗
[Equation]
, : and average auxiliary task gradient g aux at iteration t are
[Equation]
p prim and p aux will agree on some directions (same sign on those components).
We use the operator [ x ] + to mark these directions of agreement.
This operator preserves components that agree and sets those that disagree to zero.
As an example given p prim =
[1 , 1 , − 1] and p aux =
[1 , 3 , 10] , [ p prim ] + = [1 , 1 , 0] and [ p aux ] + = [1 , 3 , 0] .
For directions that disagree (different signs of the respective components), we introduce the operator [ x ] − .
In the above example [ p prim ] − =
[0 , 0 , − 1] and [ p aux ] − =
Note that our operators are deﬁned by comparing two vectors x 1 and x 2 ,
Our operators have the following properties by deﬁnition :
[Equation]
and
[Equation]
± ⊥ [ x 2 ] ∓
From Equation 2 : ˜ g
[Equation]
=
[Equation]
+
[Equation]
+
[Equation]
−
[Equation]
⊥
[Equation]
We can re-write this in terms of [ x ] ± as :
[Equation]
We now proceed to show the effect of the gradient descent update below on L aux ( θ t +1 ) and
L prim ( θ t +1 ) .
[Equation]
How does this update affect the loss on the primary task loss
L prim ( θ t +1 ) ?
[Equation]
p
[Equation]
[Equation]
[ p
[Equation]
[Equation]
[ p
[Equation]
] T + [ p
[Equation]
[Equation]
+
[Equation]
(cid:18) [ p aux ] T − [ p prim ] + + [ p aux ] T − [ p prim ] − (cid:19)(cid:19) =
L prim
[Equation]
θ
[Equation]
[Equation]
(cid:18) η
+ [ p aux ] T + [ p prim ] + + η − [ p aux ] T − [ p prim ] −
[Equation]
≤ L prim ( θ t ) (if η − ≤ 0 , η ⊥ , η + ≥ 0 )
Note that in going from line 3 to 4 in the proof above, we use the fact that (cid:0) g ⊥ aux (cid:1) T g prim = 0 since g ⊥ aux lies outside the subspace and g prim lies inside it.
For the last step of the proof, we use the observations below : [ p aux ] + [ p prim ] + ≥ 0 since these directions agree in sign [ p aux ] − [ p prim ] − ≤ 0 since these directions disagree in sign [ p aux ] + [ p prim ] − = 0 by the property of the [ x ] ± operator [ p aux ] − [ p prim ] + = 0 same motivation as above
How does Equation 3 affect the auxiliary task loss L aux ( θ t +1 ) ?
L aux ( θ t +1 ) =
L aux ( θ t − α · ˜ g aux )
≈ L aux ( θ t ) − α (cid:0) ˜
g aux (cid:1) T g aux
(First order Taylor Expansion)
=
L aux ( θ t ) − α (cid:0) η ⊥ g ⊥ aux + η + g + aux + η − g − aux
cid:1) T (cid:0) g ⊥ aux + g + aux + g − aux (cid:1)
=
{+L aux ( θ t ) − α (cid:0) η ⊥ (cid:107) g ⊥ aux (cid:107) 
+}
+ η + (cid:107) g
+ aux (cid:107) 2 + η − (cid:107) g − aux (cid:107) 2 (cid:1)
(Cross terms cancel due to orthogonality) ≤ L aux ( θ t )
(If η − , η ⊥ , η + ≥ 0 )
Thus, choosing η − = 0 ensures that we are minimizing both L aux ( θ t ) and L prim ( θ t ) .
We can combine this with the constraint on α ≤ 1 L to derive convergence guarantees after some T steps as in optimization literature.
B R ANDOMIZED M ATRIX T HEORY
Algorithm
Require : J ∈ R m × D : Input Matrix
Require : k : Rank of subspace Π ∼ N (0 ,
I
[Equation]
[Equation]
The Gram Schmidt procedure orthogonalizes the rows of an input matrix.
C M ORE E XPERIMENTAL D ETAILS
Image Classiﬁcation For MultiCifar100, unlike Rosenbaum et al.
Yu et al.
(2020) who use a 500-100 train-test split for examples under each ﬁne-grained CIFAR 100 label, we include a validation set and therefore opt for a 400-100-100 train-validation-test split.
We test on all 1000 test examples per class.
For Cat-vs-Dog, we use 100 examples from the training set as validation and test on all 1000 test examples per-class.
For Image Classiﬁcation experiments, we perform pre-training with a learning rate of 1e-4 for all experiments and ﬁnetuning learning rate of 5e-4.
These values were selected after coarse hyperparameter search.
In both pre-training and ﬁnetuning settings, we decay the learning rate by 0.5 if the validation loss has not improved over 4 epochs, up till a minimum learning rate of 1e-5.
we use the Adam Optimizer (Kingma & Ba, 2014) with β = (0 . 9 , 0 . 999) .
We clip all gradient norms to 1.
before performing gradient descent.
We cross-validated dropout rates within the set { 0.05, 0.1, 0.2, 0.3 } for both pre-training and ﬁnetuning steps.
We cross validate η prim based on the relative sizes of primary and auxilary task datasets.
All experiments are averaged over 5 random seeds.
For all our
Vision experiments, we either recompute our subspace basis every n = 5 or n = 10 times.
Due to the fact that Yue et al (PCGrad) treat all tasks symmetrically, which is different from our primary-auxiliary setting, we introduced an extra parameter, α prim , for PCGrad to account for weighting the primary task.
We cross validated values of α prim ∈ { 0 . 1 , 0 .
Medical Imaging Transfer Table 4 presents a more detailed breakdown of the ChexPert task.
For 50k examples from Imagenet, our best performing conﬁguration was η aux =
(1 . 0 , 0 . 0 , − 1 . 0) .
{+We did not use the primary task gradient directly for pre-training so η prim = 0 . 
+}
for all cases.
For ATTITUD, we use the same learning rates as in the Image classiﬁcation setup above.
For the No-Pretraining and
Vanilla pretraining we cross-validated the learning rates for both ﬁnetuning and pre-training from the set { 1e-3, 1e-4 } .
We cross-validated the same list of dropout values above.
[Table]
Text Classiﬁcation For our NLP experiments, we tried limiting the number of layers we applied
ATTITUD to.
We achieved good performance without applying ATTITUD to the word embedding layers (these were updated with untouched auxiliary task gradients).
We cross-validated η prim = { 0 . 01 , 0 . 05 , 0 . 0025 } .
For all our NLP experiments, we either recompute our subspace basis every n = 1 or n = 4 times
For all experiments involving ATTITUD, We cross-validate the following choices of the subspace size k ∈ { 5 , 10 , 20 } from
J ∗ ∈
R m × D using m ∈ { 32 , 64 } .
We recompute the subspace everysteps for vision experiments and every 4 steps for NLP experiments.
We run all experiments for a maximum of 150 pretraining epochs and 500 ﬁnetuning epochs.
We performed early stopping for all experiments if no improvement after 10 consecutive epochs.
Ablation of Fraction of Norm within Subspace The left pane of Figure 3 reinforces our intuition and conﬁrms that our choice of the top-k singular vectors ( randomized svd ) gives the best accuracy as averaged across 5 seeds.
random is the basis spanned by k randomly chosen orthogonal vectors in R D , unit avg grad is the basis spanned by the average primary task gradient whilst canonical uses the per-parameter basis.
Note that k
= 5 for random and randomized svd whilst for unit avg grad and canonical , k = 1 and k = D respectively.
We use the fraction of the norm of sample gradients within a subspace as indicators of how semantically meaningful that choice of subspace is.
We expect that a semantically meaningful choice of basis will achieve better generalization performance because it captures the essential parts of the gradient with k (cid:28) D .
canonical trivially captures all the norm of the sampled gradient vectors but because k = D , it generalizes poorly.
Notice that only small fractions of the norms of sample primary and auxiliary task average gradients lie in the subspace for random and unit avg grad , whilst signiﬁcant fractions lie in randomized svd .