Bayesian Optimization on Large Graphs via a Graph Convolutional Generative Model: Application in Cardiac Model Personalization

TL;DR

This paper introduces a graph convolutional VAE for Bayesian optimization on large non-Euclidean heart models, improving tissue property estimation by capturing complex geometrical features.

Contribution

It proposes a novel graph convolutional VAE that models non-Euclidean data, enabling effective Bayesian optimization on large cardiac graphs with transferability across geometries.

Findings

Enhanced modeling of heart geometry improves tissue property estimation.

The method successfully transfers learned features across different heart geometries.

Experiments demonstrate superior performance over traditional Euclidean VAEs.

Abstract

Personalization of cardiac models involves the optimization of organ tissue properties that vary spatially over the non-Euclidean geometry model of the heart. To represent the high-dimensional (HD) unknown of tissue properties, most existing works rely on a low-dimensional (LD) partitioning of the geometrical model. While this exploits the geometry of the heart, it is of limited expressiveness to allow partitioning that is small enough for effective optimization. Recently, a variational auto-encoder (VAE) was utilized as a more expressive generative model to embed the HD optimization into the LD latent space. Its Euclidean nature, however, neglects the rich geometrical information in the heart. In this paper, we present a novel graph convolutional VAE to allow generative modeling of non-Euclidean data, and utilize it to embed Bayesian optimization of large graphs into a small latent space. This approach bridges the gap of previous works by introducing an expressive generative model that is able to incorporate the knowledge of spatial proximity and hierarchical compositionality of the underlying geometry. It further allows transferring of the learned features across different geometries, which was not possible with a regular VAE. We demonstrate these benefits of the presented method in synthetic and real data experiments of estimating tissue excitability in a cardiac electrophysiological model.

Tables

Table 1. Table 1: Comparison of reconstruction accuracy with the presented gVAE, PCA, and fVAE with various depths  [ 6 ] in five different geometries.

	SSE					DC					Trainable parameters
Anatomy	1	2	3	4	5	1	2	3	4	5
PCA	12.73	14.18	12.35	23.85	23.26	39.80	39.87	41.31	49.17	54.42	NA
fVAE-3h [6]	8.03	8.45	8.44	13.07	13.70	61.77	66.20	60.45	70.30	70.72	2,087,822
fVAE-4h	7.97	8.29	8.33	12.47	13.99	61.76	64.60	61.58	71.51	69.84	2,613,134
fVAE-5h	7.42	8.01	8.19	13.96	12.41	64.59	65.72	62.21	68.04	74.50	3,138,446
gVAE	6.66	6.89	6.79	11.28	11.43	68.43	70.92	70.70	75.10	76.86	2,778,069

Equations

$$\begin{array}{rl}\partial u/\partial t& =\nabla (\mathbf{D}\nabla u)-cu(u-\theta )(u-1)-uv,\\ \partial v/\partial t& =\epsilon (u,v)(-v-cu(u-\theta -1)),\end{array}$$

$$\hat{\mathbf{\theta }}=\underset{\mathbf{\theta }}{\mathrm{argmax}}\{-\sum _t||{\mathbf{y}}_d(t)-M(\mathbf{\theta })|{|}^2\}.$$

$$g_l(\mathbf{\upsilon })=\sum _{\mathbf{p}\in \mathcal{P}}{w}_{\mathbf{p},l}\prod _{i=1}^d{N}_{i,p_i}^m({\upsilon }_i),$$

$$(\mathbf{f}\ast \mathbf{g})(i)=\frac{1}{|\mathcal{N}(i)|}\sum _{l=1}^M\sum _{j\in \mathcal{N}(i)}{f}_l(j)g_l(\mathbf{\upsilon }(i,j)).$$

$$\displaystyle\hat{\mathbf{z}}=\operatorname*{\arg\max}_{\mathbf{z}}\{{-\sum_{t}||\mathbf{y}_{d}(t)-M\big{(}\mathrm{E}[p_{\boldsymbol{\beta}}(\boldsymbol{\theta}|\textbf{z})]\big{)}||^{2}}\}.$$

