CardioComposer: Leveraging Differentiable Geometry for Compositional Control of Anatomical Diffusion Models

TL;DR

This work proposes CardioComposer: a programmable, inference-time framework for generating multi-class anatomical label maps based on interpretable ellipsoidal primitives, and develops differentiable measurement functions based on voxel-wise geometric moments, enabling loss-based gradient guidance during diffusion model sampling.

Abstract

Generative models of 3D cardiovascular anatomy can synthesize informative structures for clinical research and medical device evaluation, but face a trade-off between geometric controllability and realism. We propose CardioComposer: a programmable, inference-time framework for generating multi-class anatomical label maps based on interpretable ellipsoidal primitives. These primitives represent geometric attributes such as the size, shape, and position of discrete substructures. We specifically develop differentiable measurement functions based on voxel-wise geometric moments, enabling loss-based gradient guidance during diffusion model sampling. We demonstrate that these losses can constrain individual geometric attributes in a disentangled manner and provide compositional control over multiple substructures. Finally, we show that our method is compatible with a wide array of anatomical systems containing non-convex substructures, spanning cardiac, vascular, and skeletal organs.

Figures (16)

• Figure 1: We present a guidance framework to constrain diffusion models of multi-label anatomical segmentations based on simple geometric features. Such features include size, shape, and position, and can be represented as ellipsoids in 3D space (panel A). Our inference-time approach enables generation based on independent geometric features (panels B-D), and supports multi-component compositional generation (panels E-G). Gray and blue voxels represent components that are unconstrained and constrained, respectively. Purple ellipsoids indicate a strong overlap between target and sample ellipsoids .
• Figure 2: Our method involves applying a geometric guidance correction step for every denoising iteration. Left: The noised latent $\mathbf{z}_\sigma$ is passed through the diffusion model and VAE decoder to produce a clean voxel space prediction $\mathbf{\hat{x}}_0$ (\ref{['seg:segmentation_denoise']}). Middle: The segmentation is parsed for relevant substructures $\mathbf{\Omega}$, and geometric moments $\mathcal{G}$ are extracted for each substructure (\ref{['sec:extract_geo_moments']}). Right: Measured geometric moments $\mathcal{G}$ are compared to target moments $\bar{\mathcal{G}}$ through geometric moment losses. Bottom: The gradient derived from the aggregate loss corrects the denoising step.
• Figure 3: Geometric guidance can generate synthetic anatomy with geometric constraints. Grid shows example synthetic label maps where constraints are applied to the myocardium voxels . Rows: baseline conditioning and guidance methods (CFG = classifier-free guidance, ANG = adaptive null guidance).
• Figure 4: Geometric guidance can enforce conditional fidelity while maintaining realism. Line plots compare conditioning and guidance mechanisms based on the geometric properties of the myocardium. MMD values are multiplied by $10^3$.
• Figure 5: Geometric guidance enables independent control of size, shape, and position.Columns show synthetic label maps generated by geometric guidance applied to the right ventricle voxels using various geometric losses. Rows represent which geometric feature is being independently controlled.