Contrastive Diffusion Guidance for Spatial Inverse Problems

This paper introduces CoGuide, a novel diffusion-based framework that addresses spatial inverse problems with non-differentiable and partially specified forward operators by reformulating likelihood guidance in a contrastively learned embedding space to achieve robust and consistent reconstructions.

The Gist

Imagine you walk into a house you've never seen before. You don't have a map, and you can't see the walls because it's pitch black. However, you have a friend who is walking around with you, and you can hear their footsteps and see the path they take on a radar screen.

Your goal? To draw the floorplan of the house (where the walls are, where the rooms are) just by watching that path.

This is the core problem the paper "Contrastive Diffusion Guidance for Spatial Inverse Problems" tries to solve. The authors call their solution CoGuide.

Here is the breakdown in simple terms, using some creative analogies.

1. The Problem: The "Bumpy" Road

Usually, when computers try to solve puzzles like this, they use a method called Diffusion Models. Think of a diffusion model like a sculptor starting with a block of noisy static (like TV snow) and slowly chipping away the noise to reveal a statue (the floorplan).

To chip away the noise correctly, the sculptor needs a "guide" or a compass. This guide tells the sculptor: "Hey, the path you just drew doesn't match the footsteps we heard. Move the wall here."

The Catch: In this specific puzzle, the "guide" is broken.
The relationship between the floorplan and the walking path is like a light switch.

• If you move a wall by a tiny, invisible amount, the person walking might hit a dead end and have to take a completely different route.
• If you move the wall a tiny bit the other way, they walk straight through.

In math terms, this is called non-differentiable. It's like trying to roll a ball down a staircase; the ball doesn't slide smoothly—it just falls off the edge. Because the "guide" is so bumpy and unstable, standard computer methods get confused and fail to draw the right floorplan.

2. The Solution: The "Vibe Check" (Contrastive Learning)

The authors realized they couldn't fix the bumpy road, so they decided to build a new road entirely.

Instead of trying to calculate exactly how the wall moves the path, they taught the computer to recognize compatibility.

• The Old Way: "If I move the wall 1mm left, the path changes by 5 meters." (Too complicated, too bumpy).
• The New Way (CoGuide): "Does this floorplan feel right for this path?"

They created a special embedding space. Imagine this as a giant, smooth dance floor.

• They take a floorplan and a walking path and ask the computer to put them on the dance floor.
• If the floorplan and the path match (the path makes sense for that house), the computer pulls them close together, like two people who just met and hit it off.
• If they don't match (the path goes through a wall), the computer pushes them far apart, like strangers at a party who have nothing in common.

This is called Contrastive Learning. It's like training a bouncer at a club: "If the outfit matches the vibe, let them in (close). If not, keep them out (far)."

3. How It Works: The Smooth Glide

Now, when the computer is trying to draw the floorplan (the sculptor chipping away the noise), it doesn't look at the bumpy math of the walls anymore. Instead, it looks at the dance floor.

It asks: "Is the current sketch of the house close to the walking path on the dance floor?"

• If the sketch is far away, the computer gently nudges it closer.
• Because the dance floor is smooth (mathematically speaking), the computer can glide smoothly toward the correct answer without getting stuck on the "bumps" of the real world.

4. Why It's a Big Deal

The authors tested this on a dataset of thousands of house layouts.

• The Competitors: Other methods tried to use "differentiable" path planners (math tricks to make the bumpy road smooth). They failed often, creating floorplans with walls in weird places or paths that went through solid objects.
• CoGuide: Because it used the "Vibe Check" (the smooth dance floor), it produced floorplans that were much more accurate and consistent. It could even handle real-world data where the walking path was noisy or sparse (like a few footsteps here and there).

5. Beyond Houses: The "Blind" Fix

The coolest part? The authors showed this trick works for other things too, like restoring old, scratchy audio recordings.

