Image-Conditional Diffusion Transformer for Underwater Image Enhancement

arXiv:2407.05389 · cs.CV, cs.AI · Submitted 2026-08-24 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Image-Conditional Diffusion Transformer for Underwater Image Enhancement".

Tom: , extracted directly from its abstract and introduction. Summary:

Jane: First, who's behind it and why it matters.

Title and authors: Tom: Let's talk about the title itself, "Image-Conditional Diffusion Transformer for Underwater Image Enhancement." It tells us immediately what the paper is focused on: using a specific kind of AI model to fix underwater pictures.

Jane: The authors are Xingyang Nie and his colleagues, and they’ve really defined their method around conditioning the diffusion process on the degraded underwater image itself. That conditioning seems central to their whole idea.

Lu: I think the combination of Diffusion Transformers and image-conditional input is quite sophisticated; it shows a deep understanding of how generative models can be steered effectively for restoration tasks.

Meng: Steering generative models is important because if you can condition them this precisely, you control the output quality much more than with just a random generation process.

Lalam: For me, the title suggests a system that isn't just doing generic enhancement but is specifically tailored to the unique challenges of underwater visuals through its conditional nature.

The paper's summary: Tom: So, what’s actually going on in terms of what they did? Essentially, they propose ICDT, which swaps out the usual U-Net structure inside a diffusion model for a transformer architecture.

Jane: That’s correct; it replaces the conventional U-Net backbone within a denoising diffusion probabilistic model with this new transformer structure, which brings in those scalability advantages from transformers.

Lu: The core mechanism involves taking the degraded underwater image and converting both that image and a reference image into latent space using pre-trained Variational Autoencoders before applying ICDT there.

Meng: Converting everything to latent space first sounds smart for efficiency, but I wonder how robust that initial VAE conversion is when dealing with the extreme degradation found in deep underwater scenes.

Lalam: The paper summarizes this as a hybrid architecture that uses this conditional latent space, which they call computationally efficient, meaning it’s designed to be practical for real-world use rather than just a theoretical exercise.

The paper's improvements: Tom: Now we look at the specific upgrades they introduced to make ICDT better than what came before. They focused on replacing the U-Net with a transformer, which gives it that long-range dependency capability we talked about earlier.

Jane: And on top of that, they incorporated a hybrid loss function involving variances during training, which helps guide the model to sample much faster without sacrificing the quality of what it produces.

Lu: That hybrid loss function is significant because incorporating learned variances allows the DDPM to estimate the covariance structure better than standard methods did previously, which directly impacts how well it learns the underlying data distribution.

Meng: Faster sampling is a big deal for engineering; if you can get results in fewer steps, you save a lot of computational time when deploying these kinds of systems in an operational environment.

Lalam: That acceleration means we can move from slow, iterative processing to something that feels much closer to real-time performance for image enhancement tasks.

Conclusion: Tom: So, wrapping things up on the Image-Conditional Diffusion Transformer for Underwater Image Enhancement paper, the main point is that ICDT provides a scalable and conditional framework for UIE by using a transformer backbone and smart training techniques.

Jane: It seems they’ve shown that this approach can achieve state-of-the-art results on benchmarks like Underwater ImageNet while being much more efficient in terms of inference speed thanks to the sampling improvements.

Lu: The implication here is that we are moving toward generative models that are not just good at generating images, but specifically capable of high-fidelity, context-aware restoration based on complex input conditions.

Meng: From an engineering standpoint, the scalability they demonstrated through FLOPs correlation suggests we have a clear path for selecting the right model size for any given deployment scenario.

Lalam: It really shows that these AI advances can become universal frameworks for image-to-image tasks because of how well they handle conditioning and latent space manipulation.

Tom: It’s been a really insightful look at ICDT, and it sets a new standard for conditional generative models in this area. We’ll keep an eye on what comes next in this research area.

M. Chen

cs.CV, cs.AI

Submitted: 2026-08-24

Updated: 2026-08-25

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 74/100

The gist: * Underwater image enhancement (UIE) is a critical field due to its importance for "underwater operation and marine engineering." However, images degraded by underwater conditions—caused by

Key concepts

Image-Conditional Diffusion Transformer (ICDT)
ICDT is a proposed method that replaces the standard U-Net structure in a denoising diffusion probabilistic model with a transformer architecture. It conditions the diffusion process on the degraded underwater image itself to perform image enhancement.
Latent Space Conversion
The core mechanism involves converting both the degraded underwater image and a reference image into latent space using pre-trained Variational Autoencoders before applying ICDT. This step is used for computational efficiency in the hybrid architecture.
Hybrid Loss Function with Variances
The model uses a hybrid loss function that incorporates variances during training. This helps guide the diffusion model to sample much faster without sacrificing output quality by allowing it to estimate the covariance structure better than standard methods.

Terminology

Summary

Underwater image enhancement (UIE) is a critical field due to its importance for underwater operation and marine engineering. However, images degraded by underwater conditions—caused by wavelength- and distance-dependent light attenuation and scattering—suffer from issues of low contrast and color deviation. These degradations negatively affect performance in computer vision tasks such as object detection, image classification, and semantic segmentation.

Limitations of Existing Methods:

