Learning Single Index Models with Diffusion Priors

Anqi Tang, Youming Chen, Shuchen Xue, Zhaoqiang Liu

Abstract: Diffusion models (DMs) have demonstrated remarkable ability to generate diverse and high-quality images by efficiently modeling complex data distributions. They have also been explored as powerful generative priors for signal recovery, resulting in a substantial improvement in the quality of reconstructed signals. However, existing research on signal recovery with diffusion models either focuses on specific reconstruction problems or is unable to handle nonlinear measurement models with discontinuous or unknown link functions. In this work, we focus on using DMs to achieve accurate recovery from semi-parametric single index models, which encompass a variety of popular nonlinear models that may have {\em discontinuous} and {\em unknown} link functions. We propose an efficient reconstruction method that only requires one round of unconditional sampling and (partial) inversion of DMs. Theoretical analysis on the effectiveness of the proposed methods has been established under appropriate conditions. We perform numerical experiments on image datasets for different nonlinear measurement models. We observe that compared to competing methods, our approach can yield more accurate reconstructions while utilizing significantly fewer neural function evaluations.

Lay Summary:

In many real-world applications (such as imaging or signal processing), we often need to recover an image or signal from measurements that are incomplete, indirect, or distorted in a complex way. Sometimes, the way these measurements are taken is nonlinear and hard to describe with a simple formula, especially when the relationship between the true signal and the observed data is unknown or contains abrupt changes.
Our work explores how to use powerful image generation tools, known as diffusion models, to solve this problem. Diffusion models have recently gained attention for their ability to generate highly realistic images, but we show they can also help recover original signals from challenging observations. Our focus is on two representative types of measurements: one where only the sign of each pixel is observed (known as 1-bit measurements), and another where the measurement is a cubic function of the original signal.
Our method works in two simple steps. First, it estimates the noise through a partial inversion of the diffusion process based on the observed measurements. Then, it applies a pre-trained diffusion model to denoise this estimate and recover the original image. This avoids reliance on detailed knowledge of the measurement process, making the method both fast and efficient.
Experiments on standard image datasets show that our method achieves more accurate reconstructions than previous techniques, using significantly fewer neural function evaluations. Our work opens the door to better signal recovery in practical settings where measurement systems are imperfect, unknown, or too complex to model directly.

Primary Area: General Machine Learning->Representation Learning

Keywords: Single index models, diffusion priors, inversion of diffusion models

Submission Number: 10286