Adapting Auxiliary Losses Using Gradient Similarity

Introduction

Neural networks are powerful function approximators that have excelled on a wide range of tasks (Simonyan and Zisserman, 2015; Mnih et al., 2015; He et al., 2016a; Silver et al., 2016; Vaswani et al., 2017). Despite the state of the art results across domains, they remain data-inefficient and expensive to train. In supervised learning, large deep learning benchmarks with millions of examples are needed for training (Russakovsky et al., 2015) and the additional implication of requiring human intervention to label a large dataset can be prohibitively expensive. In reinforcement learning (RL), agents typically consume millions of frames of experiences before learning to act reasonably in complex environments (Silver et al., 2016; Espeholt et al., 2018), which not only puts pressure on computing power but also makes particular domains (e.g., robotics) impractical.

Different techniques have been studied for improving data efficiency, from data augmentation (Krizhevsky et al., 2012; Simonyan and Zisserman, 2015; Hauberg et al., 2016) to multi-task learning (Teh et al., 2017; Kendall et al., 2018; Chen et al., 2018; Hessel et al., 2018; Sener and Koltun, 2018) to transfer learning (Taylor and Stone, 2009; Pan et al., 2010). In this work, we focus on a particular type of transfer learning: transferring knowledge of an auxiliary task to a main task. We assume that besides the main task, one has access to one or more additional auxiliary tasks that share some unknown structures with the main task. To improve data efficiency, these additional tasks can be used as auxiliary losses to transfer knowledge to the main task. Note that only the performance of the main task is of interest, even though the model is trained simultaneously on all tasks; any improvement on the auxiliary losses is useful only to the extent that it helps learning features or behaviors for the main task.

Auxiliary tasks have been shown to work well for reinforcement learning (e.g., Zhang et al., 2016; Jaderberg et al., 2017; Mirowski et al., 2017; Papoudakis et al., 2018; Li et al., 2019). The success of these approaches depend on how well aligned the auxiliary losses are with the main task. Knowing this apriori is typically non-trivial and the usefulness of an auxiliary task can change through the course of training. In this work, we explore a simple yet effective heuristic for measuring the similarity between an auxiliary task and the main task of interest (given the value of their parameters) using gradient information.

Notation and Problem Description

Assume we have a main task $\mathcal{T}_{{main}}$ and some update rule induced by an auxiliary task $\mathcal{T}_{{aux}}$. This can be given by the gradients of an auxiliary loss $\mathcal{L}_{{aux}}$ (though our derivation holds also when there is no such loss, as will be discussed later). Let us assume the main task induces a loss $\mathcal{L}_{{main}}$. We care only about maximizing performance on $\mathcal{T}_{{main}}$ while $\mathcal{T}_{{aux}}$ is an auxiliary task which is not of direct interest. The goal is to devise an algorithm that/ can automatically:

leverage $\mathcal{T}_{{aux}}$ as long as it is helpful to $\mathcal{T}_{{main}}$,

detect when $\mathcal{T}_{{aux}}$ becomes harmful to $\mathcal{T}_{{main}}$ and block negative transfer to recover the performance of training only on $\mathcal{T}_{{main}}$.

We propose to parameterize the solution for $\mathcal{T}_{{main}}$ and $\mathcal{T}_{{aux}}$ by two neural networks, $f(\cdot,{\bm{\theta}},{\bm{\phi}}_{{main}})$ and $g(\cdot,{\bm{\theta}},{\bm{\phi}}_{{aux}})$, where they share a subset of parameters (denoted by ${\bm{\theta}}$) and have their own set of parameters (denoted by ${\bm{\phi}}_{{main}}$ and ${\bm{\phi}}_{{aux}}$ respectively). Generally, the auxiliary loss literature proposes to minimize weighted losses of the form

under the intuition that modifying ${\bm{\theta}}$ to minimize $\mathcal{L}_{{aux}}$ will improve $\mathcal{L}_{{main}}$ if the two tasks are sufficiently related. There are two potential problems with minimizing equation 1: (i) $\lambda$ is usually kept as a constant which is not adaptable where the usefulness of $\mathcal{T}_{{aux}}$ changes during training (e.g., helps initially but hurts later), and (ii) the relatedness of the two tasks are often ambiguous to quantify.

Ideally one would want to modulate the weight $\lambda$ at each learning iteration $t$ such that either final performance on main task is maximized or convergence speed of the main task is maximized.Note that both quantities might be ill-defined without considering a fixed budget of iterations. Also, while they are correlated, the two maximization problems can lead to different behaviors.

A greedy variant of this objective is to modulate $\lambda$ at each learning iteration $t$ by how useful $\mathcal{T}_{{aux}}$ is for $\mathcal{T}_{{main}}$ given ${\bm{\theta}}^{(t)},{\bm{\phi}}_{{main}}^{(t)},{\bm{\phi}}_{{aux}}^{(t)}$. That is, at each optimization iteration, we want to efficiently approximate the solution to

