Contrastive Diffusion Alignment: Learning Structured Latents for Controllable Generation

AI summary

Overview

Research area: Generative machine learning — specifically latent-space organization and controllability in diffusion models, applied across scientific and clinical spatiotemporal datasets.

Technical level: Advanced. The paper assumes familiarity with diffusion models, DDIM inversion, contrastive learning (InfoNCE), and latent embedding methods.

Scope: The paper introduces ConDA (Contrastive Diffusion Alignment), a "plug-and-play geometry layer" that reorganizes pretrained diffusion latents into a compact, contrastively learned embedding so that simple nonlinear operators can drive controllable, interpretable trajectory generation across five spatiotemporal domains.

What This Paper Is About

Diffusion models generate high-fidelity data, but their latent spaces are high-dimensional and not explicitly organized for interpretation or control, so simple linear interpolations often produce implausible intermediate states. The authors propose ConDA, a framework that separates editing (performed in a low-dimensional contrastively learned embedding) from rendering (performed in the original diffusion latent space), allowing smooth interpolation, extrapolation, and counterfactual editing along time, stimulation parameters, or facial action units.

Key Contributions

1. ConDA framework: A dynamics-aware diffusion approach that separates editing in a compact, contrastively learned embedding space from rendering in the original diffusion latent space.
2. Nonlinear traversal operators: Demonstrates that standard operators such as splines, finite differences (TEX-1/TEX-2), and LSTMs become effective when applied in the structured embedding space, enabling traversal beyond linear interpolation or direct conditioning.
3. Multi-domain validation: Validates ConDA across five spatiotemporal domains — fluid dynamics, neurostimulation, neural calcium imaging, facial expression dynamics, and monkey motor control.
4. Practical pipeline components: Introduces a neighborhood-preserving kNN decoder for lifting edited embeddings back to diffusion latents, and shows robustness across inversion solvers (DDIM, Rex-RK4) and diffusion models (cLDM, pretrained IsoDiff).

Main Findings

• Nonlinear methods improve in the structured space: In the diffusion latent space Z, linear baselines (Lerp, Slerp) underperform and nonlinear methods such as LSTM and TEX-1 achieve limited accuracy. In the ConDA embedding space C, all methods improve, with Spline and TEX-2 performing best (Table 1).
• Fluid dynamics interpolation: TEX-2 in Z-space reaches 35.29 ± 0.60 PSNR / 0.94 ± 0.01 SSIM / 0.26 RMSE; in ConDA C-space, TEX-2 and Spline reach 35.70 ± 0.36 PSNR / 0.94 ± 0.01 SSIM, with RMSE 0.02 (TEX-2) and 0.00 (Spline). Lerp and Slerp in C-space reach 28.27 PSNR. The introduction summarizes this as 35.7 PSNR vs. 28.3 for linear baselines.
• Calcium imaging (Ca2+): Spline and TEX-2 in C-space reach 38.58 ± 0.59 PSNR / 0.93 ± 0.01 SSIM / 0.00 RMSE, versus Lerp at 32.92 ± 2.86 PSNR / 0.82 ± 0.04 SSIM / 43.87 RMSE.
• DISFA facial expressions: Spline and TEX-2 in C-space reach 38.99 ± 0.23 PSNR / 0.96 ± 0.00 SSIM / 0.00 RMSE; Lerp in C-space reaches 33.08 PSNR / 0.83 SSIM / 13.61 RMSE.
• n-step-ahead prediction stays stable: Generalization is not degraded by approximation sources such as kNN lifting, diffusion inversion, and VAE reconstruction (Table 2). Spline in C-space: Fluid 34.51 ± 2.13 PSNR / 0.92 ± 0.07 SSIM / 1.48 RMSE; Ca2+ 35.00 ± 2.57 / 0.86 ± 0.06 / 1.09; DISFA 37.17 ± 2.17 / 0.95 ± 0.04 / 0.52.
• Classification improves in the embedding space: Classifiers trained in C-space outperformed those trained in Z-space across Accuracy, F1, and ROC-AUC regardless of kernel. On the E-Field responder vs. non-responder task, an RBF-SVM reached F1 of 0.63 in C-space versus 0.30 in Z-space. On steady vs. unsteady flow, SVM-RBF reached 0.87 Accuracy / 0.85 F1 / 0.95 AUC in C-space versus 0.68 / 0.52 / 0.59 in Z-space.
• Ablation against PCA and β-VAE: ConDA showed higher PSNR and SSIM on all three image datasets. Fluid: ConDA 33.90 PSNR / 0.90 SSIM / 7.60 RMSE versus PCA 31.82 / 0.87 / 10.43 and β-VAE 32.20 / 0.85 / 7.47. Ca2+: ConDA 35.69 / 0.87 / 11.02 versus PCA 34.73 / 0.86 / 12.40 and β-VAE 34.62 / 0.85 / 11.58. DISFA: ConDA 36.26 / 0.94 / 5.50 versus PCA 34.82 / 0.92 / 7.60 and β-VAE 35.76 / 0.94 / 6.59. On the Fluid RMSE metric specifically, β-VAE (7.47) is lower than ConDA (7.60).
• cLDM beats a GAN: The conditional latent diffusion model substantially outperformed a continuous-label conditioned GAN in both image generation and reconstruction (protocols and full results in Appendix C.10).
• kNN lifting beats MLP decoding: Neighborhood-preserving kNN lifting consistently outperformed a tuned MLP decoder across datasets (Appendix C.4).
• Interpretability on neural data: On monkey S4 latents, the paper compares ConDA embeddings to PCA in the same diffusion latents and evaluates prediction on 216 test samples using RMSE, total absolute error (mean/sd), and Procrustes distance; the specific values are not reported in the provided content.
• Clinical interpretability: KDE-based interpolation between class-conditional modes in E-field data produced sequences morphing a non-responder E-field toward a responder pattern, with color-coded per-pixel changes highlighting cortical shifts.

