Deep learning cardiac motion analysis for human survival prediction

In plain words

The heart's walls move through a repeating cycle, and the way the right side of the heart moves may say something about how long a patient is likely to live. This paper films hearts with MRI, follows a set of points on the right ventricle through the cycle, and turns the way each point moves into one long list of numbers per patient.

A neural network squeezes that list into a short summary code. It learns that code with two goals at once: rebuild the original motion from a partly scrambled copy, and predict how soon a patient is likely to die. Scrambling the input pushes the code toward sturdy features rather than noise. In the authors' own internal testing, the summary code ranks patients by risk better than a model built from the standard measurements that doctors take from the same scans.

TL;DR

• Author-reported. The right ventricle is tracked on 202 points over 20 frames, giving 11,514 displacement values per patient. A denoising autoencoder compresses them to a latent code, and its training loss adds a Cox partial-likelihood survival term § Network Design and Training eq. 7.
• In 302 pulmonary hypertension patients from one centre, the optimism-corrected Harrell's C-index is 0.73 (95% CI 0.68–0.78), against 0.59 (0.53–0.65) for a Cox model on manual RV volumetrics § Predictive performance.
• Validation is internal only (bootstrap on the same cohort). The tuned hyperparameter values are not reported in the paper § Hyperparameter Tuning.
• The public repository's default configuration differs from the paper in several places. Its PyTorch backend has no L1 term, computes Cox risk sets inside each mini-batch, and ships only simulated demo data [code: survival4D/nn/torch/init.py#L10-L22].

Problem

Pulmonary hypertension (PH) has high mortality, and the choice of treatment depends on individual risk stratification § Introduction. Global measures of heart contraction from imaging are only moderately reproducible and do not capture the underlying disturbances well § Discussion. The paper asks whether features learned from 3D cardiac motion predict survival better than manual volumetric indices § Introduction.

Survival outcomes are right-censored. A patient still alive at the last follow-up contributes a censoring time rather than a death, and patients who had surgery are censored at the date of surgery § Network Design and Training eq. 4. Standard survival models handle this; the paper's contribution is to learn the input representation jointly with that model.

Method

The full pipeline has three stages. Only the last two are in the public repository.

1. Segmentation and motion tracking (not in the repository). A 2.5D multi-task fully convolutional network predicts ventricle label maps and six landmark locations on each short-axis cine stack. Multiple atlas shapes are then registered to it to enforce a smooth anatomical shape § MR Image Processing § MR Image Acquisition, Processing and Computational Image Analysis. Motion is tracked with a 4D spatio-temporal B-spline registration § MR Image Acquisition, Processing and Computational Image Analysis. The paper cites earlier work for the segmentation network and the atlas § MR Image Acquisition, Processing and Computational Image Analysis.
2. Motion representation. Each subject's meshes are sampled at 202 vertices (about 90-fold fewer than the 18,028 per frame), and the same vertices are kept across frames and subjects § Characterization of right ventricular motion. For each vertex and frames 2 to 20, the coordinate-wise displacement from frame 1 is computed. This gives a vector of length 3 × 19 × 202 = 11,514 eq. 2 § Characterization of right ventricular motion. The repository's setup script computes the same displacements from tracked point files inputdata_setup.py.
3. Encoder. The input passes through a stochastic masking filter that sets a fraction m of the entries to zero, redrawn at every backward pass § Network Design and Training. Then a hidden layer of h1 units with ReLU and a latent code of d_h units with ReLU. The code is undercomplete, with d_h much smaller than the input size. The hidden width and d_h are treated as tuned hyperparameters § Network Design and Training.
4. Decoder. One intermediate hidden layer of h1 units with ReLU, then a linear output of 11,514 units. It mirrors the encoder § Network Design and Training.
5. Risk head. A single linear unit computes W'φ(x), read as the log hazard ratio of the patient eq. 6 § Network Design and Training.
6. Loss. The hybrid loss is L = α·L_r + γ·L_s with γ = 1 − α eq. 7 § Network Design and Training. L_r is the mean over patients of the squared reconstruction error on the uncorrupted input eq. 3. L_s is the negated Cox log partial likelihood summed over patients with an event, where the risk set R(t_i) is the set of patients with a time greater than t_i eq. 5 eq. 6. The text also says L1 regularisation is used for sparsity § Network Design and Training § Discussion, but no equation gives it.
7. Training. Adam, 100 epochs, mini-batches of 16 patients, ReLU in all layers except the prediction output, which is linear § Network Design and Training. The activation of the reconstruction output is not stated; the repository uses a linear one models.py. Learning rate, α, h1, h2, m and the L1 penalty are searched over the ranges in tab. 2.
8. Hyperparameter search. Particle swarm optimisation runs for 50 iterations. Each candidate is scored by 6-fold cross-validated Harrell's C. The paper says the values at the final iteration are chosen § Hyperparameter Tuning. The search uses the Optunity library § Hyperparameter Tuning.
9. Internal validation. The optimism-corrected C is the apparent C of the model trained on all patients, minus the mean optimism over B = 50 bootstrap samples. Each bootstrap repeats the full training and hyperparameter search. The optimism for one sample is its bootstrap performance minus its performance on the full sample § Internal Validation eq. 9. The paper does not state how its confidence intervals were computed § Predictive performance.
10. Benchmark model. A Cox model with L2 penalty λ, chosen by cross-validation. Its inputs are RV end-diastolic volume, RV end-systolic volume, their percentage difference, and RV ejection fraction § Conventional Parameter model eq. 10. It is fitted with lifelines § Conventional Parameter model.
11. Interpretation. The latent code is projected to two dimensions with Laplacian eigenmaps § Model Interpretation. For each of the 202 vertices, a univariate linear regression of predicted risk on the vertex's mean displacement magnitude over 19 frames gives a coefficient. The absolute coefficients are mapped onto a template RV mesh § Model Interpretation fig. 3.