Note that the input space of $\mathcal{T}_{{main}}$ and $\mathcal{T}_{{aux}}$ do not have to match. In particular, $\mathcal{T}_{{aux}}$ does not need to be defined for an input of $\mathcal{T}_{{main}}$ or the other way around.In the supervised learning case when the input features are shared, this resembles Quadrianto et al. (2010), multi-task learning without label correspondences setting. Solving the general problem, or even the greedy variant given by equation 2 is expensive, as it requires at each step to solve a non-linear maximization problem. Instead, we look for a cheap heuristic to approximate $\lambda^{(t)}$ which does not require hyperparameter tuning or computing derivatives through the learning process—and which will outperform a constant $\lambda^{(t)}$.

Note that our setup is in sharp contrast to multi-objective optimization where both the tasks are of interest. We are only interested in minimizing the $\mathcal{L}_{{main}}$ with the help of the $\mathcal{L}_{{aux}}$ where possible.

Gradient Cosine Similarity

We propose to use the cosine similarity of gradients between tasks to quantify the relatedness between tasks and hence for approximating $\lambda^{(t)}$. Consider the example in Figure 1 where the main function to minimize is $\mathcal{L}_{{main}}=(\theta-10)^{2}$ and the auxiliary function is $\mathcal{L}_{{aux}}=\theta^{2}$, their gradients are $\nabla_{\theta}\mathcal{L}_{{main}}=2(\theta-10)$ and $\nabla_{\theta}\mathcal{L}_{{aux}}=2\theta$ respectively. When $\theta=-20$, the gradients of the main and auxiliary functions point in the same direction and the cosine similarity is $1$; minimizing the auxiliary loss is beneficial for minimizing the main loss. However, at a different point, $\theta=5$, the two gradients point in different directions and the cosine similarity is $-1$; minimizing the auxiliary loss would hinder minimizing the main loss. This example suggests a natural strategy for approximating $\lambda^{(t)}$: minimize the auxiliary loss as long as its gradient has non-negative cosine similarity with the main gradient; otherwise, the auxiliary loss should be ignored. This follows the well-known intuition that if a vector is in the same half-space as the gradient of a function $f$, then it is a descent direction for $f$. This reduces our strategy to ask if the gradient of the auxiliary loss is also a descent direction for the main loss of interest.

We start by first introducing two propositions in this section and prove that our proposed method ensures convergence on the main task—although one can not guarantee the speed of convergence (we discuss this in more details later in this section). The proof holds even when adding to the gradient vectors that might not be the gradient of any function, which we argue is a practical scenario in reinforcement learning. Later in Section 4, we show empirically that our method is a good heuristic for blocking negative transfer and in practice often leads to positive transfer.

Given any gradient vector field $G({\bm{\theta}})=\nabla_{\bm{\theta}}\mathcal{L}({\bm{\theta}})$ and any vector field $V({\bm{\theta}})$ (such as the gradient of another loss function, or an arbitrary set of updates), an update rule of the form

converges to the local minimum of $\mathcal{L}$ given small enough $\alpha^{(t)}$.

We point out three important properties of the above statement. First, it guarantees only lack of divergence, but does not guarantee any improvement of convergence. That is, cosine similarity is not a silver bullet that guarantees positive transfer, but it can drop the “worst-case scenarios” thus preventing negative transfer. In principle, the convergence speed of the main loss could be affected both positively or negatively, though we notice empirically that the effect tends to be positive in realistic scenarios. Second, it is worth noting that simply adding an arbitrary vector field to $\nabla\mathcal{L}$ does not have the convergence property and hence Proposition 1 has practical utility—there are realistic scenarios where one does not always form a gradient field. For example, the update rule of the Q-learning algorithm in RL does not form a gradient field. Another example is when adding another gradient field to the update rule of REINFORCE (which is a Monte-Carlo estimate of a gradient), depending on the setting, the sum might not be a gradient field anymore. Lastly, it should be noted that gradient similarity is a local quantity on which we can not technically rely to draw global conclusions on learning dynamics. That is, having negative cosine similarity at step $t$ does not guarantee divergence of the main loss, nor does it guarantee that using the auxiliary loss (with fixed weight) will lead to slower convergence than not having it at all.

We illustrate both positive and negative scenarios of our proposed method by running a set of experiments using steepest descent method. There are two main losses to be minimized: $L_{1}(\theta_{1},\theta_{2})=\theta_{1}^{2}+\theta_{2}^{2}$ and $L_{2}(\theta_{1},\theta_{2})=(\theta_{1}<0)(\theta_{1}^{2}+\theta_{2}^{2})+(\theta_{1}>0)(1-\exp(-2(\theta_{1}^{2}+\theta_{2}^{2})))$. Two auxiliary losses are: $L_{3}(\theta_{1},\theta_{2})=(\theta_{1}-1)^{2}+(\theta_{2}-1)^{2}$ and $L_{4}(\theta_{1},\theta_{2})=(\theta_{1}-2)^{2}+(\theta_{2}-0.5)^{2}$. We also consider an arbitrary vector field $V(\theta_{1},\theta_{2})=[-\tfrac{\theta_{2}}{\theta_{1}^{2}+\theta_{2}^{2}}-2\theta_{1},\tfrac{\theta_{1}}{\theta_{1}^{2}+\theta_{2}^{2}}-2\theta_{2}]$. Each experiment runs for $600$ iterations with a constant step size of $0.01$. Convergence time is defined as the number of steps needed to get below value $0.1$ of the main loss. Each colored curve shows an example trajectory.

