Finding Local Diffusion Schr\"odinger Bridge using Kolmogorov-Arnold Network

TL;DR

This paper introduces a novel method called LDSB that leverages Kolmogorov-Arnold Networks to find local diffusion Schr"odinger bridges, significantly enhancing image generation quality and efficiency over traditional global SB methods.

Contribution

It proposes the first approach to find local diffusion Schr"odinger bridges in the diffusion path subspace using KAN, improving image generation performance and reducing computational costs.

Findings

FID reduced by over 15%, up to 48.50% on CelebA.

LDSB improves image quality with minimal additional model size.

Method is efficient, using less than 0.1MB for optimization.

Abstract

In image generation, Schr\"odinger 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)\epsilon$). To address the limitations of SB-based methods, this paper proposes for the first time to find local Diffusion Schr\"odinger 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.