Contrastive Diffusion Alignment: Learning Structured Latents for Controllable Generation

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.

1. Introduction

Controlling the generation of complex system dynamics, such as fluid flows around an obstacle, neural responses to therapeutic stimulation, or evolving facial expressions, requires generative models whose latent spaces have a geometry that captures trajectories and allows movement along time and other system variables in a principled way. Despite the success of diffusion models in generative modeling, a fundamental limitation remains: their latent spaces are not explicitly organized to represent temporal or other auxiliary-variable–dependent structure, leaving no principled mechanism for traversing complex dynamics. Addressing this gap is crucial for applications where controllability is as important as fidelity.

Diffusion models provide unmatched fidelity, stable training, and reliable inversion (Ho et al., 2020; Dhariwal and Nichol, 2021; Rombach et al., 2022; Mokady et al., 2023). Yet their latents are not dynamics-aware: linear interpolations often produce implausible intermediate states (Wang and Golland, 2023; Hahm et al., 2024), and existing mechanisms that incorporate auxiliary information such as ControlNet or InstructPix2Pix (Zhang et al., 2023; Brooks et al., 2023) guide samples but do not yield consistent traversal directions across time or other variables. As shown in Sec. 5, this result in blurry or entangled transitions when interpolating fluid dynamics trajectories and inconsistent progression in neural recordings and facial expressions (Fig. 2).

Figure 2. Visual comparison of generated images from interpolations in $𝓒$-space.

Linear methods (Lerp, Slerp) show inferior perceptual quality, whereas Spline and TEX-2 yield smoother, more dynamics-consistent reconstructions.

Recent advances attempt to address this challenge. Geometry-preserving traversals (Hahm et al., 2024), physics-informed priors (Shu et al., 2023), and video diffusion models such as MCVD (Voleti et al., 2022), Imagen Video (Ho et al., 2022a), and Make-A-Video (Singer et al., 2023) improve interpolation or temporal realism, but they do not organize diffusion latents for controllable dynamics. The missing piece is a general recipe for dynamics-aware diffusion: a way to align latent geometry with system variables while preserving the generative fidelity of diffusion models.

We introduce ConDA, a flexible framework that addresses this limitation by separating control from synthesis. It organizes pretrained diffusion latents into a compact, contrastively-learned embedding space using auxiliary variables such as time, stimulation parameters, or expression intensities. Within this structured space, standard nonlinear operators (e.g., splines, finite differences, and LSTMs) enable smooth and interpretable trajectory traversal, while rendering is performed in the the original diffusion latent space to preserve generative fidelity. This separation of editing and rendering (see Fig. 1 and Sec. 3) transforms diffusion models into controllable generators of nonlinear spatiotemporal dynamics.

Figure 1. Method overview of ConDA framework.

We encode a high-dimensional sequence of frames into a diffusion model feature latent trajectory via diffusion inversion (e.g., DDIM/Rex-RK4), then map these latents into a compact, contrastively learned embedding space that organizes local geometry to reflect nonlinear dynamics. An editing operator modifies the trajectory in the embedding space under new conditions (interpolation, local traversal, or class transfer). The edited embedding trajectory is lifted back to the feature latent space (via a learned or neighborhood-preserving kNN decoder) and decoded frame-by-frame by the diffusion sampler, yielding a new output sequence with controlled dynamics and preserved diversity and fidelity.

Across five different domains, ConDA consistently improves controllability while preserving fidelity. In fluid dynamics, ConDA achieves higher reconstruction quality (35.7 PSNR vs. 28.3 for linear baselines) and smoother flow-field interpolations (Fig. 2). In neural calcium imaging, it produces smoother temporal progression across activity states (Fig. 2). In neurostimulation, it captures class-dependent transitions that vary systematically with a given stimulation coil angle (Fig. 3). In facial expression dynamics, it preserves subject identity while enabling smooth traversal of expressions (Fig. 4). In monkey motor control, it organizes condition-dependent dynamics in delayed reaching tasks (Fig. 5). Together, these results demonstrate that contrastive organization reshapes high-dimensional diffusion latents into an interpretable, low-dimensional space in which simple nonlinear operators model dynamics effectively, improving classification and revealing neural trajectories aligned with velocity reaches.

Figure 3. KDE-based class interpolation.

Trajectory morphing a non-responder E-Field (•) toward a responder (×) by interpolating between class-conditional KDE peaks; colors indicates per-pixel E-field changes from the source, highlighting cortical shifts needed for responder-like patterns and suggesting a path to personalized targeting.

Figure 4. Comparison of facial expressions at test frames.

ConDA captures facial expressions (e.g., lip-corner pull/stretch) at test frames with higher fidelity than PCA and $\beta$-VAE.

Figure 5. Interpretability of $𝓒$ on monkey reach spiking data.

Using LDNS (Kapoor et al., 2024) with structured state-space (S4) layers, we map neural spikes to time-aligned latents and train diffusion models conditioned on reach velocity. a) True test reaches, along with predicted test embeddings of neural spikes in 2D PCA and C spaces. b) Performance metrics for PCA and ConDA on test data prediction. c) A true and ConDA generated sample of neural activity. d) Mean firing rate per neuron, averaged across all trials and time, showing a close match between ConDA predicted firing rates and the real monkey data.

Our main contributions are:

