PaperlayerView exact revision on arXiv

Source-linked AI summary

MMDiff: Extending Diffusion Transformers for Multi-Modal Generation

Yagmur Akarken, Orest Kupyn, Christian Rupprecht

arXiv:2606.16673v2cs.CV

TL;DR

Diffusion transformers compute rich perceptual representations during generation, but these intermediate features are discarded and downstream modalities are re-extracted from rendered pixels. MMDiff reuses a frozen diffusion transformer with lightweight decoder heads and multi-timestep feature fusion, improving semantic segmentation by +28.7% mIoU over single-timestep extraction while supporting dense perceptual prediction and synthetic data generation.

  • Problem

    Diffusion pipelines discard intermediate perceptual representations and rely on separate models to re-extract semantic, geometric, and structural information from generated pixels.

  • Method

    MMDiff decodes combinations of dense perceptual modalities from a frozen diffusion transformer using lightweight heads, multi-timestep features, and concept-driven attention.

  • Results

    +28.7% mIoU improves semantic segmentation over single-timestep extraction, while MMDiff achieves strong performance across semantic segmentation, salient object detection, and depth estimation.

  • Takeaways & Limitations

    MMDiff shows that recycling diffusion features across the denoising trajectory can support joint image generation and dense perceptual annotation, enabling effective synthetic data generation.

Abstract

from arXiv · show

Diffusion transformers have demonstrated remarkable generative capabilities, yet the rich perceptual representations computed across their denoising trajectory are discarded once the content is rendered. We present MMDiff, a framework that transforms a frozen diffusion transformer into a multi-modal generative system that jointly produces images alongside any combination of dense perceptual modalities using lightweight decoder heads. Our central finding is that perceptual information is temporally distributed along the denoising trajectory, and that multi-timestep feature fusion with spatially varying aggregation weights is essential, improving semantic segmentation results by up to 28.7% mIoU over single-timestep extraction. We further adopt concept-driven attention extraction for interpretable spatial guidance, and show that frozen diffusion features are competitive with and complementary to state-of-the-art encoders such as DINOv3. By training only lightweight decoder heads on a frozen backbone, we achieve strong performance in semantic segmentation, salient object detection, and depth estimation, and demonstrate that this framework enables effective synthetic data generation at scale.

1 Introduction

MMDiff repurposes a frozen diffusion transformer to generate images and dense perceptual modalities through lightweight decoder heads. It shows that multi-timestep fusion, concept-driven attention, and complementary generative-discriminative features improve downstream prediction and synthetic annotation generation.

  • Motivation: Diffusion transformers discard intermediate representations after rendering, requiring separate task-specific models to extract information such as segmentation or depth.These models process thousands of tokens across dozens of layers and denoising timesteps for one image or video.
  • Framework: MMDiff decodes any combination of dense perceptual modalities from a frozen diffusion transformer using lightweight decoder heads without modifying the generative backbone.The framework extends a frozen DiT for dense perceptual prediction.
  • Feature extraction: +28.7% mIoU: Multi-timestep feature fusion improves segmentation over single-timestep extraction by aggregating richer representations across the denoising trajectory.The method addresses information discarded by single-timestep extraction, used by prior works.
  • Feature extraction: Concept-driven attention extraction enriches features with richer spatial guidance than averaged cross-attention.This attention mechanism is used alongside multi-timestep feature fusion for dense prediction.
  • Results and applications: Frozen diffusion transformer features achieve comparable performance to DINO representations and are complementary, while MMDiff supports synthetic images with high-quality annotations.A single frozen model simultaneously generates images and extracts annotations, substantially outperforming prior U-Net-based methods.

2 Related Work

