Simplifying Multi-Task Architectures Through Task-Specific Normalization

TL;DR

The paper argues that multi-task learning challenges can be addressed with minimal architectural changes by making normalization task-specific. It introduces Task-Specific sigmaBatchNorm (TSσBN), a sigmoid-gated BN variant that softly allocates capacity while preserving a shared backbone. TSσBN not only improves stability and performance across CNNs and transformers but also yields an interpretable framework to analyze capacity allocation, filter specialization, and task relationships. Extensive experiments on standard MTL benchmarks show competitive or superior results with substantially fewer parameters, challenging the need for heavier task-specific modules or routing schemes.

Abstract

Multi-task learning (MTL) aims to leverage shared knowledge across tasks to improve generalization and parameter efficiency, yet balancing resources and mitigating interference remain open challenges. Architectural solutions often introduce elaborate task-specific modules or routing schemes, increasing complexity and overhead. In this work, we show that normalization layers alone are sufficient to address many of these challenges. Simply replacing shared normalization with task-specific variants already yields competitive performance, questioning the need for complex designs. Building on this insight, we propose Task-Specific Sigmoid Batch Normalization (TS$σ$BN), a lightweight mechanism that enables tasks to softly allocate network capacity while fully sharing feature extractors. TS$σ$BN improves stability across CNNs and Transformers, matching or exceeding performance on NYUv2, Cityscapes, CelebA, and PascalContext, while remaining highly parameter-efficient. Moreover, its learned gates provide a natural framework for analyzing MTL dynamics, offering interpretable insights into capacity allocation, filter specialization, and task relationships. Our findings suggest that complex MTL architectures may be unnecessary and that task-specific normalization offers a simple, interpretable, and efficient alternative.

Figures (10)

• Figure 1: Illustration of soft parameter sharing architectures in a two-task setting. Cross-Stitch Networks CrossStitch and MTAN MTAN incorporate additional feature extractors, which lead to scalability challenges as the number of tasks increases. Task-Specific $\sigma$BN Networks introduce only task-specific normalization layers, offering a highly parameter-efficient solution.
• Figure 2: Left: Distribution of cosine similarities between the gradients of NYUv2 tasks over the shared convolutions in the early stages of training. Middle: t-SNE visualization of the encoder representations for the first five CelebA tasks. Right: Encoder parameter count for various numbers of tasks relative to a ResNet50 backbone. Overall, TS$\sigma$BN has a greater concentration of orthogonal gradients, produces well-separated task representations and has a negligible parameter growth.
• Figure 3: Decomposed task capacity into shared and independent components using the TS$\sigma$BN framework. In all standard scenarios, tasks share most capacity without signs of dominance.
• Figure 4: Left: Percentage of specialized filters per layer in a TS$\sigma$BN SegNet. Specialization increases in the latter layers. Right: Performance drop across tasks (columns) after pruning filters based on their primary specialization (rows).
• Figure 5: Effect of BN-specific learning rate multipliers on the $\sigma(\gamma_t)$ filter importances distribution (left) and relative performance of models under loss scale perturbations (right).