Align & Invert: Solving Inverse Problems with Diffusion and Flow-based Models via Representation Alignment

TL;DR

This paper introduces REPA, a representation alignment method that enhances inverse problem solutions by aligning diffusion or flow models with pretrained encoders, leading to improved reconstruction quality and efficiency.

Contribution

We extend representation alignment to inverse problems, integrating it with diffusion and flow models and providing theoretical insights into its regularization effects.

Findings

REPA improves reconstruction quality across multiple inverse tasks.

Aligning model representations enhances perceptual realism.

REPA reduces the number of steps needed for high-quality reconstructions.

Abstract

Enforcing alignment between the internal representations of diffusion or flow-based generative models and those of pretrained self-supervised encoders has recently been shown to provide a powerful inductive bias, improving both convergence and sample quality. In this work, we extend this idea to inverse problems, where pretrained generative models are employed as priors. We propose applying representation alignment (REPA) between diffusion or flow-based models and a DINOv2 visual encoder, to guide the reconstruction process at inference time. Although ground-truth signals are unavailable in inverse problems, we empirically show that aligning model representations of approximate target features can substantially enhance reconstruction quality and perceptual realism. We provide theoretical results showing (a) that REPA regularization can be viewed as a variational approach for minimizing a divergence measure in the DINOv2 embedding space, and (b) how under certain regularity assumptions REPA updates steer the latent diffusion states toward those of the clean image. These results offer insights into the role of REPA in improving perceptual fidelity. Finally, we demonstrate the generality of our approach by We integrate REPA into multiple state-of-the-art inverse problem solvers, and provide extensive experiments on super-resolution, box inpainting, Gaussian deblurring, and motion deblurring confirming that our method consistently improves reconstruction quality, while also providing efficiency gains reducing the number of required discretization steps.