Shown in Figure 2, the first column depicts the loss when following the update given by $\nabla L_{1}$ and $\nabla L_{2}$ respectively. The second column shows the performance on the main loss when using an update that sums between the gradient of the main loss and the vector fields produced by $\nabla L_{3},V,$ and $\nabla L_{4}$ respectively. The last column shows the performance on the main loss when using our proposed cosine similarity method. Figure 2 shows our method ensures convergence in all cases. In particular, the top two rows show the positive case where our method speeds up learning, even when combined with $V$. Observe that simply adding $\nabla L_{1}+V$ is not a gradient vector field and following it leads to divergence of $L_{1}$. The last row shows a negative case where our method slows down convergence. Note, however, that our method still converges, while simply adding $\nabla L_{2}$ and $\nabla L_{4}$ (the second column) diverges. More details are provided in Appendix B to help intuitively understand the kind of scenarios for which our approach could help.

Proposition 1 refers to losses with the same set of parameters ${\bm{\theta}}$, while equation 2 refers to the scenario when each loss has task specific parameters (e.g. ${\bm{\phi}}_{{main}}$ and ${\bm{\phi}}_{{aux}}$). The following proposition extends to this scenario:

Given two losses parametrized with $\mathbf{\Theta}$ (some of which are shared ${\bm{\theta}}$ and some unique to each loss ${\bm{\phi}}_{{main}}$ and ${\bm{\phi}}_{{aux}}$), learning rule:

leads to convergence to local minimum of $\mathcal{L}_{{main}}$ w.r.t. $({\bm{\theta}},{\bm{\phi}}_{{main}})$ given small enough $\alpha^{(t)}$.

Comes directly from the previous proposition that $G=\nabla_{{\bm{\theta}}}\mathcal{L}_{{main}}$ and $V=\nabla_{{\bm{\theta}}}\mathcal{L}_{{aux}}$. For any vector fields $A,B,C$, we have $\langle A,B\rangle\geq 0$ and $\langle C,B\rangle\geq 0$ implies $\langle A+C,B\rangle\geq 0$. ∎

Analogous guarantees also hold for the unweighted version of this algorithm. Instead of weighting by $\cos(G,V)$, we use a binary weight $(\operatorname{sign}(\cos(G,V))+1)/2$ which is equivalent to using $V$ iff $\cos(G,V)>0$. When training with mini-batches, accurately estimating $\cos(G,V)$ can be difficult due to noise; the unweighted variant only requires $\operatorname{sign}(\cos(G,V))$ which can be estimated more robustly. Hence, we use the unweighted variant in our experiments unless otherwise specified.

Despite its simplicity, our proposed update rule can give rise to interesting phenomena. We can show that the emerging vector field could be non-conservative, which means there does not exist a loss function for which it is a gradient. While this might seem problematic (for gradient-descent-based optimizers), it describes only the global structure—typically used optimizers are local in nature and they do local, linear or quadratic approximations of the function (Shwartz-Ziv and Tishby, 2017). Consequently, in practice, one should not expect any negative effects from this phenomena, as it simply shows that our proposed technique is in fact qualitatively changing the nature of the update rules for training.

In general, the proposed update rule does not have to create a conservative vector field.

Applications

We now demonstrate how to use the gradient cosine similarity to decide when to leverage the auxiliary. We aim to answer two questions: (i) is gradient cosine similarity capable of detecting negative interference, which can be changing during the course of training between tasks, and (ii) compared with a constant weight, can our proposed heuristic block negative transfer when $\mathcal{T}_{{aux}}$ starts to hinder $\mathcal{T}_{{main}}$. All experiments (unless otherwise stated) follow the unweighted version of our method, summarized in Algorithm 1. The weighted version of our method is summarized in Algorithm 2, Appendix C.

First, we apply our method in supervised learning and design a multi-task binary classification problem on ImageNet (Russakovsky et al., 2015). We take a pair of classes from ImageNet, refer to these as class $A$ and class $B$; all the other 998 classes (except $A$ and $B$) are referred to as the $background$. Our tasks $\mathcal{T}_{{main}}$ and $\mathcal{T}_{{aux}}$ are then formed as a binary classification of if an image is class $A$ (otherwise $background$) and if an image is class $B$ (otherwise $background$) respectively.

Our goal is to show that cosine similarity can automatically detect when the auxiliary task becomes unhelpful and block negative transfer thus we consider two scenarios: (i) auxiliary task helps and (ii) auxiliary task hurts. It is natural to hypothesis that a pair of near classes would fit scenario (i) and a pair of far classes fit scenario (ii). However, there are no ground truth labels for such a class distance measure thus it is difficult to decided if two classes are similar or not. To combat this, we instead use two distance measures to serve as an estimated class similarity for selecting class pair $A$ and $B$: lowest common ancestor (LCA) is the ImageNet label hierarchy, and Frechet Inception Distance (FID) (Heusel et al., 2017) is the image embedding from a pre-trained model. Based on these measures, we picked three pair of classes for near, class $871$ (trimaran) vs. $484$ (catamaran), $250$ (Siberian husky) vs. $249$ (malamute), and $238$ (Greater Swiss Mountain dog) vs. $241$ (Entleucher); and for far, class $920$ (traffic light) vs. $62$ (rock python), $926$ (hotpot) vs. $800$ (slot), and $48$ (Komodo dragon) vs. $920$ (traffic light). See Appendix D.1 for details on class pair selection.

