VLD-MMDiT: Variable Layers Decomposition

• The paper presents VLD-MMDiT, which decomposes an RGB image into a variable number of semantically disentangled RGBA layers using a flow-matching diffusion process.
• Its architecture integrates multi-modal attention and a novel 3D rotary positional encoding to achieve efficient, end-to-end layer decomposition without recursive peeling.
• Empirical results demonstrate state-of-the-art performance, with significant improvements in both quantitative metrics (e.g., soft-IoU) and qualitative layer separability.

Variable Layers Decomposition MMDiT (VLD-MMDiT) is a key architectural component within Qwen-Image-Layered, a diffusion-based image generation system designed for decomposing a single RGB image into a stack of multiple, semantically disentangled RGBA layers. Unlike traditional raster representations, which entangle all visual content into a single canvas, VLD-MMDiT leverages a conditional diffusion process with a shared RGBA-VAE latent space, providing variable-length layer decomposition in an efficient, end-to-end manner. This architecture supports inherent editability, with each output layer independently manipulable, and achieves state-of-the-art results in quantitative and qualitative image decomposition metrics (Yin et al., 17 Dec 2025).

1. Objectives and Problem Formulation

VLD-MMDiT addresses the task of converting an input image $I \in \mathbb{R}^{H \times W \times 3}$ into a set of $N$ RGBA layers, $L = [L_1, \ldots, L_N]$, where each $L_i \in \mathbb{R}^{H \times W \times 4}$. The decomposition goal is to achieve semantic disentanglement such that edits to one layer do not compromise the integrity of the others. The core challenge is modeling this as a single-shot process, capable of supporting a variable number $N$ of layers, without relying on recursive foreground–background peeling or error-prone stepwise methods.

The flow-matching diffusion approach adopted in VLD-MMDiT operates in the unified latent space of an RGBA-VAE encoder $\mathcal{E}$ and decoder $\mathcal{D}$. The architecture conditions on both a latent vector $z_I = \mathcal{E}(I)$ and an optional text descriptor $h$; these inputs inform the decomposition by guiding the denoising trajectory from isotropic Gaussian noise $x_1 \sim \mathcal{N}(0, I)$ towards the true layer latents $N$0.

2. Architecture and Design Modifications

The VLD-MMDiT extends the standard Multi-Modal Diffusion Transformer (MMDiT) by incorporating innovations tailored for multi-layer decomposition:

A. Multi-Modal Attention over Layer Stacks:

The input image latent $N$1 and the current noisy layer latent $N$2 are patchified by a factor of 2× along each spatial dimension, converting them into tractable sequences. These, along with projected text prompt features ($N$3), are concatenated to form a multimodal sequence which is processed via a single MultiHeadAttention block:

$N$4

This mechanism jointly fuses intra-layer, inter-layer, and cross-modal interactions, supporting dense conditioning and semantic alignment across layers.

B. Layer3D Rotary Positional Encoding (RoPE):

A third positional axis $N$5 is introduced, supplementing standard $N$6 positional encodings, where $N$7 corresponds to the conditioning image latent and $N$8 index each output layer. The 3D RoPE equips the attention module with the capability to attend dynamically over arbitrary layer counts, enabling sequence-length agnostic processing and robust generalization to variable $N$9.

3. Mathematical Foundations

The layer decomposition process in VLD-MMDiT utilizes the flow-matching paradigm, parametrized as follows:

• Latent Initialization: $L = [L_1, \ldots, L_N]$0
• Noise Sampling: $L = [L_1, \ldots, L_N]$1, $L = [L_1, \ldots, L_N]$2
• Time-t Interpolation and Velocity:

$L = [L_1, \ldots, L_N]$3

• Prediction: The network $L = [L_1, \ldots, L_N]$4 predicts $L = [L_1, \ldots, L_N]$5 given $L = [L_1, \ldots, L_N]$6.

Loss Function:

$L = [L_1, \ldots, L_N]$7

No adversarial or perceptual losses are used within VLD-MMDiT; the training signal is exclusively driven by flow-matching reconstruction.

Attention Update Example:

$L = [L_1, \ldots, L_N]$8

Here, $L = [L_1, \ldots, L_N]$9 and $L_i \in \mathbb{R}^{H \times W \times 4}$0 carry the 3D RoPE embeddings for $L_i \in \mathbb{R}^{H \times W \times 4}$1.

4. Integration with RGBA-VAE

The RGBA-VAE encoder-decoder pair defines the latent space for both input conditioning and output prediction:

