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

1. Problem Statement

Generative models for 3D cardiovascular anatomy are essential for creating virtual cohorts for clinical trials, medical device evaluation, and biophysical simulations. However, existing approaches face a critical trade-off between geometric controllability and anatomical realism:

• Simple Primitives: Parametric models (e.g., cylinders, ellipsoids) offer precise control over size and position but fail to capture complex anatomical realism (e.g., non-convex shapes, branching structures).
• Data-Driven Models: Deep generative models (e.g., VAEs, Diffusion Models) produce realistic anatomy but often lack fine-grained, interpretable control over specific geometric attributes (size, shape, position) without retraining.
• Conditional Training Limitations: Existing conditional diffusion models often require retraining for specific constraints and struggle with compositional control (manipulating multiple substructures simultaneously) or topological changes.
• Scale Sensitivity: In biophysical simulations, millimeter-scale geometric variations can drastically alter physiological outcomes, necessitating precise control over substructure geometry.

The core challenge is to develop a framework that can independently and compositionally control the size, shape, and position of multiple anatomical substructures within a realistic 3D diffusion model without retraining.

2. Methodology: CardioComposer

The authors propose CardioComposer, an inference-time guidance framework that controls unconditional latent diffusion models (LDMs) using differentiable geometric losses.

A. Core Architecture

• Base Model: An unconditional Latent Diffusion Model (LDM) trained on 3D multi-class anatomical segmentation maps (voxel grids).
• Inference-Time Guidance: Instead of conditioning the model during training, the authors apply energy-based guidance during the sampling (denoising) process.

B. Differentiable Geometric Measurement

The method introduces a set of differentiable functions to compute geometric moments from the predicted segmentation at each denoising step:

1. Substructure Selection: Specific tissue labels (or unions of labels) are selected as target substructures ($\Omega$).
2. Geometric Moments ($G$): For each substructure, the system computes:
   • Mass ($M$): Zeroth-order moment (volume/size).
   • Centroid ($C$): First-order moment (position).
   • Covariance ($S$): Second-order moment (shape and orientation).
   • Note: The covariance is scale-normalized to isolate shape and orientation from size.
3. Differentiability: These moments are computed using soft binarization (softmax) and differentiable operations (matrix multiplication, division with noise stabilization), allowing gradients to flow back to the latent space.

C. Geometric Guidance Loss

The sampling process is guided by minimizing a composite loss function ($L_{geom}$) comparing the measured moments ($G$) against target moments ($\bar{G}$):
$$L_{geom} = \lambda_0 L_{size} + \lambda_1 L_{pos} + \lambda_2 L_{shape}$$
Where each term is a Mean Squared Error (MSE) between measured and target values.

The guidance is applied via a gradient correction step to the denoising prediction ($D_\theta$):
$$D^w_\theta(z_\sigma; \sigma) = D_\theta(z_\sigma; \sigma) - \sigma^2 \cdot w \cdot \nabla_{z_\sigma} L_{geom}$$
Here, $w$ is the guidance weight, and the gradient $\nabla_{z_\sigma} L_{geom}$ steers the generation toward the target geometry.

D. Compositional Control

The framework supports disentangled control (controlling one attribute while leaving others free) and compositional control (simultaneously constraining multiple substructures, e.g., Left Ventricle + Right Ventricle + Aorta) by summing the losses of individual substructures.

3. Key Contributions

1. Differentiable Geometry for Anatomy: Introduction of voxel-wise geometric moment functions (mass, centroid, covariance) that are fully differentiable, enabling gradient-based guidance for 3D anatomical label maps.
2. Inference-Time Control without Retraining: Demonstration that unconditional diffusion models can be precisely controlled for size, shape, and position using geometric losses, eliminating the need for conditional retraining.
3. Disentangled and Compositional Control: The ability to independently manipulate specific geometric attributes of specific substructures (e.g., changing the size of the Right Ventricle without altering its position or shape) and to compose constraints across multiple complex substructures simultaneously.
4. Generalizability: Validation of the method across diverse anatomical systems (cardiac, vascular, skeletal) and complex geometries (non-convex, branching, and Boolean unions of tissues).

4. Experimental Results

The method was evaluated on the TotalSegmentator dataset (cardiac) and additional datasets for the aorta, spine, and knee.

• Conditional Fidelity vs. Realism:
  • CardioComposer significantly outperformed conditional baselines (Explicit Concatenation, Implicit Concatenation, Cross-Attention) in matching target geometric constraints (Size, Position, Shape).
  • Unlike baselines, CardioComposer maintained high morphological quality (Precision, Recall, Fréchet Distance) and pointcloud metrics (MMD, Coverage) even at high guidance weights, avoiding the "over-smoothing" or artifacts common in other guidance methods.
• Disentanglement:
  • Experiments showed that applying a loss for "Mass" only affected the size distribution while leaving position and shape distributions broad and unchanged.
  • Conversely, applying all losses collapsed all distributions to the target values.
• Multi-Part Composition:
  • The method successfully generated anatomies with simultaneous constraints on 1, 3, and 6 substructures.
  • Baseline conditional models with dropout failed to generate high-quality anatomy when all conditioning channels were present (a rare training configuration), whereas CardioComposer handled arbitrary combinations seamlessly.
• Complex Geometries:
  • The framework successfully controlled non-convex structures (e.g., star-shaped chambers, branching vessels) and Boolean unions (e.g., combining both vena cavae into a single constrained substructure).
• Biophysical Simulation Application:
  • The authors demonstrated "geometric inpainting" to edit patient-specific anatomy (e.g., enlarging the Right Ventricle) and subsequently ran finite element method (FEM) simulations. The edited models showed physically plausible changes in wall displacement, validating the utility for counterfactual clinical trials.

5. Significance and Impact

• Clinical Translation: CardioComposer bridges the gap between generative AI and physics-based simulation. It allows researchers to generate "digital siblings" with specific, controlled anatomical variations to test medical devices or drug responses in silico.
• Interpretability: By using ellipsoidal primitives and geometric moments as control knobs, the method provides an interpretable interface for clinicians and engineers, moving away from "black box" latent vectors.
• Efficiency: The inference-time approach avoids the computational cost and data requirements of training separate conditional models for every new geometric constraint.
• Scalability: The method's success across cardiac, vascular, and skeletal systems suggests it is a generalizable solution for 3D anatomical generation, potentially accelerating the development of virtual clinical trials and personalized medicine.

In summary, CardioComposer represents a significant advancement in anatomical generative modeling by enabling precise, interpretable, and compositional geometric control over complex 3D structures without sacrificing realism or requiring retraining.