DiffusionDet, A Diffusion Model for Object Detection

Abstract

DiffusionDet is a new framework that formulates object detection as a denoising diffusion process.

• Training: The model diffuses ground truth boxes into a random distribution and learns to reverse this noise process.
• Inference: The model progressively refines a set of randomly generated boxes into final output results.

This approach offers significant flexibility and demonstrates favorable performance, achieving gains of 5.3 AP and 4.8 AP on the COCO dataset under zero-shot transfer settings.

---

Introduction

Object detection aims to predict a set of bounding boxes and associated labels for targets within an image. As a fundamental visual recognition task, it serves as the basis for various applications such as image segmentation and pose estimation.

Evolution of Object Detection

Modern approaches have evolved from using heuristics to learnable object queries:

• Traditional Detectors: Use surrogate regression or classification tasks on empirical methods like sliding windows, region proposals, and anchor boxes.
• DETR: Introduced learnable object queries to eliminate empirical components, creating an end-to-end pipeline. However, DETR remains dependent on a fixed set of learnable queries, limiting the number of detectable boxes.

The DiffusionDet Approach

This paper explores a simpler approach: Can we detect objects directly from random boxes without surrogate learnable queries?

DiffusionDet treats detection as a generative process, inspired by Denoising Diffusion Probabilistic Models (DDPM). By starting with random boxes that require no optimization during training, the model gradually refines their positions and sizes to cover target objects perfectly.

The Training Stage

1. Ground Truth Boxes: The model identifies the actual bounding boxes (position and size).
2. Add Gaussian Noise: Boxes are corrupted with Gaussian noise, making them uncertain/blurry.
3. Feature Extraction: An image backbone (such as ResNet or Swin Transformer) extracts features.
4. Region of Interest (RoI) Crop: Noisy boxes are used to crop features from the feature maps.
5. Detection Decoder: A specialized decoder is trained to predict the original denoised boxes from the noisy versions.

The Inference Stage

During inference, the model starts with completely random bounding boxes and reverses the learned process step-by-step to align the boxes with actual object locations.

Key Advantages

Being a probabilistic model, DiffusionDet provides immense flexibility:

• Dynamic Number of Boxes: Training and evaluation are decoupled. You can use $N_{train}$ boxes during training and a different $N_{eval}$ during evaluation.
• Iterative Evaluation: The detection head can be reused iteratively to further improve performance through successive denoising steps.

Related Work

• Object Detection: Most methods rely on heuristic anchors or fixed learnable queries (DETR).
• Diffusion Models: A class of deep generative models that recover data from random distributions via gradual denoising.
• Perception Tasks: While successful in image generation and segmentation, diffusion hasn't been widely used in detection because detection is a set-prediction problem rather than an image-to-image task.

Approach

Object Detection: The learning objective is defined as $(x, b, c)$, where $x$ is an image, $b$ is the set of bounding boxes, and $c$ is the label. Specifically, for the $i$-th box: $b_i = (c_x^i, c_y^i, w^i, h^i)$

Diffusion Models: These are likelihood-based probabilistic models inspired by non-equilibrium thermodynamics. They work by gradually adding noise to data and learning to reverse the process to generate new samples.

Training Objective and Loss Function

The diffusion process is defined by the following transition distribution:

$q(z_t | z_0) = \mathcal{N}(z_t; \sqrt{\bar{\alpha}_t}z_0, (1 - \bar{\alpha}_t)I)$

Where:

• $z_0$ is the original data.
• $z_t$ is the noisy data at time $t$.
• $\bar{\alpha}_t$ is the cumulative product of the noise scheduling terms.
• $I$ is the identity matrix.

During training, the model learns to predict the original data ($z_0$) from the noisy versions ($z_t$).

The loss function used for training is the Mean Squared Error (MSE) or $L_2$ loss:

$L_{train} = \frac{1}{2} \| f_{\theta}(z_t, t) - z_0 \|^2$

Architecture

Since diffusion models generate data samples iteratively, the model would typically need to run $f_\theta$ multiple times during inference. To make this computationally tractable—avoiding applying the full model to the image iteratively—the authors split the system into two distinct components:

1. Image Encoder: Runs once to extract deep feature maps from the raw image $x$.
2. Detection Decoder: Takes these feature maps and progressively refines box predictions.

Image Encoder

• Backbones: Both Swin-Transformers and ResNet are supported as the backbone.
• Feature Pyramid Network (FPN): Used to generate multi-scale feature maps for both types of backbones.

Detection Decoder

The decoder is inspired by Sparse R-CNN. It takes a set of proposal boxes to calculate the Region of Interest (RoI) from the feature maps.

• Input: While Sparse R-CNN uses fixed learned boxes, DiffusionDet uses proposal boxes disturbed from ground truth (training) or sampled from a Gaussian distribution (inference).
• Differences from Sparse R-CNN:
  1. Randomness: Starts with random boxes rather than a fixed set of learned queries.
  2. Input format: DiffusionDet only requires proposal boxes (Sparse R-CNN requires box-feature pairs).
  3. Iterative Evaluation: The parameters are shared across steps and specific to the diffusion process via timestamp embeddings.

Training

The training process involves constructing a diffusion process (Ground Truth $\rightarrow$ Noise) and learning to reverse it.

• Ground Truth Box Padding: Since the number of objects varies per image, boxes are padded to reach $N_{train}$. Box concatenation was found to be the most effective padding method.
• Box Corruption: Gaussian noise is added to padded boxes, controlled by a scaling factor $\alpha_t$ based on a cosine schedule.
• Training Losses: The model applies a set prediction loss over all $N_{train}$ predictions. Only the top $k$ predictions are selected using an optimal transport assignment** method.

Inference

Inference is a denoising sampling process moving from noise to refined object boxes.

1. Sampling Step: Random boxes are sampled from a Gaussian distribution. In each step, the boxes are sent to the decoder. DDIM is used to generate the estimate for the next step.
2. Box Renewal: Predictions are classified as "desired" or "undesired." Undesired predictions (those below a certain confidence threshold) are replaced with new random Gaussian boxes to better align the inference with the training distribution.

Flexible Usage: "One training for all inference cases." Once trained, DiffusionDet can be evaluated with an arbitrary number of boxes or iterations ($N_{eval} \neq N_{train}$) to trade speed for accuracy without retraining.

Experiments and Benchmarking

Datasets

• COCO: 118k training / 5k validation images.
• LVIS: COCO-based but with more diverse categories.
• CrowdHuman: Used to test performance in highly crowded scenarios (15k training / 4.4k validation).

Implementation Details

• Initialization: Xavier Init for the decoder.
• Optimizer: AdamW (Learning rate: $2.5 \times 10^{-5}$, Weight decay: $10^{-4}$).
• Hardware: Trained on 8 GPUs with a mini-batch size of 16.

Results

• Dynamic Boxes: Unlike DETR, which suffers when the number of queries changes, DiffusionDet's performance improves as more boxes are added.
• Zero-Shot Transfer: When trained on COCO and tested on CrowdHuman, DiffusionDet shows a significant advantage in dense scenes by simply increasing the iteration count or box count at test time.