EchoGNN: Explainable Ejection Fraction Estimation with Graph Neural Networks

The original echo videos are high-dimensional and must be mapped into lower-dimensional embeddings to reduce memory footprint and remove redundant information.

The Video Encoder is used to learn a mapping $f_{\text{ve}}:R^{T\times H\times W}\rightarrow R^{T\times d}$ from input echo videos $x^{i}\in\mathbb{R}^{T\times W\times H}$ to $d$-dimensional embeddings $h^{i}_{j}\in\mathbb{R}^{d}$, where $j\in[T]$ is the frame number. The temporal dimension is preserved because the Attention Encoder requires embeddings for all frames to produce interpretable weights over them. We use a custom network consisting of 3D convolutions and residual connections to use both the spatial and temporal information in the video in generating the embeddings. This network’s architecture is provided in the supp. material. Lastly, following [25], periodic positional encodings are added to the generated frame embeddings to encode the sequential nature of video data.

For each patient, we construct an echo-graph, which is a complete graph where each node corresponds to a frame in the echo video, and the edges show the non-Euclidean relationship between these frames. Formally, we denote the echo-graph with $G_{\text{echo}}(V,E)$ where V is the set of nodes corresponding to echo frames such that $|V|=T$, and E is the set of edges between the nodes to show the relationship between video frames such that if $v_{1},v_{2}\in V$ are connected, then $e_{v_{1},v_{2}}\in E$. We use the frame embeddings from our Video Encoder as node features of $G_{\text{echo}}$. That is, $\{h^{i}_{1},h^{i}_{2},...,h^{i}_{T}\}$ are the set of features for $\{v_{1},v_{2},...,v_{T}\}$. These embeddings can be represented as a matrix $H^{i}\in\mathbb{R}^{T\times d}$ such that each row is the embedding for a frame in the echo video for patient $i$.

Inspired by [14], we propose using GNNs to learn and assign weights to both edges and nodes of the echo-graph. The edge and node weights are learned to encode the importance of each frame (node weights) and the relationships among frames (edge weights) for the final EF estimation.

The Attention Encoder infers weights over edges and nodes of the echo-graph using message passing based GNNs [7]. A single message passing step is enough for each node to capture information from all other nodes due to echo-graph being a complete graph. More specifically, the following operations are used to obtain weights over each edge $e_{v_{k},v_{s}}$:

where $\sigma$ is the Sigmoid function, $[.\|.]$ is the concatenation operator, and $a_{k,s}\in[0,1]$ is the inferred weight for the directed edge from $v_{k}$ to $v_{s}$. Similarly, weights for each node $w_{s}\in[0,1]$ are generated by inserting another $\text{edge}\rightarrow\text{node}$ operation after Eq. 3. All MLPs use two fully connected linear layers with ELU [4] activation and batch normalization.

Our Regressor network uses GNN layers with the learned weighted echo-graph to perform EF estimation. Specifically, for each patient, the output of the Attention Encoder can be represented as a weighted adjacency matrix $A\in[0,1]^{T\times T}$ and a node weight vector $w\in[0,1]^{T}$. The Regressor uses $A$ to generate embeddings over frames of the echo video:

where $H^{l}\in\mathbb{R}^{T*d_{g}}$ is the matrix of learned node embeddings at layer l, $H^{0}$ is the matrix of frame embeddings from the Video Encoder, and $g^{l}$ is composed of a Graph Convolutional Network (GCN) layer followed by batch normalization and ELU activation [15]. To represent the whole graph with a single vector embedding, the node embeddings are averaged using the frame weights $w$ generated by the Attention Encoder:

where $H^{l}_{j}\in\mathbb{R}^{d}$ is the $j$th row of $H^{l}$, and $w_{j}$ is the $j$th scalar weight in the frame weight vector. $h^{i}_{\text{graph}}$ is mapped into an EF estimate using an MLP with two fully connected linear layers, ELU activation and batch normalization.

The model is differentiable in an end-to-end manner. Therefore, we use gradient descent with Mean-Absolute-Error (MAE) between predicted EF estimates $\tilde{y}^{i}$ and ground truth EF values $y^{i}\in Y$ as the optimization objective, which is computed as $L=\frac{1}{N}\sum_{i=1}^{N}|\tilde{y}^{i}-y^{i}|$.

The key advantage of EchoGNN over prior work is the explainability it provides through the learned weights on the echo-graph. As shown in Fig. 2, the learned weights can indicate when human intervention is required. We observe two different scenarios: (1) the model learns the periodic nature of echo videos and assigns larger weights to frames and edges that are in between ES and ED phases before performing EF estimation. This means that the location of ES and ED can be approximated using these weights as illustrated in Fig. 2. (2) The model cannot detect the location of ES and ED frames and distributes weights more evenly. We see that in these cases, we have either an atypical zoomed-in AP4 echo or an echo where the LV is not entirely visible and is cropped. In such cases, an expert can evaluate the video and determine if new videos must be obtained. More explainability examples are provided in the supp. material.

To quantitatively measure the explainability of EchoGNN, for the cases where the model learns the periodic nature of the data (1173 samples out of 1277), we use the average frame distance (aFD) as in [21], which is computed as $\text{aFD}=\frac{1}{N}\sum_{i=1}^{N}|j_{i}-\tilde{j}_{i}|$ with $j_{i}$ and $\tilde{j}_{i}$ being the true and approximated indices, respectively, for sample $i$. As shown in Table 4.3.2, our model achieves better ED aFD and comparable ES aFD without using ground-truth ES/ED locations for training, whereas Reynaud et al. [21] uses such supervision. This shows the explainability power of EchoGNN. aFD computation details are provided in the supp. material.

To evaluate the error in predicted EF values, we use Mean-Absolute-Error (MAE). Additionally, as a measure of the amount of explained variance in the data, we report the model’s $\text{R}^{2}$ score. Moreover, we report the $\text{F}_{1}$ score for the task of indicating whether EF values are lower than 40%, which is a strong indicator of heart failure [11].

As shown in Table 4.3.2, our model significantly outperforms [21] without direct supervision of ES and ED frame locations during training. Our model has similar predictive performances as [12] with a much lower number of parameters and the added benefit of explainability through the learned latent graph structures. EchoNet (AF) [18] requires large amounts of RAM due to sampling all 32-frame clips in a video, making us unable to train and evaluate the model. Because of this we only report results from the paper and cannot produce additional metrics such as $\text{F}_{1}$ score which is not originally reported, and hence we show this with N/A in Table 4.3.2. This model’s weak performance compared to our model shows the sensitivity of EchoNet (AF) to frame locations in a clip. Lastly, our model has a significantly lower number of parameters, making it desirable for deployment on mobile clinical devices. Our model’s EF scatter plot and confusion matrix are provided in the supp. material.