$$\mathrm{EI}(\mathbf{z})$$

Full text

1 Introduction

Personalized computer models of the heart have shown promise in various clinical tasks such as risk stratification [2, 13] and predicting treatment response [14]. With advances in medical imaging technologies, personalization of high-resolution anatomical models is now feasible. By contrast, organ tissue properties vary spatially over the three-dimensional anatomical domain but have to be estimated from indirect, sparse and noisy data, resulting in a difficult ill-posed optimization problem with a high-dimensional (HD) unknown.

Numerous works have been presented in estimating these spatially-varying tissue properties in the form of high-dimensional model parameters. Most existing methods choose to represent spatially-varying tissue properties via a low-dimensional (LD) partitioning of the underlying geometrical model, either as pre-defined segments [17], or iteratively optimized in a coarse-to-fine fashion [5, 14]. This LD-to-HD definition directly exploits the spatial proximity and hierarchical composition of the underlying geometry. However, it is of such limited expressiveness that the number of partitioning is either too low to faithfully represent high-resolution tissue properties, or too high to allow effective optimization. In contrast, recent work presented the use of a data-driven generative model of HD tissue properties, via a variational auto-encoder (VAE), to embed the optimization into a LD latent space [6]. Being more expressive, this VAE-based generative model is able to represent high-resolution tissue properties with a latent code sufficiently small for effective optimization. However, as the regular VAE is defined over Euclidean data, it does not take into account the valuable geometry information in the data, nor does it allow transferring among different geometry without first establishing point-by-point correspondence.

If we view organ tissue properties over a 3D geometrical model as an image, convolutional neural networks (CNNs) are a natural choice to incorporate knowledge of the spatial proximity and hierarchical composition of the image  [4]. However, standard CNNs have been most successful on data with an underlying Euclidean structure (i.e., image grids). Generalizing CNNs to non-Euclidean domains is an emerging area of research [4], where significant efforts have been presented on addressing the challenges of defining convolution [8], pooling, and up-sampling operations [18]. However, most developments to non-Euclidean CNNs are focused on supervised discriminative networks. To date, very limited work have been presented to enable generative modeling of non-Euclidean data.

In this paper, we present a novel VAE architecture that allows generative modeling of data over non-Euclidean domains, and utilize this generative model to embed Bayesian optimization of large graphs into a LD manifold. The presented approach bridges the gap of previous works by introducing an expressive generative model that is able to represent high-resolution tissue properties with a small latent code, while incorporating the geometrical knowledge in the data and being transferable across geometries. We evaluate the presented method in synthetic and real-data experiments of estimating tissue excitability in a cardiac electrophysiological model, where we compare the expressiveness of the generative model and the accuracy of the subsequent parameter optimization with those obtained by using a linear reconstruction model based on principal component analysis (PCA) and a regular fully-connected VAE  [6]. We further demonstrate the feasibility of transferring the presented non-Euclidean VAE across patients. To our knowledge, this is the first introduction of a graph convolutional VAE and its use to enable Bayesian optimization over large graphs. Our code is available at https://github.com/jwaladhamala/BO-GVAE.

2 Background: Models of Cardiac Electrophysiology

Cardiac Electrophysiological Model: Among various computational models of cardiac electrophysiology, phenomenological models are widely used in parameter personalization as they are able to express key macroscopic properties of cardiac excitation with a small number of parameters [14, 5]. We thus adopt the phenomenological two-variable Aliev-Panfilov (AP) model [1]:

[TABLE]

where $u$ is the normalized transmembrane potential, and $v$ is the recovery current. Parameter $\mathit{\varepsilon}$ controls the coupling between $u$ and $v$, $\mathbf{D}$ is the diffusion tensor, $\mathit{c}$ controls the repolarization, and $\mathit{\theta}$ controls tissue excitability. As $u$ is most sensitive to parameter $\theta$ in the AP model (1) [5], we focus on its estimation and use standard literature values for the remaining parameters [1]. Solving the AP model on a 3D discrete cardiac anatomy of meshfree nodes [16], we obtain a 3D electrophysiological model of the heart that describes the temporal evoluation of 3D transmembrane potential $\mathbf{u}(t,\boldsymbol{\theta})$.