Methodology in Plain English

The approach works in three stages.

Stage 1 — Encode to a diffusion latent. The authors train a conditional latent diffusion model (cLDM) to work in the latent space Z of a pretrained VAE rather than pixel space. For images, the LDM autoencoder maps a 256×256×3 frame to a 32×32×4 latent, a UNet2DModel is trained with 1000 diffusion steps, and DDIM inversion maps each real frame (with its auxiliary variable, such as time, coil angle, or action units) to a diffusion latent z_s. Rex-RK4 is also tested as a reversible solver baseline with coupling parameter 0.9999 and 8 steps.

Stage 2 — Learn a compact structured embedding. A mapping is trained with a supervised InfoNCE contrastive loss so that latents sharing the same label or condition cluster together and differently labeled latents separate. The resulting embedding is low-dimensional — dim(C) < 10 for spline-based interpolation, with d = 8 used for fluid and d = 3 otherwise — and its local geometry is aligned with the auxiliary variables. The authors use CEBRA with the offset10-model-mse configuration for this step.

Stage 3 — Edit and render. Trajectory editing happens in the embedding space using one of three operators: (1) a C²-continuous parametric spline fit with a curvature penalty, where edits are small moves along the curve's phase parameter; (2) Taylor extrapolation (TEX-1 uses the first derivative, TEX-2 adds the second-order term estimated via finite differences); (3) a class-transfer method that trains an RBF-kernel SVM, estimates class-conditional densities with KDE, finds the peaks of the density difference, and traverses along the line connecting them. The edited embeddings are then lifted back to the diffusion latent space using a local-neighborhood-preserving kNN decoder, and frames are rendered by the diffusion decoder.

Data and evaluation. Five domains are used: flow past a cylinder (2D incompressible Navier–Stokes simulated in FEniCS, 246,000 RGB images, 82 trajectories of 3000 time points, y = τ ∈ [0, 1.5], blocked splits); two-photon calcium imaging from mouse prefrontal cortex (706,452 RGB images, y = session and within-session frame); DISFA facial expressions (261,576 frames of 27 subjects, 12D action units y = AU ∈ {0,…,5}¹²); TMS-induced electric fields simulated with SimNIBS (569,520 images from 121 patients, coil angles y = Θ ∈ [10°, 360°]); and monkey reach spiking data (2008 training trials, 140 time bins, 182 neurons, mapped to 16-D S4 latents with 2D reach velocity as y). Held-out segment splits: Fluid M = 50, n = 12; Ca²⁺ M = 100, n = 7; DISFA M = 88, n = 11. Metrics are PSNR, SSIM, RMSE, F1/Accuracy/AUC for classification, and Procrustes distance for interpretability.