• Imagine you have a recording of a piano that is full of static and pops. You don't know exactly what caused the noise (the "forward operator" is unknown).
• CoGuide can still fix it! It learns to recognize what a "clean piano" and a "noisy piano" look like in its "dance floor" space, and it guides the restoration process to bring the noisy sound closer to the clean sound.

Summary Analogy

Imagine you are trying to guess a secret code.

• Old Method: You try to guess the code by calculating the exact mathematical difference between your guess and the real code. But the calculator is broken and gives you wild, jumping numbers. You get lost.
• CoGuide Method: You have a friend who knows the code. You show them your guess. They don't give you a number; they just say, "Warm!" (You're close) or "Cold!" (You're far). You keep guessing until you are "Hot."

CoGuide teaches the computer to be that friend, using a smooth, intuitive sense of "match" rather than a broken, bumpy calculator.

Technical Summary

Here is a detailed technical summary of the paper "Contrastive Diffusion Guidance for Spatial Inverse Problems" (CoGuide), published as a conference paper at ICLR 2026.

1. Problem Definition

The paper addresses a specific class of inverse problems (IPs) where the forward operator $A$ is partially specified, non-smooth, and non-differentiable.

• Specific Application: Reconstructing 2D spatial layouts (floorplans, $x$) from human movement trajectories ($y$).
• The Forward Process: $y = A(x) + n$, where $A$ represents a human path-planning policy (e.g., walking from point A to B avoiding walls).
• The Challenge:
  • Non-differentiability: Path planning algorithms (like A*) involve discrete decisions (argmin operations) and local searches. Small changes in the floorplan (e.g., a tiny hole in a wall) can cause drastic, discontinuous changes in the resulting path.
  • Instability: Existing diffusion-based inverse solvers (like Diffusion Posterior Sampling - DPS) rely on gradient-based guidance using the likelihood score $\nabla_x \log p(y|x)$. When $A$ is non-smooth, the Jacobian $J_A$ is unstable or undefined, causing optimization to diverge or produce artifacts.
  • Blindness: In many real-world scenarios, the exact forward operator $A$ is unknown or only partially known.

2. Methodology: CoGuide

The authors propose CoGuide, a framework that replaces the unstable, direct likelihood calculation with a surrogate likelihood score learned in a smooth embedding space using contrastive learning.

A. Core Insight: Embedding Space Surrogate

Instead of computing $\nabla_x \|y - A(x)\|^2$ directly (which is unstable), CoGuide projects both the floorplan $x$ and the trajectory $y$ into a shared, smooth embedding space $\mathcal{E}$ using encoders $f_\phi$ and $g_\psi$.

• Goal: Learn an embedding space where compatible $(x, y)$ pairs are close, and incompatible pairs are far apart.
• Surrogate Likelihood: The likelihood score is approximated by the gradient of the squared distance in the embedding space:
  $$\nabla_x \log p(y|x) \approx -\frac{1}{2\tau} \nabla_x \| f_\phi(\hat{x}_0) - g_\psi(y) \|^2_2$$
  where $\hat{x}_0$ is the Tweedie estimate of the clean signal from the diffusion process.

B. Contrastive Learning Objective

The encoders are trained using a Supervised Contrastive Loss (InfoNCE-style) on synthetic data (floorplans and their corresponding A*-generated trajectories).

• Mechanism: The loss pulls matching floorplan-trajectory pairs closer together in the embedding space while pushing mismatched pairs apart.
• Theoretical Justification: The paper proves that at the optimum of the InfoNCE loss, the inner product in the embedding space is proportional to the log-likelihood ratio $\log p(y|x) - \log p(y)$. Thus, the gradient of the embedding distance serves as a valid, smooth approximation of the true likelihood gradient.
• Multi-Positive Handling: The method supports multiple compatible trajectories for a single floorplan, utilizing a symmetric loss that anchors on both the floorplan and the trajectory.

C. Inference Pipeline (Algorithm 1)

CoGuide integrates this guidance into the diffusion sampling process (DDIM):

1. Denoising: A standard diffusion prior $s_\theta(x_t, t)$ generates a candidate floorplan $\hat{x}_0$.
2. Contrastive Guidance: The model computes the gradient of the embedding distance between $\hat{x}_0$ and the observed trajectory $y$.
3. Intersection Penalty: An additional term $\|y \odot (1-\hat{x}_0)\|_1$ is added to penalize trajectories that cross walls, ensuring physical consistency.
4. Optimization: The authors replace standard Gradient Descent (GD) with Adam during the reverse diffusion steps. This is crucial because the embedding-based gradients are high-order and require robust integration to converge effectively, especially with the fewer steps used in DDIM.
5. Learning Rate Annealing: A cosine annealing schedule is applied to the guidance strength to prevent late-stage instability.

3. Key Contributions

1. Novel Framework for Non-Differentiable IPs: CoGuide is the first diffusion-based solver to handle forward operators that are inherently non-smooth and non-differentiable (like path planning) by bypassing direct operator modeling.
2. Contrastive Likelihood Surrogate: The paper establishes a theoretical link between InfoNCE contrastive learning and likelihood estimation, demonstrating that a learned embedding space can provide a stable, valid surrogate for intractable likelihood scores.
3. Robustness to Blindness: The method does not require an analytical form of the forward operator $A$ during inference, making it applicable to "blind" inverse problems where $A$ is unknown.
4. Generalization: The approach is shown to extend beyond spatial mapping to blind audio restoration, where the degradation operator is completely unknown.

4. Experimental Results

The authors evaluated CoGuide on the HouseExpo dataset (35k+ floorplans) and real-world UWB sensor data.

• Baselines: Compared against 6 baselines, including:
  • DPS variants with differentiable path planners (Neural A*, TransPath, DiPPeR).
  • Established inverse solvers (DiffPIR, DMPlug).
  • Classifier-Free Guidance (CFG) diffusion models.
• Quantitative Performance:
  • Sparse/Moderate Trajectories: CoGuide significantly outperformed all baselines in F1 score and Intersection over Union (IoU). For example, in the sparse regime, CoGuide achieved 0.91 F1 vs. 0.86 for the best baseline (CFG).
  • Dense Trajectories: CoGuide remained competitive, though CFG performed slightly better in metrics (likely due to overfitting to the synthetic distribution).
  • Combination: The combination of CFG + CoGuide yielded the highest overall performance (0.99 F1 in dense regimes).
• Qualitative Results:
  • Baselines using differentiable planners (DPS+NA*, etc.) often produced floorplans with artifacts or paths that did not align with the trajectory due to optimization instability.
  • CoGuide produced structurally consistent floorplans with fewer visual artifacts.
• Real-World Validation: On real UWB data, CoGuide successfully reconstructed wall segments and room structures that CFG missed or misaligned, demonstrating superior generalization to non-synthetic data distributions.
• Audio Extension: In blind audio restoration, CoGuide reduced Fréchet Audio Distance (FAD) significantly compared to the LTAS baseline, proving the method's applicability to non-spatial domains.

5. Significance and Impact

• Solving the "Non-Differentiable" Bottleneck: This work opens a new avenue for applying diffusion models to inverse problems where the physics or logic of the forward process is discrete, combinatorial, or unknown. It moves beyond the limitation of requiring differentiable forward operators.
• Privacy-Preserving Mapping: By inferring floorplans from movement trajectories (which can be collected via IMUs or UWB without cameras), the method offers a privacy-preserving alternative to vision-based layout reconstruction.
• Generalizable Paradigm: The "Contrastive Guidance" paradigm suggests a general recipe for solving blind inverse problems: learn a compatible embedding space between measurements and signals, then use the distance in that space to guide the generative prior. This could apply to molecular structure synthesis, network topology inference, and other complex systems.

In summary, CoGuide represents a significant shift in diffusion-based inverse solving, trading direct gradient computation of complex operators for a learned, smooth, and robust contrastive guidance mechanism.