Auxiliary Learning by Implicit Differentiation

1 INTRODUCTION . The performance of deep neural networks can significantly improve by training the main task of interest with additional auxiliary tasks ( Goyal et al. , 2019 ; Jaderberg et al. , 2016 ; Mirowski , 2019 ) . For example , learning to segment an image into objects can be more accurate when the model is simultaneously trained to predict other properties of the image like pixel depth or 3D structure ( Standley et al. , 2019 ) . In the low data regime , models trained with the main task only are prone to overfit and generalize poorly to unseen data ( Vinyals et al. , 2016 ) . In this case , the benefits of learning with multiple tasks are amplified ( Zhang and Yang , 2017 ) . Training with auxiliary tasks adds an inductive bias that pushes learned models to capture meaningful representations and avoid overfitting to spurious correlations . In some domains , it may be easy to design beneficial auxiliary tasks and collect supervised data . For example , numerous tasks were proposed for self-supervised learning in image classification , including masking ( Doersch et al. , 2015 ) , rotation ( Gidaris et al. , 2018 ) and patch shuffling ( Doersch and Zisserman , 2017 ; Noroozi and Favaro , 2016 ) . In these cases , it is not clear what would be the best way to combine all auxiliary tasks into a single loss ( Doersch and Zisserman , 2017 ) . The common practice is to compute a weighted combination of pretext losses by tuning the weights of individual losses using hyperparameter grid search . This approach , however , limits the potential of learning with auxiliary tasks because the run time of grid search grows exponentially with the number of tasks . In other domains , obtaining good auxiliaries in the first place may be challenging or may require expert knowledge . For example , for point cloud classification , few self-supervised tasks have been proposed ; however , their benefits so far are limited ( Achituve et al. , 2020 ; Hassani and Haley , 2019 ; ∗Equal contributor †Equal contributor Sauder and Sievers , 2019 ; Tang et al. , 2020 ) . For these cases , it would be beneficial to automate the process of generating auxiliary tasks without domain expertise . Our work takes a step forward in automating the use and design of auxiliary learning tasks . We name our approach AuxiLearn . AuxiLearn leverages recent progress made in implicit differentiation for optimizing hyperparameters ( Liao et al. , 2018 ; Lorraine et al. , 2020 ) . We demonstrate the effectiveness of AuxiLearn in two types of problems . First , in combining auxiliaries , for cases where auxiliary tasks are predefined . We describe how to train a deep neural network ( NN ) on top of auxiliary losses and combine them non-linearly into a unified loss . For instance , we combine per-pixel losses in image segmentation tasks using a convolutional NN ( CNN ) . Second , designing auxiliaries , for cases where predefined auxiliary tasks are not available . We present an approach for learning such auxiliary tasks without domain knowledge and from input data alone . This is achieved by training an auxiliary network to generate auxiliary labels while training another , primary network to learn both the original task and the auxiliary task . One important distinction from previous works , such as ( Kendall et al. , 2018 ; Liu et al. , 2019a ) , is that we do not optimize the auxiliary parameters using the training loss but rather on a separate ( small ) auxiliary set , allocated from the training data . This is a key difference since the goal of auxiliary learning is to improve generalization rather than help optimization on the training data . To validate our proposed solution , we extensively evaluate AuxiLearn in several tasks in the low-data regime . In this regime , the models suffer from severe overfitting and auxiliary learning can provide the largest benefits . Our results demonstrate that using AuxiLearn leads to improved loss functions and auxiliary tasks , in terms of the performance of the resulting model on the main task . We complement our experimental section with two interesting theoretical insights regarding our model . The first shows that a relatively simple auxiliary hypothesis class may overfit . The second aims to understand which auxiliaries benefit the main task . To summarize , we propose a novel general approach for learning with auxiliaries using implicit differentiation . We make the following novel contributions : ( a ) We describe a unified approach for combining multiple loss terms and for learning novel auxiliary tasks from the data alone ; ( b ) We provide a theoretical observation on the capacity of auxiliary learning ; ( c ) We show that the key quantity for determining beneficial auxiliaries is the Newton update ; ( d ) We provide new results on a variety of auxiliary learning tasks with a focus on the low data regime . We conclude that implicit differentiation can play a significant role in automating the design of auxiliary learning setups . 2 RELATED WORK . Learning with multiple tasks . Multitask Learning ( MTL ) aims at simultaneously solving multiple learning problems while sharing information across tasks . In some cases , MTL benefits the optimization process and improves task-specific generalization performance compared to single-task learning ( Standley et al. , 2019 ) . In contrast to MTL , auxiliary learning aims at solving a single , main task , and the purpose of all other tasks is to facilitate the learning of the primary task . At test time , only the main task is considered . This approach has been successfully applied in multiple domains , including computer vision ( Zhang et al. , 2014 ) , natural language processing ( Fan et al. , 2017 ; Trinh et al. , 2018 ) , and reinforcement learning ( Jaderberg et al. , 2016 ; Lin et al. , 2019 ) . Dynamic task weighting . When learning a set of tasks , the task-specific losses are combined into an overall loss . The way individual losses are combined is crucial because MTL-based models are sensitive to the relative weightings of the tasks ( Kendall et al. , 2018 ) . A common approach for combining task losses is in a linear fashion . When the number of tasks is small , task weights are commonly tuned with a simple grid search . However , this approach does not extend to a large number of tasks , or a more complex weighting scheme . Several recent studies proposed scaling task weights using gradient magnitude ( Chen et al. , 2018 ) , task uncertainty ( Kendall et al. , 2018 ) , or the rate of loss change ( Liu et al. , 2019b ) . Sener and Koltun ( 2018 ) proposed casting the multitask learning problem as a multi-objective optimization . These methods assume that all tasks are equally important , and are less suited for auxiliary learning . Du et al . ( 2018 ) and Lin et al . ( 2019 ) proposed to weight auxiliary losses using gradient similarity . However , these methods do not scale well with the number of auxiliaries and do not take into account interactions between auxiliaries . In contrast , we propose to learn from data how to combine auxiliaries , possibly in a non-linear manner . Devising auxiliaries . Designing an auxiliary task for a given main task is challenging because it may require domain expertise and additional labeling effort . For self-supervised learning ( SSL ) , many approaches have been proposed ( see Jing and Tian ( 2020 ) for a recent survey ) , but the joint representation learned through SSL may suffer from negative transfer and hurt the main task ( Standley et al. , 2019 ) . Liu et al . ( 2019a ) proposed learning a helpful auxiliary in a meta-learning fashion , removing the need for handcrafted auxiliaries . However , their system is optimized for the training data , which may lead to degenerate auxiliaries . To address this issue , an entropy term is introduced to force the auxiliary network to spread the probability mass across classes . Implicit differentiation based optimization . Our formulation gives rise to a bi-level optimization problem . Such problems naturally arise in the context of meta-learning ( Finn et al. , 2017 ; Rajeswaran et al. , 2019 ) and hyperparameter optimization ( Bengio , 2000 ; Foo et al. , 2008 ; Larsen et al. , 1996 ; Liao et al. , 2018 ; Lorraine et al. , 2020 ; Pedregosa , 2016 ) . The Implicit Function Theorem ( IFT ) is often used for computing gradients of the upper-level function , this operation requires calculating a vector-inverse Hessian product . However , for modern neural networks , it is infeasible to calculate it explicitly , and an approximation must be devised . Luketina et al . ( 2016 ) proposed approximating the Hessian with the identity matrix , whereas Foo et al . ( 2008 ) ; Pedregosa ( 2016 ) ; Rajeswaran et al . ( 2019 ) used conjugate gradient ( CG ) to approximate the product . Following Liao et al . ( 2018 ) ; Lorraine et al . ( 2020 ) , we use a truncated Neumann series and efficient vector-Jacobian products , as it was empirically shown to be more stable than CG . 3 OUR METHOD . We now describe the general AuxiLearn framework for learning with auxiliary tasks . For that purpose , we use two networks , a primary network that is optimized on all tasks and an auxiliary network that is optimized on the main task only . First , we introduce our notations and formulate the general objective . Then , we describe two instances of this framework : combining auxiliaries and learning new auxiliaries . Finally , we present our optimization approach for both instances . 3.1 PROBLEM DEFINITION . Let { ( xti , yti ) } i be the training set and { ( xai , yai ) } i be a distinct independent set which we term auxiliary set . Let f ( · ; W ) denote the primary network , and let g ( · ; φ ) denote the auxiliary network . Here , W are the parameters of the model optimized on the training set , and φ are the auxiliary parameters trained on the auxiliary set . The training loss is defined as : LT = LT ( W , φ ) = ∑ i ` main ( x t i , y t i ; W ) + h ( x t i , y t i , W ; φ ) , ( 1 ) where ` main denotes the loss of the main task and h is the overall auxiliary loss , controlled by φ . In Sections 3.2 & 3.3 we will describe two instances of h. We note that h has access to both W and φ . The loss on the auxiliary set is defined as LA = ∑ i ` main ( x a i , y a i ; W ) , since we are interested in the generalization performance of the main task . We wish to find auxiliary parameters ( φ ) such that the primary parameters ( W ) , trained with the combined objective , generalize well . More formally , we seek φ∗ = arg min φ LA ( W ∗ ( φ ) ) , s.t . W ∗ ( φ ) = arg min W LT ( W , φ ) . ( 2 )

Summary: The paper proposes AuxiLearn, a framework that can be used to combine losses from multiple auxiliary tasks (if present) into a single combined, loss function that does not require expensive grid search over possible linear combination. It uses an implicit differentiation-based approach to train a (deep) non-linear network that weighs the various auxiliary losses to optimize the generalization capabilities of the network. In the absence of such pre-defined tasks, a variation of the approach, using teacher-student networks, helps create relevant tasks to improve the performance of the network. Experiments across tasks such as classification (both few-shot and with limited labels) and segmentation show that the approach helps improve the performance of the model on the main task.