Improving Generalization of Deep Networks for Inverse Reconstruction of Image Sequences

TL;DR

This paper proposes that incorporating stochastic latent spaces and the information bottleneck principle into deep networks enhances their ability to generalize in inverse image sequence reconstruction tasks, demonstrated on cardiac potential data.

Contribution

The paper introduces a novel variational approach combining stochastic latent spaces and the information bottleneck to improve generalization in inverse image reconstruction networks.

Findings

Stochastic latent spaces improve test data generalization.

Information bottleneck minimizes irrelevant input information.

The proposed method outperforms deterministic models in cardiac potential reconstruction.

Abstract

Deep learning networks have shown state-of-the-art performance in many image reconstruction problems. However, it is not well understood what properties of representation and learning may improve the generalization ability of the network. In this paper, we propose that the generalization ability of an encoder-decoder network for inverse reconstruction can be improved in two means. First, drawing from analytical learning theory, we theoretically show that a stochastic latent space will improve the ability of a network to generalize to test data outside the training distribution. Second, following the information bottleneck principle, we show that a latent representation minimally informative of the input data will help a network generalize to unseen input variations that are irrelevant to the output reconstruction. Therefore, we present a sequence image reconstruction network optimized by a variational approximation of the information bottleneck principle with stochastic latent space. In the application setting of reconstructing the sequence of cardiac transmembrane potential from bodysurface potential, we assess the two types of generalization abilities of the presented network against its deterministic counterpart. The results demonstrate that the generalization ability of an inverse reconstruction network can be improved by stochasticity as well as the information bottleneck.

Figures

Tables

Table 1. Table 1: Accuracy of different architectures at reconstructing unseen pathological conditions

o p2.5cm X[c] X[c] X[c] X[r] Method \Metric	MSE	TMP Corr.	AT Corr.	Dice Coeff.
svs stochastic	$0.037\pm 0.021$	$0.885\pm 0.061$	$0.885\pm 0.072$	$0.645\pm 0.181$
svs deterministic	$0.075\pm 0.013$	$0.77\pm 0.038$	$0.12\pm 0.13$	$0.01\pm 0.006$
svs-L stochastic	$0.068\pm 0.023$	$0.838\pm 0.053$	$0.601\pm 0.074$	$0.28\pm 0.154$
svs-L deterministic	$0.067\pm 0.02$	$0.84\pm 0.053$	$0.57\pm 0.052$	$0.165\pm 0.092$
Greensite	–	–	$0.514\pm 0.006$	$0.138\pm 0.005$

Equations

$$los{s}_{IB}=-I(\mathbf{x};\mathbf{w})+\beta I(\mathbf{w};\mathbf{y})$$

$$\text{vspace}-.1cmp(\mathbf{x},\mathbf{y},\mathbf{w})=p(\mathbf{x})p(\mathbf{y}|\mathbf{x})p(\mathbf{w}|\mathbf{x},\mathbf{y})=p(\mathbf{x},\mathbf{y})p(\mathbf{w}|\mathbf{y})$$

$$I(\mathbf{x};\mathbf{w})=\int p(\mathbf{x},\mathbf{w})\log (\frac{p(\mathbf{x}|\mathbf{w})}{p(\mathbf{x})})d\mathbf{x}d\mathbf{w}=H(\mathbf{x})+\int p(\mathbf{x},\mathbf{w})\log (p(\mathbf{x}|\mathbf{w}))d\mathbf{x}d\mathbf{w}$$

$$\int p(\mathbf{x},$$

$$=\int p(\mathbf{w})D_{KL}(p(\mathbf{x}|\mathbf{w})||q(\mathbf{x}|\mathbf{w}))d\mathbf{w}+\int p(\mathbf{x},\mathbf{w})\log q(\mathbf{x}|\mathbf{w})d\mathbf{x}d\mathbf{w}$$

$$I(\mathbf{x};\mathbf{w})\ge \int p(\mathbf{x},\mathbf{y},\mathbf{w})\log q(\mathbf{x}|\mathbf{w})]d\mathbf{x}d\mathbf{y}d\mathbf{w}=E_{p(\mathbf{x},\mathbf{y})}[E_{p(\mathbf{w}|\mathbf{y})}[\log q(\mathbf{x}|\mathbf{w})]]$$