1. We propose ConDA, a framework for dynamics-aware diffusion that separates editing in a compact, contrastively learned embedding space from rendering in the original diffusion latent space. 2. We show that standard nonlinear traversal operators such as splines, finite differences, and LSTMs become effective when applied in this structured space, enabling smooth trajectory traversal beyond linear interpolation or direct conditioning. 3. We validate ConDA across five spatiotemporal domains (fluid dynamics, neurostimulation, neural calcium imaging, facial expression dynamics, and monkey motor control) demonstrating improved temporal consistency, controllable condition-dependent transitions, and interpretable latent trajectories aligned with underlying system dynamics.

2. Related work

Diffusion.

Diffusion models are now the leading framework for high-fidelity image and video generation (Ho et al., 2020; Song et al., 2021; Dhariwal and Nichol, 2021; Rombach et al., 2022), offering photorealistic outputs, stable training, and practical deterministic inversion via DDIM (Mokady et al., 2023). Recent work has further improved diffusion inversion by constructing algebraically reversible solvers such as Rex-RK4 (Blasingame and Liu, 2026). However, their latent spaces are not explicitly structured for dynamics: linear interpolations often yield unrealistic transitions (Wang and Golland, 2023; Hahm et al., 2024), and conditioning methods such as ControlNet or InstructPix2Pix guide edits without ensuring consistent traversal directions (Zhang et al., 2023; Brooks et al., 2023).

Controllability.

Efforts to improve controllability include geometry-preserving traversals (Hahm et al., 2024), physics-informed priors (Shu et al., 2023), and video diffusion models (Ho et al., 2022b; Voleti et al., 2022; Singer et al., 2023) that emphasize realism over latent geometry. Related contrastive approaches such as NoiseCLR (Dalva and Yanardag, 2024) focus on discovering unsupervised, interpretable semantic directions for static image editing, while DRCT (Chen et al., 2024) leverages diffusion reconstruction and contrastive training for robust detection of diffusion-generated images. Time-series approaches such as Diffusion-TS (Yuan and Qiao, 2024), as well as label-efficient segmentation using diffusion (Baranchuk et al., 2022), remain task-specific and do not aim to learn a general-purpose, structured latent manifold for controllable spatiotemporal generation.

GAN editing.

GANs revealed interpretable latent directions (Radford et al., 2015; Karras et al., 2020; Shen et al., 2020), inspiring disentanglement strategies such as Info-GAN (Chen et al., 2016) and $\beta$-VAE (Higgins et al., 2017). Yet GAN inversion is imperfect (Abdal et al., 2019; Xia et al., 2022), latent directions are often linear, and editing typically remains prompt- or attention-based (Meng et al., 2022; Avrahami et al., 2022; Hertz et al., 2022) rather than organizing latent space into interpretable trajectories.

Dynamics models.

Sequence models such as the Kalman filter (Kalman, 1960), SINDy (Brunton et al., 2016), RNNs e.g., LSTM/GRU (Hochreiter and Schmidhuber, 1997; Chung et al., 2014), structured SSMs (HiPPO, S4) (Gu et al., 2020; 2022; Smith et al., 2023), and neural ODE/SDE approaches (Rubanova et al., 2019) provide principled dynamics. Diffusion models have been combined with S4 and RNNs for forecasting and imputation (Alcaraz and Strodthoff, 2022; Rasul et al., 2021), while LDNS (Kapoor et al., 2024) focuses on learning low-dimensional latent dynamics with structured S4 models. We show that ConDA further improves interpretability by organizing LDNS diffusion latents into a dynamics-aware geometry.

Disentanglement and contrastive learning.

Contrastive objectives structure latent embeddings by enforcing similarity and separation constraints (Oord et al., 2018; Chen et al., 2020; Poole et al., 2018; Wang and Isola, 2020), with successful applications in neuroscience (Schneider et al., 2023) and identifiability guarantees (Lyu and Fu, 2022; Cui et al., 2022). In contrast, diffusion disentanglement methods (Hahm et al., 2024) mainly address static factors. ConDA extends this line of work by incorporating auxiliary variables (e.g., time, stimulation, expressions) to align diffusion latents with spatiotemporal structure, enabling smooth nonlinear trajectory traversal.

Summary.

Prior work advances interpolation, video realism, or disentanglement, but lacks a principled framework for organizing pretrained diffusion latents for nonlinear, interpretable control. ConDA fills this gap by (1) contrastively structuring embeddings with auxiliary variables, (2) enabling nonlinear traversal in this space, and (3) separating compact embedding edits from rendering in diffusion latent space. Our experiments show improved temporal consistency, smooth controllable trajectories, and faithful reconstructions across spatiotemporal data.

3. Method

3.1. Notation and Problem Statement

