🤖 machine learning

Contrastive Diffusion Alignment: Learning Structured Latents for Controllable Generation

The paper introduces ConDA, a plug-and-play contrastive learning layer that aligns pretrained diffusion latents with auxiliary variables to create a low-dimensional, interpretable embedding enabling smooth interpolation, extrapolation, and counterfactual editing across diverse dynamical systems.

Original authors: Ruchi Sandilya, Sumaira Perez, Charles Lynch, Lindsay Victoria, Benjamin Zebley, Derrick Matthew Buchanan, Mahendra T. Bhati, Nolan Williams, Timothy J. Spellman, Faith M. Gunning, Conor Liston, Logan

Published 2026-08-06

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Ruchi Sandilya, Sumaira Perez, Charles Lynch, Lindsay Victoria, Benjamin Zebley, Derrick Matthew Buchanan, Mahendra T. Bhati, Nolan Williams, Timothy J. Spellman, Faith M. Gunning, Conor Liston, Logan Grosenick

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

Imagine you have a magical paintbrush that can create incredibly realistic pictures of anything you can imagine: a swirling storm, a smiling face, or a neuron firing in a brain. This is what modern "diffusion models" do. They are like master chefs who have tasted every dish in the world and can now cook up a perfect replica of any meal just by describing it. But here's the catch: while these chefs are amazing at cooking, they don't really understand the recipe. If you ask them to slowly turn a frown into a smile, or to show a river flowing from calm to turbulent, they often get confused. They might just mash the two pictures together, creating a blurry, weird mess in the middle, because their internal "kitchen" (the math space where they store ideas) is a giant, messy warehouse where everything is jumbled up.

Scientists have been trying to fix this by giving the chefs a better map. They want to organize that messy warehouse so that moving in one direction always means "getting happier" or "getting faster." The big question is: Can we teach these AI chefs to understand the story behind the picture, not just the picture itself? This paper tackles that problem by introducing a new way to tidy up the chef's kitchen, turning a chaotic storage room into a neat, organized library where every book is placed exactly where it belongs based on how it moves and changes.

The Problem: The Messy Attic

Think of a standard AI image generator as a giant, high-tech attic filled with millions of boxes. Inside each box is a picture. The problem is that the boxes are stacked randomly. If you want to find a picture of a cat that is slowly turning into a dog, you can't just walk down a straight aisle. You might have to jump over boxes, and when you try to mix two pictures together, the result is often a weird, blurry monster. The AI knows what a cat looks like and what a dog looks like, but it doesn't know how to smoothly transition between them because its internal map is too messy and high-dimensional to navigate easily.

The Solution: ConDA (The Magic Organizer)

The authors of this paper introduce a new tool called ConDA (Contrastive Diffusion Alignment). Imagine ConDA as a super-smart librarian who comes into that messy attic and reorganizes everything.

Instead of leaving the boxes in a giant pile, ConDA takes the pictures and sorts them into a small, neat, low-dimensional room. In this new room, the boxes are arranged by a specific rule: how they change over time or under different conditions.

If you have a video of a river, ConDA lines up the boxes so that moving one step forward in the room means the water flows a little faster.
If you have a video of a face, moving to the right means the smile gets wider, and moving up means the eyebrows go higher.

The magic trick is that ConDA uses "contrastive learning." Think of this as a game of "hot and cold." The librarian looks at two pictures: one where a person is smiling and one where they are frowning. It learns that these two boxes should be far apart. Then it looks at two pictures of the same smile, and it learns those boxes should be close together. By playing this game millions of times, it builds a perfect map where the distance between boxes tells you exactly how different the images are.

How It Works: The Two-Step Dance

The paper describes a clever two-step process to make this work without ruining the beautiful pictures the AI creates:

