[Paper Review] Auxiliary Learning by Implicit Differentiation

Aviv Navon, Idan Achituve+3 more|arXiv (Cornell University)|Jun 22, 2020

Domain Adaptation and Few-Shot Learning51 references4 citations

TL;DR

This paper proposes AuxiLearn, a novel framework for automated auxiliary learning in deep learning using implicit differentiation. It unifies the learning of non-linear loss combinations and the discovery of novel auxiliary tasks from data, significantly improving performance in low-data regimes across image segmentation and few-shot classification tasks, with consistent gains over baselines like MAXL and STL.

ABSTRACT

Training neural networks with auxiliary tasks is a common practice for improving the performance on a main task of interest. Two main challenges arise in this multi-task learning setting: (i) designing useful auxiliary tasks; and (ii) combining auxiliary tasks into a single coherent loss. Here, we propose a novel framework, AuxiLearn, that targets both challenges based on implicit differentiation. First, when useful auxiliaries are known, we propose learning a network that combines all losses into a single coherent objective function. This network can learn non-linear interactions between tasks. Second, when no useful auxiliary task is known, we describe how to learn a network that generates a meaningful, novel auxiliary task. We evaluate AuxiLearn in a series of tasks and domains, including image segmentation and learning with attributes in the low data regime, and find that it consistently outperforms competing methods.

Motivation & Objective

• To address the challenge of designing effective auxiliary tasks for multi-task learning, especially in low-data regimes where overfitting is prevalent.
• To develop a unified framework that learns non-linear combinations of predefined auxiliary losses, improving generalization beyond linear weighting.
• To automate the discovery of meaningful auxiliary tasks without domain expertise, using only input data and implicit differentiation.
• To improve model generalization by learning auxiliary tasks that are beneficial for the main task, rather than optimizing for training loss alone.

Proposed method

• Uses implicit differentiation to jointly optimize the primary network and an auxiliary network that generates novel labels, avoiding direct optimization of auxiliary parameters on the main training set.
• Employs a bi-level optimization setup: the primary network is trained on the main task and the auxiliary task, while the auxiliary network is optimized on a separate, small auxiliary dataset to improve generalization.
• Introduces a deep neural network (e.g., CNN) to learn a non-linear, task-aware combination of multiple auxiliary losses, replacing simple weighted sums.
• Applies the implicit function theorem (IFT) to backpropagate through the inner optimization problem, enabling end-to-end training of the loss-combination network.
• Uses smooth activation functions (e.g., Softplus) and weight normalization to ensure differentiability and stability in the implicit differentiation process.
• Employs approximations such as the Neumann series and limited inner optimization steps to make the bi-level optimization computationally feasible.

Experimental results

Research questions

• RQ1 Can implicit differentiation enable effective, non-linear combination of multiple auxiliary losses in a way that improves generalization beyond linear weighting?
• RQ2 Can a neural network be trained to automatically generate meaningful, data-driven auxiliary tasks without domain expertise?
• RQ3 Does learning auxiliary tasks on a separate auxiliary set—rather than the main training set—lead to better generalization and reduced overfitting?
• RQ4 How does the proposed framework perform in low-data regimes, where auxiliary learning is most beneficial?
• RQ5 What theoretical properties govern the effectiveness of auxiliary tasks, and which components are most predictive of performance gain?

Key findings

• AuxiLearn achieves a test accuracy of 26.1% ± 0.7 on the CUB-200-2011 dataset with 15-shot learning, significantly outperforming MAXL (24.2% ± 0.8) and STL (22.6% ± 0.2).
• On the STL-10 dataset with 15% of training data, AuxiLearn reaches 81.42% ± 0.30 accuracy, surpassing MAXL (81.37% ± 0.26) and STL (75.85% ± 0.32).
• In CIFAR-10 with 10% of training data, AuxiLearn achieves 76.75% ± 0.08 accuracy, exceeding STL (72.63% ± 2.14) and MAXL (75.85% ± 0.32).
• The method demonstrates stable optimization and generalization even when using non-smooth components like ReLU, due to the zero-measure nature of non-smoothness points.
• Theoretical analysis shows that the Newton update is a key predictor of which auxiliary tasks benefit the main task, suggesting a principled basis for auxiliary design.
• AuxiLearn consistently improves performance across diverse tasks, including image segmentation and few-shot classification, especially in low-data regimes where benefits are most pronounced.