$$I(\mathbf{y};\mathbf{w})$$

$$=\int p(\mathbf{y})p(\mathbf{w}|\mathbf{y})\log (\frac{p(\mathbf{w}|\mathbf{y})}{r(\mathbf{w})})d\mathbf{y}d\mathbf{w}-D_{KL}(p(\mathbf{w})||r(\mathbf{w}))$$

$$\le \int p(\mathbf{y})p(\mathbf{w}|\mathbf{y})\log (\frac{p(\mathbf{w}|\mathbf{y})}{r(\mathbf{w})})d\mathbf{y}d\mathbf{w}=E_{p(\mathbf{y})}[D_{KL}(p(\mathbf{w}|\mathbf{y})||r(\mathbf{w}))]$$

$$\displaystyle loss_{IB}\leq E_{p(\boldsymbol{x},\boldsymbol{y})}[-E_{p(\boldsymbol{w}|\boldsymbol{y})}\big{[}\log q(\boldsymbol{x}|\boldsymbol{w})]+\beta D_{KL}(p(\boldsymbol{w}|\boldsymbol{y})||r(\boldsymbol{w}))\big{]}=\mathcal{L}_{IB}$$

$$p_{{\mathbf{\theta }}_1}(\mathbf{w}|\mathbf{y})=\mathcal{N}(\mathbf{w}|{\mathbf{t}}_{{\theta }_1}(y),{{\mathbf{\sigma }}_{\mathbf{t}}}^2(\mathbf{y}))q_{{\mathbf{\theta }}_2}(\mathbf{x}|\mathbf{w})=\mathcal{N}(\mathbf{x}|{\mathbf{g}}_{{\theta }_2}(\mathbf{w}),{{\mathbf{\sigma }}_{\mathbf{x}}}^2(\mathbf{w}))$$

$${\mathcal{L}}_{IB}(\mathbf{\theta })=E_{p(\mathbf{x},\mathbf{y})}[-E_{p_{{\mathbf{\theta }}_1}(\mathbf{w}|\mathbf{y})}[\log q_{{\mathbf{\theta }}_2}(\mathbf{x}|\mathbf{w})]+\beta .D_{KL}(p_{{\mathbf{\theta }}_1}(\mathbf{w}|\mathbf{y})||r(\mathbf{w}))]$$

$${\mathcal{L}}_{IB}(\mathbf{\theta })$$

$$\displaystyle+\beta.D_{KL}(p_{\boldsymbol{\theta}_{1}}(\boldsymbol{w}|\boldsymbol{y})||\mathcal{N}(\boldsymbol{w}|\boldsymbol{0},\boldsymbol{I}))\Big{]}$$

$${\Delta }_g=E_{\mu }\ell (\mathbf{x},h_{\mathcal{A}(D_n)}(\mathbf{y}))-E_{Z_m}\ell (\mathbf{x},h_{\mathcal{A}(D_n)}(\mathbf{y}))\le V[f]{\mathcal{D}}^{\ast }[{\mathcal{T}}_{\ast }\mu ,\mathcal{T}(Z_m)]$$

$$\ell (\mathbf{x},h(\mathbf{y}))=||\mathbf{x}-{\mathbf{g}}_{{\mathbf{\theta }}_2}(\mathbf{t})|{|}_F^2=\sum _i({\mathbf{x}}_i-g_i(\mathbf{t}){)}^2=\sum _i{\ell }_i$$

$$T_1$$

$$=\sum _i\frac{1}{{\mathbf{\sigma }}_{xi}^2}({\mathbf{x}}_i-g_i(\mathbf{t}+\mathbf{\eta }){)}^2=\sum _i\frac{1}{{\mathbf{\sigma }}_{xi}^2}{E}_{\mathbf{ϵ}}[{\ell }_i({\mathbf{x}}_i,\mathbf{t}+\mathbf{\eta })]$$

$$T_1$$

$$\displaystyle+\frac{1}{2}\langle[\boldsymbol{\sigma}_{t}\otimes\boldsymbol{\sigma}_{t}]\odot E_{\epsilon}[\boldsymbol{\epsilon}\otimes\boldsymbol{\epsilon}],\Big{[}\frac{\partial^{2}}{\partial\boldsymbol{t}_{j_{1}},\partial\boldsymbol{t}_{j_{2}}}\ell_{i}(\boldsymbol{x}_{i},\boldsymbol{t})\Big{]}\rangle$$