Measurement Model: $\mathbf{u}(t,\boldsymbol{\theta})$ is measured on the body surface following the quasi-static electromagnetic theory, solving which on a discrete subject-specific heart-thorax mesh gives a linear relationship between $\mathbf{u}(t,\boldsymbol{\theta})$ and its surface potential measurement $\mathbf{y}(t)$ as: $\mathbf{y}(t)=\mathbf{H}\mathbf{u}(t,\boldsymbol{\theta})$ [16].

3 Personalizing HD Parameters on Unstructured Meshes

We seek parameter $\boldsymbol{\theta}$ that minimizes the sum of squared errors between model output $M(\boldsymbol{\theta})=\mathbf{H}\mathbf{u}(t,\boldsymbol{\theta})$ and patient’s measurements $\mathbf{y}_{d}(t)$ as:

[TABLE]

Directly solving (2) for the spatially-distributed $\boldsymbol{\theta}$ is difficult [5]. Below we describe how we learn a LD-to-HD generation of $\boldsymbol{\theta}$ that accounts for the underlying geometry, and embed the HD optimization into the expressive LD manifold.

3.1 Graph Convolutional VAE

We model the generation of spatially-distributed $\boldsymbol{\theta}$ with a VAE [11]. A VAE consists a probabilistic encoder network with parameters $\boldsymbol{\alpha}$ that approximates the intractable true posterior density as $q_{\boldsymbol{\alpha}}(\textbf{z}|\boldsymbol{\theta})$; and a probabilistic decoder network with parameters $\boldsymbol{\beta}$ that represents the likelihood as $p_{\boldsymbol{\beta}}(\boldsymbol{\theta}|\textbf{z})$. For data defined on a Euclidean grid, structural information is incorporated in VAE through CNNs. Here, data resides over the 3D heart geometry, on which standard convolution and pooling operations are not applicable. Below we present a novel VAE architecture that enables convolution, pooling, and unpooling over non-Euclidean geometry of the heart.

Local Connectivity & Graph Convolution: We model the cardiac mesh as a graph: $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{U})$, where vertices $\mathcal{V}$ consist of all $N$ meshfree nodes and edges $\mathcal{E}$ exist between each meshfree node and its $k$ nearest neighbors. $\mathbf{U}\in[0,1]^{N\times N\times 3}$ consists of edge attributes $\boldsymbol{\upsilon}(i,j)$, calculated as the normalized $(x_{1}-x_{2},y_{1}-y_{2},z_{1}-z_{2})$ if an edge $(i,j)\in\mathcal{E}$ exists between vertices at $(x_{1},y_{1},z_{1})$ and $(x_{2},y_{2},z_{2})$ and 0 otherwise. On this graph, we use a convolution operator based on spatial continuous convolution kernels because it was shown to allow better generalization to similar graphs [8]. In specific, given the graph $\mathcal{G}$ and $M$-dimensional input features $\{\mathbf{f}(i)|i\in\mathcal{V}\}$, the $l$-th convolution kernel is:

[TABLE]

where $\small{((N_{1,i}^{m})_{1\leq i\leq k_{1}},\dots,(N_{d,i}^{m})_{1\leq i\leq k_{d}})}$ denotes $d$ open B-spline bases of degree $m$ based on equidistant knot vectors with $d$-dimensional kernel size of $\small{\mathbf{k}=(k_{1},\dots,k_{d})}$, $\small{\mathcal{P}=(N_{1,i}^{m})_{i}\times\dots\times N(_{d,i}^{m})}$ is the Cartesian product of the B-spline bases, and $w_{\mathbf{p},l}$ are the trainable parameters. Given kernel functions $\mathbf{g}=(g_{1},\dots,g_{M})$ and input features $\mathbf{f}\in\mathcal{R}^{M}$, the spatial convolution operator for each vertex $i\in\mathcal{V}$ with a neighborhood $\mathcal{N}(i)$ based on its edge connectivity is then defined as:

