Technical Summary: UNIGP – Taming Diffusion Transformer for Prior-Preserved Unified Generation and Perception

Problem Statement

Recent advances in diffusion models have achieved remarkable success in both controllable image generation (e.g., Text-to-Image guided by depth or edges) and dense prediction tasks (e.g., monocular depth and surface normal estimation). However, existing approaches typically treat these as separate domains. Controllable generation methods (like ControlNet) are limited to generation tasks and often require massive datasets for precise control, while dense prediction methods (like Marigold or regression-based models) are often fine-tuned from generative backbones, leading to "catastrophic forgetting" of the model's versatile generative priors.

Current unified attempts (e.g., JointNet, UniCon) often fail to fully exploit the complementary relationship between generation and perception. They either incur massive computational overhead by duplicating entire backbones or compromise generative capabilities by directly fine-tuning the backbone for perception. The core challenge is to unify these heterogeneous tasks within a single framework that preserves the pre-trained generative priors while enabling accurate dense prediction and controllable generation, without requiring complex task-specific architectures or losses.

Methodology: UNIGP Framework

The authors propose UNIGP, a framework built upon the Multi-Modal Diffusion Transformer (MMDiT) architecture. The method consists of two primary innovations: a disentangled architectural branch and a unified training strategy.

1. Disentangled Unified Generation and Perception (DUGP)

Following the principle of Occam's razor, UNIGP avoids duplicating the entire MMDiT backbone. Instead, it introduces the DUGP branch, which copies only the image branch ($\mathcal{F}_i$) of the pre-trained MMDiT.

• Architecture: The DUGP branch is initialized by copying the image branch parameters. It is divided into a Condition Patcher, Stacked Control Layers, and Stacked Perception Layers.
• Prior Preservation (Gradient Isolation): A critical design choice is the "bypass" of the additive connection in the perception layers. In the control layers, the output of the DUGP branch is added to the backbone's output via a zero-initialized linear layer ($\mathcal{L}_{zero}$), allowing the condition to guide generation. However, in the perception layers, this additive path is severed. This ensures that gradients from the dense prediction loss flow exclusively through the DUGP branch, preventing them from altering the pre-trained weights of the backbone. This effectively quarantines the generative priors while allowing the DUGP branch to learn perception as a non-invasive "reader."
• Multi-Condition Handling: For scenarios requiring multiple conditions, the DUGP condition patcher is duplicated for each condition, with interactions handled via sequence concatenation to maintain inference efficiency.

2. Unified Dataset and Training Strategy

To bridge the gap between generation and perception, UNIGP employs a unified training strategy on a heterogeneous dataset combining filtered generation data and synthetic perception data.

• Binary Loss Weighting: Instead of complex loss balancing, the model uses a simple binary routing mechanism. For each sample in a mixed batch, the loss weights ($\lambda_g, \lambda_p$) are set to 1 or 0 based on the data source. If the sample is from the generation dataset, the generation loss is active; if from the perception dataset, the perception loss is active.
• Objective: The network learns to predict the velocity field for RGB, depth, and normal latents simultaneously, optimizing for both generative quality and geometric accuracy within a single forward pass.

Key Contributions

1. UNIGP Framework: A unified Diffusion Transformer model that integrates controllable generation and dense prediction. It utilizes a DUGP branch that copies only the image branch of MMDiT, adhering to Occam's razor to minimize computational overhead while preserving the backbone's versatile priors.
2. Unified Training Strategy: A novel dataset and training approach that jointly learns heterogeneous tasks using a simple binary loss weighting scheme, eliminating the need for task-specific architectural designs or complex loss balancing.
3. Demonstrated Synergy: The paper provides evidence that joint training yields complementary benefits: generative priors enrich perceptual details (capturing fine textures), while perceptual learning enforces strict spatial alignment in generation.

Experimental Results

Extensive experiments demonstrate that UNIGP surpasses prior unified approaches and performs on par with specialized "expert" methods.

• Controllable Generation: UNIGP outperforms single-control (ControlNet) and multi-control (UniControlNet) models in FID and CLIP scores across various conditions (depth, normal, sketch, pose, Canny). Notably, it achieves the lowest RMSE in depth-to-image and normal-to-image tasks, indicating superior structural alignment.
• Dense Prediction: In zero-shot depth and surface normal estimation, UNIGP achieves performance comparable to state-of-the-art regression models (e.g., DepthAnything V2, DSINE) despite training on significantly fewer images (59K vs. 62.6M for DepthAnything). It significantly outperforms other generative perception methods (Marigold, GeoWizard) and unified models (JointNet, UniCon).
• Joint Generation: UNIGP supports the simultaneous generation of RGB, depth, and surface normal maps from text prompts, outperforming baselines like LDM3D and JointNet in both image quality (FID) and geometric consistency (GD/GN metrics).
• Ablation Studies:
  • Synergy: Removing perception training degrades the structural alignment of generated images (higher RMSE), while removing generation training results in overly smoothed depth maps lacking fine details.
  • Design Choices: Replacing DUGP with a full backbone duplication (JointNet-style) or full fine-tuning (Marigold-style) results in significantly inferior performance, confirming that the disentangled gradient flow is the key to success.
  • Layer Balance: An equal allocation of stacked control and perception layers (12:12) yields the optimal balance between generation control and perception accuracy.

Significance and Claims

The paper claims that UNIGP successfully bridges the separated fields of diffusion-based synthesis and dense geometry prediction. Its primary significance lies in demonstrating that generative priors and perceptual constraints are not mutually exclusive but mutually reinforcing.

• Prior Preservation: By isolating the gradient flow, UNIGP retains the high-quality generative capabilities of the pre-trained MMDiT backbone, which is often lost in perception-specific fine-tuning.
• Efficiency: The method achieves state-of-the-art results with a relatively small parameter increase (from 2B to 3B) compared to duplicating the entire backbone, making it a computationally efficient unified solution.
• Versatility: The framework supports a wide range of tasks within a single model, including text-to-image, depth/normal estimation, and any-condition generation, offering a more versatile alternative to maintaining multiple specialized models.

The authors acknowledge that while UNIGP introduces additional computational overhead compared to the base model (increasing inference time from ~5s to ~9.5s), this trade-off is justified by the capability to replace multiple specialized architectures with a single, highly versatile unified model.