Finding Local Diffusion Schrödinger Bridge using Kolmogorov-Arnold Network

TL;DR

This work addresses the inefficiency and gap to global optima in Schrödinger Bridge-based diffusion models by introducing Local Diffusion Schrödinger Bridge (LDSB), which restricts optimization to the diffusion path subspace $\\mathcal{P}_{A,B}$ generated by $f_A(t)$ and $f_B(t)$. It couples a Kolmogorov-Arnold Network (KAN) with an IPF-inspired training scheme to iteratively refine forward and reverse diffusion paths while keeping the pre-trained denoising network fixed, achieving a reparameterized, continuous diffusion path with reduced memoryFootprint. The authors demonstrate substantial improvements in image generation quality and sampling efficiency across DDIM and Flow Matching frameworks on CIFAR10, CelebA, and CelebA-HQ, with FID reductions up to nearly 49% at low step counts and a KAN size under 0.1 MB. This approach tightens the link between SB and diffusion models, offering a practical pathway to faster, higher-quality image synthesis without retraining large denoisers.

Abstract

In image generation, Schrödinger Bridge (SB)-based methods theoretically enhance the efficiency and quality compared to the diffusion models by finding the least costly path between two distributions. However, they are computationally expensive and time-consuming when applied to complex image data. The reason is that they focus on fitting globally optimal paths in high-dimensional spaces, directly generating images as next step on the path using complex networks through self-supervised training, which typically results in a gap with the global optimum. Meanwhile, most diffusion models are in the same path subspace generated by weights $f_A(t)$ and $f_B(t)$, as they follow the paradigm ($x_t = f_A(t)x_{Img} + f_B(t)ε$). To address the limitations of SB-based methods, this paper proposes for the first time to find local Diffusion Schrödinger Bridges (LDSB) in the diffusion path subspace, which strengthens the connection between the SB problem and diffusion models. Specifically, our method optimizes the diffusion paths using Kolmogorov-Arnold Network (KAN), which has the advantage of resistance to forgetting and continuous output. The experiment shows that our LDSB significantly improves the quality and efficiency of image generation using the same pre-trained denoising network and the KAN for optimising is only less than 0.1MB. The FID metric is reduced by more than 15\%, especially with a reduction of 48.50\% when NFE of DDIM is $5$ for the CelebA dataset. Code is available at https://github.com/PerceptionComputingLab/LDSB.

Figures (10)

• Figure 1: Adjusting the weights of image and noise ($f_A(t)$ and $f_B(t)$) in diffusion enhances image generation quality, as the presence of perturbation reduces the FID. This paper achieves optimization by solving a local diffusion Schrödinger bridge in the diffusion path subspace generated by $f_A(t)$ and $f_B(t)$.
• Figure 2: Visualisation of training objectives. Blue dotted lines and arrow represent the optimization of the inverse path $\pi^{2n+1}$ and green dotted lines and arrow for the forward path $\pi^{2n+2}$.
• Figure 3: Comparison of different methods on 2D checkerboard with NFE of $20$. Optimisation with LDSB generates the checkerboard more accurately and with less error points. And LDSB in combination with DDIM or FM outperforms DSB.
• Figure 4: Comparison of diffusion paths of LDSB and the generated image on CIFAR10. LDSB makes the images more realistic and clearer on same NFE. $f(t)$ is the original diffusion path, and $\tilde{f}(t)$ is the optimised path using LDSB.
• Figure 5: Comparison of diffusion paths and the generated image on CelebA. LDSB accelerates the image generation, as shown that images generated by DDIM transform along a fixed trajectory as the NFE increases, such as from front face to side face in the first column. Our LDSB accelerates this transformation.