3D Vessel Graph Generation Using Denoising Diffusion

TL;DR

This work tackles the generation of realistic 3D vessel graphs that include cycles, such as capillaries and the Circle of Willis, by introducing the first application of denoising diffusion models to vascular graph generation. It proposes a novel two-stage approach that first performs continuous diffusion to generate node coordinates $\boldsymbol{X} \in \mathbb{R}^{n \times 3}$ and then applies discrete diffusion to determine edge attributes in $\mathbf{E} \in \mathbb{R}^{n \times n \times c}$ while keeping node positions fixed. The node-denoising component uses a lightweight neural network with time embeddings and cross-attention, yielding a loss $\mathcal{L}_{\mathrm{Node}}$, and the edge-denoising component uses a graph transformer with cross-entropy loss plus a node-degree KL term $\mathcal{L}_{\circ}$ via Gumbel-softmax, for a total $\mathcal{L}_{\mathrm{Edge}} = \mathcal{L}_{\mathrm{CE}} + \mathcal{L}_{\circ}$. Evaluations on VesSAP capillary graphs and Circle of Willis datasets show the method produces diverse, anatomically plausible graphs with accurate topology and geometry, outperforming baselines like Congress and MiDi. This diffusion-based vessel graph generator enables realistic data augmentation and synthetic image generation, with potential conditioning on disease labels for clinically relevant analyses.

Abstract

Blood vessel networks, represented as 3D graphs, help predict disease biomarkers, simulate blood flow, and aid in synthetic image generation, relevant in both clinical and pre-clinical settings. However, generating realistic vessel graphs that correspond to an anatomy of interest is challenging. Previous methods aimed at generating vessel trees mostly in an autoregressive style and could not be applied to vessel graphs with cycles such as capillaries or specific anatomical structures such as the Circle of Willis. Addressing this gap, we introduce the first application of \textit{denoising diffusion models} in 3D vessel graph generation. Our contributions include a novel, two-stage generation method that sequentially denoises node coordinates and edges. We experiment with two real-world vessel datasets, consisting of microscopic capillaries and major cerebral vessels, and demonstrate the generalizability of our method for producing diverse, novel, and anatomically plausible vessel graphs.

Figures (5)

• Figure 1: Overview of our method. Left. First, we use a continuous diffusion model to generate plausible node locations. Right. Subsequently, we apply discrete diffusion to generate a plausible edge configuration given the node coordinates. The node coordinates remain unchanged during edge learning. The model is optimized by $\mathcal{L}_{\hbox{Edge}}$ and focuses on the conditional edge distribution given node locations.
• Figure 2: Two real-world vessel datasets. Left. An example from VesSAP and the capillary level vessel graph patch. Right.The extracted metric graph for the blood vessel in CoW. (Best viewed when zoomed in)
• Figure 3: VesSAP and CoW graphs generated by our model in comparison to ground truth samples. Our model is able to learn complex structures, such as loops and orientation characteristics for both datasets.
• Figure 4: Comparison of our method against the MiDi for the VesSAP and Circle of Willis (CoW) datasets. We show the results for learning coordinates ($x,y,z$), edge angles ($\mathcal{E} ~\angle$), edge orientation ($\theta,\phi,\psi$) and edge length ($l_\mathcal{E}$). We also compare the statistics for node degree(deg($\mathcal{V}$)), number of edges (|$\mathcal{E}$|), and the Betti values ($\beta_0$$\beta_1$). The ground truth distribution is in purple, the distribution learned by MiDi is in orange, and our method is in green. As observed, our method emulates the statistics of the ground truth graphs more faithfully than MiDi. The degree of distribution of MiDi vs. our method on the VesSAP dataset is especially interesting. While the VesSAP graphs contain no degree 2 nodes, MiDi generates graphs with a large number of degree 2 nodes. Our method overcomes this shortcoming and generates minimal degree 2 nodes.
• Figure 5: VesSAP and CoW graphs generated by our model and MiDi, respectively. Note that, in the case of VesSAP, MiDi failed to capture the node coordinate distribution, which is the driving property for correct edge distribution, and hence produces an overconnected graph, resulting in a high Betti number error. However, MiDi fares relatively well in the CoW configuration, which resembles a molecular layout. In contrast, our model is able to generate diverse and valid novel graphs for both datasets.