🤖 machine learning

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

This paper proposes a representation alignment (REPA) framework that leverages pretrained DINOv2 encoders to guide diffusion and flow-based models in solving inverse problems, theoretically demonstrating that this approach minimizes embedding-space divergence to improve reconstruction quality and perceptual realism while reducing inference steps across various tasks.

Original authors: Loukas Sfountouris, Giannis Daras, Paris Giampouras

Published 2026-05-27

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Loukas Sfountouris, Giannis Daras, Paris Giampouras

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Fixing Broken Photos with a "Smart Guide"

Imagine you have a beautiful, high-resolution photograph, but it gets damaged. Maybe it's blurry, maybe a chunk of it is missing (like a black box covering a face), or maybe it's been shrunk down to a tiny, pixelated speck. Your goal is to fix it and get the original image back.

In the world of AI, we have powerful tools called Diffusion Models that are great at "hallucinating" or imagining what a missing or damaged image should look like. They work like a sculptor slowly chipping away noise from a block of stone until a statue appears.

However, sometimes these AI sculptors get a little lost. They might fix the blur but make the texture look like plastic, or they might fill in a missing face with the wrong features. They are good at the structure of the image, but they sometimes struggle with the vibe or the fine details that make a photo look real to a human eye.

The Problem: The "Blind" Sculptor

The paper identifies a specific problem: When the AI tries to fix a damaged photo, it doesn't have the original "blueprint" (the ground truth) to compare against. It's trying to guess the answer without knowing the question.

To help, researchers usually tell the AI, "Make sure your guess matches the blurry pixels you were given." But this isn't enough. The AI needs a better sense of what "real" looks like.

The Solution: The "Art Critic" (REPA)

The authors introduce a new method called REPA (Representation Alignment). To understand this, imagine the AI sculptor is working in a dark room.

The Sculptor (Diffusion Model): This is the AI trying to fix the photo. It has its own internal "feelings" about what the image looks like as it works.
The Art Critic (DINOv2): This is a pre-trained AI that is an expert at recognizing visual features. It's like a seasoned art critic who has seen millions of photos and knows exactly what a "real" tree, a "real" eye, or a "real" texture looks like. It doesn't care about the pixels; it cares about the meaning and structure of the image.

The Magic Trick:
Usually, the Sculptor and the Art Critic speak different languages. The Sculptor thinks in "latent codes" (abstract math), and the Critic thinks in "visual features."

The paper's innovation is to force the Sculptor to listen to the Critic while it's working. They don't just look at the final picture; they check in at every step of the process.

The Sculptor says, "I think this part of the image is a tree."
The Critic checks its internal database and says, "Yes, but the texture of that tree in your current draft doesn't match a real tree. Adjust your internal representation to look more like a real tree."

By constantly aligning the Sculptor's internal thoughts with the Critic's expert knowledge, the final result is much more realistic and detailed.

How They Handle the "Missing Blueprint"

You might ask: "But wait, if the original photo is damaged, how does the Critic know what a 'real' tree looks like in this specific photo?"

The paper has a clever workaround. Since they don't have the original photo, they use a Proxy (a stand-in).

Early in the process: The AI uses the damaged, blurry photo as a rough guide for the Critic.
Later in the process: As the AI starts to clean up the image, it uses its own current best guess of the clean image as the guide.

It's like trying to restore an old painting. At first, you only have the dirty, scratched version to look at. But as you clean it, you start using the newly revealed clean parts to guide your restoration of the remaining dirty parts. The paper shows that the "Art Critic" (DINOv2) is so robust that it can still recognize the subject matter even if the image is blurry or noisy, making this stand-in strategy work perfectly.

The Results: Sharper, Faster, and More Real

The authors tested this on four types of image damage:

Super-Resolution: Making a small, blurry image big and sharp.
Deblurring: Fixing motion blur (like a camera shaking).
Inpainting: Filling in a missing square block of the image.
Gaussian Deblurring: Fixing general fuzziness.

What happened?

Better Quality: The images looked much more realistic. The textures (like skin pores or grass) were sharper and less "plastic."
Faster Work: Surprisingly, using this "Art Critic" guidance allowed the AI to reach high-quality results in fewer steps. It's like the sculptor needed fewer chisel strikes because the Critic kept them on the right path, preventing them from wandering off into weird, unrealistic shapes.
The Trade-off: The paper notes that while the images looked better to human eyes (perceptual quality), the standard math scores (which measure pixel-by-pixel accuracy) didn't always go up. This is common: a photo can look more "real" even if it's not mathematically identical to the original pixel-for-pixel.

The Theory: Why It Works

The paper also provides a mathematical proof (a "theory") to explain why this works.

The Divergence: They show that by aligning with the Critic, the AI is essentially minimizing the "distance" between its guess and the true nature of the image in a feature space.
The Contraction: They prove that as the AI gets closer to the solution, this alignment acts like a magnet, pulling the AI's internal state closer to the "clean" state and pushing away errors.

Summary

In short, this paper teaches a powerful image-restoration AI to listen to an expert visual AI (DINOv2) while it works. Even without seeing the original perfect photo, this "alignment" helps the AI fix damaged images faster and with much more realistic details, turning a blurry or broken mess into a sharp, lifelike picture.

Technical Summary: Align & Invert

Problem Statement