Prior work has shifted diffusion models toward scalable transformer architectures and extended them to generate perceptual modalities alongside images. Existing approaches commonly rely on separate decoders, while image-based methods face limitations in spatial resolution or concept coverage.

  • Diffusion Transformer Architectures: Diffusion models have transitioned from U-Nets to transformer architectures, with latent diffusion enabling efficient high-resolution synthesis in compressed latent spaces.DiT showed that transformers scale more efficiently than U-Nets by processing image patches with equal capacity through self-attention rather than hierarchical downsampling.
  • Multi-Model Diffusion: Image and video generators increasingly predict modalities such as segmentation, depth, geometry, point maps, and camera poses alongside generated content.Many existing approaches train decoder networks or geometric branches separately from the generation process.
  • Multi-Model Diffusion: For images, DatasetDM decodes segmentation masks and depth from Stable Diffusion features, but coarse U-Net spatial resolution limits annotation quality.DiffuMask is similarly limited to single-concept segmentation because it relies on cross-attention maps.

3 Method

MMDiff extends a frozen diffusion transformer into a multi-modal generator by fusing internal features across layers and denoising timesteps with lightweight decoder heads. It supports joint image generation and dense prediction, including synthetic annotation generation, real-image analysis, and concept-driven spatial guidance.

  • Framework Overview: MMDiff extracts DiT representations to jointly generate images and dense modalities such as semantic segmentation, depth, and saliency using lightweight decoder heads.The framework also generates synthetic datasets with automatically extracted annotations during image generation.
  • Multi-Timestep Extraction: Early denoising timesteps capture coarse layouts and scene semantics, whereas late timesteps preserve fine details and local structure, motivating multi-timestep feature extraction.The same extraction mechanism operates on generated images and real images noised through the forward diffusion process.
  • Token Aggregation: The weighted feature combination is refined with CBAM and a residual connection from the cleanest timestep, allowing smooth regions and boundaries to favor different temporal information.Per-pixel weights sum to 1, while the residual from t = 0 preserves fine details.
  • Concept-Driven Attention: Concept-driven attention adds a separate concept stream that gathers information from image patches without affecting generation, producing task-specific spatial guidance at inference time.Concepts include “object” and “background” for segmentation, “near” and “far” for depth, and “salient” and “contour” for salient object detection.
  • Feature Complementarity: DINO-v3 features can optionally complement diffusion representations because generative and contrastive training capture distinct visual characteristics.The integration uses features from DINO-v3-ViT-B/16.

4 Experimental Evaluation

MMDiff is evaluated across dense prediction, consistency, synthetic-data, and controlled feature-quality settings. It outperforms prior methods, with multi-timestep aggregation and concept-driven attention providing major gains.

  • Evaluation setup: MMDiff evaluates salient object detection, semantic segmentation, and depth estimation using task-specific metrics, while comparing diffusion extractors, discriminative encoders, and synthetic-data methods.Baselines are retrained under identical conditions, including 20 epochs, optimizer, learning rate, and decoder architectures.
  • Multi-modal consistency: MMDiff achieves substantially higher multi-modal consistency than Stable Diffusion-based baselines across all reported metrics, indicating more accurate visual-content capture from DiT features.Consistency is measured on 200 images per task using shared prompts, pseudo-ground-truth labels, and manually verified annotations.
  • Synthetic data evaluation: 78.9% mIoU with pure synthetic training surpasses DatasetDM (65.9%), DiffuMask (54.7%), Dataset Diffusion (68.2%), and the DINO-v3 baseline (75.2%).MMDiff combines diffusion and DINO-v3 features for synthetic-data evaluation using EoMT as the downstream model.
  • Synthetic data evaluation: MMDiff outperforms the DINO-v3 standalone baseline on salient object detection and depth estimation, while synthetic pretraining plus real-data fine-tuning produces the best results across all metrics.This shows that labels extracted by MMDiff complement real data across DUTS and NYU Depth v2.
  • Feature quality validation: +2.08% mIoU on segmentation and +0.011 S-measure over VPD demonstrate that frozen MMDiff features outperform U-Net-based diffusion methods.The controlled evaluation isolates prediction quality on fixed real images and tests multi-timestep aggregation, concept-driven attention, and discriminative feature integration.
  • Ablation studies: 50.20% to 78.90% mIoU (+28.7%) is achieved by aggregating features over four timesteps, the largest ablation improvement.Concept-driven attention adds +13.57% mIoU with single-timestep features and reaches 63.77% mIoU versus 58.31% for averaged cross-attention maps.

