Source-linked AI summary
Branched Multi-Task Networks: Deciding What Layers To Share
Simon Vandenhende, Stamatios Georgoulis, Bert De Brabandere, Luc Van Gool
TL;DR
Multi-task networks must choose which layers to share and where tasks should branch, but ad hoc decisions are suboptimal and full NAS is expensive. The paper measures task affinities with RSA and constructs branched networks under a parameter budget. Across diverse multi-tasking scenarios, the resulting architectures consistently improve performance relative to parameter count.
Problem
Choosing shared layers and branch points is difficult because configurations grow combinatorially, while full NAS jointly optimizing sharing and architecture is expensive.
Method
The method pre-trains single-task networks, measures feature-based task affinities with RSA at encoder locations, and clusters tasks into branches within a computational budget.
Results
The method outperforms existing approaches on multi-tasking performance versus number of parameters across diverse scenarios and datasets.
Takeaways & Limitations
The approach automatically selects layer sharing and separates dissimilar tasks into branches, supporting branched networks for heterogeneous task sets when negative transfer is limited.
Takeaways & Limitations
Loss weighting is outside the paper’s scope, and all experiments use a simple uniform loss-weighting scheme.
Abstract
from arXiv · showhide
In the context of multi-task learning, neural networks with branched architectures have often been employed to jointly tackle the tasks at hand. Such ramified networks typically start with a number of shared layers, after which different tasks branch out into their own sequence of layers. Understandably, as the number of possible network configurations is combinatorially large, deciding what layers to share and where to branch out becomes cumbersome. Prior works have either relied on ad hoc methods to determine the level of layer sharing, which is suboptimal, or utilized neural architecture search techniques to establish the network design, which is considerably expensive. In this paper, we go beyond these limitations and propose an approach to automatically construct branched multi-task networks, by leveraging the employed tasks' affinities. Given a specific budget, i.e. number of learnable parameters, the proposed approach generates architectures, in which shallow layers are task-agnostic, whereas deeper ones gradually grow more task-specific. Extensive experimental analysis across numerous, diverse multi-tasking datasets shows that, for a given budget, our method consistently yields networks with the highest performance, while for a certain performance threshold it requires the least amount of learnable parameters.
1 Introduction
Multi-task networks can reduce memory and inference costs by sharing layers, but choosing where tasks should branch is difficult as task configurations grow combinatorially. The paper proposes using task affinities to automate this design instead of relying on ad hoc branching or expensive NAS.
- Layer sharing can reduce memory use and inference computation, while joint learning may improve performance over separate single-task models.
- Hard parameter sharing creates many possible configurations as the number of tasks increases, making trial-and-error architecture design unwieldy.
- Neural architecture search is expensive here because it must optimize layer sharing together with layer types and connectivity.
- The proposed approach measures task affinity with RSA and uses those scores to construct branched multi-task networks automatically.
2 Related work
Prior multi-task learning work includes soft and hard parameter sharing, task grouping, adaptive sharing, and neural architecture search. These approaches differ in how they share features, choose task groupings, and control design cost.
- Multi-task learning: Soft-sharing methods include cross-stitch networks, while hard-sharing methods commonly use a shared encoder followed by task-specific decoders.
- Multi-task learning: Hard parameter sharing divides parameters into shared and task-specific sets, whereas soft parameter sharing gives each task its own parameters and uses feature-sharing mechanisms.
- Adaptive feature sharing: FAFS grows a model greedily layer by layer and groups tasks using the probability of concurrent easy or difficult examples.
- Neural architecture search: NAS automates architecture construction, but applying it to multi-task learning requires jointly optimizing sharing, layer types, and connectivity.
- Task affinity and clustering: The proposed task-affinity approach instead uses RSA-based feature relationships and a clustering procedure under a computational budget.
- Scope: Loss weighting is outside this work’s scope because all experiments use uniform loss weighting.
3 Method
The method measures task affinities from single-task representations at multiple encoder locations, then selects a branched architecture that groups tasks under a computational budget. It uses global dissimilarity minimization to separate tasks with dissimilar representations.
- Problem formulation: The method targets N visual recognition tasks under a computational budget and represents layer sharing as a tree whose shallow layers are shared before later task-specific splits.
- Calculate task affinity scores: It trains one single-task encoder-decoder model per task under identical conditions before measuring affinity.
- Calculate task affinity scores: RSA compares representation dissimilarity matrices at D encoder locations using held-out images, producing a D × N × N task-affinity tensor.
- Calculate task affinity scores: The affinity computation focuses on features used to solve each task rather than example difficulty, with negligible overhead beyond pre-training the single-task networks.
- Construct a branched multi-task network: Task dissimilarity at location d is defined as 1 − A_d,i,j, and the network separates dissimilar tasks by assigning them to different branches.
- Construct a branched multi-task network: The selected tree minimizes clustering dissimilarity across encoder depths among all trees that fit the computational budget.
4 Experiments
Experiments across Cityscapes, Taskonomy, and CelebA evaluate affinity-based branched networks under parameter budgets. The method generally improves the performance–parameter trade-off and remains effective across diverse task settings.
- Experimental setup: Experiments span diverse datasets, tasks, and settings, with parameter counts reported for fair model comparisons.The evaluation includes dense prediction, classification, few-task, and many-task scenarios.
- Cityscapes: Task affinity decreases after deeper ResNet-50 blocks on Cityscapes, indicating increasingly task-specific features.Affinity is measured after each of the four encoder blocks.
- Cityscapes: Within a given computational budget, the proposed task groupings achieve higher Cityscapes performance than FAFS and usually select the best performance-versus-parameters grouping.Figure 5 compares all groupings formed by branching in the last three ResNet blocks.
- Cityscapes: Branched networks offer a better performance–parameter trade-off than computationally expensive soft-sharing models such as cross-stitch and NDDR-CNN.They can sample architectures between a baseline multi-task model and soft-sharing architectures.
- Taskonomy: On Taskonomy, the proposed models consistently outperform corresponding FAFS models, while task grouping becomes more influential as task diversity increases.The authors associate the greater sensitivity to grouping with possible negative transfer.
- Cross-dataset findings: Across diverse multi-tasking scenarios and datasets, the approach shows stable performance while separating dissimilar tasks to limit negative transfer.The reported scope includes heterogeneous task dictionaries that challenge soft parameter sharing.
5 Conclusion
The paper concludes that task affinities can guide automatic construction of branched multi-task networks under a computational budget. Across diverse scenarios and datasets, the resulting architectures improve the performance-versus-parameter trade-off consistently.
- Conclusion: The method automatically constructs branched multi-task networks for a given computational budget using task affinities to guide layer sharing.It optimizes layer sharing without jointly searching layer types and connectivity.
- Conclusion: Across diverse multi-tasking scenarios and datasets, the method outperforms existing approaches on multi-task performance versus number of parameters.The conclusion reports consistent results across the evaluated settings.
A.1 Cityscapes
Cityscapes experiments compare single-task, baseline multi-task, branched, cross-stitch/NDDR-CNN, and MTAN models under specified training setups. The branched architectures evaluated quantitatively are shown in Fig. S4.
- Experimental setup: The encoder is an ImageNet-pretrained ResNet-50 with dilated convolutions, while PSP modules provide task-specific decoders.Inputs are rescaled to 512 x 256 pixels, with outputs upsampled during training and testing.
- Experimental setup: Semantic segmentation uses weighted pixel-wise cross-entropy, instance segmentation regresses pixels to instance centers, and depth estimation uses L1 loss.Losses are normalized so one task does not overwhelm the others during training.
- Baselines: Single-task and baseline multi-task models are evaluated over 48 hyperparameter settings spanning batch sizes, learning-rate schedules, optimizers, and initial learning rates.The single-task search is repeated separately for semantic segmentation, instance segmentation, and monocular depth estimation.
- Evaluation: Multi-task performance is measured as the average per-task performance increase or drop relative to the corresponding single-task models.The sign is adjusted according to whether lower or higher values indicate better performance.
- Compared architectures: Branched multi-task architectures use the best baseline multi-task hyperparameter setting, with the evaluated Cityscapes designs shown in Fig. S4.Cross-stitch/NDDR-CNN and MTAN are also included as comparison approaches.
A.2 Taskonomy
Taskonomy experiments reuse a ResNet-50-based pixel-to-pixel setup and compare baseline, branched, cross-stitch/NDDR-CNN, and MTAN models. Figures show task groupings, generated architectures, and qualitative predictions for the evaluated approaches.
- Experimental setup: Taskonomy uses a ResNet-50 encoder with a 15-layer fully-convolutional decoder for pixel-to-pixel prediction tasks.Inputs are rescaled to 256 x 256 pixels, and the final stride-2 convolution is replaced by stride 1.
- Experimental setup: Task losses are weighted as ws = 1, wd = 1, wk = 10, we = 10, and wc = 1.Heatmaps are rescaled to [0, 1], and depth maps are normalized by their standard deviation.
- Baselines: Single-task and baseline multi-task models use Adam with initial learning rate 1e-4, 120000 training iterations, batch size 32, and learning-rate decay after 80000 iterations.The baseline multi-task performance is calculated using Eq. 1.
- Architectures and qualitative results: The proposed branched architectures are shown in Fig. S8, while architectures from the task-grouping method of [31] are shown in Fig. S12.Qualitative predictions from the third proposed branched network appear in Fig. S16.
- Comparison methods: Cross-stitch/NDDR-CNN models reuse Cityscapes-optimal hyperparameters, and MTAN likewise reuses the settings found optimal on Cityscapes.Cross-stitch/NDDR units use initialization weights α = 0.8 and β = 0.05.
A.3 CelebA
CelebA experiments use a thin-ω VGG-16-based architecture to group 40 person-attribute classification tasks. The associated grouping is shown in Fig. S17.
- Training: The branched multi-task network is trained with SGD with momentum 0.9, initial learning rate 0.05, batch size 32, and weight decay 0.0001.Training lasts 120000 iterations, with the learning rate divided by 10 every 40000 iterations.
- Tasks and grouping: The loss is sigmoid cross-entropy with uniform weighting across the 40 person-attribute classification tasks.The task grouping is presented in Fig. S17.
A.4 Computational Analysis
The computational analysis decomposes task-affinity computation into training single-task networks, computing representational dissimilarity matrices, and calculating task correlations. It concludes that the method requires training N single-task networks plus minor overhead.
- Cost decomposition: The analysis identifies three computational steps for calculating task-affinity scores.These steps are training single-task networks, computing RDM matrices, and calculating the RSA matrix.
- Single-task training: Training N single-task networks can use a subset of the available data; 500 randomly selected Cityscapes training images produced the same task groupings.This provides a way to reduce training time in the first step.
- Representation analysis: Computing RDM matrices requires features from K images at D predetermined layers across all N networks; the experiments use K = 500.The features are stored before calculating the RDM matrices.
- Representation analysis: Computing the RSA matrix requires D × N × N correlations between lower-triangle RDM entries, and is negligible compared with training the single-task networks.The computation is performed at D locations for all N tasks.
- Overall cost: The method's computational cost is training N single-task networks plus overhead, with only minor overhead relative to cross-stitch networks and NDDR-CNNs.Those comparison methods also pre-train single-task networks before adding soft parameter sharing.