We train modified ResNetV2-18 model (He et al., 2016b) with all parameters in the convolutional layers shared (denote as ${\bm{\theta}}$) between tasks, followed by task-specific parameters ${\bm{\phi}}_{{main}}$ and ${\bm{\phi}}_{{aux}}$. First, we use a multi-task learning setup and minimize $\mathcal{L}_{{main}}+\mathcal{L}_{{aux}}$ (i.e., $\lambda=1$), and measure cosine similarity on ${\bm{\theta}}$ through the course of training. Figure 3(a) shows that cosine similarity is higher for near pairs (blue lines) and lower (mostly negative) for far pairs (red lines), indicating that cosine similarity captures task relatedness well. Next, we compare single-task training (i.e., $\lambda=0$), multi-task training, and our proposed method on near and far class pairs to illustrate scenario (i) and (ii) respectfully. Note that, the multi-task baseline here is not minimizing the desired objective, $\mathcal{L}_{{main}}$ and is only used for demonstrating what would happen if the auxiliary loss is weighted as a constant without any adaptions based on its usefulness.

Figure 3(b) shows scenario (i) where auxiliary task helps. While all variants perform similarly in terms of final performance, our method obtains an initial boost compared to the single task learning and thereafter achieves similar performance as the multi-task learning. This behavior indicates that when gradient cosine similarity is high, our method automatically “switches on” $\mathcal{L}_{{aux}}$ and trains as multi-tasking. Figure 3(c) shows scenario (ii) where the auxiliary task could hinder. Here, the multi-task baseline suffers from poorer performance than single-task learning throughout the training process due to the class distance is far and lack of transferability. In contrast, our method blocks negative transfer by “switching off” the unhelpful auxiliary task and later recovers the performance of single-task. We note that the blocking did not happen until later in training (at around step $15,000$). This observation potentially reflects two aspects of our method: (i) estimating the cosine similarity between gradients can be noisy and unreliable, and (ii) our method does not guarantee maximizing transfer but only “drop the worst.” Nevertheless, this experiment shows that our method can eventually notice that the auxiliary task hinders the main task and is able to block negative transfer then recover the performance of the single task learning.

2 Multi Class Classification Tasks

Second, we evaluate our method on the multi-class classification task of rotated MNIST digits, a common benchmark in multi-task continual learning (Lopez-Paz and Ranzato, 2017; Farajtabar et al., 2020) and multi-task generalization (Ghifary et al., 2015) problems. In this setting, the main task is classifying the original MNIST images to one of the 10 digit classes. The auxiliary task consists of all images rotated by a certain degree. For this task we use a simple multi-layer fully connected network with 3 hidden layers each with 100 neurons and ReLU nonlinearity on the top. Each task is then is equipped with a separate 10 neuron head for the associated 10 dimensional logit. We used all the 60K training inputs and train the network for 50 epochs with batch size set to 128. For the optimizer we used RMSprop with learning rate $0.001$.

Table 1 shows the results. We reported the mean and standard deviation of the classification error on 10k held-out test examples over 10 runs with random network initialization and data order. Overall, the proposed method shows a strong performance compared to single- and multi-task baselines. More specifically, in the case of 0 rotation (i.e. the two tasks are similar) the auxiliary task has positive influence in both cases over the single-task with some marginal advantage of the proposed method compared to the multi-task baseline. Interestingly, when the digits are rotated (45, 90, etc) the auxiliary task is negatively interfering with the main task. Again, the proposed method succeeds to mitigate the negative behavior significantly. For the case of 45 degrees it not only outperform the multi-task baseline but also beats the single task again by benefiting from positive transfer while cleverly avoiding negative interference. In the case of 90 and 135 degrees the negative effect of the auxiliary task compromised its positive effect and the performance degrades, however, as expected the proposed method is avoiding a significant amount of negative interference compared to the multi-task baseline.

3 Gridworld Tasks

We design a simple RL experiment to show that our method is capable of finding the strategy of combining the main loss and the auxiliary loss. We define a distribution over a set of $15\times 15$ gridworlds, where an agent observes its surrounding (up to four pixels away) and can move in four directions, up, down, left, and right. We randomly place two types of positive rewards, $+5$ and $+10$ points, both terminating an episode. To guarantee a finite length of episodes, we add a fixed probability of $0.01$ of transitioning to a non-rewarding terminal state. See Appendix D.2 for more experiment details and a visualization of the environment.

