Research map / Image generation / Fast diffusion sampling
Denoising Task Routing for Diffusion Models
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 Denoising Task Routing (DTR), a parameter-free add-on that routes timestep-specific channel masks through DiT and ADM blocks, with sliding windows for task affinity and timestep-dependent allocation for task weights.
Problem
Diffusion models implicitly treat denoising at each timestep as a multi-task problem with negative transfer, but existing work only addresses it via timestep conditioning or loss/optimization methods, not via explicit architectural design.
Approach
Apply task routing (binary channel masks per task) to residual blocks of diffusion networks, with masks defined by a sliding window over channels that shifts with timestep (task affinity) and a nonlinear exponent alpha that allocates more task-dedicated channels to higher timesteps (task weights); no extra parameters are added.
Key results
- DiT-L/2 + DTR reduces ImageNet 256x256 FID from 12.59 to 8.90 (IS 134.60 to 156.48) with no extra parameters.
- ADM + DTR improves ImageNet 64x64 FID from 6.34 to 4.91 and FFHQ FID from 3.12 to 2.65.
- DiT-L/2 + DTR with ANT-UW trained 2M iterations reaches FID 2.33, close to DiT-XL/2 at 7M iterations (FID 2.27).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add DTR routing to DiT | DiT · FFHQ 256×256 unconditional | FID: 10.99 -> 7.32 | Helped |
| Add DTR routing to DiT | DiT · ImageNet 256×256 class-conditional | FID: 12.59 -> 8.90 | Helped |
| Add DTR routing to DiT | DiT · ImageNet 256×256 class-conditional | Recall: 0.49 -> 0.51 | Helped |
| Add DTR routing to DiT | DiT · MS-COCO 256×256 text-conditional | FID: 7.70 -> 7.04 | Helped |
| Add DTR routing to ADM | ADM · ImageNet 64×64 class-conditional | FID: 6.34 -> 4.91 | Helped |
| Add DTR routing to ADM | ADM · ImageNet 64×64 class-conditional | Recall: 0.46 -> 0.46 | No clear effect |
| Replace DTR with random routing (R-TR) (ablation) | DiT and ADM · FFHQ, ImageNet, and MS-COCO | FID | Hurt |
| Add DTR to ANT-UW loss weighting | DiT-L/2 · ImageNet 256×256 class-conditional | FID: 5.85 -> 4.61 | Helped |
| Add DTR to Min-SNR loss weighting | DiT-L/2 · ImageNet 256×256 class-conditional | FID: 9.58 -> 8.24 | Helped |
| Set DTR α to 4 rather than the other tested values (ablation) | DiT-B/2 + DTR · FFHQ 256×256 unconditional | FID: 7.78 -> 7.32; 7.64 -> 7.32; 7.77 -> 7.32; 7.40 -> 7.32 | Helped |
| Increase DTR α from 4 to 5 (ablation) | DiT-B/2 + DTR · FFHQ 256×256 unconditional | FID: 7.32 -> 7.40 | Hurt |
| Increase DTR α from 4 to 5 (ablation) | DiT-L/2 + DTR · ImageNet 256×256 class-conditional | FID: 8.90 -> 9.66 | Hurt |
Surprising findings
- Naive random task routing (R-TR) degrades performance relative to the baseline on most metrics, e.g., DiT ImageNet FID rises from 12.59 to 16.18.
- On unconditional FFHQ, DTR alone nearly matches the best loss-weighting-plus-DTR combination, so DTR largely substitutes for loss weighting there.
Limitations
- Masks are fixed and hand-designed, not learned during training.
- Hyperparameters alpha and beta must be tuned (alpha=4, beta=0.8 chosen empirically).
- Unconditional and text-conditional DiT variants are extensions made by the authors and not fully optimized.
Experimental setup
Datasets: FFHQ ImageNet MS-COCO
Benchmarks: FFHQ 256x256 unconditional ImageNet 256x256 class-conditional ImageNet 64x64 MS-COCO text-to-image
Metrics: FID IS Precision Recall CKA similarity
Baselines: DiT ADM R-TR (random task routing) Min-SNR ANT-UW P2 weighting Multi-experts (DiT-B/2 x 4) MaxRoaming CDTR ERCDT
Builds on: task routing DiT ADM Min-SNR ANT
Closest papers by meaning
- DiffusionMTL: Learning Multi-Task Denoising Diffusion Model from Partially Annotated Data (CVPR 2024 · similarity 0.92)
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.92)
- Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning (ICLR 2025 · similarity 0.89)
- Diffusion Models for Multi-Task Generative Modeling (ICLR 2024 · similarity 0.88)
- Decouple-Then-Merge: Finetune Diffusion Models as Multi-Task Learning (CVPR 2025 · similarity 0.88)
- Why DiT Models Underperform as Representation Learners without Long Skip Connections (NeurIPS 2026 · similarity 0.88)
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.