$$\displaystyle+..+\frac{1}{k!}\langle[\boldsymbol{\sigma}_{t}\otimes^{k}\boldsymbol{\sigma}_{t}]\odot E_{\epsilon}[\boldsymbol{\epsilon}\otimes^{k}\boldsymbol{\epsilon}],\Big{[}\frac{\partial^{k}}{\partial\boldsymbol{t}_{j_{1}},..,\partial\boldsymbol{t}_{j_{k}}}\ell_{i}(\boldsymbol{x}_{i},\boldsymbol{t})\Big{]}\rangle+..\Big{]}$$

$$E_{ϵ}[{\ell }_i({\mathbf{x}}_i,\mathbf{t}+\mathbf{\eta })]$$

$$\displaystyle+...+\frac{1}{k!}\langle[\boldsymbol{\eta}\otimes^{k}\boldsymbol{\eta}],\Big{[}\frac{\partial^{k}}{\partial\boldsymbol{t}_{j_{1}},..,\partial\boldsymbol{t}_{j_{k}}}\ell_{i}(\boldsymbol{x}_{i},\boldsymbol{t})\Big{]}\rangle+..\Big{]}$$

Full text

1 Introduction

There has been an upsurge of deep learning approaches for traditional image reconstruction problems in computer vision and medical imaging [11]. Examples include image denoising [13], inpainting [14], and medical image reconstructions across a variety of modalities such as magnetic resonance imaging [19] and computed tomography [6]. Despite state-of-the-art performances brought by these deep neural networks, their ability to reconstruct from data not seen in the training distribution is not well understood. To date, very limited work has investigated the generalization ability of these image reconstruction networks from a theoretical perspective, or provided insights into what aspects of representation and learning may improve the ability of these networks to generalize outside the training data.

In this paper, we take an information theoretic perspective – along with analytical learning theory – to investigate and improve the generalization ability of deep image reconstruction networks. Let $\boldsymbol{x}$ be the original image and $\boldsymbol{y}$ be the measurement obtained from $\boldsymbol{x}$ by some transformation process. To reconstruct $\boldsymbol{x}$ from $\boldsymbol{y}$, we adopt a common deep encoder-decoder architecture [14, 19] where a latent representation $\boldsymbol{w}$ is first inferred from $\boldsymbol{y}$ before being used for the reconstruction of $\boldsymbol{x}$. Our objective is to learn transformations that are general, possibly learning the underlying generative process rather than focusing on every detail in training examples. To this end, we propose that the generalization ability of a deep reconstruction network can be improved from two means:

1. the ability to generalize to data $\boldsymbol{y}$ that are generated from $\boldsymbol{x}$ (and thereby $\boldsymbol{w}$) outside the training distribution; and 2) the ability to generalize to unseen variations in data $\boldsymbol{y}$ that are introduced during the measurement process but irrelevant to $\boldsymbol{x}$.

For the first type of generalization ability, we hypothesize that it can be improved by using stochastic instead of deterministic latent representations. We support this hypothesis by the analytical learning theory [8], showing that stochastic latent space helps to learn a decoder that is less sensitive to perturbations in the latent space and thereby leads to better generalization. For the second type of generalization ability, we hypothesize that it can be improved if the encoder compresses the input measurement into a minimal latent representation (codes in information theory), containing only the necessary information for $\boldsymbol{x}$ to be reconstructed. To obtain a minimal representation from $\boldsymbol{y}$ that is maximally informative of $\boldsymbol{x}$, we adopt the information bottleneck theory formulated in [17] to maximize the mutual information between the latent code $\boldsymbol{w}$ and $\boldsymbol{x}$, $I(\boldsymbol{x},\boldsymbol{w})$, while putting a constraint on the mutual information between $\boldsymbol{y}$ and $\boldsymbol{w}$ such that $I(\boldsymbol{w},\boldsymbol{y})<I_{0}$. This can be achieved by minimizing the following objective:

[TABLE]

where $\beta$ is the Lagrange multiplier. Based on these two primary hypotheses, we present a deep image reconstruction network optimized by a variational approximation of the information bottleneck principle with stochastic latent space.