First, we train a Q-learning agent to navigate the gridworld which gives us a teacher policy $\pi^{\text{Q}}$. Then, we create tasks to which there is a possible positive knowledge transfer by keeping the environment with the same layout but removing the $+10$ rewards (and corresponding states are no longer terminating). Consequently, we have two tasks: the auxiliary task $\mathcal{T}_{{aux}}$ where we have a strong teacher policy $\pi^{\text{Q}}$, and the main task $\mathcal{T}_{{main}}$ where the $+10$ rewards are removed. That is, we know $\mathcal{T}_{{aux}}$ is helpful initially but hurts later since $\mathcal{T}_{{main}}$ no longer has the $+10$ reward; our proposed method should adapt to this change during training. We sample $1,000$ such environment pairs and report expected returns obtained ($100$ evaluation episodes per evaluation point) using five training methods: 1) reward using only policy gradient (Williams, 1992) in the new task, this is the baseline; 2) distill using only distillation (Rusu et al., 2015) cost towards the teacher; 3) add simply adding the two above; 4) cos-weighted using the weighted version of our method (Algorithm 2); 5) cos-unweighted using the unweighted version of our method (Algorithm 1).

Results are shown in the top row of Figure 4. The baseline reward gives a score of slightly above $1$ point after $10,000$ steps of training. To leverage teacher policies, we define the auxiliary loss to be a distillation loss (i.e., a per-state cross-entropy between the teacher’s and student’s distributions over actions). First, we test using solely the distillation loss while sampling trajectories from the student (distill). This recovers a subset of teacher’s behaviors and end up with point—an expected negative transfer as the teacher is guiding us to states that are no longer rewarding. Then, we test simply adding gradients estimated by policy gradient and distillation (add). The resulting policy learns quickly but saturates at a return of $1$ point, showing very limited positive transfer. Lastly, when using our proposed gradient cosine similarity as the measure of transferability (cos-weighted and cos-unweighted), we get a significant performance boost. The learned policies reach baseline performance after just one-third of steps taken by the baseline, and on average obtain $3$ points after $10,000$ steps.Note that we compute cosine similarity between a distillation gradient and a single sample of the policy gradient estimator, meaning that we are using a high variance estimate of the similarity. For larger experiments, one would need to compute running means for reliable statistics.

This experiment shows that gradient cosine similarity allows using knowledge from other related tasks in an automatic fashion. The agent is simply ignoring the teacher signal when it disagrees with its policy gradient estimator. If they do agree in terms of which actions to reinforce, the teacher’s logits are used for better replication of useful policies. In addition, in the bottom row of Figure 4, we present an experiment of transferring between the same task $\mathcal{T}_{{main}}$. Here, the teacher is perfect and it is the optimal to follow it everywhere. We see that the cosine similarity methods underperformed that of simply adding the two losses. This is expected as the noise in the gradients makes it hard to measure if the two tasks are a good fit.

4 Single and Multi-task Atari Games

Finally, we consider a similar RL setup in the more complex Atari domain (Bellemare et al., 2013). We use a deep RL agent the same as previous works (Mnih et al., 2015; 2016; Espeholt et al., 2018; Hessel et al., 2018) and train using the batched actor-critic with V-trace algorithm (Espeholt et al., 2018). See Appendix D.3 for more experiment details.

First, we look at training an agent on a main task (here, to play Breakout) given a sub-optimal teacher (obtained by stop training before it converges) to the task. Analogous to the previous experiment, we define the auxiliary loss as the distillation (i.e., the KL divergence) between the sub-optimal teacher and the training model. Figure 5) shows as expected that solely relying on distillation loss (Only KL) leads to lower performance. Training with both distillation and RL losses without adaptation (RL+KL(Baseline)) leads to slightly better but also sub-optimal performance. While both approaches learn very quickly initially, they plateau much lower than the pure RL approach (RL(Baseline), minimizing just the main task) due to the potential negative effect from the imperfect teacher and the lack of ability to detect and prevent such an effect. In our method (RL+KL(Our Method)), the KL penalty is scaled at every time-step by the cosine similarity between the policy gradient and distillation losses; once this falls below a fixed threshold, the loss is “turned off” thus preventing negative transfer. This experiment shows our approach is able to learn quickly at the beginning then continue fine-tuning with pure RL loss once the distillation loss is zeroed out. That is, when the teacher does not provide useful information anymore, our method simply encourages the agent to learn on its own.

Then, we consider a setting where the main task $\mathcal{T}_{{main}}$ is to train an agent to play two Atari games, Breakout and Ms. PacMan, such that the performance on both games are jointly maximized (i.e., $\mathcal{L}_{{main}}=\mathcal{L}_{Breakout}+\mathcal{L}_{PacMan}$). Similar to previous experiments, we have access to a teacher trained on just Breakout as the auxiliary task. Note that $\mathcal{T}_{{main}}$ itself is chosen to be multi-task to illustrate a complex scenario where $\mathcal{T}_{{aux}}$ helps with only part of $\mathcal{T}_{{main}}$ (here, Breakout).

Figure 6 shows that, compared to the baseline method which trains only on the main loss (Multitask) and the simple addition of the updates of the two tasks method (Multitask RL + Distillation) where the agent learns one task at the expense of the other, our method (Multitask RL + Distillation (Our Method)) of scaling the auxiliary loss by gradient cosine similarity is able to correct itself by “turning off” the auxiliary distillation when the teacher is no longer helpful and is able to learn Ms. PacMan without forgetting Breakout. The evolution of the gradient cosine similarity between the auxiliary teacher and the main Breakout and Ms.PacMan task in Figure 6 provides a meaningful cue for the usefulness of $\mathcal{L}_{{aux}}$.

