Source-linked AI summary
Learning to Branch for Multi-Task Learning
Pengsheng Guo, Chen-Yu Lee, Daniel Ulbricht
TL;DR
Multi-task learning needs efficient parameter sharing without forcing unrelated tasks into one representation, while prior topology choices rely on manual design or pre-computed relatedness. The paper learns a tree-structured sharing topology with differentiable Gumbel-softmax branching and validates it across synthetic data, CelebA, and Taskonomy. On CelebA, its models outperform counterpart baselines with similar capacity while using fewer parameters.
Problem
Choosing which parameters to share across tasks remains unclear, while over-sharing can cause negative knowledge transfer and exhaustive topology search is combinatorially complex.
Method
The method jointly optimizes network weights and branching distributions in a hard-sharing topology, using Gumbel-softmax sampling to learn where branches split.
Results
LearnToBranch-VGG and LearnToBranch-Deep-Wide outperform counterpart CelebA baselines under similar network capacity, while using fewer total parameters.
Takeaways & Limitations
The framework provides an automated way to learn task grouping and network branching directly from multi-task objectives across synthetic, CelebA, and Taskonomy settings.
Takeaways & Limitations
The controlled synthetic evaluation represents task relatedness through shared activation functions and noise differences, and future work targets multi-modality and partial annotations.
Abstract
from arXiv · showhide
Training multiple tasks jointly in one deep network yields reduced latency during inference and better performance over the single-task counterpart by sharing certain layers of a network. However, over-sharing a network could erroneously enforce over-generalization, causing negative knowledge transfer across tasks. Prior works rely on human intuition or pre-computed task relatedness scores for ad hoc branching structures. They provide sub-optimal end results and often require huge efforts for the trial-and-error process. In this work, we present an automated multi-task learning algorithm that learns where to share or branch within a network, designing an effective network topology that is directly optimized for multiple objectives across tasks. Specifically, we propose a novel tree-structured design space that casts a tree branching operation as a gumbel-softmax sampling procedure. This enables differentiable network splitting that is end-to-end trainable. We validate the proposed method on controlled synthetic data, CelebA, and Taskonomy.
1. Introduction
Multi-task networks can reduce inference latency and improve generalization, but choosing where to share or branch is difficult because over-sharing can cause negative transfer. The paper proposes learning a tree-structured topology through differentiable Gumbel-softmax branching.
- Motivation: Sharing one network across tasks can reduce latency by producing multiple predictions in one forward propagation.This benefit is particularly relevant for portable devices with limited computational budgets.
- Motivation: Other tasks can act as regularizers by encouraging representations that explain more than a single task requires.The passage presents this as a potential route to improved generalization.
- Challenges: Multi-task training must handle differences in task difficulty, label noise, and loss-function scale.Examples include the contrast between L1 and cross-entropy losses.
- Challenges: Exhaustively searching for an optimal parameter-sharing schema has combinatorial complexity as the number of tasks grows.The paper therefore asks whether topology can be searched from back-propagated multi-task signals.
- Challenges: Earlier layers are expected to share general visual features, while later layers may need to be untied across tasks.The paper identifies deciding the sharing boundary as a key challenge.
- Proposed direction: The proposed method learns where to branch by sampling categorical latent variables with Gumbel-softmax, minimizing overall multi-task loss without task-relatedness scores.It also avoids relying on human intuition about which layers should capture task-specific features.
2. Related Work
Prior multi-task methods use soft or hard parameter sharing, loss-balancing strategies, and task-oriented architecture allocation, but effectively deciding what to share remains unresolved. This work instead directly learns sharing and branching from gradients while optimizing the end outputs.
- Soft parameter sharing: Soft parameter sharing gives each task separate backbone parameters and regularizes their distance or connects task activations.Neural Network Parser and Cross-Stitch Networks exemplify these strategies.
- Hard parameter sharing: Hard parameter sharing uses common backbone parameters with task-specific branches or output layers.Examples include Deep Relationship Networks, Fully-Adaptive Feature Sharing, UberNet, and Meta Multi-Task Learning.
- Alternative strategies: Other approaches balance task losses or dynamically allocate parameters using uncertainty, gradient magnitudes, or task-specific mechanisms.The related work discusses uncertainty weighting, GradNorm, and architecture methods for distributing network parameters.
- This work: The paper introduces an end-to-end algorithm that automatically designs hard-sharing networks by sharing or splitting branches according to back-propagated update gradients.It directly optimizes end outputs and bypasses pre-computed task-relatedness scores.
3. Method
The method jointly learns a tree-structured multi-task network architecture and its weights by differentiably sampling branching connections, then selects and retrains the converged architecture.
- Network Topological Space: The design space represents a tree-structured network as a DAG whose nodes perform operations and whose edges define data flows.Each branching block contains parent and child nodes connected by sampled paths.
- Network Topological Space: Each child node samples one parent path, and the selected connectivities define the network’s tree structure.A one-hot indicator assigns one parent node to each child node, while a child may receive propagated activations from selected parents.
- Differentiable Branching Operation: The sampling distributions are updated through the final loss, increasing the likelihood of network configurations that reduce the overall multi-task loss.Temperature is gradually decayed so training explores early and exploits the converged topology distribution later.
- Network Topological Space: Branching blocks can be stacked into deeper tree-structured networks, with parent and child counts adjusted to control model capacity.The architecture and network weights are jointly optimized end-to-end rather than through greedy layer-wise optimization.
- Final Architecture Selection: After validation loss convergence, each child selects the parent with the highest probability, unselected paths are removed, and the resulting architecture is retrained from scratch.Final selection uses the converged categorical distributions without sampling noise.
- Differentiable Branching Operation: The branching operation uses categorical sampling with gumbel-softmax relaxation so architecture probabilities can receive gradients during backpropagation.Discrete parent selection is used during feedforward computation, while relaxed samples update the probability distribution during the backward pass.
4. Experiments
Experiments show that the learned branching structure captures task relatedness in controlled regression settings and improves multi-task performance on CelebA and Taskonomy.
- 4.1. Controlled Synthetic Data: Synthetic experiments grouped tasks with the same activation function while separating tasks with different activations.Tasks differing only by scale remained grouped, while square branched earlier than similar cos and sinc tasks.
- 4.2. CelebA: CelebA evaluates whether automatically learned task grouping outperforms grouping based on human intuition or pre-computed task relatedness.The benchmark contains 40 binary attribute classification tasks over more than 200K face images.
- 4.2. CelebA: Architecture search takes 10 hours plus 11 hours for weight training on LearnToBranch-VGG, and 4 hours plus 10 hours on LearnToBranch-Deep-Wide using one 16GB Tesla GPU.The method samples discrete actions for child nodes without exponential tree growth.
- 4.2. CelebA: LearnToBranch-VGG and LearnToBranch-Deep-Wide outperform their comparable CelebA baselines under similar capacity, using fewer parameters.The models differ from their baselines in the branching operation while preserving kernel sizes and channel counts.
- 4.3. Taskonomy: On Taskonomy, the method achieves the best performance on all five tasks compared with recent baselines, including AdaShare.It also uses fewer parameters than the single-task setting and is more efficient than Cross-Stitch and MTAN.
- 4.3. Taskonomy: Four sampled Taskonomy architectures converged to similar task-grouping strategies, with edge and keypoint branching from normal, depth, and segmentation near the network’s end.One architecture branched at the last layer, while the other three shared the same tree structure.
5. Conclusion
The framework automatically learns task grouping by sharing and branching a neural network, using differentiable optimization and evaluation across synthetic and real-world datasets.
- The method learns task grouping strategies by sharing and branching a neural network.
- Its topological space jointly optimizes network weights and branching distributions through gumbel-softmax sampling.
- The framework is validated on controlled synthetic data, CelebA, and the large-scale Taskonomy dataset.
- Future work includes extending the approach to multi-modality inputs and tasks with partial annotations.
A. Implementation Details
The experiments use staged topology learning followed by training the sampled architecture from scratch, with dataset-specific optimization settings and task weightings.
- CelebA: CelebA experiments train the network topological distribution for 30 epochs, sample the final architecture, then train it from scratch for 30 epochs.LearnToBranch-VGG and LearnToBranch-Deep-Wide use different learning rates for global weights and branching operations.
- CelebA: The CelebA architectures use exponential learning-rate decay, weight decay of 5e−4, and momentum of 0.9 during final training.
- Taskonomy: Taskonomy topology learning runs for 30 epochs with global learning rate 10−3 and branching-operation learning rate 10−1.The final sampled architecture is then trained from scratch for 30 epochs with learning rate 5e−4.
- Taskonomy: Taskonomy combines task weightings of 1.0 for semantic segmentation, 3.0 for surface normals, 2.0 for depth, and 7.0 each for keypoints and edges.The authors note that adaptive task-weighting methods could be combined with their method, leaving that investigation for future work.
- Taskonomy: Preprocessing ignores specified uncertain, background, and invalid pixels while normalizing depth disparities, surface normals, and other task inputs.
B. Learned Branching Features
Network Dissection shows that learned branches develop feature profiles aligned with their task groupings, separating higher-level and lower-level visual representations.
- The SDN {segmentation, depth, normal} branch shows a 35% increase in high-level features and a 20% decrease in low-level features after splitting.The comparison is against the shared layer before branching at layer 13.
- The EK {edge, keypoint} branch continues focusing on low-level features, with no increase in high-level features.The passage attributes this to edge and keypoint tasks generally being considered low-level tasks.
- Table 3 reports detector counts for different input-image categories at different network layers using Network Dissection.
C. Generalizability of the Learned Branching
The learned task grouping transfers from Taskonomy to NYUv2 with a small overall performance difference, while CelebA visualizations show intuitive groupings among related attributes.
- 1.23% is the overall performance difference when transferring Taskonomy-learned grouping to NYUv2 instead of learning grouping on NYUv2 training data.
- On NYUv2, Taskonomy-learned grouping achieves {1.611, 0.739, 0.058} for segmentation, normal, and depth, respectively.
- NYUv2-trained grouping achieves {1.572, 0.748, 0.058} on the same three tasks.
- CelebA architectures group Eyeglasses with Narrow Eyes, Mustasche with No Beard, Black Hair with Gray Hair, and Bald with Receding Hairline.