While the presented network applies for general reconstruction problems, we test it on the sequence reconstruction of cardiac transmembrane potential (TMP) from high-density body-surface electrocardiograms (ECGs) [18]. Given the sequential nature of the problem, we use long short-term memory (LSTM) networks in both the encoder and decoder, with two alternative architectures to compress the temporal information into vector latent space. We tackle two specific challenges regarding the generalization of the reconstruction. First, because the problem is ill-posed, it has been important to constrain the reconstruction with prior physiological knowledge of TMP dynamics [5, 18, 4]. This however made it difficult to generalize to physiological conditions outside those specified by the prior knowledge. By using the stochastic latent space, we demonstrate the ability of the presented method to generalize outside the physiological knowledge provided in the training data. Second, because the generation of ECGs depends on heart-torso geometry, it has been difficult for existing methods to generalize beyond a patient-specific setting. By the use of the information bottleneck principle, we demonstrate the robustness of the presented network to geometrical variations in ECG data and therefore a unique ability to generalize to unseen subjects. These generalization abilities are tested in two controlled synthetic datasets as well as a real-data feasibility study. We hope that these findings may initiate more theoretical and systematic investigations of the generalization ability of deep networks in image reconstruction problems.

2 Related Work

Deep neural networks have become popular in medical image reconstructions across different modalities such as computed tomography [6], magnetic resonance imaging [19], and ultrasound [12]. Some of these inverse reconstruction networks are based on an encoder-decoder structure [6, 19], similar to that investigated in this paper. Among these, the presented work is the closest to Automap [19] in that the output image is reconstructed directly from the input measurements without any intermediate domain-specific transformations. However, these existing works have not investigated either the use of stochastic architectures or the information bottleneck principle to improve the ability of the network to generalize outside the training distributions.

The presented theoretical analysis of stochasticity in generalization utilizes analytical learning theory [8], which is fundamentally different from classic statistical learning theory in that it is strongly instance-dependent. While statistical learning theory deals with data-independent generalization bounds or data-dependent bounds for certain hypothesis space, analytical learning theory provides the bound on how well a model learned from a dataset should perform on true (unknown) measures of variable of interest. This makes it aptly suitable for measuring the generalization ability of stochastic latent space for the given problem and data.

The presented variational formulation of the information bottleneck principle is closely related to that presented in [1]. However, our work differs in three primary aspects. First, we investigate image reconstruction tasks in which the role of information bottleneck has not been clearly understood. Second, we define generalization ability in two different categories, and provide theoretical as well as empirical evidence on how stochastic latent space can improve the network’s generalization ability in a way different from the information bottleneck. Finally, we extend the setting of static image classification to image sequences, in which the latent representation needs to be compressed from temporal information within the whole sequence.

To learn temporal relationship in ECG/TMP sequences, we consider two sequence encoder-decoder architectures. One is commonly used in language translation [16], where the code from the last unit of the last LSTM encoder layer is used as the latent vector representation to reconstruct $\boldsymbol{x}$. We also present a second architecture where fully connected layers are used to compress all the hidden codes of the last LSTM layer into a latent vector representation. This is in concept similar to the attention mechanism [3] to selectively use information from all the hidden LSTM codes for decoding. We experimentally compare the generalization ability of using stochastic versus deterministic latent vectors in both architectures, which has not been studied before.

In the application area of cardiac TMP reconstruction, most related to this paper are works constraining the reconstruction with prior temporal knowledge in the form of physics-based simulation models of TMP [18] and, more recently, generative models learned from physics-based TMP simulation [4]. This however to our knowledge is the first work that investigated the use of deep learning for the direct inference of TMP from ECG. This method will also have the unique potential to generalize outside the patient-specific settings and outside pathological conditions included in the prior knowledge.

3 Methodology

