PaperlayerView exact revision on arXiv

Source-linked AI summary

Auto-Lambda: Disentangling Dynamic Task Relationships

Shikun Liu, Stephen James, Andrew J. Davison, Edward Johns

arXiv:2202.03091v2cs.LGcs.AIcs.CV

TL;DR

When task relationships are unknown, multi-task learning can require costly searches over task groupings and may perform worse than single-task learning. Auto-λ learns dynamic task weightings from primary-task validation loss, and the paper reports state-of-the-art performance across computer vision and robotics benchmarks for multi-task and auxiliary learning.

  • Problem

    Unknown task relationships can make multi-task learning underperform single-task learning, while exhaustive grouping search requires training 2|T| −1 networks.

    Source: §1 · pp. 1–2

  • Method

    Auto-λ uses gradient-based meta-learning to update task-specific weightings from primary-task validation loss during joint network training.

    Source: §1 · p. 3, §4 · p. 5

  • Results

    Auto-λ achieves state-of-the-art performance in computer vision and robotics benchmarks for both multi-task and auxiliary learning, including against methods designed for each setting.

    Source: §8 · p. 15

  • Takeaways & Limitations

    The learned task relationships remain consistent across multi-task architectures and converge in alignment with relationships found by brute-force search.

    Source: §1 · p. 3

    Computing second-order gradients is expensive, and the finite-difference approximation may still make Auto-λ slower than alternative optimization methods.

    Source: §8 · p. 15

Abstract

from arXiv · show

Understanding the structure of multiple related tasks allows for multi-task learning to improve the generalisation ability of one or all of them. However, it usually requires training each pairwise combination of tasks together in order to capture task relationships, at an extremely high computational cost. In this work, we learn task relationships via an automated weighting framework, named Auto-Lambda. Unlike previous methods where task relationships are assumed to be fixed, Auto-Lambda is a gradient-based meta learning framework which explores continuous, dynamic task relationships via task-specific weightings, and can optimise any choice of combination of tasks through the formulation of a meta-loss; where the validation loss automatically influences task weightings throughout training. We apply the proposed framework to both multi-task and auxiliary learning problems in computer vision and robotics, and show that Auto-Lambda achieves state-of-the-art performance, even when compared to optimisation strategies designed specifically for each problem and data domain. Finally, we observe that Auto-Lambda can discover interesting learning behaviors, leading to new insights in multi-task learning. Code is available at https://github.com/lorenmt/auto-lambda.

1 Introduction

Multi-task learning can help related tasks but may underperform single-task learning when task relationships are unknown, while exhaustive task-group searches are costly. Auto-λ learns dynamic task weightings from primary-task validation loss and is reported to outperform optimization strategies and static groupings across computer vision and robotics.

1 Introduction

  • Unknown task relationships can make multi-task learning perform worse than single-task learning because irrelevant tasks may produce conflicting gradients and complicate optimization.The authors define task relationship by how much two tasks should be trained together, illustrated by comparing task A's performance with B versus C.Source: §1 · p. 1
  • Exhaustive task-group search requires training 2|T| −1 networks, while prior grouping methods require a separate search and retraining stage and approximate higher-order relationships.As the number of tasks grows, even evaluating combinations of lower-order relationships may become prohibitively costly.Source: §1 · p. 2
  • Auto-λ treats task relationships as dynamic properties of the network during training and aims to unify multi-task learning with auxiliary learning for selected primary tasks.In auxiliary learning, the primary tasks are the subset being optimized for; other training tasks are included to assist them.Source: §1 · pp. 2–3
  • Validation loss on primary tasks guides task-specific weightings, which are jointly updated with the multi-task network in an end-to-end meta-learning framework.The weightings are intended to improve primary-task performance in the next iteration.Source: §1 · p. 3
  • Auto-λ outperforms multi-task and auxiliary optimization strategies and selected static task groupings across computer vision and robotics evaluations.The authors report that learned task relationships remain consistent across multi-task architectures and align at convergence with relationships found by brute-force search.Source: §1 · p. 3

2 Related Work

Prior work addresses task sharing through architecture design, optimization, task grouping, and meta-learning, often with approaches tailored to particular settings. Auto-λ is presented as a weighting-based framework that unifies multi-task and auxiliary learning while learning dynamic task relationships.