We consider controlled sequence generation for high-dimensional spatiotemporal data (e.g., videos of fluid dynamics, calcium imaging, neurostimulation E-fields, facial expression sequences, monkey electrophysiology). Each sequence $x_{1:S}=\left(x_1,\dots ,x_S\right)$ consists of a series of frames $x_s\in 𝓧$ with associated per-frame auxiliary variables $y_s\in 𝓨$ (e.g., encoding time, neurostimulation coil angle, facial action units, etc.). Our goal is to generate a target sequence ${\widehat{x}}_{j+1:j+n}=\left({\widehat{x}}_{j+1},\dots ,{\widehat{x}}_{j+n}\right)$ frame-by-frame under a user-specified auxiliary-variable trajectory $y_{j+1:j+n}^{\prime }$ (e.g., a single target image is the special case $n=1$). We study three evaluation settings: (i) Interpolation, where we regenerate intermediate frames under observed conditions $y_{j+1:j+n}^{\prime }=y_{j+1:j+n}$ to assess the interpolation fidelity of the learned embedding space. (ii) Held-out segment prediction (n-step-ahead), where given observed frames $\left(x_{1:j},y_{1:j}\right)$ we extrapolate the next $n$ frames by performing nonlinear traversal under target conditions $y_{j+1:j+n}^{\prime }=y_{j+1:j+n}$. (iii) Counterfactual editing, where given two classes $\left(x_j^{(0)},y_j^{(0)}\right)$ and $\left(x_j^{(1)},y_j^{(1)}\right)$, we generate class-conditional interpolated trajectories to transfer from class 0 to class 1, with condition $y_j^{\prime }=y_j^{(1)}$.

ConDA separates control from synthesis by introducing two latent spaces and a local-neighborhood lifting step. The pipeline proceeds in three stages. First, we train a conditional diffusion model (e.g., cLDM or pretrained IsoDiff) and apply diffusion inversion (e.g., DDIM or Rex-RK4) $g_{\varphi }:𝓧\times 𝓨\to 𝓩$ to map each observation–auxiliary-variable pair $\left(x_s,y_s\right)$ to a diffusion feature latent space $z_s=g_{\varphi }\left(x_s,y_s\right)\in 𝓩$. This feature latent space is a high-dimensional, near-lossless representation that supports high-fidelity reconstruction through the diffusion decoder $f_{\theta }:𝓩\times 𝓨\to 𝓧$.

Second, we learn a compact structured embedding $h_{\psi }:𝓩\times 𝓨\to 𝓒$ such that $c_s=h_{\psi }\left(z_s,y_s\right)\in 𝓒$, designed so that the local geometry of $𝓒$ aligns with the auxiliary variables and supports smooth, interpretable traversal. While $𝓩$ preserves reconstruction fidelity, it remains high-dimensional (e.g., $\dim (𝓩)=32\times 32\times 4$ when using cLDM on an image frame with $\dim (𝓧)=256\times 256\times 3$) and its geometry is not explicitly organized for stable temporal or condition-dependent editing; direct manipulation in $𝓩$ is therefore computationally expensive and brittle. In contrast, the structured embedding space $𝓒$ is low-dimensional (e.g. $\dim (𝓒)<10$ for spline-based interpolation), whose local geometry aligns with auxiliary variables and explicitly designed for smooth and interpretable trajectory manipulation, enabling efficient and stable control operations.

Third, we perform trajectory editing by traversing the embedding sequence $c_{1:S}$ using an editing operator ${\widehat{c}}_{s^{\prime }}={𝓣}_{\eta }\left(c_s;y_s\to y_{s^{\prime }}\right)$ to interpolate, extrapolate, or transfer between conditions. The edited embeddings are then lifted back to the feature latent space using a local-neighborhood–preserving kNN decoder $\ell :𝓒\to 𝓩$, which reconstructs each edited latent as a weighted combination of nearby training latents in $𝓒$. Finally, we render the edited sequence by decoding ${\widehat{x}}_{s^{\prime }}=f_{\theta }\left({\widehat{z}}_{s^{\prime }},y_{s^{\prime }}\right)$ with the diffusion model decoder.

This principled separation of editing and rendering enables stable, interpretable control of spatiotemporal sequences while preserving the high generative fidelity of diffusion models. It supports faithful interpolation and extrapolation in $𝓒$, including missing-frame imputation and structured condition control, with high-quality synthesis achieved by lifting edited embeddings back to $𝓩$ for decoding.

3.2. Encoding/Decoding to Diffusion Latent space

We train a conditional latent diffusion model (cLDM) to encode or decode to feature latent space, which operates diffusion process in the latent space $𝓩$ of a pretrained VAE instead of pixel space $𝓧$ which is very high-dimensional. The forward diffusion process progressively adds Gaussian noise to the data and a reverse diffusion iteratively denoises it, using a U-Net architecture. (Rombach et al., 2022) demonstrated that this approach achieves state-of-the-art performance in tasks such as text-to-image synthesis while significantly reducing computational overhead. We denote $E:𝓧\to 𝓩$ as encoder and $D:𝓩\to 𝓧$ as decoder of the pretrained VAE. We then learn the cLDM by minimizing the following objective

$${𝓛}_{\text{cLDM}}:={\mathbb{E}}_{E\left(x_s\right),y_s,ϵ\sim 𝓝(0,1),t}\left[{‖ϵ-{ϵ}_{\theta }\left(z_{s,t},t,y_s\right)‖}_2^2\right].$$

Here, ${ϵ}_{\theta }$ denote the UNet model (Ronneberger et al., 2015), $t$ is a diffusion timestep, $z_{s,t}$ is a noisy version of the clean input $E\left(x_s\right)$, and $y_s\in 𝓨$ denotes conditional label. This reduces compute cost and helps the model focus on semantic features rather than pixel-level noise.

