MMDiff: Extending Diffusion Transformers for Multi-Modal Generation

Conference: ECCV 2026
Paper: ECCV 2026 Publication
Area: 3D Vision
Keywords: diffusion transformer, multimodal generation, feature aggregation, dense perceptual prediction, concept attention

TL;DR

Addressing the issue where diffusion transformers discard rich perceptual representations during rendering, MMDiff transforms a frozen DiT into a multi-modal generator using adaptive multi-timestep feature fusion, concept-driven attention, and complementary discriminative features to accurately predict semantic segmentation, saliency, and depth alongside generated images with lightweight decoder heads.

Background & Motivation

Modern image and video generative models are rarely deployed in isolation to produce raw RGB visual data alone. In real-world robotic planning, autonomous driving simulation, and 3D reconstruction pipelines, downstream tasks simultaneously require detailed scene geometry and fine-grained semantics, such as dense depth maps, semantic segmentation masks, and salient object contours. State-of-the-art diffusion transformers (DiTs), such as FLUX and Stable Diffusion 3, process thousands of visual tokens across dozens of transformer layers and iterative denoising steps. Yet, once the final RGB pixels are rendered, all the high-dimensional intermediate perceptual representations computed during generation are completely discarded. As a consequence, existing systems must rely on external task-specific perceptual networks to re-extract semantic and structural attributes from the generated pixels, incurring substantial computational redundancy and compounding upstream synthesis artifacts.

Prior efforts exploring multimodal annotation from generative models, such as DatasetDM and DiffuMask, have predominantly built upon convolutional U-Net diffusion architectures. These earlier approaches either extract features exclusively from the final denoised latent or from a single terminal denoising step. While U-Nets naturally construct semantic hierarchies through hierarchical spatial downsampling, modern DiTs operate on a flat architecture where all patch tokens maintain a uniform spatial resolution throughout every layer. In the absence of spatial hierarchy, semantic information is not confined to deeper layers; rather, it is temporally dispersed across the entire denoising trajectory. Early denoising timesteps possess high feature entropy and govern global layout and coarse scene semantics; intermediate timesteps refine object boundaries; and terminal timesteps preserve high-frequency textures and local details. Re-applying naive single-timestep extraction to flat DiT architectures fails to capture this coarse-to-fine trajectory, precipitating severe performance drops in dense prediction tasks.

The key insight of this paper is to exploit the temporal dynamics of the diffusion process itself, enabling a frozen generative backbone to produce aligned dense perceptual outputs without modifying its weights. The core idea is to perform adaptive multi-timestep feature aggregation with spatially varying weights and extract targeted spatial guidance via decoupled one-directional concept attention, thereby turning a frozen DiT into an efficient multi-modal generator with lightweight decoder heads.

Method

Overall Architecture

The pipeline accepts a text prompt (or an image with forward diffusion noise) and leverages a completely frozen 12B-parameter FLUX.1-dev backbone for image synthesis. During generation, intermediate token features are extracted across sparse layers and multiple denoising timesteps, routed to an adaptive token aggregation module, paired with decoupled concept attention maps, optionally enriched with frozen discriminative representations (e.g., DINO-v3), and finally decoded into target modalities via task-specific decoders totaling only ~36M trainable parameters.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Text Prompt / Noisy Image Latent"] --> B["Multi-Timestep Sparse Feature Extraction<br/>Sample FLUX single-stream layers and denoising timesteps"]
    B --> C["Adaptive Spatiotemporal Token Aggregation<br/>3-layer Transformer predicts dynamic weights + CBAM"]
    C --> D["Concept-Driven One-Directional Attention<br/>Decoupled concept tokens yield dot-product spatial maps"]
    D --> E["Discriminative Feature Complementarity<br/>Concatenate frozen DINO-v3 contrastive embeddings"]
    E --> F["Lightweight Task-Specific Decoders<br/>DeepLabV3+ / DPT heads output dense modalities"]

Key Designs