2 Related Work

  • Multi-task architectures include hard-parameter sharing, soft-parameter sharing, and neural architecture search, while optimization work studies task weighting and gradient modification.Weighting criteria include uncertainty, task prioritization, and loss magnitude.Source: §2 · p. 3
  • Auxiliary-task methods have used cosine similarity to assess auxiliary usefulness and neural networks to combine auxiliary losses nonlinearly.Source: §2 · p. 3
  • Auto-λ learns task weightings to unify multi-task and auxiliary optimization, while prior task-grouping methods compute fixed relationships and approximate higher-order relationships from lower-order ones.Earlier grouping work improved training efficiency, including a method using a single training run.Source: §2 · pp. 3–4
  • Related meta-learning work includes task scheduling, multi-objective optimization, and balancing worst-performing tasks; Auto-λ instead learns task relationships for primary-task performance.It can also be viewed as meta-learning-based hyperparameter optimization with a bi-level objective, but it includes training strategies tailored to multi-task learning and studying learned relationships.Source: §2 · p. 4

3 Background

The background formalizes task-specific network parameters and weights, distinguishes multi-task from auxiliary learning by which tasks are primary, and reviews optimization approaches. Auto-λ is positioned as exploring dynamic relationships alongside existing fixed groupings.

3 Background

  • A multi-task network has shared and task-specific parameters, assigns each task a weighting, and represents tasks by their input-output spaces.Source: §3 · p. 4
  • When all tasks are primary, the setting is multi-task learning; when only a selected subset is primary, it is auxiliary learning.The task spaces may share one input domain or use different input domains.Source: §3 · p. 4
  • Single-objective optimization linearly combines task losses using task-specific weights, while gradient-based methods can directly balance task gradients.Weights may be manually chosen or learned using heuristics.Source: §3 · p. 5
  • Task-grouping methods use fixed binary weights to indicate which tasks train together, whereas Auto-λ explores dynamic task relationships.Source: §3 · p. 5
  • Multi-objective optimization uses a vector-valued loss and seeks Pareto optimality, where no common gradient update decreases every task-specific loss.This strategy is limited to multi-task learning settings.Source: §3 · p. 5

4 Auto-λ: Exploring Dynamic Task Relationships

Auto-λ learns dynamic task relationships through task-specific weights, using primary-task validation loss to guide their updates. Its bi-level optimization is made practical through gradient approximations and training-data reuse.

The Design Philosophy

  • Auto-λ learns dynamic task relationships with task-specific weights, optimized from primary-task validation loss in a standalone meta-learning framework.It is designed for multi-task and auxiliary learning, using any combination of primary tasks.Source: §4 · p. 5

Bi-level Optimisation

  • The bi-level objective chooses weights to minimize primary-task validation loss while network parameters minimize the weighted training loss across tasks.The validation loss serves as a measure of generalization.Source: §4 · p. 6

Approximation via Finite Difference

  • Finite differences approximate the weight gradient to reduce the memory and training-time costs of computing second-order gradients.The method iteratively updates network parameters and task weights; the forward model uses one gradient step, and the approximation uses a small constant ϵ.Source: §4 · p. 6
  • Task weights apply only to training loss; weighting validation loss can yield the trivial solution λ = 0, while assuming θ′ = θ* makes the weight gradient zero.The primary tasks are assigned equal importance by default, though different constants can encode prior knowledge.Source: §4 · p. 6

Swapping Training Data

  • Swapping training and validation batches from the same dataset produced similar learned weights to sampling batches from different datasets, enabling single-stage end-to-end optimization.The batches serve as training and validation data without splitting the dataset into separate sets.Source: §4 · p. 7

Stochastic Task Sampling

  • Stochastically sampling K′ ≪ K tasks for weight updates allows constant memory independent of the number of training tasks.In practice, the largest K′ that fits on the GPU is chosen; performance was consistent across a wide range of K′ values.Source: §4 · p. 7

5 Experiments

Across dense prediction, multi-domain classification, and robot manipulation, Auto-λ outperformed the reported baselines in multi-task and auxiliary learning settings. Experiments also examine auxiliary-task noise, challenging domains, and performance across 10 RLBench tasks.

5 Experiments

  • Experiments span single-domain and multi-domain vision and robotics, comparing weighting-based baselines in multi-task learning and GCS in auxiliary learning.The methods use the same optimization strategies within each dataset and data domain; auxiliary experiments generally treat each task individually as primary.Source: §5 · p. 7