For diffusion inversion, we employ DDIM (Song et al., 2021; Mokady et al., 2023) to invert a real encoded image $E\left(x_s\right)$ to a noisy latent code $z_s={\tilde{g}}_{\varphi }\left(E\left(x_s\right),y_s\right)$ with $g_{\varphi }\left(\cdot ,y_s\right)={\tilde{g}}_{\varphi }\left(\cdot ,y_s\right)\circ E$. When used as the starting point in the sampling process, this latent code enables reconstruction of the original image as ${\widehat{x}}_s=\left(D\circ {\tilde{f}}_{\theta }\right)\left(z_s,y_s\right)=f_{\theta }\left(z_s,y_s\right)$ (see Appendix B.1 for details). DDIM inversion provides a deterministic mechanism to compute inverted latent $z_s=z_{s,T}$, yielding a latent trajectory ${\left\{z_{s,t}\right\}}_{t=0}^T$ that maps cyclically to and from the data sample. This cycle-consistency property ensures that latent edits remain faithful when decoded back into data space.

3.3. Structured Embedding and Controlled Generation via ConDA

The complete ConDA pipeline, including diffusion model training (or use of a pretrained model), diffusion inversion, contrastive embedding, trajectory traversal, kNN lifting, and diffusion sampling is summarized in Algorithm 1 (Training) and Algorithm 2 (Generation) in Appendix B.1.

Structured Embedding of Diffusion Latents.

We introduce a supervised contrastive learning framework ConDA by learning a map $h_{\psi }$ that organizes high dimensional diffusion latents $𝓩$ into a low-dimensional structured space $𝓒$ that encodes ordered trajectories. We train $h_{\psi }$ using a supervised InfoNCE objective, where positives are defined as $\left(z_s,z_p\right)$ belonging to the same label or condition $y_s=y_p$, while negatives are drawn from samples with differing labels. The supervised contrastive loss is given by

$$\begin{array}{l}{𝓛}_{\text{InfoNCE}}=\\ -{\mathbb{E}}_{(s)}\left[\frac{1}{|P(s)|}\sum _{p\in P(s)}\log \frac{\exp \left(\text{sim}\left(c_s,c_p\right)/\tau \right)}{\sum _{a\ne s}\exp \left(\text{sim}\left(c_s,c_a\right)/\tau \right)}\right],\end{array}$$ (1)

where $c_s=h_{\psi }\left(z_s,y_s\right)$, $P(s)$ denotes the set of positives for anchor $s$ (same label), $\text{sim}\left(c_s,c_p\right)=-{‖c_s-c_p‖}_2^2$ is the similarity measure, and $\tau$ is a temperature parameter. This loss encourages embeddings with the same condition to cluster together while separating those with different conditions, enforcing that $𝓒$ encodes dynamics in a condition-aware, low-dimensional structure.

Trajectory Modeling and Controlled Generation.

Within the space $𝓒$ we introduce approaches that enable nonlinear dynamic modeling and controlled generation. The goal is to identify directions of movement in latent space that correspond to meaningful changes in observed data, either along the trajectory of a sequence or across distinct classes. We define a map ${𝓣}_{\eta }:𝓒\times 𝓨\times 𝓨\to 𝓒$ that shifts a source embedding $\left(c_{\text{s}},y_s\right)$ towards a target embedding $\left(c_{s^{\prime }},y_{s^{\prime }}\right)$, while leaving unrelated factors invariant. This is obtained by solving

$${\widehat{c}}_{s^{\prime }}=\underset{𝓣}{\text{arg min}}𝓛\left({𝓣}_{\eta }\left(c_s;y_s\to y_{s^{\prime }}\right);c_{s^{\prime }}\right)+\lambda \text{Ω}\left({𝓣}_{\eta }\right),$$ (2)

where $𝓛$ is a task-specific discrepency loss that ensures that the generated sample under the edited latent aligns with the target representation, while a regularization term $\text{Ω}(\cdot )$, weighted by $\lambda$, enforces edits to be small. The solution, ${\widehat{c}}_{s^{\prime }}$ thus reflects the minimal, structured change in the latent space $𝓒$ needed to achieve faithful, controllable generation ${\widehat{x}}_{s^{\prime }}=f_{\theta }\left(l\left({\widehat{c}}_{s^{\prime }}\right),y_{s^{\prime }}\right)$ aligned with the target condition.

1. Spline Interpolation. For sequential data, we assume that embeddings $c_{1:S}$ lie on or near a smooth manifold that captures the intrinsic phase of the sequence. To recover this structure, we fit a a $C^2$-continuous parametric curve $\gamma :ℝ\to 𝓒$, $\gamma (\alpha )\approx c_s$ for $\alpha ={\alpha }_s$, to the ordered pairs ${\left\{\left({\alpha }_s,c_s\right)\right\}}_{s=1}^S$, where ${\alpha }_s$ is a monotone ordering parameter (e.g., normalized index or latent phase). Concretely, parametric spline curve $\gamma$ is obtained by minimizing a regularized least-squares objective that balances data fidelity and curvature: ${\min }_{\gamma }{\sum }_{s=1}^S{‖\gamma \left({\alpha }_s\right)-c_s‖}^2+\lambda \int {‖{\gamma }^{\prime \prime }(\alpha )‖}^{2}d\alpha$, subject to appropriate boundary conditions. This smooth parameterization lets us define edits as motion along the curve. Given a point $c_s$ at coordinate ${\alpha }_s$, the edit operator moves by a small phase increment $\text{Δ}\alpha$:

   $${𝓣}_{\alpha }\left(c_s;y_s\to y_{s+\text{Δ}s}\right)=\gamma \left({\alpha }_s+\text{Δ}\alpha \right).$$ (3)

   Choosing $\text{Δ}\alpha >0$ or $\text{Δ}\alpha <0$ traverses forward or backward along the trajectory, enabling interpolation of missing states and extrapolation to unobserved phases.