1. Multi-Timestep Sparse Feature Extraction: Eliminating Spatial and Temporal Redundancy DiT layers process over 1,000 spatial tokens each spanning 3,072 channels, making naive multi-layer, multi-timestep concatenation prohibitively expensive in GPU memory, while simple averaging destroys complementary temporal cues. Empirical analysis reveals that activation magnitudes in FLUX grow monotonically across single-stream blocks with high inter-layer correlation, while feature entropy varies substantially across denoising steps. Consequently, MMDiff adopts a sparse sampling scheme: selecting four representative single-stream layers \(\mathcal{L} = \{4, 15, 26, 37\}\) to span network depth, and four timesteps \(\mathcal{T} = \{0.78, 0.52, 0.26, 0.001\}\) across the noise schedule. To resolve the extreme activation magnitudes characteristic of DiTs, extracted features undergo linear projection down to dimension \(C'\) followed by Layer Normalization, establishing cross-layer numerical stability.

2. Adaptive Spatiotemporal Token Aggregation: Dynamic Content-Aware Weighting Smooth object interiors benefit from the global semantic coherence formed in early denoising steps, whereas complex object boundaries demand the sharp high-frequency localization present in late steps. To allow aggregation weights to adapt dynamically to local image content, MMDiff deploys a 3-layer Transformer (8 heads, hidden dimension 768, temporal positional embeddings, and LayerScale initialization). The module predicts per-pixel normalized temporal weights \(\alpha \in \mathbb{R}^{|\mathcal{T}|}\) satisfying \(\sum_{t} \alpha_t = 1\). The combined representation is further enhanced via a Convolutional Block Attention Module (CBAM) incorporating spatial and channel attention, alongside a residual skip connection from the cleanest timestep (\(t=0\)):

\[F_{\text{fused}} = \text{CBAM}\left(\sum_{t \in \mathcal{T}} \alpha_t \cdot F_t^{\text{proj}}\right) + F_0^{\text{proj}}\]

This learned spatially varying aggregation is the central driver of performance, yielding a +28.7% mIoU surge over single-timestep extraction on PASCAL VOC.

3. Concept-Driven One-Directional Attention: Non-Invasive Spatial Guidance Standard cross-attention maps between text and image tokens suffer dilution from stop-words and punctuation, and DiT prompt embeddings drift across layers, hindering the extraction of arbitrary target masks without altering image synthesis. MMDiff adopts concept-driven attention by introducing decoupled concept tokens (e.g., "object/background" for segmentation, "near/far" for depth, and "salient/contour" for saliency). Concept embeddings are updated via a strict one-directional attention mechanism:

\[o_c = \text{softmax}(q_c k_{xc}^T) v_{xc}\]

Concept queries attend to concatenated image and concept keys and values, whereas image tokens never attend to concepts, completely preserving the generative trajectory. Spatial guidance maps are computed via dot-product similarities in the attention output space \(\phi(o_x, o_c) = \text{softmax}(o_x o_c^T)\) and averaged across multi-modal layers, providing sharp, interpretable spatial cues to downstream heads.

4. Generative-Discriminative Feature Complementarity and Lightweight Decoding Diffusion features encode rich semantic context and compositional structures from generative objectives, whereas contrastive discriminative representations (e.g., DINO-v3) capture complementary localized boundary representations. MMDiff concatenates fused DiT features, concept maps, and frozen DINO-v3-ViT-B/16 features into a unified perceptual pool. Downstream decoding uses compact architectures: DeepLabV3+ with ASPP for semantic segmentation, and DPT-style hierarchical reassembly heads for depth and saliency. With only ~36M trainable parameters against the 12B frozen generative backbone, the entire decoding system operates at minimal additional training overhead.

Loss & Training

The 12B FLUX generative backbone remains completely frozen. The lightweight decoders and aggregation module are trained on a small set of annotated real images using AdamW with an initial learning rate of \(10^{-4}\) (\(5 \times 10^{-5}\) for decoders) and an Exponential Moving Average (EMA) factor of \(\beta = 0.999\).

Key Experimental Results

Main Results

Evaluating feature quality on real benchmark images demonstrates that MMDiff surpasses existing diffusion-based extractors and discriminative baselines across multiple dense perception tasks.

Method	Frozen Encoder	PASCAL VOC 2012 (mIoU ↑)	DUTS Saliency (\(S_m\) ↑ / \(F_m\) ↑ / MAE ↓)	NYU Depth v2 (AbsRel ↓ / RMSE ↓ / \(\delta_1\) ↑)
DatasetDM (U-Net)	Yes	41.19%	0.845 / 0.909 / 0.077	0.1536 / 0.612 / 0.772
Diffusion Hyperfeatures (U-Net)	Yes	67.57%	0.892 / 0.910 / 0.021	0.1348 / 0.503 / 0.844
VPD (U-Net Finetuned)	No	82.36%	0.912 / 0.916 / 0.024	0.1244 / 0.402 / 0.867
MMDiff (DiT Features Only)	Yes	78.90%	0.918 / 0.929 / 0.020	0.1175 / 0.370 / 0.951
DINO-v3 (Discriminative)	No	83.09%	0.920 / 0.922 / 0.021	0.1288 / 0.389 / 0.945
VPD + DINO-v3	No	83.77%	0.918 / 0.914 / 0.020	0.1223 / 0.387 / 0.946
MMDiff + DINO-v3 (Full Model)	Yes	84.95%	0.934 / 0.947 / 0.018	0.1164 / 0.365 / 0.952

Ablation Study

The ablation analysis isolates the performance contribution of each core component.

Configuration	VOC'12 mIoU (%) ↑	DUTS \(S_m\) ↑	DUTS MAE ↓	NYU AbsRel ↓	NYU \(\delta_1\) (%) ↑	Note
FLUX double-stream block (single step \(t=0\))	40.52	0.884	0.033	0.622	74.0	Early text-image interaction features
FLUX single-stream block (single step \(t=0\))	50.20	0.890	0.032	0.478	78.0	Unified single-stream representation
+ Cross-attention maps	58.31	0.869	0.043	0.455	80.0	Standard text prompt attention
+ Concept-driven attention	63.77	0.903	0.024	0.431	82.0	Decoupled concept tokens (+13.57% mIoU)
Multi-timestep + Concept attention	78.90	0.918	0.020	0.118	95.1	4-timestep fusion (+28.70% mIoU)
Full model (w/o adaptive aggregation)	78.21	0.918	0.017	0.130	81.4	Without Transformer weights and CBAM
Full model (4T + Concepts + DINO-v3)	84.95	0.934	0.018	0.116	95.2	Full framework with discriminative fusion

Key Findings

• Multi-timestep fusion drives representation richness: Expanding from single-step \(t=0\) extraction (50.20%) to multi-timestep aggregation (78.90%) yields a remarkable +28.7% mIoU improvement on PASCAL VOC, verifying that DiT perceptual knowledge is temporal rather than hierarchical.
• Concept attention outclasses averaged prompt attention: Task-specific concept tokens improve mIoU from 58.31% to 63.77% over prompt cross-attention, eliminating noise introduced by irrelevant sentence tokens.
• Effective synthetic data generation at scale: When training downstream segmentation models purely on synthetic data generated with extracted labels, MMDiff achieves 78.9% mIoU, substantially beating prior generative pipelines (DatasetDM 65.9%, Dataset Diffusion 68.2%).
• Generalization to Video-to-3D synthesis: When applied to the video-to-3DGS pipeline (Lyra) on RealEstate10K, replacing the terminal denoised latent with MMDiff aggregated features improves novel-view PSNR from 21.37 to 21.61 and SSIM from 0.705 to 0.747.

Highlights & Insights

• Computational recycling: Repurposes intermediate activations calculated during generation as high-value perceptual signals, eliminating redundant forward passes through separate vision models.
• Leak-free concept injection: The one-directional attention mechanism extracts arbitrary concept guidance maps without altering the appearance or identity of generated imagery.
• Bridging generative and discriminative features: Proves that frozen DiT representations rival supervised/contrastive models on dense geometry and semantics, and that fusing generative and discriminative representations achieves synergistic performance gains.

Limitations & Future Work

• Inference memory footprint: Caching intermediate features across 4 timesteps and 4 layers increases intermediate activation memory consumption by ~4x during generation.
• Manual concept definitions: Concept tokens for attention guidance currently require explicit task-level specification (e.g., "object", "background", "near", "far") rather than being learned automatically.
• Future directions: Integrating temporal aggregation into parameter-efficient adapters (e.g., LoRA) and expanding into dynamic online perception within 4D world models.

Related Work & Insights

• vs DatasetDM / DiffuMask: Previous methods extract single-timestep representations from U-Net downsampled feature maps; MMDiff extracts temporally distributed representations across a flat DiT architecture with adaptive spatial weighting.
• vs VPD: VPD requires computationally heavy full-network fine-tuning of the diffusion backbone; MMDiff keeps the 12B FLUX generative foundation frozen and only trains 36M decoder parameters.
• vs DINO-v3: While DINO-v3 offers strong contrastive localized representations, MMDiff demonstrates that generative diffusion trajectory features provide complementary structural and semantic context.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Formulates multi-timestep feature extraction and decoupled concept guidance specifically tailored to the architectural dynamics of diffusion transformers.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Rigorous three-protocol validation across semantic segmentation, saliency detection, depth estimation, and 3D Gaussian Splatting transfer.
• Writing Quality: ⭐⭐⭐⭐⭐ Clear progression, insightful layer/entropy visualizations, and solid analytical grounding.
• Value: ⭐⭐⭐⭐⭐ Provides a practical, highly scalable blueprint for zero-human-labor synthetic dataset curation and multimodal generation.