5.1 Results on Dense Prediction Tasks

  • Dense prediction uses three tasks on each of NYUv2 and CityScapes, evaluated with Split and MTAN architectures and task-specific metrics.Auxiliary experiments add a random-noise prediction task as a sanity check for methods facing useless gradients.Source: §5.1 · pp. 7–8
  • Auto-λ outperformed all baselines in multi-task and auxiliary learning across NYUv2 and CityScapes, using both Split and MTAN architectures.Table 1 covers auxiliary learning with an added noise-prediction task and reports results averaged over two independent runs.Source: Table 1 · p. 8, §5.1 · p. 8

5.2 Results on Multi-domain Classification Tasks

  • Auto-λ improved overall performance on 20 CIFAR-100 domains and achieved around 5–7% absolute test-accuracy gains in challenging domains under auxiliary learning.Table 2 reports the five domains with lowest single-task performance and the average across all 20 domains.Source: Table 2 · p. 9, §5.2 · p. 9

5.1 Results on Dense Prediction Tasks

  • Auto-λ [3 Tasks] performed nearly identically to multi-task Auto-λ, while the best multi-task baseline performed notably worse with noise prediction as an auxiliary task.Auto-λ [1 Task] further improved performance over Auto-λ [3 Tasks], at the cost of task-specific training for each task.Source: §5.1 · p. 9

5.2 Results on Multi-domain Classification Tasks

  • The CIFAR-100 classification experiment treats each of 20 coarse classes as a domain and trains a shared VGG-16 model with task-specific prediction layers.Each task predicts five fine classes.Source: §5.2 · p. 9

5.3 Results on Robot Manipulation Tasks

  • The RLBench evaluation uses 10 tasks with 100 demonstrations per task, converted into behavioral-cloning data through keyframe discovery.The network receives RGB and point-cloud inputs from three cameras and predicts a continuous 6D pose and discrete gripper action.Source: §5.3 · p. 10, Figure 2 · p. 10
  • Auto-λ achieved the best reported performance in both robot multi-task and auxiliary learning, improving success rates by up to 30–40% on some multi-stage tasks versus single-task learning.Table 3 reports per-task and average success rates over 10 RLBench tasks.Source: Table 3 · p. 10, §5.3 · p. 10

6 Intriguing Learning Strategies in Auto-λ

Auto-λ learns task relationships that vary with primary-task choice and training. Its learned weights correspond to relative task utility and can avoid negative transfer.

6.1 Understanding The Structure of Tasks

  • Across NYUv2 and CityScapes, learned task relationships were consistent across Split and MTAN architectures; higher weights indicate stronger relationships.Source: Figure 3 · p. 11
  • Auto-λ achieved higher per-task performance than every fixed task grouping in NYUv2 and CityScapes, with learned weights aligning with those relationships.For semantic segmentation, normal prediction yielded +6.6% versus depth prediction at −6.0%; their weights were 0.84 and 0.52, respectively.Source: §6.1 · p. 12
  • In NYUv2, a weighting crossover near the end of training reflects a dynamic learning strategy resembling an automated curriculum.Source: §6.1 · p. 12
  • Unlike Uncertainty’s constant noise-prediction weighting, Auto-λ’s dynamic relationships can avoid negative transfer from that task.Source: Figure 4 · p. 11, §6.1 · p. 12

7 Additional Analysis

The additional analyses compare Auto-λ with fixed task groupings, alternative gradient strategies, and methods using strong regularisation, while testing training choices and computational cost.

7 Additional Analysis

  • Auto-λ achieved the best per-task performance against every fixed task grouping in NYUv2 and CityScapes with the Split architecture.Source: Figure 5 · p. 12

7.1 Robustness on Training Strategies

  • First-order gradient approximation produced task weightings within < 0.05 averaged difference and near-identical performance, while potentially speeding training about ×2.3.Figure 6 reports the mean and range of weighting differences for direct and approximated gradients.Source: §7.1 · p. 12, Figure 6 · p. 13
  • Small initial task weightings and a suitable learning rate support good performance; larger learning rates saturate weights and destabilize optimization, while larger initializations fail to avoid negative transfer.Using different data to optimize network parameters and task weightings is also essential for measuring generalization; otherwise performance decreases slightly.Source: §7.1 · p. 12, Table 4 · p. 13

7.2 Comparison to Gradient-based Methods

  • Auto-λ remained the best weighting-based method against gradient-based alternatives in the vanilla setting, and combining it with CAGrad improved performance further.The comparison used NYUv2’s standard three-task multi-task learning setup.Source: §7.2 · p. 13, Table 5 · p. 13