Inverse problems aim to recover a clean signal $x_0$ from noisy or degraded measurements $y$ . While pretrained diffusion and flow-based models have become powerful priors for these tasks, they often struggle to produce highly detailed images in scenarios involving severe degradation (e.g., extreme super-resolution, motion blur) or complex scenes. A specific limitation arises with Latent Diffusion Models (LDMs), where the non-linearities of the encoder-decoder architecture can introduce artifacts during the reverse sampling process. Furthermore, existing methods often require a large number of discretization steps to achieve high-quality reconstructions, and they lack strong semantic constraints during inference to guide the trajectory toward perceptually realistic solutions.

Methodology

The paper proposes Align & Invert, a framework that integrates Representation Alignment (REPA) into the inference process of diffusion and flow-based solvers for inverse problems.

Core Concept

The method enforces alignment between the internal representations of the generative model (diffusion/flow) and the features of a pretrained, self-supervised visual encoder (specifically DINOv2). While REPA was previously used as a training regularizer to improve generative quality, this work adapts it for inference in inverse problems where ground-truth signals are unavailable.

Technical Implementation

Proxy Construction: Since the ground truth $x_0$ is unknown, the method constructs a proxy representation $c_{proxy}$ to stand in for the missing target features.
- Initially, $c_{proxy}$ is derived from the DINOv2 features of the available measurement $y$ .
- As the reverse diffusion process progresses, the proxy is gradually updated to the DINOv2 features of the model's current denoised estimate, $f_{DINOv2}(D(E[z_0 | z_t]))$ .
- This strategy relies on the observed robustness of DINOv2 features to common image degradations.
Tilted Distribution: The sampling process is guided by a tilted distribution that incorporates the alignment term as a reward:
$\tilde{p}(x_t | y) \propto p(y | x_t) p_t(x_t) \exp(\lambda \cdot \text{REPA}(\bar{x}_0, h_t))$
Here, $p(y | x_t)$ enforces measurement consistency, $p_t(x_t)$ is the unconditional prior, and the REPA term encourages alignment between the diffusion model's intermediate hidden states $h_t$ and the proxy DINOv2 features.
Algorithm Integration: The approach is integrated into existing state-of-the-art solvers (e.g., Latent DPS, ReSample) by adding a gradient step based on the REPA objective during the reverse sampling loop. This step aligns the latent diffusion state with the semantic features of the target.
Adaptability: For models not originally pretrained with REPA, the framework learns a lightweight post-hoc projection head ( $g_\phi$ ) to map the diffusion model's internal representations to the DINOv2 feature space. This mapping is optimized while keeping the diffusion model frozen, making the adaptation computationally inexpensive.

Theoretical Insights

The paper provides two main theoretical results characterizing the REPA regularizer:

Divergence Minimization: Maximizing the REPA objective is shown to be equivalent to minimizing the Maximum Mean Discrepancy (MMD) between the DINOv2 feature distributions of the reconstructed image and the ground truth, up to terms accounting for feature misalignment and proxy approximation errors.
Latent State Contraction: Under regularity assumptions regarding the Jacobians of the diffusion encoder and the projection head, a single REPA update induces a contraction in the diffusion state space. Specifically, it steers the current latent state $z_t$ closer to the latent representation of the clean image $z^*$ , provided the current state is within a local neighborhood of the solution.

Experimental Results

The authors evaluated the method on four inverse problems: Super-Resolution (4×, 8×, 16×), Box Inpainting, Gaussian Deblurring, and Motion Deblurring, using ImageNet and FFHQ datasets.

Perceptual Quality: Integrating REPA consistently improved perceptual metrics (LPIPS and FID) across all tasks and solvers (Latent DPS, ReSample, Latent DAPS). For instance, in 4× Super-Resolution on ImageNet, Latent DPS + REPA reduced LPIPS from 0.238 to 0.217 and FID from 123.82 to 86.88.
Pixel Metrics: Standard pixel-wise metrics (PSNR, SSIM) remained comparable to baselines, reflecting the known perception-distortion trade-off where REPA prioritizes semantic realism over pixel-perfect accuracy.
Efficiency: A significant finding is that REPA allows solvers to reach the same level of perceptual quality with significantly fewer discretization steps. In super-resolution and Gaussian deblurring, the method achieved comparable LPIPS scores with roughly 4× fewer steps than the baseline, leading to reduced overall runtime despite the additional computation per step.
Robustness: The method remained effective under severe degradations (e.g., 16× upscaling, high-intensity blur) and when applied to independently pretrained models via the learned projection head.

Significance and Claims

The paper claims that representation alignment constitutes an effective tool for improving generative inverse problem solvers at inference time.

Semantic Guidance: By aligning internal model representations with robust visual features (DINOv2), the method provides a strong semantic prior that guides the reconstruction toward perceptually meaningful solutions, even without ground-truth supervision.
Theoretical Justification: The work elucidates why alignment works, showing it acts as a mechanism to minimize feature-space divergence and contract the latent state toward the clean image solution.
General Applicability: The framework is general, compatible with various diffusion/flow-based solvers, and adaptable to models not originally trained with alignment via a lightweight post-hoc mapping.
Efficiency: It addresses the computational bottleneck of diffusion sampling by accelerating convergence, allowing high-quality reconstruction with fewer sampling steps.

The authors conclude that this approach bridges the gap between the generative capabilities of diffusion models and the semantic understanding of self-supervised encoders, offering a principled way to enhance reconstruction quality in challenging inverse problem settings.