Related Work

Our work is related to the literature on identifying task similarity in transfer learning. It is generally believed that positive transfer can be achieved when source task(s) and target task(s) are related. However, it is usually assumed that this relatedness mapping is provided by human experts (Taylor and Stone, 2009; Pan et al., 2010); few works have addressed the problem of finding a general measure of similarity to predict transferability between tasks. In RL, methods have been proposed to use the Markov Decision Process (MDP) similarity as a measure of task relatedness (Carroll and Seppi, 2005; Ammar et al., 2014). In image classification, Yosinski et al. (2014) defined image similarity in ImageNet by manually splitting classes into man-made versus natural objects. However, none of these works have explicitly used the learned similarity metric to quantify the transferability from one task to another. In our work, we propose to use cosine similarity of gradients as a generalizable measure across domains and show it can be directly leveraged to detect when unwanted interference occurs and block negative transfer. One important aspect of task similarity for transfer is that it is highly dependent on the parametrization of the model and the current value of the parameters. We exploit this property by providing a heuristic similarity measure for the current parameters, resulting in an approach that relies on an adaptive weight over the updates of the model.

Auxiliary tasks have shown to be beneficial in facilitating learning across domains. In RL, several work have studied using different tasks as auxiliaries. The UNREAL framework (Jaderberg et al., 2017) incorporated unsupervised control along with reward prediction learning as auxiliary tasks; Mirowski et al. (2017) studied auxiliary tasks in the context of navigation; Kartal et al. (2019) used terminal prediction as an auxiliary task; Hernandez-Leal et al. (2019) considered multiagent system where the other agents’ policy can be used as auxiliary tasks for the main agent. In image classification, Zhang et al. (2016) used unsupervised reconstruction tasks. Papoudakis et al. (2018) also explored auxiliary losses for VizDoom. However, these works rely on empirical results and do not address how the auxiliary tasks were selected. In our work, we propose a simple yet effective heuristic that provides explicit guidance on the important question of how to select a good auxiliary task.

Our work is also related to, but different from multi-task learning (Caruana, 1997), particularly the line of work on using adaptive scaling techniques for multi-objective learning. GradNorm (Chen et al., 2018) uses gradient magnitude to scale loss function for each task, aiming to learn well for all tasks. Kendall et al. (2018) proposed a weighting mechanism by considering the homoscedastic uncertainty of each task. Sener and Koltun (2018) propose formulating multi-task learning as multi-objective optimization and leverage multiple gradient descent algorithm (MGDA) (Désidéri, 2012) to find a common descent direction among multiple tasks. Our work distinct in two ways: (i) we care only about the performance of the main task but not all tasks; our optimization goal is not the traditional multi-objective optimization, and (ii) the aforementioned work scale the losses individually without looking at their interaction (which can lead to poor performance in our problem setup when the auxiliary task hurts the main task), whereas we look for alignments in the vector field between the main and the auxiliary task, and the auxiliary task is used only when it is well-aligned with the main task. Nevertheless, our approaches are complementary since they are solving different problems and it would be interesting to combine them to further improve the performance.

Shortcomings

We discuss here a few shortcomings of our proposed method and present some initial analyses. First, estimating the cosine similarity between the gradients of tasks could be expensive or noisy and we currently use a fixed threshold for turning off the auxiliary loss. These could be addressed by calculating a running average of the cosine similarity to get a smoother result and potentially hyper-tune the threshold instead of setting it as a fixed constant. We point out also that even for a perfect alignment—a cosine similarity of $1$—is not informative. In this case, using the auxiliary loss would be equivalent to increasing the learning rate. The positive transfer comes mostly from the fact that the gradients are not perfectly aligned, meaning that the auxiliary loss is learning features that could become useful to the main loss later on. This is reflective on the fact that the gradient descent is not necessarily the optimal descent direction. In addition, one might argue that our approach would fail in high-dimensional spaces since random vectors in such spaces tend to be orthogonal, so that cosine similarity will be naturally driven to . In fact, this is not the case; if two gradients are meant to be co-linear, the noise components cancel each other thus will not affect the cosine similarity estimation. We empirically explore this in Appendix E in the supplementary material.

Second, the new loss surface might be less smooth which can be problematic when using optimizers that rely on statistics of the gradients or second order information (e.g. Adam or RMSprop). In these cases, the transition from just the gradient of the main task to the sum of the gradients can affect the statistics of the optimizer in unexpected ways. While this can be technically true, we have not observed this behaviour in practice.

Thirdly, although the proposed approach works well empirically on complex and noisy tasks like Atari games, as discussed in Section 3, it guarantees only the main task’s convergence, but not how fast it is. While removing the worst case scenarios is important and a good first step, one might care more for data efficiency when using auxiliary losses (i.e., faster convergence). In Appendix B we provide counter examples where the proposed update rule slows down learning, compared to optimizing the main task alone.

Finally, similar to other concurrent works (Yu et al., 2020; Schaul et al., 2019) we are making the assumption that the auxiliary loss harming learning of main loss locally (at time $t$) implies it will affect the convergence speed globally. In theory, however, this is not necessarily the case. By moving in an ascent direction on the main loss, one might learn better representations that could lead to better performance or faster learning. Empirically we do not notice this to happen and we hypothesize that overparametrization helps making gradient similarity reliable as a signal for task interference. We provide a longer discussion in the Appendix F.