Hybrid loss and corrected C-index, as written in the paper:

L_{hybrid} = \alpha L_r + \gamma L_s, \quad \gamma = 1 - \alpha eq. 7

C_{corrected} = C_{full}^{full} - \frac{1}{B}\sum_{b=1}^{B}\bigl(C_b^b - C_b^{full}\bigr) eq. 9

The training step, in pseudocode. Names follow the paper; the risk-set rule is the paper's, not the repository's (see Paper vs code).

# Forward pass of the 4Dsurvival network  [§ Network Design and Training]
# Input:  motion (B, 11514), mask_rate m
# Output: reconstruction (B, 11514), log hazard ratio (B,)

def forward(motion, mask_rate):
    # Zero a random fraction of the input entries, redrawn each pass
    # [§ Network Design and Training]
    corrupted = motion * random_keep_mask(mask_rate)
    # Encoder: hidden layer of h1 units, then latent code of d_h units,
    # both ReLU  [§ Network Design and Training]
    hidden = relu(linear(corrupted, h1))            # (B, h1)
    code = relu(linear(hidden, d_h))                # (B, d_h)
    # Decoder: one hidden layer, then linear output of size 11514
    # [§ Network Design and Training]; output activation not stated
    decoded_hidden = relu(linear(code, h1))         # (B, h1)
    reconstruction = linear(decoded_hidden, 11514)  # (B, 11514)
    # Linear predictor W' phi(x): log hazard ratio  [eq. 6]
    log_hazard = linear(code, 1)                    # (B,)
    return reconstruction, log_hazard

# Hybrid loss for one mini-batch  [eq. 3] [eq. 5] [eq. 7]
# Risk set of i uses the full training sample in the paper  [eq. 5]

def hybrid_loss(motion, reconstruction, log_hazard, event, time, alpha):
    # Reconstruction term: mean over patients of squared error  [eq. 3]
    reconstruction_loss = mean(sum((motion - reconstruction) ** 2))
    # Negated Cox log partial likelihood over patients with an event
    # [eq. 6]
    cox_loss = 0
    for i in patients:
        if event[i] == 1:
            risk_set = [j for j in patients if time[j] > time[i]]
            risk_sum = sum(exp(log_hazard[j]) for j in risk_set)
            cox_loss -= log_hazard[i] - log(risk_sum)
    # Weighted sum with gamma = 1 - alpha  [eq. 7]
    return alpha * reconstruction_loss + (1 - alpha) * cox_loss