[TABLE]

Hierarchical Composition & Pooling: To define pooling and unpooling operations necessary for the encoding-decoding architecture, a hierarchical representation of the graph is needed. We obtain this by an efficient multilevel graph clustering method based on minimizing the normalized cuts (Graclus) [7], which reduces the graph size by half the number of vertices at each coarsening.

We store hierarchical graph representation in matrices which reduces pooling/unpooling operations to efficient matrix multiplications. In specific, if $\mathcal{G}$ is a graph with $N_{1}$ vertices and $\mathcal{G}_{c}$ is its coarsened graph with $N_{2}<N_{1}$ vertices, we populate a binary matrix $\mathbf{P}^{N_{1}\times N_{2}}$, where $\mathbf{P}_{ij}=1$ if the $i^{th}$ vertex in $\mathcal{G}$ was grouped to the $j^{th}$ vertex in $\mathcal{G}_{c}$ and $\mathbf{P}_{ij}=0$ otherwise. Given $M$ feature maps $\mathbf{F}\in\mathcal{R}^{N_{1}\times M}$ over the vertices of graph $\mathcal{G}$ and $\mathbf{F}_{c}\in\mathcal{R}^{N_{2}\times M}$ over graph $\mathcal{G}_{c}$, the average pooling in the encoder can be obtained by $\mathbf{F}_{c}=\mathbf{P}_{n}^{\mathrm{T}}\mathbf{F}$ and unpooling in the decoder by $\mathbf{F}=\mathbf{P}\mathbf{F}_{c}$, where $\mathbf{P}_{n}$ is column normalized from $\mathbf{P}$.

Graph Convolutional VAE: Using these building blocks along with global pooling and standard convolution layers, we construct a VAE architecture as shown in Fig. 1. It is trained by optimizing the variational lower bound on the marginal likelihood of the training data $\boldsymbol{\Theta}=\{\boldsymbol{\theta}^{(i)}\}_{i=1}^{N}$: $\mathcal{L}(\boldsymbol{\alpha};\boldsymbol{\beta};\boldsymbol{\theta}^{(i)})=-D_{\mathrm{KL}}(q_{\boldsymbol{\alpha}}(\textbf{z}|\boldsymbol{\theta}^{(i)})||p(\textbf{z}))+E_{q_{\alpha}(\textbf{z}|\boldsymbol{\theta}^{(i)})}[\mathrm{log}p_{\boldsymbol{\beta}}(\boldsymbol{\theta}^{(i)}|\textbf{z})]$. We set $q_{\boldsymbol{\alpha}}(\textbf{z}|\boldsymbol{\theta})$ and $p_{\boldsymbol{\beta}}(\boldsymbol{\theta}|\textbf{z})$ to be Gaussian parameterized by the graph convolutional networks. The prior $p(\textbf{z})\sim\mathcal{N}(0,1)$ is set to be an isotropic Gaussian, producing an analytical form for the KL divergence. Using the re-parameterization trick [11], to express the random $\mathbf{z}$ as a deterministic variable, standard stochastic gradient methods can be used to optimize $\mathcal{L}(\boldsymbol{\alpha};\boldsymbol{\beta};\boldsymbol{\theta}^{(i)})$.

3.2 Bayesian Optimization on Large Graphs

Bayesian optimization is a popular choice in optimizing complex objective functions such as (2)  [3]. It begins by defining a surrogate over the objective function. The optimization then consists of two iterative steps: 1) actively find a point that optimizes a utility function based on the surrogate, and 2) update the surrogate with the newly-selected point. Direct Bayesian optimization over HD space is difficult and its use over large graphs has not been reported [3, 9]. To enable this, we reformulate the original objective function in (2) by replacing the unknown $\boldsymbol{\theta}$ with the expectation of its non-Euclidean generative model:

[TABLE]