Conclusions

In this work, we explored a simple yet efficient technique to ensure that an auxiliary loss does not hurt the learning on the main task. The proposed approach reduces to applying gradients of the auxiliary task only if they are a descent direction of the main task. We have empirically shown the potential of using the proposed hypothesis as an elegant way of picking a suitable auxiliary task. While we have mostly considered scenarios where the auxiliary task helps initially but hurts later, it would be interesting to explore settings where the auxiliary task hurts initially but helps in the end. Examples of such are annealing $\beta$ in $\beta$-VAE (Higgins et al., 2017) and annealing the confidence penalty in Pereyra et al. (2017).

References

Appendix A Proofs

Given any gradient vector field $G({\bm{\theta}})=\nabla_{\bm{\theta}}\mathcal{L}({\bm{\theta}})$ and any vector field $V({\bm{\theta}})$ (such as gradient of another loss function, but could be arbitrary set of updates), an update rule of the form

converges to the local minimum of $\mathcal{L}$ given small enough $\alpha^{(t)}$.

Our update rule is simply ${\bm{\theta}}^{(t+1)}:={\bm{\theta}}^{(t)}-\alpha^{(t)}\Delta{\bm{\theta}}^{(t)}$ and we have

And it can be 0 if and only if $\|\nabla\mathcal{L}^{(t)}\|=0$ (since sum of two non-negative terms is zero iff both are zero, and step from (4) to (5) is only possible if this is not true), thus it is 0 only when we are at the critical point of $\mathcal{L}$. Thus the method converges due to convergence of steepest descent methods, see “Cauchy’s method of minimization” [Goldstein, 1962]. ∎

A.2 Proof for Proposition 3

In general, the proposed update rule does not have to create a conservative vector field.

Proof comes from a counterexample, let us define in 2D space:

for some fixed $a\neq 0$. Let us now define two paths (parametrized by $s$) between points $(0,0)$ and $(2,2)$, path $A$ which is a concatenation of a line from $(0,0)$ to $(0,2)$ (we call it $U$, since it goes up) and line from $(0,2)$ to $(2,2)$ (which we call $R$ as it goes right), and path $B$ which first goes right and then up. Let $V_{\text{cos}}$ denote the update rule we follow, then:

At the same time, since gradient of $\mathcal{L}_{{main}}$ is conservative by definition:

where $C$ is a part of $B$ that goes through $\times$. We conclude that $\int_{A}V_{\text{cos}}ds\neq\int_{B}V_{\text{cos}}ds$, so our vector field is not path invariant, thus by Green’s Theorem it is not conservative, which concludes the proof. See Figure 7 for visualization. ∎

Appendix B Positive and Negative Examples of Auxiliary Losses

We show in figure 8 positive examples of where an auxiliary loss can help a main loss to converge faster. In addition, as mentioned in Section 3 that an auxiliary task does not guarantee faster convergence, we discuss here a few potential issues of using cosine similarity of gradients to measure task similarity and show in Figure 9 a negative example on where the auxiliary loss could slow-down the convergence of the main task. First, the method depends on being able to compute cosine between gradients. However, in deep learning we rarely are able to compute exact gradients in practice, we instead depend on their high variance estimators (mini-batches in supervised learning, or Monte Carlo estimators in RL). Consequently, estimating the cosine similarity might require additional tricks such as keeping moving averages of estimates. Second, adding additional task gradient in selected subset of iterates can lead to very bumpy surface from the perspective of optimizer, causing methods which keep track of gradient statistics/estimate higher order derivatives, can be less efficient. Finally, one can construct specific functions, where despite still minimizing the loss, one significantly slows down optimization process.

Appendix C Weighted Version of Our Method

Algorithm 2 describes the weighted version of our method.

Appendix D Experimental Details

We present in this section experimental details for the ImageNet classification task, the RL gridworld task, and the RL Atari game task.

As a ground truth for class similarity, we identify pairs of ImageNet classes to be near or far using, lowest common ancestor (LCA) and Frechet Inception Distance (FID) [Heusel et al., 2017].

ImageNet follows a tree hierarchy where each class is a leaf node. We define the distance between a pair of classes as at which tree level their LCA is found. In particular, there are 19 levels in the class tree, each leaf node (i.e. class) is considered to be level 0 while the root node is considered to be level 19. We perform bottom-up search for one pair of random sampled classes and find their LCA node—the class distance is then defined as the level number of this node. For example, class $871$ (“trimaran”) and class $484$ (“catamaran”) has class distance 1 because their LCA is one level up.

FID is used as a second measure of similarity. We obtain the image embedding of a pair of classes using the penultimate layer of a pre-trained ResNetV2-50 model [He et al., 2016b] and then compute the embedding distance using FID, defined in Heusel et al. as:

where $m_{k},C_{k}$ denote the mean and covariance of the embeddings from class $k$.