2. Taylor Extrapolation (TEX). In this local extrapolation approach, we estimate the direction of change in $𝓒$ that moves a source embedding forward along the trajectory by a step size $\text{Δ}s$ using a second-order Taylor expansion around a local coordinate $s$:

   $$\begin{array}{l}𝓣\left(c_s;y_s\to y_{s+\text{Δ}s}\right)=c(s+\text{Δ}s)\approx c(s)+\dot{c}(s)\text{Δ}s\\ +\frac{1}{2}\ddot{c}(s){(\text{Δ}s)}^2,\end{array}$$ (4)

   where $\dot{c}$, $\ddot{c}$ are the first and second derivatives of the local trajectory estimated using finite differences. The first order TEX-1 uses only the linear term, while the second order TEX-2 adds $\frac{1}{2}\ddot{c}(s){(\text{Δ}s)}^2$, which captures local nonlinearities of the trajectory. This allows us to navigate the latent space along locally directed trajectories, enabling generation of sequences consistent with the estimated dynamics.
3. Classification+KDE-Based Class Traversal. Given sequences of latent embeddings ${\left\{c_{1:S}^{(j)}\right\}}_{j=1}^N$, we first train an SVM with an RBF kernel on $\left\{c^{(j)}\right\}$ to separate classes (e.g., responder vs. non-responder). To identify navigation directions between classes, we estimate class-conditional densities via KDE: $f_k(u)=\frac{1}{N_k}{\sum }_{i=1}^{N_k}K\left(\left(u-c_i^{(k)}\right)/h\right)$, $k\in \{0,1\}$, where $K(\cdot )$ is a Gaussian kernel, $h$ is the bandwidth, and $\left\{c_i^{(k)}\right\}$ are samples from class $k$. We then form the density difference $\text{Δ}f(u)=f_1(u)-f_0(u)$, and detect class-specific peaks by $m_{\text{class}0}={\text{arg max}}_u(-\text{Δ}f(u)),m_{\text{class}1}={\text{arg max}}_u(\text{Δ}f(u))$. The transformation function here is defined as a traversal from the source to the target class along the line connecting these peaks:

   $$\begin{array}{l}{𝓣}_{\eta }\left(c_s^{\text{(class}0)};y_s^{\text{class}0}\to y_s^{\text{class}1}\right)=\\ m_{\text{class}0}+\eta \left(m_{\text{class}1}-m_{\text{class}0}\right),\eta \in [0,1].\end{array}$$ (5)

   This procedure provides interpretability by linking movements in latent space to density maxima of the two classes, enabling smooth and explainable transitions between non-responder and responders.

4. Experiments

Datasets.

We evaluate on five spatiotemporal domains: fluid dynamics, neural calcium imaging, facial expressions, therapeutic neurostimulation, and neural spike dynamics; see Appendix A for acquisition, simulation, and split details.

1. Flow past a cylinder. 2D incompressible Navier–Stokes with cylinder interactions following the benchmark of Schäfer et al. (1996); simulated in FEniCS (Logg et al., 2012). We render total $N=\mathrm{246,000}$ RGB images (velocity magnitude) of 82 different flow trajectories, with 3000 time points as $y=\tau \in [0,1.5]$. To reduce temporal leakage, we, use blocked splits (Roberts et al., 2017).

2. Two-photon calcium imaging (PFC). High-resolution recordings from mouse prefrontal cortex during cognitive flexibility tasks (Spellman et al., 2021). We convert videos to $N=\mathrm{706,452}$ RGB images with $y=\left(v_n,v_t\right)$ (session and within-session frame).

3. DISFA facial expressions. Standard AU dataset (Mavadati et al., 2013). We use $N=\mathrm{261,576}$ frames of 27 subjects with 12D facial action units as $y=\text{AU}\in {\{0,\dots ,5\}}^{12}$.

4. TMS-induced electric fields. Individualized head models and E-field simulations with SimNIBS (Thielscher et al., 2015); $N=\mathrm{569,520}$ images from 121 patients with coil angles as $y=\text{Θ}\in \left[{10}^{°},{360}^{°}\right]$. Used for classification and class-transfer analyses.

5. Monkey reach neural spiking. Premotor population activity during delayed reaches (Churchland and Kaufman, 2022; Pei et al., 2021). The number of training trials for dataset (non-image dynamics) is 2008 with trial length of 140 time bins and recordings from 182 neurons. Following Kapoor et al. (2024), spikes are mapped to 16-D S4 latents; we train conditional diffusion directly on these latents with 2D reach velocity as $y$.

Experimental setup.