7.3 Comparison to Strong Regularisation Methods

  • With strong regularisation on CelebA, Auto-λ improved performance marginally in multi-task learning and significantly in auxiliary learning, unlike the other methods shown.The reported gains had clear statistical significance; Figure 7 summarizes the comparison.Source: §7.3 · pp. 13–14, Figure 7 · p. 14
  • Auto-λ required three times the training time of Unit. Scal. but remained an order of magnitude faster than PCGrad and CAGrad.The extra time reflects two additional forward and two backward passes for second-order-gradient approximation.Source: §7.3 · p. 14

8 Conclusions, Limitations and Discussion

Auto-λ unifies multi-task and auxiliary learning through dynamically changing task weights and achieves state-of-the-art results across the reported computer vision and robotics benchmarks. The authors identify hyperparameter search and computational expense as limitations, and suggest further directions for task decomposition and open-ended learning.

8 Conclusions, Limitations and Discussion

  • Auto-λ dynamically adjusts task weightings in a unified framework for multi-task and auxiliary learning, achieving state-of-the-art performance across computer vision and robotics benchmarks.Source: §8 · p. 15
  • Optimal performance still requires hyperparameter search, although it is primarily sensitive to the learning rate.Weighting decay or bounded task weights might help find hyperparameters that work across datasets.Source: §8 · p. 15
  • Second-order gradients are computationally expensive; finite-difference approximation adds two forward and two backward passes but may still be slower than alternatives.Source: §8 · p. 15
  • Applying Auto-λ to sub-tasks created by decomposing a single task might reveal learning behaviors that improve single-task learning efficiency.Source: §8 · p. 15
  • The authors propose studying Auto-λ in open-ended learning systems where tasks are continually added during training.They suggest its dynamic task-relationship optimization may facilitate this without manual hyperparameter selection for each new task.Source: §8 · p. 15

A Detailed Training Strategies

The appendix reports training configurations for dense prediction, multi-domain classification, and robot manipulation tasks.

A Detailed Training Strategies

  • Dense prediction used MTAN training settings from the authors’ public code, with Auto-λ learning rates of 10^-4 for NYUv2 and 3 × 10^-5 for CityScapes.Source: A Detailed Training Strategies · p. 21
  • Multi-domain classification used SGD with momentum 0.9, initial learning rate 0.1, weight decay 5 × 10^-4, cosine annealing, 200 epochs, batch size 32, and Auto-λ learning rate 3 × 10^-4.Source: A Detailed Training Strategies · p. 21
  • Robot manipulation used Adam at a constant learning rate of 10^-3 for 8000 iterations, with batch size 32 and Auto-λ learning rate 3 × 10^-5.Source: A Detailed Training Strategies · p. 21

B Detailed Experimental Setting for Robotic Manipulation Tasks

For robotic manipulation, demonstrations are converted into stage-wise targets, and an encoder-decoder predicts pose and gripper actions from three-camera RGB and point-cloud inputs.

B Detailed Experimental Setting for Robotic Manipulation Tasks

  • Keyframe discovery turns 100 demonstrations per task into simple stages by selecting transitions such as gripper-state changes or near-zero velocities.The resulting dataset supplies end-effector poses and gripper actions as prediction targets for new task configurations.Source: B Detailed Experimental Setting for… · p. 21
  • The encoder-decoder uses RGB images and point clouds from three cameras to predict a continuous 6D pose and a binary gripper action.Separate decoders predict position and rotation; a learnable task embedding enters the bottleneck for multi-task and auxiliary learning.Source: B Detailed Experimental Setting for… · p. 21, Figure 8 · p. 21
  • Position is predicted through RGB attention maps and spatial soft argmax on point clouds, with a learned stage-specific offset extending positions beyond image-visible locations.Rotation and gripper action are predicted by direct regression.Source: B Detailed Experimental Setting for… · p. 21

C Auto-λ Learned Weightings for NYUv2 and CityScapes

Learned task-weighting dynamics differ across settings: NYUv2 and CityScapes relationships are usually static, RLBench relationships shift during training, and CIFAR-100 weights correlate with domain difficulty.

C Auto-λ Learned Weightings for NYUv2 and CityScapes

D Auto-λ Learned Weightings for RLBench

E Auto-λ Learned Weightings for CIFAR-100

  • For multi-task multi-domain classification, the authors report a clear correlation between task weighting and single-task performance, describing higher weights as assigned to more difficult domains with lower single-task performance.Source: E Auto-λ Learned Weightings for… · p. 24

Related papers

Loading 2202.03091v2…