Unscrambling disease progression at scale: fast inference of event permutations with optimal transport

Summary

This paper proposes to use Sinkhorn algorithm to compute optimal transport in order to speed up disease progression that was previously computationally prohibited. This method enables disease progression models with higher dimensionality in features as well as 1000x faster inference speed. Authors provide experiments on Alzheimer’s disease and age-related macular degeneration at pixel-level disease progression events.

Strengths

* Interesting experimental results showing wall-clock improvement on synthetic dataset * Idea is well motivated and easy to understand

Weaknesses

* On the Pixel-level disease experiments, it's hard to judge how realistic the simulated diseases progressions are without quantitative comparison with baseline/groundtruth, such as co-occurances of events. Since it's a scientific study on method compared with baseline, the claims become unfalsifiable if such setup is provided. * Overall, this work gives audience the impression of application of Sinkhorn algorithm. Given the lack of my domain expertise in medical science, it's hard for me to judge the novelty of the problem setting. However, method-wise, novelty is limited. * Practicality of the method in real life problems: Although the claim about "high dimension", which goes up to 200 features, however modern medical MRI machines can easily capture high resolution images with magnitude higher number features/voxels. It's unclear how this method will scale in the more relevant settings.

Questions

In both figure 4, and figure 6, it mentions, " White pixels correspond to events that have occurred; black not yet occurred". However, there doesn't seem to be any black spots?

Rating

5

Confidence

2

Soundness

2

Presentation

2

Contribution

2

Limitations

The limitation has been discussed at end of the paper.

Summary

The authors investigate the task of disease progression modeling, an area of research that learns underlying disease trajectory from temporal snapshots of individual patients. The authors claim that all previous approaches either sacrifice computational tractability for direct interpretability in the feature space or vice versa. The authors introduce the variational event-based model (vEBM) to remedy the former situation, by enabling high dimensional interpretable models through a computationally efficient approach that circumvents dimensionality reduction or manual feature selection. vEBM borrows concepts from optimal transport to directly infer a continuous probability over events. The authors further claimed a 1000x speed-up, better accuracy and improved robustness to noise.

Strengths

1. It is quite innovative to view the disease progression modeling task from the optimal transport perspective. 2. Figure 1 makes the paper slightly easier to read, as it outlines the proposed method with references to subsequent paper sections. 3. The proposed vEBM, compared to the baselines (EBM and ALPACA), show significant advantages in efficiency evaluated by wall-clock time. vEBM also shows better scaling with data dimensionality. 4. The datasets the authors used for empirical results (Alzheimer’s disease and age-related macular degeneration) are of significant clinical importance. I particularly like that the authors compare the disease progression patterns with the known changes from the literature.

Weaknesses

1. Correct me if I am wrong: the proposed method seems to be a method that models the disease progression on a population level. This might limit the method to population level studies for disease research purposes and render it unsuitable for predicting individual-level progressions which could facilitate personalized treatment plans. 2. While producing pixel-level disease progression sequences for certain diseases is fantastic, I would suspect the proposed vEBM method is not ideal for pixel-level predictions, since vEBM presumably treats different pixels as separate features, ignoring the spatial information formed by multiple nearby pixels. 3. For the results in Figure 7, while it is visually informative, it will be great if the authors can incorporate quantitative metrics. 4. Minor issue: For Figure 1, it would be great if the authors can improve the aesthetics. 5. Minor issue: For Figure 5, it would be helpful to provide the color bar.

Questions

1. May I ask the authors to explain in a little more detail what Figure 3 Bottom row is visualizing? I am unfamiliar with the positional variance diagrams and it will be helpful if the authors can explain what the ordering of the feature (vertical axis) indicates, what specific traits on the diagram tells us, etc. 2. Would the authors consider comparing with alternative deep-learning-based models, such as Transformers, neural ODE, latent ODE aka ODE-RNN, neural CDE, etc.? If not, could they provide justifications?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes, the authors adequately addressed the limitations and, if applicable, potential negative societal impact of their work.

Summary

This manuscript derives an Event-Based Model via variational inference and optimal transform. This approach significantly enhances computational efficiency, robustness to noise, and scalability, outperforming current methods by a factor of 1000 in wall-clock.

Strengths

1. The experiments have been performed with multiple datasets: synthetic data, neuroimaging, and optical coherence tomography. 2. The proposed method has been compared to two baselines: the event-based model (EBM) and the Alzheimer’s Disease Probabilistic Cascades (ALPACA) model - on synthetic data. 3. The methodology development was pretty clear and supported by relevant literature. 4. Novel formulation of Event-Based Model via variational inference and optimal transport.

Weaknesses