This allows us to embed surrogate construction and active selection of training points in a LD manifold. We initialize the Gaussian process (GP) surrogate of (5) with a zero mean function and an anisotropic Mátern 5/2 kernel.

Active Selection of Training Points: To select a training point, we maximize the expected improvement (EI) utility function that favours a point with the highest expected improvement over the current optimum $f^{+}$ [3, 15]:

[TABLE]

where $\mu(\mathbf{z})$ and $\sigma(\mathbf{z})$ are the predictive mean and standard deviation of the GP, $\Phi$ is the cumulative normal distribution, and $\phi$ is the normal density function. The first term promotes exploitation by favoring regions of high $\mu(\mathbf{z})$, while the second term promotes exploration by favoring the regions with high $\sigma(\mathbf{z})$.

GP Update: After picking a new point $\mathbf{z}^{(i)}$, the value of the optimization objective (5) is evaluated at $\mathbf{z}^{(i)}$ as $\mathcal{J}^{(i)}$. The GP is then updated by including the new input-output pair of $(\mathbf{z}^{(i)},\mathcal{J}^{(i)})$ and maximizing the log marginal likelihood for kernel hyperparameters: length-scales and co-variance amplitude.

4 Synthetic experiments

We evaluate the presented graph convolutional VAE (termed as gVAE) by: 1) its reconstruction accuracy, and 2) optimization accuracy of the gVAE-based Bayesian optimization, both in comparison to existing methods. Accuracy is evaluated by the sum of squared error (SSE) in $\boldsymbol{\theta}$, and the dice coefficient (DC) of the abnormal region obtained by thresholding $\boldsymbol{\theta}$ with Otsu’s method [12].

We synthetically generate data of heterogeneous tissue excitability via random region growing in each cardiac model. Beginning with a single meshfree node as abnormal, we randomly grow the abnormal region by adding one of the nearest neighbors of the abnormal nodes, until the abnormal region reaches a desired size ($2\%$ to $40\%$ of the heart volume). On average, we generated $78,208\pm 12,541$ data for training and $13,545\pm 7654$ data each for validation and testing. The various layers and sizes of feature maps in each layer of the presented gVAE architecture are detailed in Fig. 1. We use B-spline basis degree of $m=1$ with kernel size of $k_{1}=k_{2}=k_{3}=5$ in all graph convolution layers. All models are trained with a learning rate of $0.001$ with Adam optimizer [10].

gVAE as a Generative Model: On five different human heart models constructed from CT images, we first evaluate the ability of the presented gVAE model to reconstruct tissue excitability in comparison to: 1) PCA-based linear reconstruction, and 2) fully-connected VAE (termed as fVAE) [6] with three to five hidden layers (termed as fVAE-3h, -4h, and -5h, respectively). As summarized in Table 1, PCA being a linear model has the lowest accuracy. Compared to fVAE, the reconstruction accuracy of gVAE is consistently higher in both DC and SSE, even when its number of trainable parameters is similar to or lower than fVAE. However, gVAE is more expensive to train: 37.21 hrs vs. 7.51 mins for fVAE-3h in TITAN Xp GPU. Nevertheless, note that the number of trainable parameters for gVAE does not increase for larger meshes.

We further compare gVAE with two-dimensional (2d) latent codes vs. fVAE and PCA with various latent dimensions. As shown in Fig. 2(a), to achieve the reconstruction accuracy of gVAE with 2d latent codes, at least 13 principles components are required with PCA: this increase in dimension will make the subsequent optimization difficult. With fVAE, a similar SSE is attained with four latent dimensions, while a similar DC could not be attained with even 50 latent dimensions. This may be because fVAE does not consider the geometry underlying the spatial distribution of tissue excitability. Note that this increase in the number of latent dimensions will increase both the difficulty and the computational cost of the subsequent optimization. Fig. 2(b) and (c) show that the latent code learned with gVAE are clustered by the location of the abnormal tissue and its radial direction encodes the size of the abnormal tissue.

