Research map / Architectures and efficiency / RNNs and normalization
DMTG: One-Shot Differentiable Multi-Task Grouping
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DMTG, a one-shot differentiable multi-task grouping method that formulates grouping as pruning of KN task heads via a relaxed Categorical (Gumbel-softmax) distribution, jointly training group identification and grouped weights.
Problem
Multi-task grouping with many tasks requires searching up to 2^N task-group candidates; existing methods identify groups heuristically in a separate stage from grouped training, creating objective bias and inefficiency.
Approach
Start with K encoder branches each connected to all N task heads (KN heads), then jointly learn a relaxed Categorical assignment matrix via Gumbel-softmax that masks task losses so heads are pruned to one group per task, all trained in one shot with task losses only.
Key results
- On Taskonomy-5 with K=3, total loss 0.173 with NormGain_L +63.85% vs Naive MTL 0.223 and best prior HOA 0.190 (+47.78%).
- On CelebA-9 with K=2, total error 52.97 (NormGain_E +5.75%), outperforming TAG (+4.38%) and HOA (+3.27%).
- Ablation shows one-shot training beats retraining from scratch with the same groups (0.173 vs 0.183 total loss at K=3).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use one-shot joint group identification and grouped-task training instead of retraining grouped tasks from scratch (ablation) | DMTG on Taskonomy-5 · Taskonomy-5, K=3 | total loss: 0.183 -> 0.173 | Helped |
| use one-shot joint group identification and grouped-task training instead of retraining from Naive MTL initialization (ablation) | DMTG on Taskonomy-5 · Taskonomy-5, K=3 | total loss: 0.194 -> 0.173 | Helped |
| replace Xception variant with ViT-Base (ablation) | DMTG · Taskonomy-5, K=3 | total loss: 0.379 -> 0.326 | Helped |
| use 3 shared backbone blocks instead of 0 (ablation) | DMTG on Taskonomy-5 · Taskonomy-5, K=3 | total loss: 0.174 -> 0.166 | Helped |
| increase shared backbone blocks from 3 to 9 (ablation) | DMTG on Taskonomy-5 · Taskonomy-5, K=3 | total loss: 0.166 -> 0.169 | Hurt |
| increase shared backbone blocks from 3 to 9 (ablation) | DMTG on Taskonomy-5 · Taskonomy-5, K=3 | FLOPs (G): 49.68 -> 33.36 | Helped |
| vary the Gumbel Softmax temperature schedule (ablation) | DMTG on Taskonomy-5 · Taskonomy-5, K=5 | total loss: 0.1656; 0.1665; 0.1730; 0.1710; 0.1684 | No clear effect |
| use DMTG instead of the strongest reported baseline, MTG-Net | Taskonomy-5 · Taskonomy-5, K=3 | total loss: 0.191 -> 0.173 | Helped |
| use DMTG instead of the strongest reported baseline, TAG | CelebA-9 · CelebA-9, K=2 | total classification error: 53.41 -> 52.97 | Helped |
| use DMTG instead of the strongest reported baseline, HOA | DMTG with ViT-Base · Taskonomy-5, K=3 | total loss: 0.379 -> 0.326 | Helped |
Surprising findings
- Sharing 9 of 12 backbone blocks gives comparable performance to no sharing while cutting FLOPs by ~58%.
Limitations
- Gumbel softmax is biased and sensitive to temperature.
- Efficiency is impeded when task heads dominate computation (e.g., Taskonomy-9 with Xception), since KN heads are expanded.
- Performance gain on CelebA-40 is smaller than on CelebA-9.
Experimental setup
Datasets: Taskonomy CelebA
Benchmarks: Taskonomy-5 CelebA-9 CelebA-40
Metrics: total loss NormGain_L NormGain_E total classification error FLOPs training time
Baselines: Naive MTL STL RG HOA TAG MTG-Net
Builds on: Gumbel softmax Concrete distribution HOA TAG MTG-Net
Closest papers by meaning
- Towards Principled Task Grouping for Multi-Task Learning (NeurIPS 2026 · similarity 0.90)
- Task Groupings Regularization: Data-Free Meta-Learning with Heterogeneous Pre-trained Models (ICML 2024 · similarity 0.88)
- Selective Task Group Updates for Multi-Task Optimization (ICLR 2025 · similarity 0.87)
- AdaMerging: Adaptive Model Merging for Multi-Task Learning (ICLR 2024 · similarity 0.86)
- Towards Consistent Multi-Task Learning: Unlocking the Potential of Task-Specific Parameters (CVPR 2025 · similarity 0.85)
- Multi-Task Reinforcement Learning with Mixture of Orthogonal Experts (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.