Resolving Endpoint Underfitting in Diffusion Bridges via Noise Alignment

2026-05-27

AI Summary

This work addresses a critical underfitting issue in diffusion bridge models, which arises near the target endpoint due to mismatched noise levels, manifesting as significant deviations in predicted variance and direction. To resolve this, the authors propose Noise-Aligned Diffusion Bridge (NADB), which reconstructs the diffusion process through a noise alignment mechanism. Specifically, NADB first employs a mean network to generate a cleaner conditional target and then establishes a noise-aligned mapping between the input and the regression target, effectively eliminating the noise discrepancy. This approach is the first to systematically identify and mitigate endpoint underfitting in diffusion bridges. Extensive experiments demonstrate that NADB substantially improves performance across multiple image restoration and translation tasks, effectively alleviating endpoint distortion.

Abstract

Diffusion bridge models offer a powerful framework for connecting two data distributions, such as in image restoration and translation. Many existing methods learn this bridge by mimicking the score-matching formulation of standard diffusion models. In this work, we find that this way leads to an anomalous underfitting phenomenon near the target endpoint, as the process approaches the target distribution ($t \to 0$). This underfitting, characterized by significant drift in the predicted variance and direction, results from an excessively large discrepancy in noise levels between the network's input and its regression target.To resolve this issue, we propose the Noise-Aligned Diffusion Bridge (NADB).Our approach reformulates the diffusion bridge by first employing a mean network to provide a cleaner conditional target, and then introducing a novel, noise-aligned mapping relationship. This new formulation resolves the noise mismatch and corrects the underfitting near the target endpoint. Experimental validation across multiple image restoration and image translation tasks demonstrates the effectiveness of our approach. Code is available at https://github.com/gyr02/NADB.

Problem

• diffusion bridge
• endpoint underfitting
• noise mismatch
• score matching
• image restoration

Innovation

• diffusion bridge
• noise alignment
• endpoint underfitting
• score matching
• image restoration