• The encoder $L_i \in \mathbb{R}^{H \times W \times 4}$2 (first convolution expanded from 3 to 4 channels) processes RGB and RGBA images, ensuring unified latent embedding.
• The decoder $L_i \in \mathbb{R}^{H \times W \times 4}$3 (last convolution expanded similarly) reconstructs RGBA layers from their latent representations.
• Initial weights are set to $L_i \in \mathbb{R}^{H \times W \times 4}$4 and $L_i \in \mathbb{R}^{H \times W \times 4}$5, with $L_i \in \mathbb{R}^{H \times W \times 4}$6.
• The input image $L_i \in \mathbb{R}^{H \times W \times 4}$7 is encoded as $L_i \in \mathbb{R}^{H \times W \times 4}$8, and output stack $L_i \in \mathbb{R}^{H \times W \times 4}$9 is encoded as $N$0.
• All operations within VLD-MMDiT are performed in this shared manifold, simplifying conditioning and reconstruction.

5. Multi-Stage Training Framework and Data Pipeline

The training of VLD-MMDiT involves three distinct phases:

Stage	Objective	Data Source
1	Text-to-RGB and Text-to-RGBA generation	Large text–image corpus (no layers)
2	Text-to-Multi-RGBA (T2L): composite and multi-layer outputs ($N$1)	PSD-extracted multilayer sets + generated captions
3	Image-to-Multi-RGBA (I2L): layer decomposition conditioned on $N$2	PSD-extracted multilayer sets

The PSD extraction pipeline involves parsing raw Photoshop documents, discarding irrelevant or low-quality layers, merging non-overlapping layers to manage $N$3, and generating semantic captions using Qwen2.5-VL.

6. Training and Inference Details

The core protocol for VLD-MMDiT training includes:

• Batch sampling from PSD-derived datasets $N$4.
• Encoding $N$5 and $N$6 into $N$7 and $N$8, respectively, via RGBA-VAE $N$9.
• Noise and time parameter sampling ($\mathcal{E}$0, $\mathcal{E}$1), followed by layer latent interpolation.
• Forwarding $\mathcal{E}$2 through VLD-MMDiT (with $\mathcal{E}$3, $\mathcal{E}$4, $\mathcal{E}$5 inputs).
• Computing and applying the flow-matching loss.
• Optimizer: Adam with learning rate $\mathcal{E}$6.
• Batch size: 128.
• Max layer count $\mathcal{E}$7.
• Training steps: 500 K for Stage 1, 400 K for Stage 2 and Stage 3.
• RGBA-VAE latent dimensions: $\mathcal{E}$8 channels, $\mathcal{E}$9.
• Patchification: 2× per spatial dimension.
• Attention heads: 16, with head dimension 64.
• Layer3D RoPE: $\mathcal{D}$0.

Sampling (Inference) Algorithm:

• Encode input $\mathcal{D}$1 to $\mathcal{D}$2
• Initialize $\mathcal{D}$3
• Iteratively update from $\mathcal{D}$4 to $\mathcal{D}$5:
  • Predict velocity $\mathcal{D}$6
  • Update latent $\mathcal{D}$7 using rectified flow integrator
• Decode final layer latent $\mathcal{D}$8 via $\mathcal{D}$9 to obtain $z_I = \mathcal{E}(I)$0

7. Empirical Performance and Ablation

VLD-MMDiT demonstrates superior quantitative performance on the Crello dataset relative to leading baselines, as shown:

Method	RGB $z_I = \mathcal{E}(I)$1 (weighted by $z_I = \mathcal{E}(I)$2)	$z_I = \mathcal{E}(I)$3 soft-IoU
LayerD	0.0709	0.7520
Qwen-Image-Layered-I2L	0.0594 ($z_I = \mathcal{E}(I)$4)	0.8705 ($z_I = \mathcal{E}(I)$5)

Ablation studies confirm the necessity of Layer3D RoPE, RGBA-VAE, and multi-stage training:

• Removal of Layer3D RoPE causes failure to distinguish layers (RGB $z_I = \mathcal{E}(I)$6).
• Omitting RGBA-VAE manifests a persistent gap (RGB $z_I = \mathcal{E}(I)$7, $z_I = \mathcal{E}(I)$8 IoU $z_I = \mathcal{E}(I)$9).
• Excluding multi-stage training yields suboptimal adaptation (RGB $h$0, $h$1 IoU $h$2).

Qualitative gains include elimination of recursive peeling errors, improved content separation, and efficient, robust support for decompositions ranging from one to twenty layers. The architecture is end-to-end trainable and computationally efficient, establishing a new paradigm for image editing via inherent layerwise disentanglement (Yin et al., 17 Dec 2025).