The Edit (In the Library): First, the user goes into the neat, organized library (the low-dimensional space). Here, they can easily draw a smooth path from a frown to a smile, or from a calm river to a stormy one. Because the library is so well-organized, they can use simple math tools (like drawing a curved line) to predict exactly what the next picture should look like.
The Render (Back in the Attic): Once the path is drawn in the library, ConDA takes those new coordinates and brings them back to the messy attic. It finds the closest real boxes to the new path and asks the original AI chef to paint the final picture. This ensures the result is still a high-quality, realistic image, but now it follows the smooth path the user drew.

What They Found: Smoother, Smarter, and More Real

The researchers tested this idea on five very different types of data, and the results were impressive:

Fluid Dynamics (Water Flow): When they tried to simulate water flowing around a cylinder, the old methods made the water look like it was glitching or jumping. With ConDA, the water flowed smoothly, just like in real life. The images were much sharper, with a quality score (PSNR) of 35.7 compared to 28.3 for the old methods.
Neural Activity (Brain Signals): They looked at recordings of neurons firing in a mouse's brain. The new method could predict how the brain activity would change over time much more accurately, creating a smooth movie of the brain thinking rather than a choppy slideshow.
Facial Expressions: They tested it on human faces. The old methods often made the person's face look distorted or changed their identity while they smiled. ConDA kept the person looking like themselves while smoothly transitioning their expression.
Therapeutic Stimulation: They even used it to model how magnetic fields hit the brain during therapy. The new method could show exactly how changing the angle of the magnetic coil would change the effect on the brain, helping doctors plan better treatments.

What It's Not

The paper is careful to point out what ConDA doesn't do. It doesn't invent new ways to generate images from scratch; it just organizes the ones the AI already knows how to make. It also admits that the "library" it builds is a bit of a simplification. Because it squishes a huge amount of information into a small space, it's not perfect for every single detail, but it's perfect for understanding the movement and change in the data.

Why It Matters

This work suggests that we don't need to throw away the powerful AI chefs we already have. Instead, we just need to give them a better map. By organizing the hidden space where these AI models live, we can finally control them. We can ask them to show us "what if?" scenarios—like "what if this river flowed faster?" or "what if this patient received a different type of brain stimulation?"—and get clear, smooth, and realistic answers. It turns a black box that just guesses into a tool we can actually steer.

Technical Summary: Contrastive Diffusion Alignment (ConDA)

Problem Statement

Diffusion models currently excel at generating high-fidelity images and videos; however, their latent spaces are high-dimensional and lack explicit organization for interpretation or control. While these models offer stable training and reliable inversion, their latent geometries are not "dynamics-aware." Consequently, standard linear interpolations often yield implausible intermediate states, and existing conditioning mechanisms (e.g., ControlNet, InstructPix2Pix) guide samples but fail to provide consistent traversal directions across time or auxiliary variables. This limitation hinders the ability to perform principled control over complex system dynamics, such as fluid flows, neural responses, or facial expressions, where controllability is as critical as fidelity.

Methodology: ConDA Framework

The authors introduce ConDA (Contrastive Diffusion Alignment), a plug-and-play geometry layer that separates the tasks of editing and rendering. The framework operates in three primary stages:

