RotoGrad: Gradient Homogenization in Multi-Task Learning

Multi-task learning is being increasingly adopted in applications domains like1 computer vision and reinforcement learning . However , optimally exploiting its ad-2 vantages remains a major challenge due to the effect of negative transfer . Previous3 works have tracked down this issue to the disparities in gradient magnitudes and4 directions across tasks , when optimizing the shared network parameters . While5 recent work has acknowledged that negative transfer is a two-fold problem , exist-6 ing approaches fall short as they focus only on either homogenizing the gradient7 magnitude across tasks ; or greedily change the gradient directions , overlooking8 future conflicts . In this work , we introduce RotoGrad , an algorithm that tackles9 negative transfer as a whole : it jointly homogenizes gradient magnitudes and direc-10 tions , while ensuring training convergence . We show that RotoGrad outperforms11 competing methods in complex problems , including multi-label classification in12 CelebA and computer vision tasks in the NYUv2 dataset.13 1 Introduction14 As neural network architectures get larger in order to solve increasingly more complex tasks , the15 idea of jointly learning multiple tasks ( for example , depth estimation and semantic segmentation in16 computer vision ) with a single network is becoming more and more appealing . This is precisely the17 idea of multi-task learning ( MTL ) [ 3 ] , which promises higher performance in the individual tasks18 and better generalization to unseen data , while drastically reducing the number of parameters [ 27 ] .19 Unfortunately , sharing parameters between tasks may also lead to difficulties during training as20 tasks compete for shared resources , often resulting in poorer results than solving individual tasks , a21 phenomenon known as negative transfer [ 27 ] . Previous works have tracked down this issue to the22 two types of differences between task gradients . First , differences in magnitude across tasks can make23 some tasks dominate the others during the learning process . Several methods have been proposed to24 homogenize gradient magnitudes such as MGDA [ 28 ] , GradNorm [ 6 ] , or IMTL-G [ 18 ] . However,25 little attention has been put towards the second source of the problem : conflicting directions of the26 gradients for different tasks . Due to the way gradients are added up , gradients of different tasks may27 cancel each other out if they point to opposite directions of the parameter space , thus leading to a poor28 update direction for a subset or even all tasks . Only very recently a handful of works have started to29 propose methods to mitigate the conflicting gradients problem , for example , by removing conflicting30 parts of the gradients [ 33 ] , or randomly ‘ dropping ’ some elements of the gradient vector [ 7 ] .31 In this work we propose RotoGrad , an algorithm that tackles negative transfer as a whole by ho-32 mogenizing both gradient magnitudes and directions across tasks . RotoGrad addresses the gradient33 magnitude discrepancies by re-weighting task gradients at each step of the learning , while encourag-34 ing learning those tasks that have converged the least thus far . In that way , it makes sure that no task is35 overlooked during training . Additionally , instead of directly modifying gradient directions , RotoGrad36 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . smoothly rotates the shared feature space differently for each task , seamlessly aligning gradients in37 the long run . As shown by our theoretical insights , the cooperation between gradient magnitude-38 and direction-homogenization ensures the stability of the overall learning process . Finally , we run39 extensive experiments to empirically demonstrate that RotoGrad leads to stable ( convergent ) learning,40 scales up to complex network architectures , and outperforms competing methods in multi-label41 classification settings in CIFAR10 and CelebA , as well as in computer vision tasks using the NYUv242 dataset . Alongside this paper , we will provide a simple-to-use library to include RotoGrad in any43 Pytorch pipeline with a few lines of code.44 2 Multi-task learning and negative transfer45 The goal of MTL is to simultaneously learn K different tasks , that is , finding K mappings from a46 common input datasetX ∈ RN×D to a task-specific set of labels Yk ∈ YNk . Most settings consider47 a hard-parameter sharing architecture , which is characterized by two components : the backbone and48 heads networks . The backbone uses a set of shared parameters , θ , to transform each input x ∈X49 into a shared intermediate representation z = f ( x ; θ ) ∈ Rd , where d is the dimensionality of z.50 Additionally , each task k = 1 , 2 , . . . , K has a head network hk , with exclusive parameters φk , that51 takes this intermediate feature z and outputs the prediction hk ( x ) = hk ( z ; φk ) for the corresponding52 task . This architecture is illustrated in Figure 2 , where we have added task-specific rotation matrices53 Rk that will be necessary for the proposed approach , RotoGrad . Note that the general architecture54 described above is equivalent to the one in Figure 2 when all rotations Rk correspond to identity55 matrices , such that rk = z for all k.56 r1 L1 ( h1 ( r1 ) , y1 ) x z r2 L2 ( h2 ( r2 ) , y2 ) rK LK ( hK ( rK ) , yK ) hφ1 fθ R1 R2 RK ... hφ2 ... hφK MTL aims to learn the architecture parameters57 θ , φ1 , φ2 , . . . , φK by simultaneously minimiz-58 ing all task losses , that is , Lk ( hk ( x ) , yk ) for59 k = 1 , . . . , K. Although this is a priori a multi-60 objective optimization problem [ 28 ] , in practice61 a single surrogate loss consisting of a linear com-62 bination of the task losses , L = ∑ k ωkLk , is op-63 timized . While this approach leads to a simpler64 optimization problem , it may also trigger nega-65 tive transfer between tasks , hurting the overall66 directions that do not improve any of the tasks . Figure 1 shows an example of poor direction updates73 ( left ) as well as magnitude dominance ( right ) .74 In this work , we tackle negative transfer as a whole by homogenizing tasks gradients both in magnitude75 and direction . Note that homogenizing gradients with respect to θ is equivalent to homogenizing76 gradients with respect to the shared feature z due to the chain rule , ∇θLk = ∇θz · ∇zLk . Thus,77 from now on we focus on homogenizing the feature-level task gradients∇zLk.78 3 RotoGrad79 In this section we introduce RotoGrad , a novel algorithm that addresses the negative transfer problem80 as a whole . RotoGrad consists of two building blocks which , respectively , homogenize task-gradient81 magnitudes and directions . Moreover , these blocks complement each other and provide convergence82 guarantees of the network training . Next , we detail each of these building blocks and show how they83 are combined towards an effective MTL learning process.84 3.1 Gradient-magnitude homogenization85 As discussed in Section 2 , we aim to homogenize gradient magnitudes across tasks , as large magnitude86 disparities can lead to a subset of tasks dominating the learning process . Thus , the first goal of87 RotoGrad is to homogenize the magnitude of the gradients across tasks at each step of the training.88 Let us denote the feature-level task gradient of the k-th task for the n-th datapoint , at iteration t , by89 gn , k : = ∇zLk ( hk ( xn ) , yn , k ) , and its batch versions by G > k : = [ g1 , k , g2 , k , . . . , gB , k ] , where B is90 the batch size . Then , equalizing gradient magnitudes amounts to finding weights ωk that normalize91 and scale each gradientGk , that is,92 ||ωkGk|| = ||ωiGi|| ∀i ⇐⇒ ωkGk = C ||Gk|| Gk = CUk ∀k , ( 1 ) where Uk : = Gk||Gk|| denotes the normalized task gradient and C is the target magnitude for all tasks.93 Note that , in the above expression , C is a free parameter that we need to select.94 In RotoGrad , we select C such that all tasks converge at a similar rate . We motivate this choice95 by the fact that , by scaling all gradients , we change their individual step size , interfering with the96 convergence guarantees provided by their Lipschitz-smoothness ( for an introduction to non-convex97 optimization see , for example , [ 25 ] ) . Therefore , we seek for the value of C providing the best98 step-size for those tasks that have converged the least up to iteration t. Specifically , we set C to be a99 convex combination of the task-wise gradient magnitudes , C : = ∑ k αk||Gk|| , where the weights100 α1 , α2 , . . . , αK measure the relative convergence of each task and sum up to one , that is,101 αk = ||Gk||/||G0k||∑ i ||Gi||/||G0i || , ( 2 ) withG0k being the initial gradient of the k-th task , i.e. , the gradient at iteration t = 0 of the training.102 As a result , we obtain a ( hyper ) parameter-free approach that equalizes the gradient magnitude across103 tasks to encourage learning slow-converging tasks . Note that the resulting approach resembles104 Normalized Gradient Descent ( NGD ) [ 8 ] for single-task learning , which has been proved to quickly105 escape saddle points during optimization [ 24 ] . Thus , we expect a similar behavior for RotoGrad,106 where slow-converging tasks will force quick-converging tasks to escape from saddle points.107 The resulting training algorithm may however diverge as a consequence of constantly oscillating108 between ( slow-converging ) tasks . For example , in scenarios where one task improves , there is always109 another task ( s ) that deteriorates . Fortunately , as shown in the following result ( proof in Appendix A ) ,110 such a phenomenon does not appear in the absence of conflicting gradients.111 Proposition 3.1 . LetG1 , G2 , . . . , GK be the task gradients with respect to Z as defined above . If112 K = 2 ; or cos_sim ( Gi , Gj ) ≥ 0 pairwise ; then there exists a small-enough step size ε > 0 such113 that , for all tasks , we have that Lk ( hk ( Z − ε · C ∑ kUk ; φk ) ; Yk ) < Lk ( hk ( Z ; φk ) ; Yk ) .114 In other words , Proposition 3.1 shows that , when gradients do not conflict in direction with each115 other , following the feature-level gradient C ∑ kUk improves all ( lower-bounded ) task losses for116 the given batch . This result , while restricted to the given batch and to the gradient with respect to117 the shared representation Z , still provides useful insights in favor of having as desideratum of an118 efficient MTL pipeline the absence of conflicting gradients.119 3.2 Gradient-direction homogenization120 In the previous subsection , we have shown that avoiding conflicting gradients may not only be121 necessary to avoid negative transfer , but also to ensure the stability of the training . In this section122 we introduce the second building block of RotoGrad , an algorithm that homogenizes task-gradient123 directions . The main idea of this approach is to smoothly rotate the feature-space z in order to reduce124 the gradient conflict between tasks—in following iterations—of the training by bringing ( local ) 125 optima for different tasks closer to each other ( in the parameter space ) . As a result , it complements126 the previous magnitude-scaling approach and reduces the likelihood of the training to diverge.127 In order to homogenize gradients , for each task k = 1 , . . . , K , RotoGrad introduces a matrix Rk128 so that , instead of optimizing Lk ( z ) with z being the last shared representation , we optimize an129 equivalent loss function Lk ( Rkz ) . As we are only interested in changing directions ( not the gradient130 magnitudes ) , we choose Rk ∈ SO ( d ) to be a rotation matrix1 leading to per-task representations131 rk : = Rkz . RotoGrad thus extends the standard MTL architecture by adding task-specific rotations132 before each head , as depicted in Figure 2.133 Unlike all other network parameters , matrices Rk do not seek to reduce their task ’ s loss . Instead,134 these additional parameters are optimized to reduce the direction conflict of the gradients across135 tasks . To this end , for each task we optimize Rk to maximize the batch-wise cosine similarity or,136 equivalently , to minimize137 Lkrot : = − ∑ n 〈R > k g̃n , k , vn〉 , ( 3 ) where g̃n , k : = ∇rkLk ( hk ( xn ) , yn , k ) ) ( which holds that gn , k = R > k g̃n , k ) and vn is the target vector138 that we want all task gradients pointing towards . We set the target vector vn to be the gradient we139 would have followed if all task gradients weighted the same , that is , vn : = 1K ∑ k un , k , where un , k140 is a row vector of the normalized batch gradient matrix Uk , as defined before.141 As a result , in each training step of RotoGrad we simultaneously optimize the following two problems:142 143 N etwork : minimize θ , { φ } k ∑ k ωk Lk. , Rotation : minimize { Rk } k ∑ k Lkrot ( 4 ) The above problem can be interpreted as a Stackelberg game : a two player-game in which leader144 and follower alternately make moves in order to minimize their respective losses , Ll and Lf , and the145 leader knows what will be the follower ’ s response to their moves . Such an interpretation allows us to146 derive simple guidelines to guarantee training convergence—that is , that the network loss does not147 oscillate as a result of optimizing the two different objectives in Equation 4 . Specifically , following148 Fiez et al . [ 10 ] , we can ensure that problem 4 converges as long as the rotations ’ optimizer ( leader ) 149 is a slow-learner compared with the network optimizer ( follower ) . That is , as long as we make the150 rotations ’ learning rate decrease faster than that of the network , we know that RotoGrad will converge151 to a local optimum for both objectives . A more extensive discussion can be found in Appendix B.152 3.3 RotoGrad : the full picture153 After the two main building blocks of RotoGrad , we can now summarize the overall proposed154 approach in Algorithm 1 . At each step , RotoGrad first homogenizes the gradient magnitudes such155 that there is no dominant task and the step size is set by the slow-converging tasks . Additionally,156 RotoGrad smoothly updates the rotation matrices—using the local information given by the task157 gradients—to seamlessly align task gradients in the following steps , thus reducing direction conflicts.158 3.4 Practical considerations159 In this section , we discuss the main practical considerations to account for when implementing160 RotoGrad and propose efficient solutions.161 1The special orthogonal group , SO ( d ) , denotes the set of all ( proper ) rotation matrices of dimension d. Algorithm 1 Training step with RotoGrad Input input samplesX , task labels { Yk } , network ’ s ( RotoGrad ’ s ) learning rate η ( ηroto ) Output backbone ( heads ) parameters θ ( { φk } ) , RotoGrad ’ s parameters { Rk } 1 : compute shared feature Z = f ( X ; θ ) 2 : for k = 1 , 2 , . . . , K do 3 : compute task-specific loss Lk = ∑ n Lk ( hk ( Rkzn ; φk ) , yn , k ) 4 : compute gradient of shared featureGk = ∇zLk 5 : compute gradient of task-specific feature G̃k = RkGk . Treated as constant w.r.t . Rk . 6 : compute unitary gradients Uk = Gk/||Gk|| 7 : compute relative task convergence αk = ||Gk||/||G0k|| 8 : end for 9 : make { αk } sum up to one [ α1 , α2 , . . . , αK ] = [ α1 , α2 , . . . , αK ] / ∑ k αk 10 : compute shared magnitude C = ∑ k αk||Gk|| 11 : update backbone parameters θ = θ − ηC ∑ kUk 12 : compute target vector V = 1K ∑ kUk 13 : for k = 1 , 2 , . . . , K do 14 : compute RotoGrad ’ s loss Lrotok = − ∑ n〈R > k g̃n , k , vn〉 15 : update RotoGrad ’ s parametersRk = Rk − ηroto∇RkLrotok 16 : update head ’ s parameters φk = φk − η∇φkLk 17 : end for Unconstrained optimization . As previously discussed , parameters Rk are defined as rotation162 matrices , and thus the Rotation optimization in problem 4 is a constrained problem . While this would163 typically imply using expensive algorithms like Riemannian gradient descent [ 1 ] , we can leverage164 recent work on manifold parametrization [ 5 ] and , instead , apply unconstrained optimization methods165 by automatically2 parametrizingRk via exponential maps on the Lie algebra of SO ( d ) .166 Memory efficiency and time complexity . Second , as we need one rotation matrix per task , we have167 to storeO ( Kd2 ) additional parameters . In practice , we only needKd ( d− 1 ) /2 parameters due to the168 aforementioned parametrization and , in most cases , this amounts to a small part of the total number169 of parameters . Moreover , as described by Casado et al . [ 5 ] , parametrizing Rk enables efficient170 computations compared with traditional methods , with a time complexity of O ( d3 ) independently of171 the batch size . In our case , the time complexity is of O ( Kd3 ) , which scales better with respect to the172 number of tasks than existing methods ( for example , O ( K2d ) for PCGrad [ 33 ] ) . Moreover , caching173 Rk in the forward pass and GPU parallelization can further reduce training time.174 Scaling-up RotoGrad . Even though we can efficiently compute and optimize the rotation matrixRk,175 in some application domains , like computer vision , in which the size d of the shared representation z176 is large , the time complexity for updating the rotation matrix may become comparable to the one of177 the network updates . In those cases , we propose to only rotate a subspace of the feature space , that178 is , rotate only m < < d dimensions of z . Then , we can simply apply a transformation of the form179 rk = [ Rkz1 : m , zm+1 : d ] , where za : b denotes the elements of z with indexes a , a+ 1 , . . . , b. While180 there exist other possible solutions , such as using block-diagonal rotation matricesRk , we defer them181 to future work.182 4 Illustrative examples183 In this section , we illustrate the behavior of RotoGrad in two synthetic scenarios , providing clean184 qualitative results about its effect on the optimization process . Appendix C.1 provides a detailed185 description of the experimental setups.186 To this end , we propose two different multi-task regression problems of the form187 L ( x ) = L1 ( x ) + L2 ( x ) = ϕ ( R1f ( x ; θ ) , 0 ) + ϕ ( R2f ( x ; θ ) , 1 ) , ( 5 ) where ϕ is a test function with a single global optimum whose position is parametrized by the second188 argument , that is , both tasks are identical ( and thus related ) up to a translation . We use a single input189 2For example , Geotorch [ 4 ] makes this transparent to the user . x ∈ R2 and drop task-specific network parameters . As backbone , we take a simple network of the190 form z = W2 max ( W1x+ b1 , 0 ) + b2 with b1 ∈ R10 , b2 ∈ R2 , andW1 , W > 2 ∈ R10×2.191 For the first experiment we choose a simple ( avocado-shaped ) convex objective function and , for192 the second one , we opt for a non-convex function with several local optima and a single global193 optimum . Figure 1 shows the training trajectories in the presence ( and absence ) of RotoGrad in both194 experiments , depicted as level plots in the space of z and rk , respectively . We can observe that in195 the first experiment ( Figure 1a ) , RotoGrad finds both optima—which is in stark contrast to the vanilla196 case—by rotating the feature space and matching the ( unique ) local optima of the tasks . Similarly,197 the second experiment ( Figure 1b ) shows that , as we have two symmetric tasks and a non-equidistant198 starting point , in the vanilla case the optimization is dominated by the task with an optimum closest to199 the starting point . RotoGrad avoids this behavior by equalizing gradients and , by aligning gradients,200 is able to find the optima of both functions.201 5 Related Work202 Understanding and improving the interaction between tasks is one of the most fundamental problems203 of MTL , since any improvement in this regard would translate to all MTL systems . Consequently,204 several approaches to address this problem have been adopted in the literature . Among the different205 lines of work , the one most related to the present work is gradient homogenization.206 Gradient homogenization . Since the problem is two-fold , there are two main lines of work . On207 the one hand , we have task-weighting approaches that focus on alleviating magnitude differences.208 Similar to us , GradNorm [ 6 ] attempts to learn all tasks at a similar rate , yet they propose to learn209 these weights as parameters . Instead , we provide a closed-form solution in Equation 1 , and so does210 IMTL-G [ 18 ] . However , IMTL-G scales all task gradients such that all projections ofG ontoGk are211 equal . MGDA [ 28 ] , instead , adopts an iterative method based on the Frank-Wolfe algorithm in order212 to find the set of weights { ωk } ( with ∑ k ωk = 1 ) such that ∑ k ωkGk has minimum norm . On the213 other hand , recent works have started to put attention on the conflicting direction problem . Maninis214 et al . [ 22 ] first proposed adversarial training to make task gradients statistically indistinguishable215 as part of a bigger image-tailored architecture . More recently , PCGrad [ 33 ] proposed to drop the216 projection of one task gradient onto another if they are in conflict , whereas GradDrop [ 7 ] randomly217 drops elements of the task gradients based on a sign-purity score.218 In the literature , we can also find other approaches which , while orthogonal to the gradient homoge-219 nization , are complementary to our work and thus could be used along with RotoGrad . Next , we220 provide a brief overview of them.221 A prominent approach for MTL is task clustering , that is , selecting which tasks should be learned222 together . This approach dates back to the original task-clustering algorithm [ 31 ] , but new work in223 this direction keeps coming out [ 29 , 35 ] . Alternative approaches , for example , scale the loss of each224 task differently based on different criteria such as task uncertainty [ 14 ] , task prioritization [ 11 ] , or225 similar loss magnitudes [ 18 ] . Moreover , while most models fall into the hard-parameter sharing226 umbrella , there exists other architectures in the literature . Soft-parameter sharing architectures [ 27 ] ,227 for example , do not have shared parameters but instead impose some kind of shared restrictions to the228 entire set of parameters . An interesting approach consists in letting the model itself learn which parts229 of the architecture should be used for each of the tasks [ 12 , 23 , 30 , 32 ] . Other architectures , such230 as MTAN [ 19 ] , make use of task-specific attention to select relevant features for each task . Finally,231 problems triggered by the differences between task gradients ( in magnitude and direction ) have also232 been studied in other domains like meta-learning [ 34 ] and continual learning [ 21 ] .233 6 Experiments234 In this section we assess the performance of RotoGrad on a wide range of datasets and MTL235 architectures . First , we check the effect of the learning rates of the rotation and network updates236 on the stability of the learning process of RotoGrad . Then , with the goal of applying RotoGrad237 in scenarios with extremely large sizes of z , we explore the effect of rotating a subspace of z238 instead of the whole shared representation . Finally , we compare our approach with competing MTL239 solutions in the literature , showing that RotoGrad consistently outperforms all other methods . Refer240 to Appendix C for a more detailed description of the experimental setups and additional results.241 Relative task improvement . Since MTL uses different metrics for different tasks , throughout this242 section we group results by means of the relative task improvement , first introduced in [ 22 ] . Given a243 task k , and the metrics obtained during test time by our model , Mk , and by a baseline model , Sk,244 which consists of K networks trained on each task individually , the relative task improvement for the245 k-th task is defined as246 ∆k : = 100 · ( −1 ) lk Mk − Sk Sk , ( 6 ) where lk = 1 if Mk < Sk means that our model performs better than the baseline in the k-th task , and247 lk = 0 otherwise . We depict our results using different statistics of ∆k such as its mean ( avgk ∆k ) ,248 maximum ( maxk ∆k ) , and median ( medk ∆k ) across tasks.249 6.1 Training stability250 At the end of Section 3.2 we discussed that , by casting251 problem 4 as a Stackelberg game , we have convergence252 guarantees when the rotation optimizer is the slow-learner.253 Next , we empirically show this necessary condition.254 Experimental setup . Similar to [ 28 ] , we use a multi-task255 version of MNIST [ 16 ] where each image is composed256 of a left and right digit , and use as backbone a reduced257 version of LeNet [ 17 ] with light-weight heads . Besides258 the left- and right-digit classification proposed in [ 28 ] , we259 consider three other quantities to predict : i ) sum of digits ; 260 ii ) parity of the digit product ; and iii ) number of active261 pixels . The idea here is to enforce all digit-related tasks262 to cooperate ( positive transfer ) , while the ( orthogonal ) image-related task should not disrupt these263 learning dynamics . We use negative cross-entropy and accuracy for the left- and right-digit tasks,264 binary cross-entropy and f1-score for the parity task , and mean squared error ( MSE ) as loss and265 metric for both regression tasks.266 Results . Figure 3 shows the effect averaged over ten independent runs—in terms of test error in the267 sum task , while the rest of tasks are shown in Appendix C.2—of changing the rotations ’ learning rate.268 We can observe that , the bigger the learning rate is in comparison to that of the network ’ s parameters269 ( 1e−3 ) , the higher and more noisy the test error becomes . MSE keeps decreasing as we lower the270 learning rate , reaching a sweet-spot at half the network ’ s learning rate ( 5e−4 ) . For smaller values,271 the rotations ’ learning is too slow and results start to resemble those of the vanilla case , in which no272 rotations are applied ( leftmost box in Figure 3 ) .273 6.2 Rotating a subspace274 Next , we evaluate the effect of subspace rotations as described at the end of Section 3.4 , assessing the275 trade-off between avoiding negative transfer and size of the subspace considered by RotoGrad.276 Experimental setup . We test RotoGrad on a 10-task classification problem on CIFAR10 [ 15 ] , using277 binary cross-entropy and f1-score as loss and metric , respectively , for all tasks . We use ResNet18 [ 13 ] 278 without pre-training as backbone ( d = 512 ) , and linear layers with sigmoids as task-specific heads.279 Results are summarized at the bottom part of Table 1 . We can observe that rotating the entire space280 provides the best results , and they worsen as we decrease the size ofRk . However , rotating only 64281 features ( 12.5 % of the shared feature space ) still yields better results than vanilla optimization.282 6.3 Methods comparison283 We now proceed to compare RotoGrad with the different existing approaches to gradient conflict ( for284 both magnitude and direction ) in different real-world datasets , showing how RotoGrad outperforms285 existing methods while being on par with existing methods in training time.286 Experimental setup . In order to provide fair comparisons among methods , all experiments use287 identical configurations and random initializations . For all methods we performed a hyper-parameter288 search and chose the best ones based on validation error . Our results are reported using the median and289 standard deviation computed over 5-10 random seeds . Further details can be found in Appendix C.1.290 MNIST and SVHN . We reuse291 the experimental setting from292 Section 6.1—now with multi-293 task versions of MNIST [ 16 ] and294 SVHN [ 26 ] —in order to evalu-295 ate how disruptive the orthogonal296 image-related task is for differ-297 ent methods . We can observe in298 the results from Table 2 that the299 effect of the image-related task300 is more disruptive in MNIST,301 in which MGDA utterly fails.302 Direction-aware methods ( Grad-303 Drop and PCGrad ) do not im-304 prove the vanilla results , whereas305 IMTL-G , GradNorm , and RotoGrad obtain the best results.306 CIFAR10 . We reuse the setting in Section 6.2 and compare the different MTL methods using five307 different seeds . Results are shown in Table 1 and Figure 4 . Unlike the previous setting , scaling308 gradients is not enough to solve the problem . Among existing methods , both direction-aware solutions309 ( PCGrad and GradDrop ) improve over the vanilla case on all the statistics , whereas most magnitude-310 aware solutions substantially worsen task performance . In stark contrast , RotoGrad improves task311 performance across all ten tasks , as it can be observed both in Table 1 and Figure 4 . To further show312 that this is a consequence of gradient homogenization in terms of both magnitudes and directions , we313 introduced an extra-baseline , IMTL-G+Rk , which applies IMTL-G to the extended MTL architecture314 ( Figure 2 ) , that is , with matrixRk optimizing the k-th task loss ( instead of Equation 3 ) .315 NYUv2 . Now , we test all methods using NYUv2 [ 9 ] on three different tasks : 13-class semantic316 segmentation ; depth estimation ; and surface normals . To speed up training , all images were resized317 to 288 × 384 resolution ; and data augmentation was applied to alleviate overfitting . As MTL318 architecture , we use SegNet [ 2 ] where the decoder is splitted into three convolutational heads . We use319 the same setup as Liu et al . [ 19 ] . Like in previous experiments , we observe in Table 3 that RotoGrad320 results in a consistent improvement over all tasks with respect to the vanilla case . MGDA obtains321 the best results in surface normals at the expense of overlooking the other tasks , while GradDrop322 worsens all results and PCGrad obtains minor improvements in all tasks . GradNorm finds a trade-off323 solution instead , improving results in depth estimation and surface normals , yet with worse results in324 semantic segmentation . RotoGrad obtains the best results followed by IMTL-G and , more importantly,325 RotoGrad is the only method resulting in a average positive task improvement—across the three326 tasks—over training three single-task models independently . It is worth mentioning that , with only327 three tasks , all methods trained in less than 4 hours ; and that this result consolidates RotoGrad ’ s328 scalability , as we only rotate the first 1024 dimensions of z , out of a total of 7 millions.329 CelebA . Last , we apply all methods to a 40-class multi-classification problem in CelebA [ 20 ] on330 two different settings : one using a convolutional network as backbone ( d = 512 ) ; and another using331 ResNet18 [ 13 ] as backbone ( d = 2048 ) . Similar to CIFAR10 , we use binary cross-entropy and332 f1-score as loss and metric for all tasks . Even though we face two completely different architectures,333 results in Table 4 show that RotoGrad convincingly outperforms all competing methods in all f1-score334 statistics , independently of the model . Furthermore , since this is a computationally demanding task335 with 40 tasks—in fact , we omit MGDA as it takes several days to train—we also compare methods in336 terms of training time . On the one hand , GradDrop and IMTL-G produce little overhead compared337 with the vanilla case , as expected . On the other hand , GradNorm and PCGrad take , respectively , 2.5338 and 4 times longer to train than the vanilla setting . More importantly , RotoGrad outperforms existing339 methods while staying on par with them in training time , rotating 50 % and 75 % of the shared340 feature z for the convolutional and residual backbones , respectively , which further demonstrates that341 RotoGrad can scale-up to real-world settings.342 7 Conclusions343 In this work , we have introduced RotoGrad , an algorithm that tackles negative transfer in MTL by344 homogenizing task gradients in terms of both magnitudes and directions . RotoGrad enforces a similar345 convergence rate for all tasks , while at the same time smoothly rotates the shared representation346 differently for each task in order to avoid conflicting gradients . As a result , RotoGrad leads to347 stable and accurate MTL . Our empirical results have shown the effectiveness of RotoGrad in many348 scenarios , staying on top of all competing methods in performance , while being on par in terms of349 computational complexity with those that better scale to complex networks.350 We believe our work opens up interesting venues for future work . For example , it would be interesting351 to study alternative approaches to further scale up RotoGrad using , for example , diagonal-block or352 sparse rotation matrices ; to rotate the feature space in application domains with structured features353 ( e.g. , channel-wise rotations in images ) ; and to combine different methods , for example , by scaling354 gradients using the direction-awareness of IMTL-G and the “ favor slow-learners ” policy of RotoGrad.355 References356 [ 1 ] Pierre-Antoine Absil , Robert E. Mahony , and Rodolphe Sepulchre . Optimization Algorithms357 on Matrix Manifolds . Princeton University Press , 2008.358 [ 2 ] Vijay Badrinarayanan , Alex Kendall , and Roberto Cipolla . “ SegNet : A Deep Convolutional359 Encoder-Decoder Architecture for Image Segmentation. ” In : IEEE Trans . Pattern Anal . Mach.360 Intell . 39.12 ( 2017 ) , pp . 2481–2495 . DOI : 10.1109/TPAMI.2016.2644615.361 [ 3 ] Rich Caruana . “ Multitask Learning : A Knowledge-Based Source of Inductive Bias. ” In:362 Machine Learning , Proceedings of the Tenth International Conference , University of Mas-363 sachusetts , Amherst , MA , USA , June 27-29 , 1993 . Ed . by Paul E. Utgoff . Morgan Kaufmann,364 1993 , pp . 41–48 . DOI : 10.1016/b978-1-55860-307-3.50012-5.365 [ 4 ] Mario Lezcano Casado . “ Trivializations for Gradient-Based Optimization on Manifolds. ” 366 In : Advances in Neural Information Processing Systems 32 : Annual Conference on Neural367 Information Processing Systems 2019 , NeurIPS 2019 , December 8-14 , 2019 , Vancouver , BC,368 Canada . Ed . by Hanna M. Wallach et al . 2019 , pp . 9154–9164.369 [ 5 ] Mario Lezcano Casado and David Martínez-Rubio . “ Cheap Orthogonal Constraints in Neural370 Networks : A Simple Parametrization of the Orthogonal and Unitary Group. ” In : Proceedings of371 Machine Learning Research 97 ( 2019 ) . Ed . by Kamalika Chaudhuri and Ruslan Salakhutdinov,372 pp . 3794–3803.373 [ 6 ] Zhao Chen et al . “ GradNorm : Gradient Normalization for Adaptive Loss Balancing in Deep374 Multitask Networks. ” In : Proceedings of the 35th International Conference on Machine375 Learning , ICML 2018 , Stockholmsmässan , Stockholm , Sweden , July 10-15 , 2018 . Ed . by376 Jennifer G. Dy and Andreas Krause . Vol . 80 . Proceedings of Machine Learning Research.377 PMLR , 2018 , pp . 793–802.378 [ 7 ] Zhao Chen et al . “ Just Pick a Sign : Optimizing Deep Multitask Models with Gradient Sign379 Dropout. ” In : Advances in Neural Information Processing Systems 33 : Annual Conference on380 Neural Information Processing Systems 2020 , NeurIPS 2020 , December 6-12 , 2020 , virtual.381 Ed . by Hugo Larochelle et al . 2020.382 [ 8 ] Jorge Cortés . “ Finite-time convergent gradient flows with applications to network consensus. ” 383 In : Autom . 42.11 ( 2006 ) , pp . 1993–2000 . DOI : 10.1016/j.automatica.2006.06.015.384 [ 9 ] Camille Couprie et al . “ Indoor Semantic Segmentation using depth information. ” In : 1st385 International Conference on Learning Representations , ICLR 2013 , Scottsdale , Arizona , USA,386 May 2-4 , 2013 , Conference Track Proceedings . Ed . by Yoshua Bengio and Yann LeCun . 2013.387 [ 10 ] Tanner Fiez , Benjamin Chasnov , and Lillian J. Ratliff . “ Convergence of Learning Dynamics in388 Stackelberg Games. ” In : CoRR abs/1906.01217 ( 2019 ) .389 [ 11 ] Michelle Guo et al . “ Dynamic Task Prioritization for Multitask Learning. ” In : Computer390 Vision - ECCV 2018 - 15th European Conference , Munich , Germany , September 8-14 , 2018,391 Proceedings , Part XVI . Ed . by Vittorio Ferrari et al . Vol . 11220 . Lecture Notes in Computer392 Science . Springer , 2018 , pp . 282–299 . DOI : 10.1007/978-3-030-01270-0\_17.393 [ 12 ] Pengsheng Guo , Chen-Yu Lee , and Daniel Ulbricht . “ Learning to Branch for Multi-Task394 Learning. ” In : Proceedings of the 37th International Conference on Machine Learning , ICML395 2020 , 13-18 July 2020 , Virtual Event . Vol . 119 . Proceedings of Machine Learning Research.396 PMLR , 2020 , pp . 3854–3863.397 [ 13 ] Kaiming He et al . “ Deep Residual Learning for Image Recognition. ” In : 2016 IEEE Conference398 on Computer Vision and Pattern Recognition , CVPR 2016 , Las Vegas , NV , USA , June 27-30,399 2016 . IEEE Computer Society , 2016 , pp . 770–778 . DOI : 10.1109/CVPR.2016.90.400 [ 14 ] Alex Kendall , Yarin Gal , and Roberto Cipolla . “ Multi-task learning using uncertainty to weigh401 losses for scene geometry and semantics. ” In : Proceedings of the IEEE conference on computer402 vision and pattern recognition . 2018 , pp . 7482–7491.403 [ 15 ] Alex Krizhevsky , Geoffrey Hinton , et al . “ Learning multiple layers of features from tiny404 images. ” In : ( 2009 ) .405 [ 16 ] Yann LeCun , Corinna Cortes , and CJ Burges . “ MNIST handwritten digit database. ” In : ATT406 Labs [ Online ] 2 ( 2010 ) .407 [ 17 ] Yann LeCun et al . “ Gradient-based learning applied to document recognition. ” In : Proceedings408 of the IEEE 86.11 ( 1998 ) , pp . 2278–2324.409 [ 18 ] Liyang Liu et al . “ Towards Impartial Multi-task Learning. ” In : International Conference on410 Learning Representations . 2021.411 [ 19 ] Shikun Liu , Edward Johns , and Andrew J. Davison . “ End-To-End Multi-Task Learning With412 Attention. ” In : IEEE Conference on Computer Vision and Pattern Recognition , CVPR 2019,413 Long Beach , CA , USA , June 16-20 , 2019 . Computer Vision Foundation / IEEE , 2019 , pp . 1871–414 1880 . DOI : 10.1109/CVPR.2019.00197.415 [ 20 ] Ziwei Liu et al . “ Deep Learning Face Attributes in the Wild. ” In : 2015 IEEE International416 Conference on Computer Vision , ICCV 2015 , Santiago , Chile , December 7-13 , 2015 . IEEE417 Computer Society , 2015 , pp . 3730–3738 . DOI : 10.1109/ICCV.2015.425.418 [ 21 ] David Lopez-Paz and Marc ’ Aurelio Ranzato . “ Gradient Episodic Memory for Continual419 Learning. ” In : Advances in Neural Information Processing Systems 30 : Annual Conference on420 Neural Information Processing Systems 2017 , December 4-9 , 2017 , Long Beach , CA , USA.421 Ed . by Isabelle Guyon et al . 2017 , pp . 6467–6476.422 [ 22 ] Kevis-Kokitsi Maninis , Ilija Radosavovic , and Iasonas Kokkinos . “ Attentive Single-Tasking of423 Multiple Tasks. ” In : IEEE Conference on Computer Vision and Pattern Recognition , CVPR424 2019 , Long Beach , CA , USA , June 16-20 , 2019 . Computer Vision Foundation / IEEE , 2019,425 pp . 1851–1860 . DOI : 10.1109/CVPR.2019.00195.426 [ 23 ] Ishan Misra et al . “ Cross-Stitch Networks for Multi-task Learning. ” In : 2016 IEEE Conference427 on Computer Vision and Pattern Recognition , CVPR 2016 , Las Vegas , NV , USA , June 27-30,428 2016 . IEEE Computer Society , 2016 , pp . 3994–4003 . DOI : 10.1109/CVPR.2016.433.429 [ 24 ] Ryan W. Murray , Brian Swenson , and Soummya Kar . “ Revisiting Normalized Gradient430 Descent : Fast Evasion of Saddle Points. ” In : IEEE Trans . Autom . Control . 64.11 ( 2019 ) ,431 pp . 4818–4824 . DOI : 10.1109/TAC.2019.2914998.432 [ 25 ] Yurii E. Nesterov . Introductory Lectures on Convex Optimization - A Basic Course . Vol . 87.433 Applied Optimization . Springer , 2004 . DOI : 10.1007/978-1-4419-8853-9.434 [ 26 ] Yuval Netzer et al . “ Reading digits in natural images with unsupervised feature learning. ” In:435 NeurIPS Workshop on Deep Learning and Unsupervised Feature Learning ( 2011 ) .436 [ 27 ] Sebastian Ruder . “ An Overview of Multi-Task Learning in Deep Neural Networks. ” In : CoRR437 abs/1706.05098 ( 2017 ) .438 [ 28 ] Ozan Sener and Vladlen Koltun . “ Multi-Task Learning as Multi-Objective Optimization. ” 439 In : Advances in Neural Information Processing Systems 31 : Annual Conference on Neural440 Information Processing Systems 2018 , NeurIPS 2018 , December 3-8 , 2018 , Montréal , Canada.441 Ed . by Samy Bengio et al . 2018 , pp . 525–536.442 [ 29 ] Trevor Standley et al . “ Which Tasks Should Be Learned Together in Multi-task Learning ? ” In:443 Proceedings of the 37th International Conference on Machine Learning , ICML 2020 , 13-18444 July 2020 , Virtual Event . Vol . 119 . Proceedings of Machine Learning Research . PMLR , 2020,445 pp . 9120–9132.446 [ 30 ] Ximeng Sun et al . “ AdaShare : Learning What To Share For Efficient Deep Multi-Task Learn-447 ing. ” In : Advances in Neural Information Processing Systems 33 : Annual Conference on Neural448 Information Processing Systems 2020 , NeurIPS 2020 , December 6-12 , 2020 , virtual . Ed . by449 Hugo Larochelle et al . 2020.450 [ 31 ] Sebastian Thrun and Joseph O ’ Sullivan . “ Discovering Structure in Multiple Learning Tasks:451 The TC Algorithm. ” In : Machine Learning , Proceedings of the Thirteenth International452 Conference ( ICML ’ 96 ) , Bari , Italy , July 3-6 , 1996 . Ed . by Lorenza Saitta . Morgan Kaufmann,453 1996 , pp . 489–497.454 [ 32 ] Simon Vandenhende et al . “ Branched Multi-Task Networks : Deciding what layers to share. ” 455 In : 31st British Machine Vision Conference 2020 , BMVC 2020 , Virtual Event , UK , September456 7-10 , 2020 . BMVA Press , 2020.457 [ 33 ] Tianhe Yu et al . “ Gradient Surgery for Multi-Task Learning. ” In : Advances in Neural Informa-458 tion Processing Systems . Ed . by H. Larochelle et al . Vol . 33 . Curran Associates , Inc. , 2020,459 pp . 5824–5836.460 [ 34 ] Tianhe Yu et al . “ Meta-World : A Benchmark and Evaluation for Multi-Task and Meta Re-461 inforcement Learning. ” In : 3rd Annual Conference on Robot Learning , CoRL 2019 , Osaka,462 Japan , October 30 - November 1 , 2019 , Proceedings . Ed . by Leslie Pack Kaelbling , Danica463 Kragic , and Komei Sugiura . Vol . 100 . Proceedings of Machine Learning Research . PMLR,464 2019 , pp . 1094–1100.465 [ 35 ] Amir Roshan Zamir et al . “ Taskonomy : Disentangling Task Transfer Learning. ” In : 2018 IEEE466 Conference on Computer Vision and Pattern Recognition , CVPR 2018 , Salt Lake City , UT,467 USA , June 18-22 , 2018 . IEEE Computer Society , 2018 , pp . 3712–3722 . DOI : 10.1109/CVPR.468 2018.00391.469 Checklist470 1 . For all authors ... 471 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s472 contributions and scope ? [ Yes ] RotoGrad homogenizes both magnitudes ( Equation 1 ) 473 and direction ( Equation 4 ) , and empirical results in Section 6 and Appendix C demon-474 strate our claims.475 ( b ) Did you describe the limitations of your work ? [ Yes ] In Section 6.1.476 ( c ) Did you discuss any potential negative societal impacts of your work ? [ N/A ] 477 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to478 them ? [ Yes ] 479 2 . If you are including theoretical results ... 480 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] Proposi-481 tion 3.1 ’ s assumptions are stated in Appendix A , and those regarding RotoGrad ’ s482 stability appear in Appendix B.483 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] Proof of Proposi-484 tion 3.1 appears in Appendix A . For the proofs related to Stackelberg games , which are485 not a direct contribution of our paper , please refer to [ 10 ] .486 3 . If you ran experiments ... 487 ( a ) Did you include the code , data , and instructions needed to reproduce the main exper-488 imental results ( either in the supplemental material or as a URL ) ? [ Yes ] We provide489 instructions and code to reproduce our experiments in the supplemental material.490 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they491 were chosen ) ? [ Yes ] In Appendix C.1.492 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running ex-493 periments multiple times ) ? [ Yes ] We provide statistics for most of our experiments494 computed over 5-10 independent runs . Due to time complexity required by larger495 datasets on NYUv2 and CelebA , we only report the results for a single random seed,496 but still compare the different methods using several performance metrics.497 ( d ) Did you include the total amount of compute and the type of resources used ( e.g.,498 type of GPUs , internal cluster , or cloud provider ) ? [ Yes ] All details are provided in499 Appendix C.1.500 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 501 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] 502 ( b ) Did you mention the license of the assets ? [ Yes ] We only use code from previous503 research and licence MIT , which we inherit and acknowledge in our extended version504 of the code.505 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 506 We release the code implementation to reproduce our experiments together with the507 supplementary material , and will make it publicly available after the paper acceptance.508 ( d ) Did you discuss whether and how consent was obtained from people whose data509 you ’ re using/curating ? [ N/A ] We only use publicly available datasets with no personal510 information . Moreover , our experiments only report statistics on the results.511 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable512 information or offensive content ? [ N/A ] We only use publicly available and broadly513 used image datasets.514 5 . If you used crowdsourcing or conducted research with human subjects ... 515 ( a ) Did you include the full text of instructions given to participants and screenshots , if516 applicable ? [ N/A ] 517 ( b ) Did you describe any potential participant risks , with links to Institutional Review518 Board ( IRB ) approvals , if applicable ? [ N/A ] 519 ( c ) Did you include the estimated hourly wage paid to participants and the total amount520 spent on participant compensation ? [ N/A ] 521

Summary: The work proposes RotoGrad, a multitask layer that rotates a shared feature space in a deep neural network to align gradients and mitigate negative transfer effects. The proposed method works within models with multiple prediction heads and a shared backbone, and deals both with magnitude and directional imbalance between task gradients. Authors provide multiple experiments in the computer vision space to show efficacy of the proposed method against other multitask baselines.