Body-surface electrical potential is produced by TMP in the heart. Their mathematical relation is defined by the quasi-static approximation of electromagnetic theory [15] and, when solved on patient-specific heart-torso geometry, can be derived as: $\boldsymbol{y}(t)=\boldsymbol{Hx}(t)$, where $\boldsymbol{y}(t)$ denotes the time-varying body-surface potential map, $\boldsymbol{x}(t)$ the time-varying TMP map over the 3D heart muscle, and $\boldsymbol{H}$ the measurement matrix specific to the heart-torso geometry of a subject [18]. The inverse reconstruction of $\boldsymbol{x}$ from $\boldsymbol{y}$ at each time instant is ill-posed, and a popular approach is to reconstruct TMP time sequence constrained by prior physiological knowledge of its dynamics [4, 18, 5]. This is the setting considered in this study, in which the deep network learns to reconstruct with prior knowledge from pairs of $\boldsymbol{x}(t)$ and $\boldsymbol{y}(t)$ generated by physics-based simulation. Note that it is not possible to obtain real TMP data for training, which further highlights the importance of the network to generalize. In what follows, we use $\boldsymbol{x}$ and $\boldsymbol{y}$ to represent sequence matrices with each column denoting the potential map at one time instant.

Given the joint distribution of TMP and ECG given by $p(\boldsymbol{x},\boldsymbol{y})$, the encoder gives us a conditional distribution $p(\boldsymbol{w}|\boldsymbol{y})$. These together defines a joint distribution of $(\boldsymbol{x},\boldsymbol{y}$,$\boldsymbol{w}$):

[TABLE]

The first term in $loss_{IB}$ in eq.(1) is given by

[TABLE]

where $p(\boldsymbol{x}|\boldsymbol{w})=\int\frac{p(\boldsymbol{x},\boldsymbol{w},\boldsymbol{y})}{p(\boldsymbol{w})}d\boldsymbol{y}=\int\frac{p(\boldsymbol{x},\boldsymbol{y})p(\boldsymbol{w}|\boldsymbol{y})}{p(\boldsymbol{w})}d\boldsymbol{y}$ is intractable. Letting $q(\boldsymbol{x}|\boldsymbol{w})$ to be the variational approximation of $p(\boldsymbol{x}|\boldsymbol{w})$, we have:

[TABLE]

where the KL divergence in the first term is non-negative. This gives us:

[TABLE]

The second term in $loss_{IB}$ in eq.(1) is given by

[TABLE]

Combining eq.(4) and eq.(6), we have

[TABLE]

which gives us $\mathcal{L}_{IB}$ to be minimized as an upper bound of the information bottleneck objective $loss_{IB}$ formulated in eq.(1).

3.0.1 Parameterization with neural network:

We model both $p(\boldsymbol{w}|\boldsymbol{y})$ and $q(\boldsymbol{x}|\boldsymbol{w})$ as Gaussian distributions, with mean and variance parameterized by neural networks:

[TABLE]

where $\boldsymbol{\sigma_{x}}^{2}$ denotes a matrix that consists of the variance of each corresponding element in matrix $\boldsymbol{x}$. This is based on the implicit assumption that each elements in $\boldsymbol{x}$ is independent and Gaussian, and similarly for $\boldsymbol{w}$. This gives us:

[TABLE]

where $\boldsymbol{\theta}=\{\boldsymbol{\theta}_{1},\boldsymbol{\theta}_{2}\}$. We use reparameterization $\boldsymbol{w}=\boldsymbol{t}+\boldsymbol{\sigma}_{t}\odot\boldsymbol{\epsilon}$ as described in [10] to compute the inner expectation in the first term. The KL divergence in the second term is analytically available for two Gaussian distributions. We obtain:

[TABLE]

where $g_{i}$ is the $i^{th}$ function mapping latent variable to the $i^{th}$ element of mean of $\boldsymbol{x}$, such that $\boldsymbol{g}_{\boldsymbol{\theta}_{2}}=[g_{1},g_{2}...g_{U}]$. The deep network is trained to minimize $\mathcal{L}_{IB}(\boldsymbol{\theta})$ in eq.(9) with respect to network parameters $\boldsymbol{\theta}$.

3.0.2 Network architectures:

The sequence reconstruction network is realized using long short-term memory (LSTM) neural networks in both the encoder and decoder. To compress the time sequence into a latent vector representation, we experiment with two alternative architectures. First, based on the commonly-used sequence-to-sequence language translation model [16],we consider a svs-L architecture that employs the hidden code of the last unit in the last encoding LSTM layer as the latent vector representation for reconstructing TMP sequences. Second, we propose a svs architecture where two fully connected layers are used to compress all the hidden codes of the last LSTM layer into a vector representation. In the decoder, this latent representation is expanded by two fully-connected layers before being fed into LSTM layers as shown in Fig. 1.