The paper identifies several limitations in prior UIE approaches:

  1. Conventional Methods: Traditional techniques (e.g., white balance or histogram equalization) are not adaptable to different water environment and lighting conditions.

  2. CNN-based Methods: While robust, CNN-based methods are limited by their finite receptive field and the fixed weight structure of convolutional layers, which restrict them from modeling long-range pixel dependence.

  3. GAN-based Methods: Generative Adversarial Networks (GAN) often suffer from an unstable training process and mode collapse, making it difficult to ensure the accuracy and consistency of generated results.

  4. Denoising Diffusion Probabilistic Models (DDPMs): Although DDPMs offer stable training and high-quality synthesis, a recognized challenge is their long inference time, which leads to poor real-time performance.

Proposed Solution: Image-Conditional Diffusion Transformer (ICDT)

The authors propose a novel UIE method based on the Image-Conditional Diffusion Transformer (ICDT), which utilizes a conditional DDPM framework.

Architecture and Mechanism:

  1. Conditional Input: The method uses the degraded underwater image as the conditional input to the diffusion process. This is achieved by converting both the degraded image and the corresponding reference image (x 0) into latent space via pre-trained Variational Autoencoders (VAEs).

  2. ICDT Backbone: The core innovation replaces the conventional U-Net backbone in a denoising diffusion probabilistic model (DDPM) with a transformer. This architecture, known as ICDT, inherits favorable properties such as scalability from transformers and operates on sequences of patches derived from the latent space.

  3. Hybrid Architecture: The overall system is a computationally efficient hybrid architecture that combines the VAE and a transformer-based DDPM.

Key Contributions:

The paper highlights three primary contributions:

  1. Scalability via Transformer Backbone: The use of ICDT allows for the exploration of scaling performance, demonstrating a strong correlation between the sample quality and network complexity (measured by FLOPs). This enables models to meet requirements for various downstream tasks based on available computing power.

  2. Conditional Latent Space: The the degraded image is used as an explicit conditional input in latent space, creating a computationally efficient hybrid architecture that has the potential to serve as a universal framework for image-to-image generation tasks (e.g., restoration, enhancement).

  3. Accelerated Sampling: A hybrid loss function involving variances is employed during training. This allows the DDPM to sample in much fewer [steps] without sacrificing sample quality, significantly speeding up the sampling process for practical deployment.

Experimental Results and Evaluation:

The models are tested on the Underwater ImageNet dataset, with performance assessed using several metrics: Peak Signal-to-Noise Ratio (PSNR), Structural Similarity (SSIM), Learned Perceptual Image Patch Similarity (LPIPS), and Underwater Image Quality Measure (UIQM).

  • Scalability: The experiments demonstrate that increasing model size and decreasing patch size can substantially improve the image enhancement performance. Furthermore, the results show a significant positive correlation between PSNR and model FLOPs.

  • State-of-the-Art Performance: The largest model configuration, ICDT-XL/2, achieved state-of-the-art (SOTA) quality. Quantitative comparisons show that ICDT outperforms all competing methods in terms of PSNR and SSIM while achieving the lowest LPIPS and highest UIQM score.

In conclusion, the paper presents a robust and scalable method for UIE, demonstrating that ICDT is a generic image-to-image generative model.

Improvements for AI systems

As a diligent researcher operating under high stakes, I have analyzed the core technical contributions of this paper to define specific, high-impact improvements that can be generalized across various AI systems.

The fundamental improvement is not merely applying this technique to underwater images, but replacing the traditional U-Net paradigm with a highly scalable and conditionally aware Image-Conditional Diffusion Transformer (ICDT) framework. This framework solves several critical bottlenecks in generative and restoration AI systems.

Here are the specific improvements derived from the paper, categorized by their impact on architectural design, efficiency, and performance:

Improvement: The transition from a convolutional U-Net backbone to a Transformer-based structure (ICDT).

Mechanism: By treating the image (or its latent representation) as a sequence of patches and applying Multi-Head Self-Attention, we replace the fixed, local receptive field of CNN layers with global dependency modeling. This allows the model to capture long-range spatial relationships that U-Nets inherently miss.

What the Improved AI System Can Do:

  • Solve Long-Range Dependency Problems: The system can perform complex image restoration tasks (e.g., deblurring, denoising) where artifacts or errors are spread across large areas of an input image, ensuring global coherence and consistency that local CNN models cannot achieve.

  • Achieve Superior Scalability: The system enables linear scaling of performance by allowing the user to increase the number of tokens (T) or increase the hidden dimension (d), providing a predictable path to higher quality simply by increasing computational budget.

Improvement: Implementing a two-stage, hybrid architecture utilizing pre-trained Variational Autoencoders (VAEs) combined with the diffusion transformer in latent space.

Mechanism: Instead of performing computationally prohibitive operations directly on high-resolution pixel data, the VAE compresses the image into a lower-dimensional latent representation (z). The ICDT then operates entirely within this efficient latent space.

Improvement: Integrating a hybrid loss function (L hybrid = L simple + lambda L VLB) that incorporates learned variances (theta) into the DDPM training process.

Mechanism: This modification allows the model to estimate the necessary covariance structure of the distribution, which was previously unknown or fixed. This enables a more informed and targeted noise removal process during reverse diffusion.

Improvement: Using a degraded input image as an explicit conditional input alongside the noised latent (z t) for the DDPM.

Mechanism: The system concatenates the two latents (z t and z) at the input of every transformer block. This forces every subsequent attention mechanism to be aware of both the noise distribution and the specific characteristics of the degraded input image.

Sources

Related papers