Diffusion Inversion and Latent Extraction:
The pipeline begins by training a conditional latent diffusion model (cLDM) or utilizing a pretrained model. Real-world observations (frames $x_s$ ) paired with auxiliary variables ( $y_s$ , e.g., time, stimulation parameters, facial action units) are mapped into a high-dimensional diffusion feature latent space ( $Z$ ) using diffusion inversion solvers (e.g., DDIM or Rex-RK4). This step ensures a near-lossless representation that preserves generative fidelity.
Contrastive Structured Embedding:
The core contribution is learning a compact, low-dimensional embedding space ( $C$ ) that organizes the high-dimensional latents ( $Z$ ). A supervised contrastive learning objective (InfoNCE) is applied to map $Z$ to $C$ . This loss function enforces that embeddings sharing the same auxiliary condition (e.g., same time step or class) cluster together, while those with different conditions are separated. The resulting space $C$ is designed such that its local geometry aligns with the underlying dynamical factors, enabling smooth, interpretable trajectories.
Trajectory Editing and Lifting:
Control is exerted within the structured space $C$ . Standard nonlinear operators—such as parametric splines, Taylor extrapolation (TEX), or classification-based density interpolation—are used to traverse the embedding space, enabling interpolation, extrapolation, or counterfactual editing (e.g., changing a "non-responder" neural pattern to a "responder" pattern).
Crucially, these edited embeddings are not decoded directly. Instead, they are lifted back to the high-dimensional feature latent space $Z$ using a neighborhood-preserving $k$ -Nearest Neighbors ( $k$ NN) decoder. Finally, the original diffusion model decoder renders the sequence frame-by-frame, ensuring the output maintains the high fidelity and diversity of the original diffusion model.

Key Contributions

Separation of Control and Synthesis: ConDA decouples the manipulation of dynamics from the generation of high-fidelity images. It organizes pretrained diffusion latents into a compact, contrastively learned space for editing, while retaining the original diffusion space for rendering.
Effective Nonlinear Traversal: The paper demonstrates that standard nonlinear operators (splines, finite differences, LSTMs) become highly effective for trajectory modeling when applied in the structured $C$ space, outperforming linear baselines and direct conditioning methods.
Empirical Validation Across Domains: The framework is validated across five distinct spatiotemporal domains: fluid dynamics, neural calcium imaging, therapeutic neurostimulation (TMS), facial expression dynamics, and monkey motor cortex activity.

Results

The paper reports quantitative and qualitative improvements over linear baselines (Lerp, Slerp), variational methods ( $\beta$ -VAE), and direct conditioning approaches:

Fluid Dynamics: ConDA achieved a PSNR of 35.7 compared to 28.3 for linear baselines, with significantly smoother flow-field interpolations.
Neural Imaging & Neurostimulation: In calcium imaging, ConDA produced smoother temporal progressions. In TMS-induced electric field simulations, it successfully captured class-dependent transitions, improving the F1 score for classifying treatment responders vs. non-responders from 0.30 (in raw latent space) to 0.63 (in ConDA space).
Facial Expressions: The method preserved subject identity while enabling smooth traversal of expression intensities, outperforming PCA and $\beta$ -VAE in reconstruction metrics (PSNR/SSIM) and trajectory error (RMSE).
Monkey Motor Cortex: ConDA organized condition-dependent dynamics in delayed reaching tasks, showing clearer organization by reach direction and stronger alignment with true velocity trajectories compared to PCA.
Solver Agnosticism: The method remains robust across different inversion solvers (DDIM vs. Rex-RK4) and different diffusion models (cLDM vs. IsoDiff), indicating that the contrastive alignment is a generalizable property rather than a solver-specific artifact.

Significance and Claims

The authors claim that ConDA addresses a fundamental gap in generative modeling: the lack of a principled mechanism for traversing complex dynamics in diffusion latent spaces. By demonstrating that diffusion latents encode dynamics-relevant structure that can be exploited by an explicit contrastive geometry layer, ConDA transforms diffusion models into controllable generators of nonlinear spatiotemporal dynamics.

The paper emphasizes that this approach does not require retraining the diffusion model from scratch but rather adds a lightweight alignment layer. The resulting structured latent space allows for:

Interpretability: Visualizing and manipulating specific dynamical factors (e.g., time, stimulation angle) independently.
Controllability: Enabling smooth interpolation, extrapolation, and counterfactual editing that respects the underlying physics or biology of the system.
Fidelity: Maintaining the high-quality synthesis capabilities of the underlying diffusion model through the separation of editing and rendering.

The authors conclude that while the embedding is lossy and emphasizes local geometry, the practical approximation errors are negligible, providing a viable path toward interpretable and controllable modeling of complex dynamical systems in scientific and biomedical applications.