Research map / Image generation / Fast diffusion sampling
Condition-Aware Neural Network for Controlled Image Generation
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 Condition-Aware Neural Network (CAN), which generates condition-dependent convolution/linear weights for diffusion transformers, and CaT, which combines CAN with EfficientViT for efficient class-conditional and text-to-image generation.
Problem
Adding conditional control to image generative models is typically done via feature-space manipulation (adaptive normalization, attention), while conv/linear weights stay static across conditions, limiting performance and efficiency.
Approach
Make selected conv/linear layers condition-aware by adding a weight generation module that outputs a condition-dependent weight added to a static weight; select which modules to make condition-aware, and implement per-sample kernels as a grouped convolution.
Key results
- CaT-L0 achieves 2.78 FID on ImageNet 512x512 with 10G MACs per step and 20 steps, outperforming DiT-XL/2 (3.04 FID, 525G MACs, 250 steps).
- On ImageNet 256x256 with cfg, CaT-B0 reaches 2.09 FID, outperforming DiT-XL/2 (2.27) and UViT-H/2 (2.29) at 9.9x fewer MACs than DiT-XL/2.
- On UViT-S/2 ImageNet 256x256, CAN improves FID from 28.32 to 8.82 and CLIP score from 30.09 to 31.74 with condition-aware DW conv, patch embedding, and output projection layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add CAN to UViT-S-Deep/4 | UViT-S-Deep/4 · ImageNet 512×512 | FID-50K: 4.05 -> 4.04 | Helped |
| Use CaT-L0 instead of DiT-XL/2 | ImageNet class-conditional generation · ImageNet 512×512 | FID-50K: 3.04 -> 2.78 | Helped |
| Scale CaT-L0 up to CaT-L1 | CaT · ImageNet 512×512 | FID-50K: 2.78 -> 2.48 | Helped |
| Use CaT-B0 instead of DiT-XL/2 | ImageNet class-conditional generation · ImageNet 256×256 | FID-50K: 2.27 -> 2.09 | Helped |
| Add CAN to UViT-S/2 (ablation) | UViT-S/2 · ImageNet 256×256 | FID: 28.32 -> 8.82 | Helped |
| Make the head condition-aware rather than using a static head (ablation) | UViT-S/2 · ImageNet 256×256 | FID: 8.82 -> 12.29 | Hurt |
| Make output projection condition-aware (ablation) | UViT-S/2 with condition-aware DW Conv and patch embedding · ImageNet 256×256 | FID: 10.23 -> 8.82 | Helped |
| Use class label only rather than timestep only in condition embedding (ablation) | CAN on UViT-S/2 · ImageNet 256×256 | FID: 15.16 -> 10.01 | Helped |
| Include both class label and timestep in condition embedding (ablation) | CAN on UViT-S/2 · ImageNet 256×256 | FID: 10.01 -> 8.82 | Helped |
| Use CAN alone instead of adaptive normalization (ablation) | DiT-S/2 · ImageNet 256×256 | FID: 39.44 -> 26.44 | Helped |
| Add adaptive normalization to CAN (ablation) | DiT-S/2 · ImageNet 256×256 | FID: 26.44 -> 21.76 | Helped |
| Add attention with condition as tokens to CAN (ablation) | UViT-S/2 · ImageNet 256×256 | FID: 8.79 -> 8.82 | No clear effect |
Surprising findings
- Making all modules condition-aware is not beneficial; some choices (head, QKV projection, MLP) hurt performance.
- Adaptive kernel selection (CondConv-style) does not match CAN's performance even with many base kernels.
Limitations
- Grouped-convolution implementation incurs 30-40% training overhead on GPUs.
- Evaluated only on ImageNet class-conditional and small COCO text-to-image; large-scale text-to-image, video not tested.
- Combining CAN with attention-based conditioning slightly hurts UViT performance.
Experimental setup
Datasets: ImageNet COCO 2014
Benchmarks: ImageNet 256x256 ImageNet 512x512 COCO 256x256
Metrics: FID CLIP score MACs per step latency on Jetson AGX Orin
Baselines: ADM ADM-U DiT-XL/2 UViT-H/4 LDM-4 UViT-L/2 UViT-H/2 Friro adaptive normalization attention (condition as tokens) CondConv-style adaptive kernel selection
Builds on: DiT UViT EfficientViT CondConv classifier-free guidance
Closest papers by meaning
- A Simple Approach to Unifying Diffusion-based Conditional Generation (ICLR 2025 · similarity 0.87)
- Simplifying Control Mechanism in Text-to-Image Diffusion Models (AAAI 2025 · similarity 0.87)
- CtrLoRA: An Extensible and Efficient Framework for Controllable Image Generation (ICLR 2025 · similarity 0.86)
- CoDi: Conditional Diffusion Distillation for Higher-Fidelity and Faster Image Generation (CVPR 2024 · similarity 0.86)
- FlexControl: Computation-Aware Conditional Control with Differentiable Router for Text-to-Image Generation (ICML 2025 · similarity 0.85)
- CPO: Condition Preference Optimization for Controllable Image Generation (NeurIPS 2025 · similarity 0.84)
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.