Complete hyperparameters and algorithmic steps appear in Appendix C.1 and Appendix B.1. For images, we use the LDM autoencoder (Rombach et al., 2022) to encode $x_s\in {ℝ}^{256\times 256\times 3}$ to $E\left(x_s\right)\in {ℝ}^{32\times 32\times 4}$, train a UNet2DModel (Hugging Face diffusers (Von Platen et al., 2022)) with 1000 diffusion steps, and apply DDIM inversion to obtain feature latents $z_s$. As a solver baseline, we also implement the reversible REX-RK4 method (Blasingame and Liu, 2026) using the noise-prediction parameterization, coupling parameter 0.9999, and 8 steps. We evaluate Rex-RK4 with both cLDM and the pretrained IsoDiff. Contrastive embeddings are learned with CEBRA (Schneider et al., 2023) (offset10-model-mse); we use $d=8$ for fluid and $d=3$ otherwise (Appendix Fig. 10). We lift $c_s\in {ℝ}^d$ using kNN decoder and render via diffusion sampling. The spatiotemporal datasets consist of 82 fluid flow videos (3,000 frames each), 60 calcium imaging videos (approximately 4,800–10,000 frames), and 27 facial expression videos (4,844 frames each). For held-out segment prediction, we split video sequence into $M$ contiguous train-test blocks, each containing a held-out segment of length $n$: Fluid:$M=50$, $n=12$; Ca²⁺ : $M=100$, $n=7$; DISFA : $M=88$, $n=11$.

Tasks and metrics.

Across datasets, we evaluate on the following tasks: (i) interpolation; (ii) $n$-step-ahead prediction; (iii) classification and counterfactual editing in the embedding space; and (iv) interpretability. Interpolation: parametric cubic B-splines and finite-difference TEX on latent sequences $z_{1:S}$; baselines are linear (Lerp), spherical (Slerp) (Wang and Golland, 2023; Preechakul et al., 2022), and an LSTM (Hochreiter and Schmidhuber, 1997). n-step-ahead prediction: spline extrapolation, TEX-2 and LSTM. Ablation: ConDA vs linear (PCA) and variational ($\beta$-VAE, $\beta =4$) baselines. Image fidelity: PSNR and SSIM; Trajectory error: RMSE. Classification: linear/RBF SVMs with leave-subject-out (TMS) and leave-(Re)-out (flow) (Osafo Nkansah et al., 2024); we report F1/Acc/AUC. Class transfer: KDE interpolation between class-conditional modes in the embedding (details in Appendix C.1). Interpretability: on monkey S4 latents, we compare ConDA embeddings to PCA in the same diffusion latents and evaluate prediction on 216 test samples using RMSE, total absolute error (mean/sd), and Procrustes distance.

5. Results

We benchmark cLDM against a continuous-label conditioned GAN (Ding et al., 2021) and find that cLDM substantially outperforms the GAN in both image generation and reconstruction (see Appendix C.10 for protocols and full results). In this section, we present both qualitative and quantitative results for our proposed approach ConDA. In addition to full-reference image quality metrics (PSNR, SSIM), we evaluate trajectory-level accuracy via RMSE between predicted and ground-truth embeddings.

Nonlinear Trajectory Modeling: Spline/TEX-2 vs. Baselines.

Our experiments show that the complex nonlinear dynamics in the high-dimensional diffusion latent space $𝓩$ are difficult to model directly: linear baselines underperform, and even nonlinear methods such as LSTM and TEX-1 achieve limited accuracy when operating in $𝓩$ (Tab. 1). In contrast, all methods especially Spline and TEX-2 achieve substantially better reconstruction scores when applied in the low-dimensional contrastive space $𝓒$, validating our approach. Qualitative results (Fig. 2) and generalization tests (Tab. 2, Fig. 11) further confirm that $𝓒$ captures system dynamics in a more structured, dynamics-aligned geometry, while remaining robust to approximation sources such as kNN lifting, diffusion inversion, and VAE reconstruction (Tab. 8, Appendix C.2). The near-zero RMSE achieved by Spline and TEX-2 in $𝓒$ highlights its suitability for accurate interpolation, prediction, and stable trajectory editing. We also show in Appendix C.3 that ConDA is solver-agnostic, transferring effectively across diffusion models and inversion solvers (e.g., IsoDiff, Rex-RK4, (Hahm et al., 2024; Blasingame and Liu, 2026)), and that kNN lifting which leverages neighbor-based contrastive structure, consistently outperforms a tuned MLP decoder across datasets (Appendix C.4).

Table 1. Baseline comparison of interpolations.

We compare spline and TEX-2 (second-order Taylor expansion) against Lerp, Slerp, LSTM, and TEX-1 (first-order Taylor). Linear baselines (Lerp/Slerp) consistently underperform in both $𝓩$ and $𝓒$, indicating inherently nonlinear trajectories. In $𝓩$, TEX-2 yields the best fidelity. In $𝓒$, spline and TEX-2 perform best, consistent with a dynamics-aligned geometry that supports smooth, predictable traversal. Overall, nonlinear methods (LSTM, TEX) improve notably in $𝓒$ compared to $𝓩$.

Dataset	Method	PSNR ↑	SSIM ↑	RMSE ↓	PSNR ↑	SSIM ↑	RMSE ↓
		DDIM $𝓩$-space			ConDA $𝓒$-space