Hyperparameters

Tuned values for the final model are not reported anywhere in the paper.

Results

Author-reported, on one centre's cohort of 302 incident PH patients referred between 2004 and 2017 § Baseline Characteristics. The median follow-up was 371 days, up to 27 November 2017, with 85 deaths (28%) § Baseline Characteristics. No patients were excluded § Baseline Characteristics. Validation is internal and bootstrap-based only § Predictive performance.

Baseline cohort, selected rows from the characteristics table. Values are mean ± SD, or n (%) for sex tab. 1:

Qualitative results: in the two-dimensional projection of the latent code, patients cluster by survival time (survival truncated at 7 years for display). Two patients at opposite risk ends are shown with their RV motion fig. 3A. The saliency map shows regions of the right ventricle that are spatially distant but functionally linked, all contributing to predicted risk fig. 3B § Visualization of Learned Representations.

Training cost: the full process, including the hyperparameter search and the bootstrap validation, took 76 hours on the authors' cluster, which has a TITAN Xp GPU § Network Design and Training.

Paper vs code

Code: the official repository UK-Digital-Heart-Project/4Dsurvival at commit 3329777 (see code_commit), linked from the paper's data availability statement § Data and code availability. The repository is GPL-3.0; this card cites file paths and line numbers, and quotes no code. The repository contains the survival network, the Cox baseline and the bootstrap scripts. It does not contain the segmentation network or the motion-tracking code. The README describes the bundled inputs as simulated and the demo as taking hours to run README.md.

Discrepancies between the paper and the code at this commit:

• Final hyperparameters. The paper gives only search ranges tab. 2. The repository's default config searches only the learning rate and α, and fixes the hidden width at 250, the latent size at 20 and the input dropout at 0 default_nn.conf. The README lists the wider ranges, including dropout and widths README.md.
• Search method and budget. The paper uses particle swarm for 50 iterations § Hyperparameter Tuning. The default config uses grid search with 9 evaluations default_nn.conf default_nn.conf. The demo script for the network uses particle swarm demo_hypersearch_nn.py.
• Epochs, batch and bootstrap count. The paper uses 100 epochs, batch 16 and B = 50 § Network Design and Training § Internal Validation. The README matches 100 epochs and batch 16 README.md. The default config uses 300 epochs, batch 100 and B = 100 default_nn.conf.
• Learning-rate and dropout ranges. The paper's Table 2 gives a learning-rate range up to 10^-4.5 and a dropout range of [0.1, 0.9] tab. 2. The default config searches 10^-6 to 10^-4 default_nn.conf. The README gives [0.1, 0.5] for dropout README.md.
• Framework. The paper says TensorFlow and Keras § Network Design and Training. The default config selects the PyTorch backend default_nn.conf. The TensorFlow path is not the default, and calling its wrapper as written passes one more positional argument than the TensorFlow function accepts. [code: survival4D/nn/init.py#L53] [code: survival4D/nn/tf/init.py#L38]. I did not run either backend.
• Architecture. The paper describes a linear risk head eq. 6 and no batch normalisation. The default model uses batch normalisation in the encoder, decoder and reconstruction paths, and its risk head is a small network with one hidden layer and batch normalisation models.py.
• L1 regularisation. The paper uses L1 regularisation § Discussion. The PyTorch backend applies none; its only penalty is Adam's weight decay [code: survival4D/nn/torch/init.py#L48]. The TensorFlow models apply an L1 activity penalty to the first hidden layer models.py.
• Loss scale. The paper's L_r averages a squared L2 norm over patients, and L_s sums over patients eq. 3 eq. 6. The PyTorch backend uses a mean over every element of the batch for reconstruction, so the reconstruction term is smaller by a factor of 11,514 than the paper's formula, and sums the Cox term over the batch [code: survival4D/nn/torch/init.py#L49] [code: survival4D/nn/torch/init.py#L62] [code: survival4D/nn/torch/init.py#L21]. The TensorFlow backend averages the Cox term over the batch [code: survival4D/nn/tf/init.py#L21]. A given α therefore does not mean the same thing across the paper and the two backends.
• Risk sets. The paper's risk set is over the training sample eq. 5. The PyTorch backend sorts each mini-batch and takes cumulative sums within it, so the risk set is limited to the batch [code: survival4D/nn/torch/init.py#L58-L59] [code: survival4D/nn/torch/init.py#L17-L18]. The TensorFlow backend also sorts within batches [code: survival4D/nn/tf/init.py#L47].
• Confidence intervals. The paper does not say how its confidence intervals were built. The repository computes percentiles of the apparent C minus each bootstrap optimism bootstrap_nn.py.
• Conventional model. The paper describes the penalty as chosen by cross-validation § Conventional Parameter model. The repository uses lifelines' penalised Cox fitter with the penalty set to 10^x cox_reg.py cox_reg.py. The default search range is x in [-2, 1] and only 3 evaluations default_cox.conf default_cox.conf.
• Hardware. The paper's models were trained on GPUs § Network Design and Training. The PyTorch backend moves all tensors to CUDA unconditionally and has no CPU path [code: survival4D/nn/torch/init.py#L43] [code: survival4D/nn/torch/init.py#L46].
• Input data. The README says the bundled demo data are simulated README.md. The paper says personal data are not available § Data and code availability. The reported numbers cannot be reproduced from the public repository.

Pitfalls

• Tuned values are missing. The paper reports ranges but not the chosen learning rate, α, widths, latent size, masking rate or L1 penalty. An implementer has to run the search, and the repository's defaults are not the paper's final values.
• Risk-set definition. The paper's risk set is the set of patients still at risk over the training sample. The repository's risk set is a mini-batch of 16 (or 100 in the default config), which is much smaller than the 302-patient sample. Ties are also not addressed: the paper's definition uses strict inequality, while the code's cumulative sum includes any tied patients that happen to sort earlier.
• Loss scaling. The reconstruction term's scale depends on whether it is averaged over features. With the paper's formula, α trades off a per-patient squared norm against a sum of partial-likelihood terms. Matching the paper requires the scaling in the equations, not the repository's default.
• Overflow in the Cox term. The paper's formula, and the repository's code, compute exp of the predicted log hazard ratio before taking a cumulative sum. Large predictions can overflow. A log-sum-exp form is a safer implementation; this is our suggestion, not something the paper or code states.
• Hyperparameter selection rule. The paper uses the values at the final particle-swarm iteration, not the best value seen during the search § Hyperparameter Tuning. Those can differ.
• Bootstrap cost. Each of the 50 bootstrap samples repeats the full hyperparameter search with 6-fold cross-validation. The full process took 76 hours § Network Design and Training. The optimism estimate averages over those 50 samples § Internal Validation.
• Apparent versus corrected C. The apparent C-index is 0.85 for 4Dsurvival and 0.61 for the benchmark § Predictive performance. The gap between them is inflated by optimism. The corrected values (0.73 and 0.59) are the ones to compare.
• Single centre, no external validation. All results come from one centre's cohort with internal bootstrap validation § Predictive performance. The paper says multi-centre and external validation remain future work § Discussion.
• Survival model choices. Patients who had surgery are censored at the date of surgery § Network Design and Training. This is a modelling choice that changes the risk sets and the event rate.
• Significance test. The paper reports p < .0001 for the C-index difference but does not name the test § Predictive performance.
• Equation numbers. The equation numbers above were counted from the LaTeX source in order of appearance, since the document class does not reset them. Check them against the compiled PDF before citing them in another document.
• Version. This card covers the arXiv v1 text from 2018. The repository's README cites a 2019 journal version README.md. Numbers, ranges or the code may have changed between them.
• Figures. Figures are described in words only, because the site does not yet support images.

Lineage

• Builds on: denoising autoencoders (Vincent et al., 2010) and the Cox proportional hazards model (Cox, 1972), both cited in § Network Design and Training. Neither has a card in this wiki.
• Compared against: a Cox model on manual RV volumetric indices § Conventional Parameter model.
• Superseded by: not known from this wiki's content.