5 Extension to Video Diffusion Models

MMDiff’s multi-timestep feature aggregation transfers from image-space prediction to video-to-3D Gaussian reconstruction, improving performance over denoised latents across all evaluated metrics. Video DiT analysis shows the same division of labor across timesteps and layers as image DiTs.

  • Video-to-3DGS transfer: MMDiff replaces denoised video latents with multi-timestep aggregated features in Lyra’s video-to-3DGS pipeline, training on 1,000 RealEstate10K videos and evaluating 200 held-out scenes.The decoder is trained in a low-data setting using the latent-to-3DGS reconstruction pipeline.
  • Video-to-3DGS transfer: +0.24 PSNR, +0.042 SSIM, and −0.017 LPIPS improve over the denoised latent across all three metrics.These gains indicate that multi-timestep aggregation extends beyond image-space dense prediction to video and 3D.
  • Layer/timestep analysis: Early timesteps and shallow layers carry fine detail, while later timesteps and deeper layers carry semantic structure in Wan-2.1 video DiT features.The layer/timestep behavior mirrors trends observed for image diffusion transformers.

6 Conclusion

MMDiff converts a frozen diffusion transformer into a multimodal generative system that extracts dense perceptual annotations during image generation. It recycles denoising-trajectory features instead of discarding them after rendering, addressing inefficiencies in current generative pipelines.

  • Framework: MMDiff transforms a frozen diffusion transformer into a multimodal generative system that extracts dense perceptual annotations during image generation.The framework produces annotations alongside image generation.
  • Feature reuse: MMDiff recycles rich features computed across the denoising trajectory rather than discarding them after rendering.This reuses information already computed during generation.
  • Pipeline efficiency: By reusing denoising features, MMDiff addresses the inefficiency of re-extracting semantic, geometric, and structural information from generated pixels with separate task-specific models.The approach targets redundant post-generation extraction in current generative pipelines.

Supplementary Material … A.2 Salient Object Detection

The supplementary experiments show that MMDiff’s diverse, multi-source features are especially useful with limited labeled data, while synthetic pretraining consistently improves downstream performance when combined with real-data fine-tuning. Gains are demonstrated for semantic segmentation and salient object detection across multiple supervision regimes.

  • A Learning from Limited Data: MMDiff combines multi-timestep features, concept-driven attention, and discriminative representations to reduce overfitting and enrich training signals in low-data regimes.The passage emphasizes complementary information from each sample and the risks of relying on a narrow representation.
  • A.1 Semantic Segmentation: 72.7% mIoU from 10k synthetic samples trained with 1k real samples approaches 79.3% mIoU from 1k real samples.The extracted labels are accurate enough to supervise EoMT, although limited real data retains an advantage.
  • A.1 Semantic Segmentation: 85.2% mIoU from synthetic pretraining plus real fine-tuning exceeds 79.3% mIoU from real-only training by +5.9% mIoU.This result uses 1k real samples and 10k synthetic samples with EoMT.
  • A.1 Semantic Segmentation: 87.7% mIoU from combined training exceeds 84.5% mIoU from real-only training and 77.4% mIoU from synthetic-only training with 5k real samples.Synthetic pretraining remains beneficial even with full data, improving performance by +3.0% mIoU over real-only training.
  • A.2 Salient Object Detection: 0.799 S-measure from synthetic images generated from 1k real samples approaches 0.838 S-measure from 10.5k real samples for salient object detection.IS-Net is used as the downstream model, and the synthetic-only result uses one-tenth as much real data as the comparison.
  • A.2 Salient Object Detection: 0.840 S-measure from synthetic pretraining plus 1k real fine-tuning outperforms the real-only baseline.The result is obtained with IS-Net and synthetic images generated from only 1k real samples.
  • A.2 Salient Object Detection: 0.863 S-measure and 0.049 MAE from combined training exceed 0.838 S-measure from real-only training with the full 10.5k training set.With 5k real samples, the combined approach also achieves 0.854 S-measure and 0.054 MAE, showing consistent benefits across regimes.