Why This Matters

Impact on research: The paper argues that pretrained diffusion latents already encode dynamics-relevant structure that an explicit contrastive geometry layer can exploit. If that holds generally, it means controllability in diffusion models can be added as a plug-in without retraining the generator, and that the same recipe transfers across very different data types — fluid simulations, microscopy, electrophysiology, facial video, and electromagnetic simulations.

Real-world applications:

• Therapeutic neurostimulation targeting: The E-field experiments model responder vs. non-responder patterns and produce interpretable trajectories from one to the other, suggesting a route toward personalized targeting strategies.
• Neuroscience analysis: Embeddings of calcium imaging and monkey motor cortex spiking organize condition-dependent neural trajectories, supporting analysis of how neural activity evolves during tasks such as delayed reaches.
• Scientific simulation and flow modeling: Smooth interpolation and extrapolation of fluid dynamics trajectories can fill in unobserved frames and predict future states of physical simulations.
• Facial behavior modeling: Identity-preserving traversal of facial expressions supports animation and behavioral research where action-unit intensity must be varied continuously.

Industry relevance: Applications include controllable generative media and animation pipelines, digital avatars and expression editing, scientific and engineering simulation surrogates, and medical device or treatment-planning workflows where a clinician needs to interpolate between candidate parameter settings and see the result.

Future Directions

• Broader solver and model coverage: The paper demonstrates transfer across DDIM and Rex-RK4 and across cLDM and pretrained IsoDiff; whether the same alignment holds for other diffusion families, or for pixel-space models, is left open.
• Lifting from embedding to latent: The kNN decoder outperformed a tuned MLP decoder, but the paper leaves open whether learned lifting could match or exceed neighborhood-based methods while scaling better.
• Automatic embedding dimension and operator selection: Dimension was set to d = 8 for fluid and d = 3 otherwise, and interpolation dimension was under 10; a principled way to choose dimensionality and which nonlinear operator to use per dataset is not established.
• Clinical validation of the class-transfer trajectories: The responder/non-responder and treatment-response analyses are computational; whether the interpolated E-field trajectories correspond to viable clinical targeting strategies is not reported and would require prospective study.
• Full interpretability quantification: Results for the monkey S4 latent comparison (RMSE, total absolute error, Procrustes distance) against PCA are referenced but their values are not present in the provided text, leaving the quantitative interpretability claim unanswered here.

Target Audience

This paper is most useful to machine learning researchers working on diffusion models, latent-space controllability, and disentanglement; to computational neuroscientists and clinicians interested in interpretable latent representations of neural stimulation and neural activity; and to scientific machine learning practitioners who need controllable sequence generation for simulations such as fluid dynamics. Readers should already be comfortable with diffusion model training and inversion, contrastive objectives, and standard latent-space evaluation metrics.

Authors’ abstract

Diffusion models excel at generation, but their latent spaces are high dimensional and not explicitly organized for interpretation or control. We introduce ConDA (Contrastive Diffusion Alignment), a plug-and-play geometry layer that applies contrastive learning to pretrained diffusion latents using auxiliary variables (e.g., time, stimulation parameters, facial action units). ConDA learns a low-dimensional embedding whose directions align with underlying dynamical factors, consistent with recent contrastive learning results on structured and disentangled representations. In this embedding, simple nonlinear trajectories support smooth interpolation, extrapolation, and counterfactual editing while rendering remains in the original diffusion space. ConDA separates editing and rendering by lifting embedding trajectories back to diffusion latents with a neighborhood-preserving kNN decoder and is robust across inversion solvers. Across fluid dynamics, neural calcium imaging, therapeutic neurostimulation, facial expression dynamics, and monkey motor cortex activity, ConDA yields more interpretable and controllable latent structure than linear traversals and conditioning-based baselines, indicating that diffusion latents encode dynamics-relevant structure that can be exploited by an explicit contrastive geometry layer.