gVAE-based Parameter Optimization: On 40 synthetic cases on two different heart geometries, we conduct experiments on estimating unknown tissue excitability. In each case, we set an abnormal region by using various combinations of AHA segments which are very different from the training set. Measurement data was simulated, with models described in Section 2, and then sub-sampled and corrupted with 20dB Gaussian noise. We compare the accuracy of gVAE-based Bayesian optimization with three previous approaches: 1) optimization on 17 fixed segments (termed as FS) [17, 14], 2) coarse-to-fine optimization along a fixed multi-scale mesh hierarchy (termed as FH) [5], and 3) optimization on a LD manifold obtained with fVAE-3h [6].

Result in Fig. 3(a) shows that gVAE-based Bayesian optimization is more accurate than all other approaches in both DC and SSE (paired t-test, $p<0.05$). This indicates that gVAE was better at capturing the LD generative factors necessary for accurate optimization. Fig. 3(c)-(f) shows visual comparison on a few examples. The computational cost of gVAE-based optimization is much lower than that of FH ($>22$x) and FS ($>7.5$x), and similar to that of fVAE-based optimization.

Feature Sharing across Geometries: To demonstrate the feasibility of transferring the presented gVAE across different geometries, we take a pre-trained gVAE, fix the learned features in the encoder’s graph convolution layers, and fine-tune the remaining layers for a different anatomy. We compare this training strategy to training a gVAE from scratch. Results in Fig. 4(a) show that a pre-trained model can be fined-tuned with as small as 1088 new examples. In comparison, gVAE could not be trained from scratch with $\leq 7360$ samples, shown both by the low reconstruction accuracy (Fig. 4(a): DC = 0.24; SSE = 17.68) and a flat test loss plot (Fig. 4(b)). Test loss plots in Fig. 4(b) also show that a pre-trained model starts with a lower loss and a larger size of training data leads to faster convergence. Parameter optimization via a gVAE fine-tuned with 7360 data on 20 cases achieved an average DC and SSE of 53.10 and 11.01 respectively. Fig. 4(c) shows some examples of the estimated parameters. More detailed analysis and experimentation is left to future work.

5 Real Data Experiments

We conduct real-data studies on two patients who underwent catheter ablation of ventricular tachycardia due to chronic myocardial infarction. Patient-specific heart-thorax models are obtained from axial CT images. Using 120-lead ECG as measurements, we evaluate the performance of the presented gVAE in estimating tissue excitability in comparison to the fVAE [6], FH [5], and FS [17] methods. Training dataset and network architectures are as described in Section 4. We qualitatively evaluate the results with in-vivo catheter mapping data which, as shown in Fig. 5, provides a reference for the location of the abnormal (red, voltage $\leq 0.5\mathrm{mV}$) and healthy (purple, voltage $>1.5\mathrm{mV}$) regions.

Case 1 has a large heterogeneous abnormal region in the lateral LV region (Fig. 5(a)). All methods are able to localize this region, but with varying degree of heterogeneity. Optimizations based on gVAE and fVAE are much faster requiring only 100 model evaluations, in comparison to FH and FS that required 4056 and 1058 model evaluations, respectively. By contrast, case 2 has a smaller but dense abnormal region in the lateral LV (Fig. 5(b)). While all methods identify the general location of this abnormality, gVAE more accurately differentiates the region of dense core and border. In comparison, fVAE and FH estimate a larger border region; and the abnormal region revealed by FS is less accurate. Again, in this case, gVAE and fVAE required only 100 model evaluations, whereas FH and FS required 5798 and 1501 model evaluations, respectively.

6 Conclusion

We presented a novel graph convolutional VAE model that allows generative modeling of data defined over non-Euclidean structures and integrated it with the Bayesian optimization to enable optimization on large graphs. Experiments on estimating tissue properties distributed on a 3D cardiac mesh shows higher accuracy in terms of both reconstructing the tissue excitability and estimating them from indirect measurements, both in comparison to existing baselines. In future, we will incorporate realistic data from high resolution 3D images and investigate efficient transfer learning methods for models trained on such data.