Fluid	Lerp	28.16 ± 0.49	0.56 ± 0.08	16.12	28.27 ± 0.47	0.59 ± 0.07	13.75
	Slerp	28.11 ± 0.53	0.54 ± 0.09	17.24	28.27 ± 0.48	0.59 ± 0.07	13.09
	LSTM	34.50 ± 1.81	0.92 ± 0.09	4.04	34.53 ± 2.05	0.92 ± 0.08	1.07
	TEX-1	29.28 ± 2.38	0.57 ± 0.25	14.20	32.98 ± 3.09	0.84 ± 0.15	3.16
	TEX-2	35.29 ± 0.60	0.94 ± 0.01	0.26	35.70 ± 0.36	0.94 ± 0.01	0.02
	Spline	—	—	—	35.70 ± 0.36	0.94 ± 0.01	0.00
Ca2+	Lerp	32.36 ± 1.68	0.79 ± 0.04	16.60	32.92 ± 2.86	0.82 ± 0.04	43.87
	Slerp	32.68 ± 2.54	0.80 ± 0.07	17.65	32.82 ± 2.81	0.82 ± 0.04	43.38
	LSTM	35.52 ± 2.70	0.86 ± 0.07	9.15	35.96 ± 2.91	0.87 ± 0.07	0.89
	TEX-1	30.64 ± 1.07	0.57 ± 0.08	22.58	36.13 ± 2.96	0.87 ± 0.07	0.84
	TEX-2	38.00 ± 0.51	0.92 ± 0.01	0.13	38.58 ± 0.59	0.93 ± 0.01	0.00
	Spline	—	—	—	38.58 ± 0.59	0.93 ± 0.01	0.00
DISFA	Lerp	29.98 ± 2.70	0.74 ± 0.08	13.86	33.08 ± 3.25	0.83 ± 0.09	13.61
	Slerp	31.13 ± 2.72	0.76 ± 0.09	14.97	33.13 ± 3.31	0.83 ± 0.09	13.68
	LSTM	38.13 ± 0.54	0.96 ± 0.00	3.64	38.74 ± 0.75	0.96 ± 0.01	0.34
	TEX-1	35.18 ± 4.06	0.88 ± 0.14	7.49	38.77 ± 0.70	0.96 ± 0.00	0.24
	TEX-2	38.27 ± 0.24	0.96 ± 0.00	0.07	38.99 ± 0.23	0.96 ± 0.00	0.00
	Spline	—	—	—	38.99 ± 0.23	0.96 ± 0.00	0.00

Table 2.

$n$-step-ahead prediction using nonlinear extrapolation methods (Spline, TEX-2, and LSTM). The results show that generalization remains stable and is not degraded by approximation sources such as k-NN lifting, diffusion inversion, and VAE reconstruction, indicating that these errors are negligible for downstream prediction and do not impede faithful recovery or extrapolation.

	Fluid			Ca2+			DISFA
Method	PSNR ↑	SSIM ↑	RMSE ↓	PSNR ↑	SSIM ↑	RMSE ↓	PSNR ↑	SSIM ↑	RMSE ↓
	ConDA $𝓒$-space			ConDA $𝓒$-space			ConDA $𝓒$-space
LSTM	34.18 ± 1.97	0.90 ± 0.11	1.87	35.11 ± 2.57	0.86 ± 0.06	0.90	36.44 ± 2.80	0.93 ± 0.06	0.75
TEX-2	34.32 ± 2.31	0.91 ± 0.10	1.68	34.92 ± 2.53	0.86 ± 0.06	1.05	36.59 ± 2.71	0.93 ± 0.07	0.61
Spline	34.51 ± 2.13	0.92 ± 0.07	1.48	35.00 ± 2.57	0.86 ± 0.06	1.09	37.17 ± 2.17	0.95 ± 0.04	0.52

SVM Classification and KDE-Based Class Traversal.

We evaluated the ConDA embedding space ($𝓒$) on downstream classification by training linear and RBF-kernel SVMs on two tasks: Fluid (steady vs. unsteady laminar flow) and E-Field (treatment responder vs. non-responder). As shown in Tab. 3, classifiers trained on $𝓒$-space consistently outperformed those trained directly in the diffusion latent space $𝓩$ across Accuracy, F1, and ROC-AUC, independent of the kernel choice, indicating improved separability of task-relevant structure in $𝓒$. For example, in the E-Field task, an RBF-SVM classifier trained on $𝓒$ achieved an F1 score of 0.63, a marked improvement over the 0.30 in $𝓩$. Leveraging the interpretability afforded by the low dimensionality of $𝓒$, we visualized the decision-relevant latent dynamics. We used kernel density estimation to model the class-conditional densities and then interpolated along the vector connecting the peaks of these densities. The resulting sequences, depicted in Fig. 3, illustrate the smooth trajectory of change required for a non-responder E-field pattern to morph toward that of a responder. The color-coded changes reveal specific cortical shifts, providing a tangible and interpretable path toward personalizing targeting strategies.

Table 3. SVM classification performance.

Classification of steady vs. unsteady laminar flow and responder vs. non-responder E-Fields in $𝓩$ and $𝓒$-space.

	Steady vs. Unsteady			Resp. vs. Non-Resp.
Method	Acc. ↑	F1 ↑	AUC ↑	Acc. ↑	F1 ↑	AUC ↑
	DDIM $𝓩$-Space
SVM-Linear	0.78	0.73	0.89	0.48	0.21	0.46
SVM-RBF	0.68	0.52	0.59	0.57	0.30	0.58
	ConDA $𝓒$-Space