4 Encoder-Decoder Learning from the Perspective of Analytical Learning Theory

In this section we look at the encoder-decoder inverse reconstructions using analytical learning theory [8]. We start with a general framework and then show that having a stochastic latent space with regularization helps in generalization.

Let $\boldsymbol{z}=(\boldsymbol{y},\boldsymbol{x})$ be an input-output pair, and let $D_{n}=\{\boldsymbol{z}^{(1)},\boldsymbol{z}^{(2)},...,\boldsymbol{z}^{(n)}\}$ denote the total set of training and validation data where $Z_{m}\subset D_{n}$ be the validation set. During training, a neural network learns the parameter $\boldsymbol{\theta}$ by using an algorithm $\mathcal{A}$ and dataset $D_{n}$, at the end of which we have a mapping $h_{\mathcal{A}(D_{n})}(.)$ from $\boldsymbol{y}$ to $\boldsymbol{x}$. Typically, we stop training when the model performs well in the validation set. To evaluate this performance, we define a prediction error function, $\ell(\boldsymbol{x},h_{\mathcal{A}(D_{n})}(\boldsymbol{y}))$ based on our notion of the goodness of prediction. The average validation error is given by ${E_{Z_{m}}\ell(\boldsymbol{x},h_{\mathcal{A}(D_{n})}(\boldsymbol{y}))}$. However, there exists a so-called generalization gap between how well the model performs in the validation set versus in the true distribution of the input-output pair. To be precise, let $(\mathcal{Z},\mathcal{S},\mu)$ be a measure space with $\mu$ being a measure on $(\mathcal{Z},\mathcal{S})$. Here, $\mathcal{Z}=\mathcal{Y}\times\mathcal{X}$ denotes the input-output space of all the observations and inverse solutions. The generalization gap is given by $\Delta_{g}={E_{\mu}\ell(\boldsymbol{x},h_{\mathcal{A}(D_{n})}(\boldsymbol{y}))}-{E_{Z_{m}}\ell(\boldsymbol{x},h_{\mathcal{A}(D_{n})}(\boldsymbol{y}))}$. Theorem 1 in [8] provides an upper bound on the generalization gap $\Delta_{g}$ in terms of data distribution in the latent space and properties of the decoder.

Theorem 4.1 ([8])

For any $\ell$, let $(\mathcal{T},f)$be a pair such that $\mathcal{T}:(\mathcal{Z},\mathcal{S})\rightarrow([0,1]^{d},$ $\mathcal{B}([0,1]^{d}))$ is a measurable function, $f:([0,1]^{d},\mathcal{B}([0,1]^{d}))\rightarrow(\mathbb{R},\mathcal{B}(\mathbb{R}))$ is of bounded variation as $V[f]<\infty$, and $\ell(\boldsymbol{x},h(\boldsymbol{y}))=(f\circ\mathcal{T})(\boldsymbol{z})\forall\boldsymbol{z}\in\mathcal{Z}$, where $\mathcal{B}(A)$ indicates the Borel $\sigma$- algebra on $A$. Then for any dataset pair $(D_{n},Z_{m})$ and any $\ell(\boldsymbol{x},h_{\mathcal{A}(D_{n})}(\boldsymbol{y}))$,

[TABLE]

where $\mathcal{T}_{*}\mu$ is pushforward measure of $\mu$ under the map $\mathcal{T}$.

For an encoder-decoder setup, $\mathcal{T}$ is the encoder that maps the observation to the latent space and $f$ becomes the composition of loss function and decoder that maps the latent representation to the reconstruction loss. Theorem 1 provides two ways to decrease the generalization gap in our problem: by decreasing the variation $V[f]$ or the discrepancy $\mathcal{D}^{*}[\mathcal{T}_{*}\mu,\mathcal{T}(Z_{m})]$. Here, we show that stochasticity of the latent space helps decrease the variation $V[f]$. The variation of $f$ on $[0,1]^{d}$ in the sense of Hardy and Krause [7] is defined as: $V[f]=\sum_{k=1}^{d}\sum_{1\leq j_{1}<...<j_{k}\leq d}V^{k}[f_{j_{1}...j_{k}}]$ where $V^{k}[f_{j_{1}...j_{k}}]$ is defined with following proposition.

Proposition 1** ([8])**