We randomly sampled 50 pairs of classes for each level of $LCA=\{1,2,3,4,16,17,18,19\}$ (400 pair of classes in total) and compute their FID. Figure 10 shows a plot of LCA (x-axis) verses FID (y-axis) over our sampled class pairs. It can be seen that LCA and FID are (loosely) correlated and that they reflect human intuition of task similarity for some pairs. For example, trimaran and catamaran (bottom-left) are similar both visually and conceptually, whereas rock python and traffic light (top-right) are dissimilar both visually and conceptually. However, there are contrary examples where LCA disagrees with FID; monkey pinscher and doberman pinscher (top-left) are visually dissimilar but conceptually similar, whereas bubble and sundial (bottom-right) are visually similar but conceptually dissimilar. Per the observations, in subsequent experiments we pick class pairs that are {Low LCA, Low FID} as near pairs (e.g., trimaran and catamaran), and class pairs that are {high LCA, high FID} as far pairs (e.g., rock python and traffic light).

D.2 Gridworld Experiments

We define a distribution over $15\times 15$ gridworlds, where an agent observes its surrounding (up to 4 pixels away) and can move in 4 directions (with 10% transition noise). We randomly place walls (blocking movement) as well as two types of positive rewards: $+5$ and $+10$ points, both terminating an episode. There are also some negative rewards (both terminating and non-terminating) to make problem harder. In order to guarantee (expected) finite length of episodes we add fixed probability of 0.01 of transitioning to a non-rewarding terminal state.

For this setup, the update rule for each sequence \tau=\bigl{(}(s_{1},a_{1},r_{1}),\dots(s_{N},a_{N},r_{N})\bigr{)} is thus given by

In order to make use of expert policies for $\mathcal{T}_{{aux}}$ we define auxiliary loss as a distillation loss, which is just a per-state cross-entropy between teacher’s and student’s distributions over actions. If we just add gradients estimated by policy gradient, and the ones given by distillation, the update is given by

where $V^{(t)}=\sum_{a}\pi^{\text{Q}}(a|s_{t^{\prime}})\nabla_{\bm{\theta}}\log\pi(a|s_{t^{\prime}})$ and $\text{H}^{\times}(p,q)=-\sum_{k}p_{k}\log q_{k}$ is the cross entropy.

However, if we use the proposed gradient cosine similarity, we get the following update

This get a significant boost to performance, and policies that score on average 3 points after 10,000 steps and obtain baseline performance after just one third of steps. Figure 11 shows a depiction of the task and an example solution.

D.3 Atari Experiments

For all Atari game experiments, we use a convolutional architecture as in previous work [Espeholt et al., 2018, Hessel et al., 2018, Mnih et al., 2015, 2016], trained with batched actor-critic with the V-trace algorithm [Espeholt et al., 2018]. We use a learning rate of $0.0006$ and an entropy cost of $0.01$ for all experiments, with a batch size of $32$ and $200$ parallel actors.

For the single game experiment, Breakout, we use 0.02 for the threshold on the cosine similarity and, for technical reasons we ended up computing the cosine distance on a per-layer basis and then averaged. We additionally need to do a moving average of the cosine over time ($0.999c^{(t-1)}+0.001c^{(t)}$) to ensure there are no sudden spikes in the weighting due to noisy gradients. Same setting is used for the multi-task experiment, just that the threshold is set to 0.01.

Appendix E Gradient Cosine Similarity in High Dimensions

As discussed in Section 6 of the paper, we empirically show in Figure 12 that if two gradients are meant to be co-linear (e.g., tasks share some underlying structure), the cosine similarity can still be a potential way of measuring task similarity even in high dimensions.

Appendix F Overparametrization and reliability of gradient similarity globally

As discussed in Section 6 of the paper, in theory it is not clear that moving in an ascent direction of the main loss locally (e.g. for a small finite number of steps) has negative consequences on speed of convergence when the main loss is non-convex as in the case of neural networks.

For example the auxiliary loss could drive the model outside of the convex bowl (basin of attraction) of a suboptimal minima, helping the neural network convert to a better solution for the main loss, even though locally it will move in an ascent direction on the main loss. In general if we think of the auxiliary loss as some form of regularization, one interpretation of regularization terms is that it is needed to prevent overfitting the data. That implies that it will, at some point in training, stop the model to move in the descent direction of the training error which would lead to memorization of the training set.

However we observe that in practice, as a heuristic, gradient similarity seems to work reasonably well. While we do not have an understanding of the efficiency of our heuristic, we hypothesize that overparametrization of neural networks might play an important role in the explanation. For example, it is well understood that overparametrization makes it less likely for the loss surface to exhibit suboptimal local minima. Therefore the explicit case of needing to move in ascent direction in order to escape the basing of attraction of a suboptimal solution becomes less likely. The lack of bad local minima however does not exclude that moving in an ascent direction can potentially have positive effects globally. For example, one can imagine moving in a descent direction forcing learning to traverse a plateaux until it reaches a minima, while moving in a ascent direction leads to a steep valley. The minima reached on both direction could be connected (as most minima for neural networks seem to be Garipov et al. , Draxler et al. ) and hence this is not a case of being stuck in a suboptimal solution. While such structures of the loss surface have not been investigated, and they could happen due to the non-linearity, we believe that overparametrization makes them less likely. We leave exploring this hypothesis as an open question for future work.