SVM-Linear	0.83	0.81	0.82	0.62	0.62	0.68
SVM-RBF	0.87	0.85	0.95	0.66	0.63	0.67

Ablation: ConDA vs. linear and variational baselines.

We compared ConDA against two baseline representation learning techniques: PCA, a linear method, and a $\beta$-VAE, a variational method on test-segments prediction. To ensure a fair comparison, all models used identical train/validation splits and were configured with the same latent dimension. Each representation was evaluated with its corresponding decoder. As summarized in Tab. 4, ConDA consistently achieved superior performance across all datasets, as measured by PSNR, SSIM, and RMSE. Qualitative comparisons on DISFA (Fig. 4) further illustrate ConDA’s ability to recover fine facial motions at unseen frames with higher fidelity. These results indicate that the baseline models fail to capture the dynamics-relevant variance as effectively as ConDA.

Table 4. ConDA vs linear and variational baselines.

Performance metrics for test sequence prediction comparing ConDA to PCA and $\beta$-VAE. ConDA achieves improved performance indicating closer alignment with the underlying dynamical manifold.

Dataset	Method	PSNR ↑	SSIM ↑	RMSE ↓
Fluid	PCA	31.82 ± 2.30	0.87 ± 0.11	10.43
	$\beta$-VAE	32.20 ± 2.55	0.85 ± 0.14	7.47
	ConDA	33.90 ± 2.52	0.90 ± 0.11	7.60
Ca 2+	PCA	34.73 ± 2.29	0.86 ± 0.05	12.40
	$\beta$-VAE	34.62 ± 2.40	0.85 ± 0.06	11.58
	ConDA	35.69 ± 2.54	0.87 ± 0.05	11.02
DISFA	PCA	34.82 ± 1.03	0.92 ± 0.02	7.60
	$\beta$-VAE	35.76 ± 1.23	0.94 ± 0.02	6.59
	ConDA	36.26 ± 1.03	0.94 ± 0.01	5.50

Interpretability of ConDA-space $𝓒$.

We assess interpretability on neural spiking data from a monkey performing a delayed center–out reach task (Kapoor et al., 2024). Following the LDNS pipeline, 182-channel spikes are encoded into 16-dimensional, time-aligned latents using an S4-based autoencoder. Conditional diffusion models are trained on these latents, conditioned on 2D hand velocity, with DDIM used for inversion. We visually compare predicted test embeddings in 2D PCA and $𝓒$ spaces, evaluate alignment with held-out spiking data via RMSE, total absolute error (mean/std), and geometric measures (Procrustes distance). As shown in Fig. 5, $𝓒$ embeddings exhibit clearer organization by reach direction and stronger alignment with true velocity trajectories than PCA. Quantitative metrics consistently favor $𝓒$, indicating that contrastive structuring better preserves task-relevant kinematics and improves interpretability relative to variance-only PCA embeddings. Predicted mean firing rates closely match held-out data (Pearson $r=0.95$; Fig. 5d).

The near-orthogonality of the ConDA latent subspace directions associated with temporal variations and those influenced by confounding factors (e.g., Reynolds number in the fluid dataset) is provided as evidence of successful disentanglement in Appendix C.9.

6. Conclusion and Limitations

We introduced ConDA, a contrastive alignment layer that organizes diffusion latents into compact, dynamics-aware embeddings, enabling smooth nonlinear traversals that outperform raw latent and linear baselines across physical and biological domains. ConDA is solver-agnostic and operates consistently across different inversion solvers (e.g., DDIM, Rex-RK4 (Blasingame and Liu, 2026)) as well as across different diffusion models (e.g., IsoDiff (Hahm et al., 2024) latents), see Appendix C. The method has three limitations: (1) the embedding is lossy and many-to-one, which we mitigate using a neighborhood-preserving kNN lifting (with learned decoders as future work); (2) the embedding emphasizes local rather than global geometry, which may distort long-range traversals, though locally fitted interpolators and models such as IsoDiff can help; and (3) the approach relies on DDIM inversion, adding computational cost that can be reduced through caching, multi-scale pipelines, or faster solvers such as Rex-RK4 or consistency models. Despite these trade-offs, approximation errors remain small in practice, and ConDA provides a practical path toward controllable and interpretable modeling of complex dynamical systems.

Impact Statement.

This paper proposes a method to enhance the latent space of diffusion models to support interpretable and controllable editing by aligning latent geometry with underlying dynamical factors. Our work shares the same ethical issues as generative models more broadly, including risks related to fake data generation, and malicious manipulation of images, or spatiotemporal signals. While ConDA does not introduce new generative capabilities beyond existing diffusion models, a more structured and disentangled latent representation may also lower the barrier for certain forms of misuse. At the same time, ConDA is designed primarily for scientific and biomedical applications, such as modeling physical systems, neural activity, and therapeutic neurostimulation, where interpretability and controlled counterfactual reasoning are critical. Ethical concerns related to dataset bias, representational fairness, and downstream discrimination remain unchanged by our approach; ConDA neither mitigates nor exacerbates these issues directly. We believe that responsible deployment, transparency about limitations, and domain-expert oversight are essential when applying such methods in high-stakes settings. A collective effort within the research community and society will be important to ensure that generative models and their extensions remain beneficial and are used in socially responsible ways.