Suppose that $f_{j_{1},..j_{k}}$ is a function for which $\partial_{1,...k}f_{j_{1},..j_{k}}$ exists on $[0,1]^{k}$. Then, $V^{k}[f_{j_{1}...j_{k}}]\leq\underset{\boldsymbol{t}_{j_{1}},..,\boldsymbol{t}_{j_{k}}\in[0,1]^{k}}{sup}|\partial_{1,...k}f_{j_{1},..j_{k}}(\boldsymbol{t}_{j_{1}},..,\boldsymbol{t}_{j_{k}})|$. If $\partial_{1,...k}f_{j_{1},..j_{k}}$ is also continuous on $[0,1]^{k}$, $V^{k}[f_{j_{1}...j_{k}}]=\int_{[0,1]^{k}}|\partial_{1,...k}f_{j_{1},..j_{k}}(\boldsymbol{t}_{j_{1}},..,\boldsymbol{t}_{j_{k}})|dt_{j_{1}}..dt_{j_{k}}$.

In our case, $f$ is the prediction error $\ell$ as a function of latent representations $\boldsymbol{t}$:

[TABLE]

where $||\boldsymbol{a}||_{F}$ denotes the Frobenius norm of matrix $\boldsymbol{a}$, and $\boldsymbol{g}_{\boldsymbol{\theta}_{2}}$ maps the latent space to the estimated $\bar{\boldsymbol{x}}$. Theorem 1 and Proposition 1 implies that if the cross partial derivative of the loss with respect to the latent vector at all order is low in all directions throughout the latent space, then the approximated validation loss would be closer to the actual loss over the true unknown distribution of the dataset. Intuitively, we want the loss curve as a function of latent representation to be flat if we want a good generalization.

4.0.1 Using stochastic latent space:

In our formulation, the latent vector is stochastic with the cost function given by eq.(9). Using reparameterization $\boldsymbol{\eta}=\boldsymbol{\sigma}_{t}\odot\boldsymbol{\epsilon}$, the inner expectation of the first term in the loss function $\mathcal{L}_{IB}$ is given by

[TABLE]

Result 1

[TABLE]

where $[\boldsymbol{\sigma}_{t}\otimes^{k}\boldsymbol{\sigma}_{t}]$ denotes k order tensor product of a vector $\boldsymbol{\sigma}_{t}$ by itself.

Proof

Using Taylor series expansion for $\ell_{i}(\boldsymbol{x}_{i},\boldsymbol{t}+\boldsymbol{\eta})$,

[TABLE]

We move expectation operator inside both brackets and take expectation of only the first term in the inner product. Using $\boldsymbol{\eta}=\boldsymbol{\sigma}_{t}\odot\boldsymbol{\epsilon}$, we get $E_{\epsilon}[\boldsymbol{\eta}\otimes^{k}\boldsymbol{\eta}]=[\boldsymbol{\sigma}_{t}\otimes^{k}\boldsymbol{\sigma}_{t}]\odot E_{\epsilon}[\boldsymbol{\epsilon}\otimes^{k}\boldsymbol{\epsilon}]$. Using these in eq.(11) yields the required result.

The first term of Result 1, $\ell_{i}(\boldsymbol{x}_{i},\boldsymbol{t})$ (after ignoring $\frac{1}{\boldsymbol{\sigma}_{xi}^{2}}$), would be the only term in the cost function if the latent space were deterministic. The rest of the terms are additional in stochastic training. Each of these terms is an inner product of two tensor, the first being $[\boldsymbol{\sigma}_{t}\otimes^{k}\boldsymbol{\sigma}_{t}]\odot E_{\epsilon}[\boldsymbol{\epsilon}\otimes^{k}\boldsymbol{\epsilon}]$, and the second being the $k^{th}$ order partial derivative tensor \Big{[}\frac{\partial^{k}}{\partial\boldsymbol{t}_{j_{1}},..,\partial\boldsymbol{t}_{j_{k}}}\ell_{i}(\boldsymbol{x}_{i},\boldsymbol{t})\Big{]}. We can thus consider the first tensor as providing penalizing weights to different partial derivatives in the second tensor. Since each inner product is added to the cost, we are minimizing them during optimization. This gives two important implications:

