Contrastive Diffusion Alignment: Learning Structured Latents for Controllable Generation

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.

💡 Research Summary

The paper addresses a fundamental limitation of modern diffusion models: while they excel at high‑fidelity image and video generation, their latent spaces are high‑dimensional and lack an explicit organization that reflects temporal or conditional dynamics. Consequently, tasks that require controllable generation—such as interpolating fluid flow fields, editing neural activity recordings, or smoothly transitioning facial expressions—are hampered by blurry or inconsistent intermediate states when using simple linear interpolations or existing conditioning mechanisms like ControlNet.

To solve this, the authors propose Contrastive Diffusion Alignment (ConDA), a plug‑and‑play geometry layer that sits on top of any pretrained conditional diffusion model. ConDA works in three stages. First, a conditional latent diffusion model (cLDM) together with a deterministic inversion method (DDIM or Rex‑RK4) maps each frame‑auxiliary pair ((x_s, y_s)) to a high‑dimensional diffusion feature latent (z_s). This latent space preserves reconstruction fidelity but is not suitable for direct manipulation.

Second, ConDA learns a compact contrastive embedding (c_s = h_\psi(z_s, y_s) \in \mathbb{R}^d) (typically (d<10)). Using an InfoNCE‑style loss, embeddings of samples sharing the same auxiliary variable (e.g., the same time step or stimulation parameter) are pulled together, while embeddings of different conditions are pushed apart. This forces the local geometry of the embedding space to align with the underlying dynamical factors, yielding interpretable axes that correspond to time, stimulus intensity, facial action units, etc.

Third, trajectory editing is performed entirely in the low‑dimensional space. Because the space is smooth and structured, standard nonlinear operators—cubic splines for interpolation, finite‑difference schemes for extrapolation, or recurrent networks (LSTM) for learned dynamics—can be applied to generate a modified embedding sequence (\hat c’_s). To render the edited sequence, ConDA lifts each (\hat c’_s) back to the original diffusion latent space using a neighborhood‑preserving k‑nearest‑neighbors decoder. The decoder reconstructs (\hat z’_s) as a weighted combination of nearby training latents, thereby preserving the local geometry learned during contrastive training. Finally, the diffusion decoder synthesizes the output frames (\hat x’s = f\theta(\hat z’_s, y’_s)).

The authors evaluate ConDA on five diverse spatiotemporal domains:

1. Fluid dynamics – ConDA achieves a PSNR of 35.7 dB versus 28.3 dB for linear baselines, and flow‑field interpolations are visually smooth without the vortex smearing seen in other methods.
2. Neural calcium imaging – Temporal progression of activity states is continuous and less noisy, improving downstream classification of neural events.
3. Therapeutic neurostimulation – Embeddings capture systematic changes as the stimulation coil angle varies, enabling class‑conditional transitions that respect the underlying physics.
4. Facial expression dynamics – Subject identity is preserved while facial action units transition smoothly; counterfactual edits (e.g., neutral → surprise) are realistic.
5. Monkey motor cortex recordings – Condition‑dependent reaching trajectories are disentangled in the embedding, facilitating accurate n‑step‑ahead predictions of motor intent.

Across all tasks, ConDA outperforms linear interpolation, ControlNet, InstructPix2Pix, and direct conditioning approaches in both quantitative metrics (PSNR, SSIM, LPIPS) and qualitative assessments (smoothness, realism). Ablation studies show that the method is robust to the choice of inversion solver and that a simple MLP decoder can replace k‑NN with minimal loss, confirming that the core benefit stems from the contrastively aligned geometry rather than a specific decoder architecture.

Limitations include reliance on accurate auxiliary labels (the contrastive loss degrades with noisy or missing condition information) and potential insufficiency of a very low‑dimensional embedding for highly chaotic dynamics. Memory overhead of the k‑NN decoder may also become significant for very large training sets.

The paper concludes that by explicitly aligning diffusion latents with auxiliary variables through contrastive learning, ConDA provides a general, model‑agnostic framework for controllable, interpretable generation of spatiotemporal data. Future work is suggested on self‑supervised contrastive objectives, more memory‑efficient lifting mechanisms, and integration with neural ODE/SDE dynamics to handle even more complex systems.