A.3 Depth Estimation

Depth estimation experiments show that synthetic pre-training followed by real fine-tuning consistently outperforms real-only or synthetic-only training across data regimes, including the full dataset. The full-data combination reduces relative error by 8.2%.

  • A.3 Depth Estimation: 0.121 AbsRel is achieved by synthetic pre-training with real fine-tuning using 2.4k real samples, outperforming real-only and synthetic-only training.At 2.4k samples, real-only training achieves 0.1386 AbsRel, while synthetic-only training achieves 0.1938 AbsRel.
  • A.3 Depth Estimation: 0.1192 AbsRel and 0.880 δ1 are achieved by the combined approach with 12k real samples, versus 0.1303 AbsRel for real-only training.Synthetic-only training achieves 0.1900 AbsRel in this regime.
  • A.3 Depth Estimation: 0.1185 AbsRel and 0.885 δ1 are achieved by synthetic pre-training with real fine-tuning on the full 24k training set, surpassing real-only and synthetic-only training.Real-only achieves 0.1291 AbsRel, while synthetic-only achieves 0.1880 AbsRel.
  • A.3 Depth Estimation: 8.2% relative error reduction is achieved on the full dataset, with combined training reaching 0.1185 versus 0.1291 AbsRel for real-only training.The gains indicate that synthetic data provides complementary signal that enhances generalization even when substantial real data is available.

B Unified Multi-Task Decoder · C Feature and Concept Visualizations

MMDiff supports unified decoders across semantic and depth tasks, with positive semantic transfer but substantial depth degradation under joint supervision. Its visualizations show complementary multi-timestep features and interpretable concept-specific attention maps.

  • B Unified Multi-Task Decoder: A single decoder is trained across VOC+DUTS, VOC+NYU, or all three tasks using loss weights based on early-training relative loss magnitudes.This tests whether one decoder can serve multiple tasks from the same frozen MMDiff features.
  • B Unified Multi-Task Decoder: 84.92 VOC and 0.946 DUTS Fm match per-task baselines of 84.95 and 0.947 under unified VOC+DUTS training.The paired comparisons are VOC 84.92 vs. 84.95 and DUTS Fm 0.946 vs. 0.947.
  • B Unified Multi-Task Decoder: 85.42 VOC under VOC+NYU and 84.96 under all-three training indicate positive transfer between semantic tasks.Depth instead degrades sharply when combined with semantic supervision: NYU AbsRel is 0.22 for all-three versus 0.116 single-task.
  • C Feature and Concept Visualizations: Figures 6 and 7 visualize extracted features and learned concept attention maps to explain how MMDiff leverages multi-modal representations.The visualizations connect temporal and layer-wise feature diversity with concept-driven guidance for salient object detection.
  • C Feature and Concept Visualizations: Synthetic data generated from 2.4k, 12k, and 24k real samples substantially improves limited-supervision NYU Depth v2 performance when combined with real fine-tuning.Table 9 evaluates AbsRel, RMSE, δ1, δ2, and δ3.
  • C Feature and Concept Visualizations: FLUX features across timesteps and layers capture complementary information, with L=4 preserving fine-grained texture and L=37 encoding semantic object-level information.Fig. 6 combines t = 0.52 and t = 0 features with layers L=4 and L=37, alongside Dino-v3 layers L=2 and L=11.
  • C Feature and Concept Visualizations: Learned concept-specific attention maps decompose segmentation into interpretable visual properties rather than redundant representations.The maps separately attend to salient regions, background context, object contours, and fine-grained edges.

D Effect of Concept Sets

Concept attention substantially improves dense prediction by supplying spatial localization absent from frozen diffusion features. Broad structural concepts perform comparably to exhaustive class lists, while geometric descriptors are especially effective for depth and saliency.

  • Impact of Concept Guidance: 63.8% mIoU versus 50.2% on Pascal VOC demonstrates the gain from adding concept guidance to frozen diffusion features.The concept module addresses the precise spatial localization required for dense prediction.
  • Structural vs. Categorical Concepts: Concept attention primarily improves spatial localization by guiding decoders toward precise foreground boundaries.The broad “object” concept produces a dense, high-confidence attention map covering foreground elements, while specific class tokens can be sparse or empty.
  • Geometric Priors for Depth and Saliency: AbsRel 0.431 is achieved on NYU Depth when geometric descriptors “near” and “far” augment the standard “depth” token.This outperforms semantic scene descriptions such as “furniture” or “floor.”

E Token Aggregation Module Architecture · F Extension to Other DiT Backbones

The token aggregation module fuses four timestep features with temporal modeling and spatially adaptive weighting, then refines them for downstream decoding. MMDiff transfers to Stable Diffusion 3.5 without architectural or training changes, with complementary gains from DINO-v3 features.

  • E Token Aggregation Module Architecture: The module aggregates temporal information while preserving spatial specificity for each extracted DiT layer.It operates independently on each layer’s feature maps.
  • E Token Aggregation Module Architecture: Four timestep feature maps are normalized with timestep-specific LayerNorm and projected from C = 3072 to D = 768.The two-stage preparation addresses diffusion activation variance before aggregation.
  • E Token Aggregation Module Architecture: A 3-layer temporal Transformer with 8 attention heads models nonlinear relationships using learnable temporal positional embeddings and LayerScale initialized at 10^-6.The small initialization makes residual layers initially behave as identity mappings.
  • E Token Aggregation Module Architecture: A 1 × 1 convolution predicts per-pixel timestep mixing logits, whose softmax weights sum to 1 at every spatial location.This produces content-adaptive temporal aggregation.
  • E Token Aggregation Module Architecture: The weighted features dynamically prioritize early or late timesteps, then receive a cleanest-timestep skip connection and CBAM refinement.The skip connection uses t = 0 features and a learnable scalar initialized to 0.1.
  • F Extension to Other DiT Backbones: MMDiff transfers to Stable Diffusion 3.5 without architectural or training modifications and performs strongly across semantic segmentation, saliency, and depth tasks.The transfer preserves multi-timestep aggregation and concept-driven attention transfer.
  • F Extension to Other DiT Backbones: 70.14% mIoU on PASCAL VOC is achieved by SD 3.5 features alone, while combining them with DINO-v3 substantially improves all three tasks.The results again support complementarity between generative and discriminative features.
  • F Extension to Other DiT Backbones: SD 3.5 performs worse than FLUX in the diffusion-only setting, and frozen feature quality is bounded by the generator’s internal representations.Stronger generators are described as producing richer features.

G Qualitative Results

Qualitative results show that MMDiff jointly predicts binary masks, semantic segmentation, and depth across diverse synthetic scenes. Compared with baselines, it produces coherent depth maps and handles challenging salient-object cases with more complete masks.

  • Diverse-scene predictions: MMDiff simultaneously infers binary masks, semantic segmentation, and depth across complex synthetic scenes.The examples span a wide variety of scenes.
  • Depth estimation: For depth estimation, MMDiff produces smooth, spatially coherent maps with sharp object boundaries.VPD exhibits oversmoothing, while DHF generates noisy predictions.
  • Salient object detection: For salient object detection, MMDiff accurately segments challenging multi-object scenes and low-contrast subjects.Baselines produce incomplete or fragmented masks in these cases.
  • Cross-task comparisons: Qualitative comparisons cover depth estimation, salient object detection, and semantic segmentation across NYU Depth V2, DUTS, and Pascal VOC examples.These comparisons are presented in Figures 10–12.
Loading 2606.16673v2…