1. I am not sure about the results in Section 3.3.2. You see that CDR58, MMSE, and RAVLT are pretty far in the Event order. If you check the Temporal Event-Based Model (TEBM) (Wijeratne et al., 2021; Wijeratne et al., 2023), these cognitive test scores can be found earlier in the disease timeline. However, there are differences; they used T1 images, and in your case, you used TBM images. Before making new claims, it will be essential to consider experiments with a previously explored set of features. Importantly, the proposed solution is supposed to solve the scaling question. Hence, it should work with the previous set of features. 2. The authors enable models that express progression at the pixel level compared to the previous region level. However, how the insights will compare to the region-level progression needs to be explored. I suggest exploring it via post hoc analysis by combining pixels back into regions and comparing the ordering. It needs to be clarified whether pixel-level estimation leads us to a new insight into progression. Because it might be great from a computer science perspective but not from a clinical. Furthermore, measurements based on single pixels are highly susceptible to noise. In addition, pixels are not actually pixels but features because the ADNI images were standardized to the template. Hence, it is more like the most fine-grained "atlas" to the template. Wijeratne, Peter A., Daniel C. Alexander, and Alzheimer’s Disease Neuroimaging Initiative. "Learning transition times in event sequences: The temporal event-based model of disease progression." International Conference on Information Processing in Medical Imaging. Cham: Springer International Publishing, 2021. Wijeratne, Peter A., et al. "The temporal event-based model: Learning event timelines in progressive diseases." Imaging Neuroscience 1 (2023): 1-19.

Questions

My questions were combined with weaknesses.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations.

Summary

The authors propose a method to learn a latent event sequence from cross-sectional data of modest size. Each event corresponds to a single observed feature transitioning from an initial parametric distribution (i.e., 'normal') to a second, final parametric distribution (i.e., 'abnormal'). Their inference procedure incorporates elements of variational inference and optimal transport. They demonstrate that it outperforms two available baselines both in speed / scalability and accuracy of the learned event sequence. Finally, they apply the method to learn the order in which (a) individual pixels of (carefully registered) MRI images become abnormal in Alzheimer's disease, and (b) individual OCT pixels become abnormal in macular degeneration.

Strengths

Results clearly demonstrate that their method is scalable, and it's effective in simulation where the true data match their model. The writing style is clear, and the visualizations are excellent, particularly Figures 4 and 6. Experimental results are interesting and clearly presented.

Weaknesses

- The related work section is very brief, which made me question whether it is comprehensive. - I found the second half of section 2.2 very difficult to follow. Specifically, I don't understand the relationship between S and X (see lines 138-143) and think this relationship and the optimal transport details more broadly could be presented much more clearly. - I am not convinced that this is solving a real problem. The approach learns the order in which pixel-sized image regions tend to become abnormal in AD and AMD, but I would think that the more interesting questions have to do with individual variability in the order and timing of this sequence. I may be wrong about this, but I think the authors should do more to explain how the resulting model might be useful in practice.

Questions

- I am confused about how the normal and abnormal distributions are defined in the synthetic data, and whether they are learned versus fixed during inference. Equation (3) implies to me that they are learned, but other descriptions seem to imply that they are initially learned by dividing the population into patients and controls, but then fixed during the broader inference procedure. - What is denoted by the color of the pixels in Figure 5? - The authors mention that the method requires image registration, but very few details are given. It seems to me that it would be challenging to align images from very different stages of progression. Is the method sensitive to this image registration step? - Do the authors envision applications of this method outside of medicine?

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

Somewhat, but a more comprehensive related work section is needed to understand pros and cons of this method (and the approach to disease progression modeling more generally; see Weaknesses) relative to alternatives.

Additional question

Thanks for the additional results. I am a bit confused by the definition of the y-axis, > "The horizontal axis shows the event number (from 0 – 1344), and the vertical axis shows the fraction of pixel-events that have occurred in each regional brain volume at the corresponding event number." and why as event number goes up, the fraction also becomes higher?

Response to rebuttal

Many thanks to the authors for the rebuttal. I like the general clarifications. Answering a few questions from the authors: 1. For quantitative metrics on Figure 7, I do not have specific metrics in mind. However, since you are analyzing distributions, would it be helpful to consider divergence measures (KL or JS) or earth mover distance? 2. Regarding the aesthetics, I have no issue with black and white figures. I was speaking of (1) the method figure is not quite beautiful and engaging, but rather looks a bit dull and almost like a casual sketch, and (2) some line plots are using the default color (blue, orange, green), non-obvious error bars, etc., that are not very optimized.

Response to rebuttal

I want to thank you authors for their rebuttal. I decided to increase the score to Accept.