For sufficiently large samples, $E_{\epsilon}[\boldsymbol{\epsilon}\otimes^{k}\boldsymbol{\epsilon}]$ must be close to central moments of isotropic Gaussian. However, in practice, the number of samples of $\epsilon$ remains constant. As we move to the higher order moment tensors, we can expect that they do not converge to that of the standard Gaussian. This, luckily, works in our favor. Since we are minimizing \frac{1}{k!}\langle[\boldsymbol{\sigma}_{t}\otimes^{k}\boldsymbol{\sigma}_{t}]\odot E_{\epsilon}[\boldsymbol{\epsilon}\otimes^{k}\boldsymbol{\epsilon}],\Big{[}\frac{\partial^{k}}{\partial\boldsymbol{t}_{j_{1}},..,\partial\boldsymbol{t}_{j_{k}}}\ell_{i}(\boldsymbol{x}_{i},\boldsymbol{t})\Big{]}\rangle for each order, the inner product can be vanished for arbitrary $\epsilon$ only by driving partial derivative tensors towards zero. Therefore, minimizing the sum of all the inner product for arbitrary $\epsilon$ would minimize most of the terms in the partial derivative tensor. From Proposition 1, this corresponds to minimizing the variation of function $\ell_{i}$, and consequently variation of the total error function $\ell$ according to eq.(10). Hence, additional terms in the stochastic latent space formulation contributes to decreasing the variation $V[f]$ and consequently the generalization gap. 2. 2.

Not all the partial derivatives are equally weighted in the cost function. Due to the presence of weighting tensor $[\boldsymbol{\sigma}_{t}\otimes^{k}\boldsymbol{\sigma}_{t}]$ in the first tensor of inner product, different partial derivative terms are penalized differently according to the value of $\boldsymbol{\sigma}_{t}$. Combination of the KL divergence term in eq.(9) with $T_{1}$ tries to increase standard deviation, $\boldsymbol{\sigma}_{t}$ towards 1 whenever it does not significantly increase the cost $T_{1}$: higher value of $\boldsymbol{\sigma}_{t}$ penalizes the partial derivatives of a certain direction more heavily, making the cost flatter in some directions than other.

Strictly speaking, Proposition 1 requires cross partial derivatives to be small throughout the domain of latent variable, which is not included in the above analysis. It however should not significantly affect the observation that, compared to deterministic formulation, the stochastic formulation decreases the variation $V[f]$.

5 Experiments & Results

Since it is not possible to obtain real TMP data, the reconstruction network is trained on simulated data pairs of $\boldsymbol{y}$ and $\boldsymbol{x}$. We focus on evaluating three generalization tasks of the network: to learn how to reconstruct under the prior physiological knowledge given in simulation data while generalizing to 1) unseen pathological conditions in $\boldsymbol{x}$, 2) unseen geometrical variations in $\boldsymbol{y}$ that are irrelevant to $\boldsymbol{x}$, and 3) real clinical data.

5.1 Generalizing outside the training distribution of TMP

5.1.1 Dataset and implementation details:

We simulated training and test sets using three human-torso geometry models. Spatiotemporal TMP sequences were generated using the Aliev-Panfilov (AP) model [2], and projected to the body-surface potential data with 40dB SNR noises. Two parameters were varied when simulating the TMP data: the origin of excitation and abnormal tissue properties representing myocardial scar. Training data were randomly selected with regard to these two parameters. Test data were selected such that values in these two parameters differed from those used in training in four levels:

1. Scar: Low, Exc: Low, 2) Scar: Low, Exc: High, 3) Scar:High, Exc: Low, and 4) Scar: High, Exc: High, where Scar/Exc indicates the parameter being varied and High/Low denotes the level of difference (therefore difficulty) from the training data. For example, Scar: Low, Exc: High test ECG data was simulated with region of scar similar to training but origin of excitation very different from that used in training.

For all four models being compared (svs stochastic/deterministic and svs-L stochastic/deterministic), we used ReLU activation functions in both the encoder and decoder, ADAM optimizer [9], and a learning rate of $10^{-3}$. Each neural network was trained on approximately 2500 TMP simulations on each geometry. In addition to the four neural networks, we included a classic TMP inverse reconstruction method (Greensite) designed to incorporate temporal information [5]. On each geometry, approximately 300 cases were tested for each of the four difficulty levels. We report the average and standard deviation of the results across all three geometry models.

5.1.2 Results: