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

Authors: peter a wijeratne
daniel c alexander
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 18:31:08

Abstract: Disease progression models infer group-level temporal trajectories of change in patients' features as a chronic degenerative condition plays out. They provide unique insight into disease biology and staging systems with individual-level clinical utility. Discrete models consider disease progression as a latent permutation of events, where each event corresponds to a feature becoming measurably abnormal. However, permutation inference using traditional maximum likelihood approaches becomes prohibitive due to combinatoric explosion, severely limiting model dimensionality and utility. Here we leverage ideas from optimal transport to model disease progression as a latent permutation matrix of events belonging to the Birkhoff polytope, facilitating fast inference via optimisation of the variational lower bound. This enables a factor of 1000 times faster inference than the current state of the art and, correspondingly, supports models with several orders of magnitude more features than the current state of the art can consider. Experiments demonstrate the increase in speed, accuracy and robustness to noise in simulation. Further experiments with real-world imaging data from two separate datasets, one from Alzheimer's disease patients, the other age-related macular degeneration, showcase, for the first time, pixel-level disease progression events in the brain and eye, respectively. Our method is low compute, interpretable and applicable to any progressive condition and data modality, giving it broad potential clinical utility.

Main paper: 
 introduction: 
 The main aim of disease progression modelling is to learn a hidden underlying disease trajectory from 'snapshots' (sets of observations at a single time) of individuals at hidden points along the trajectory. The classical approach is to treat the problem dynamically, using either discrete [1] [2] [3] [4] [5] [6] [7] [8] [9] or continuous [10] [11] [12] [13] [14] [15] [16] [17] [18] models with latent variables to describe the hidden disease stage or time. An abundance of such models have been proposed (see [19] for a comprehensive review) and have found extensive success in providing unique interpretability and utility across a wide range of progressive diseases, including Alzheimer's disease (AD) [1, 20, 5, 12, 21, 22 ], Huntington's disease [23] [24] [25] [26] , multiple sclerosis [27, 28] , Parkinson's disease [29] , prion disease [30] , amyotrophic lateral sclerosis [31] , and chronic obstructive pulmonary disorder [32] . However all previous approaches make a compromise: they are either i) interpretable in feature space but sacrifice computational tractability [33, 1, 34, 2, 5, 13, 35, 6, [36] [37] [38] [39] ; or ii) are made computationally tractable by encoding to a latent space but sacrifice direct interpretability [40, 41] . Models of type (i) often require preprocessing or dimensionality reduction to extract a modest number of interpretable features from high dimensional data, e.g., deriving features of anatomical regions from medical images, because computation time scales super-linearly with the number of features. The preprocessing introduces uncertainty and is often computationally burdensome in itself.Here we introduce the variational event-based model (vEBM), which enables high dimensional interpretable models through a new computationally efficient approach that avoids the need for dimensionality reduction or manual feature extraction. For example, with image-based models, it enables models that express progression at the pixel-level rather than regional level. To achieve this we reformulate disease progression modelling as the 'transport' of latent disease events to their 'optimal' location in a continuous latent permutation, unlocking benefits from recent advancements in the field of computational optimal transport [42] . Our approach generalises discrete generative models of disease progression, e.g., [1, 20, 5, [36] [37] [38] [39] 43] , which it obtains as a limit; and it directly infers a continuous probability over events, while the others require costly sampling methods. Crucially, it also facilitates variational inference of the posterior, allowing for substantial gains in computational tractability and hence larger models.Related work. The closest direct comparisons to the model we propose here are the sequence-based models proposed by [33, 1] . These models underpin both a range of direct applications [20, 5, 23, 24, 12, 21, 25, 27, 29] , as well as providing components in higher level models [5, 36, 38, 39] . Like ours, these models require data with only a single time-point per individual. They assume monotonic progression in order to learn a latent sequence of events from such cross-sectional data. However, the models in [33, 1] are severely limited by their computational tractability, and can typically only use a few 100 features at most. In contrast, our new formulation enables this type of model to include several orders of magnitude more features enabling, for example, pixel-level temporal models as we demonstrate here.Deep-learning based sequence models, using e.g. transformer architectures have recently become popular for models of high dimensional temporal sequences e.g., [41] . However, as with other deep state-space models, e.g., [44, 6] these approaches require vast amounts of data with multiple time points to train, unlike our approach which can be trained on modest datasets (order 100 subjects with observations from a single time-point). Furthermore, the computational power required to train the upstream foundation model, plus the downstream model itself, is order of magnitudes higher than our model, which can run on a single CPU in a matter of minutes.Here we address the problem of how to learn interpretable high dimensional disease progression models efficiently, which is longstanding in the machine learning community.• We leverage ideas from optimal transport to derive a new generative latent variable model of disease progression, the variational event-based model (vEBM). The vEBM characterises the disease process by a continuous latent permutation of event probabilities, permitting direct inference of event distributions and model uncertainty from mixed feature datasets.• We define a differentiable variational evidence lower bound (ELBO) and devise a suitable inference scheme to learn the vEBM efficiently from high dimensional data.• We use synthetic data to demonstrate that the vEBM achieves a factor of 1000× faster inference than baselines, provides better inference accuracy, and is robust to noise.• We use the vEBM with data from Alzheimer's disease (AD) and age-related macular degeneration (AMD) to obtain, for the first time, pixel-level disease progression events in the brain and eye, and mixed-feature models combining imaging and clinical test score data. 
variational event-based model: 
 To derive the variational event-based model (vEBM), we first derive a generative latent variable model of disease progression in terms of a latent permutation matrix of events (Section 2.1). Our key methodological contribution is reformulating the generative model in the context of optimal transport; we introduce the relevant mathematical tools to do this in Section 2.2, which we use to derive the limit relationship between the classical EBM and the vEBM in Appendix Section A.3. We then define 2), the ELBO (Section 2.3) is optimised and S is updated each iteration using the Sinkhorn-Knopp algorithm (Section 2.4). C. The resulting hard permutation, s, i.e., the disease event sequence, is obtained from S, using the Hungarian algorithm. We note that S represents the full distribution of event probabilities, which can be sampled to obtain uncertainty.our model in a variational inference setting (Section 2.3), devise a suitable inference scheme (Section 2.4), and provide a method for probabilistic individual-level staging using the trained model (Section 2.5). Figure 1 1 provides schematic overview of the vEBM.Consider a generative latent variable model with observed data Y and latent variables Z = {S, k, θ}, where S ∈ R N ×N + is a latent permutation matrix of N events; k ∈ Z N + is the latent state of an individual; and θ = {θ(1), θ(2), θ(3), ..., θ(J)} are additional model parameters. We write the joint probability as a hierarchical Bayesian model using the chain rule (see Appendix Figure 9 for the graphical model),EQUATIONEach element in the permutation matrix, S, defines an 'event', which corresponds to a feature becoming measurably abnormal with respect to a reference distribution. Following [1, 20] , we parameterise the data likelihood by probability distributions of 'abnormality' (typically from patients, p), and 'normality' (typically from controls, c) for each feature, and choose to model the distributions for feature j using univariate Gaussian mixture models with mean, µ j , standard deviation, σ j , and mixture weights, w j , so thatEQUATIONHowever any probabilistic characterisation that defines a reference group to anchor the progression is permissible. To enable S to be inferred using data from different individuals at a single time-point ('snapshots'), we make two key assumptions: i) monotonic progression of events at the group level; and ii) a consistent event permutation across the whole population. We note that assumption ii) could be relaxed to permit multiple event permutations (i.e., clusters) within the same population. Then for individuals i = {1, 2, 3, ..., I} with observed features j = {1, 2, 3, .., J}, the model likelihood can be written as (see Appendix A.2 for a full derivation),EQUATIONHere s ∈ Perm(N ) is a discrete permutation of N events, corresponding to the hard permutation obtained from S (see the next section); and θ s(j) = θ p s(j) ∪ θ c s(j) are the patient, p, and control, c, distribution parameters generating the data for feature j at position s(j) in the permutation. Note that if data are missing, the two likelihoods on the RHS of Equation 3 can be set equal and factorised, i.e., the data can be treated as missing at random. In order to impose no prior information on the permutation ordering, we chose the prior over latent stages to be uniform, P (k i |S) ∼ Unif(0, k).Under the definition of disease progression given by Equation 3 and using Bayes rule, our posterior is a probability distribution over the sequence of events. Given that permutations are factorial in N , the challenge is to make inference of S computationally tractable when the number of features -and hence the number of events -is large (of order N > 100). To address this, we propose to frame the learning problem in terms of the 'transport' of disease events to their 'optimal' location in the disease event sequence, s, defined by the permutation matrix S. Our key methodological contribution is the translation of the model likelihood (Equation 3) to the context of optimal transport; here we provide the necessary background theory to enable us to derive the relationship between s and S in our model.Optimal transport aims to identify the mass-conserving coupling between two distributions ('transport plan') that minimises the cost required to move (or transform) one into the other [45] . The minimum cost defines a distance between distributions (the Wasserstein distance) and induces a rich underlying geometry on the space of distributions, providing benefits over classical learning techniques such as maximum likelihood. While optimal transport requires solving a computationally expensive linear problem, recent advances have resolved this by substituting the original problem with an entropy regularised version [46] , paving the way for its use in learning generative models.When the couplings are restricted to be permutation matrices, we can leverage the machinery of entropy regularised optimal transport to provide computationally tractable solutions to inferring latent permutations [47] . Here we are interested in learning a latent permutation matrix, S, with a corresponding discrete permutation, s, such that,EQUATIONIn the context of permutation matrices as couplings, S belongs to the Birkhoff polytope,EQUATIONThe Birkhoff-von Neumann theorem states that B N is the convex hull of the set of doubly stochastic (soft) permutation matrices, and that its vertices are the (hard) permutation matrices [48] . The row-column normalisation equality constraints in Equation 5 demand efficient algorithms to solve for S. Following [46] , we use the Sinkhorn-Knopp algorithm with an entropy regularisation term, H(S) =i,j S i,j log(S i,j ), as an approximation to solving the optimal transport problem, K(X/τ ) = argmaxEQUATIONHere K(•) is the Sinkhorn-Knopp operator, which maps the positive orthant on to B N by iteratively normalising rows and columns [49, 50] ; X is the unnormalised assignment probability (transportation cost) matrix; and τ is a temperature parameter, analogous to the temperature-dependent softmax function for discrete categories [51] . In our context, X corresponds to the event likelihood distributions given by Equation 2; as such, we are looking to find the permutation matrix, S, that transports event probabilities to their optimal location in the event sequence, s. Alternatively, we can think of relationship as S being the transport plan that permutes event likelihoods in X to their optimal position in the latent event sequence.To obtain a hard permutation from S, we use a result from [47] , who showed that M (X), the hard permutation matrix of discrete matches (i.e., the matrix of basis vectors corresponding to the vertices of the Birkhoff polytope), can be obtained as the limit τ → 0 of the Sinkhorn-Knopp operator,EQUATIONHere e n are basis (one-hot) vectors of size N with a value of 1 in the n-th position and 0 everywhere else. In practice we compute the hard permutation matrix M (X) from S using the Hungarian algorithm [52, 53] , which solves the minimum bipartite matching problem in cubic time. We use the relation in Equation 7to show that the original EBM can be obtained as the temperature limit of the vEBM (see Appendix A.3). To facilitate inference, the value of τ must be chosen to balance between the limit of a hard permutation, where the gradients are discontinuous (and hence non-differentiable), and a uniform soft permutation, where the gradients are flat (and hence non-informative). A parametric analysis of τ , τ prior , and the number of Sinkhorn-Knopp iterations, n s , is presented in Appendix A.8.We approximate the posterior probability, obtained from applying Bayes rule to Equation 3, using variational inference [54] , and define the evidence lower bound (ELBO). To enable differentiability of the ELBO, we parameterise our variational prior and posteriors using the Gumbel-Sinkhorn distribution, G(X, τ ), with a matrix, ϵ, of i.EQUATIONThe Gumbel-Sinkhorn distribution effectively implements the reparameterisation trick [55] for permutations, and in the limit τ → 0 it has been shown to converge to the Gumbel-Matching distribution, the equivalent of the Gumbel-Sinkhorn distribution for hard matchings [47] . We choose a uniform prior over permutations, G(X = 0, τ prior ), and for the posterior, G(X, τ ; ϕ), with parameters ϕ. We seek to optimise the corresponding ELBO, logPEQUATIONThe Kullback-Leibler (KL) divergence term on the RHS of Equation 9is intractable, but can be rewritten as KL((X + ϵ)/τ ∥ϵ/τ prior ) by substituting (X + ϵ)/τ for Z and estimated using random sampling [47] . For completeness, we restate the full term derived by [47] in Appendix A.4.To optimise the ELBO we use Adam [56] with n opt = 200 iterations and a learning rate of 0.1. Temperature hyperparameters τ and τ prior were set to 1 for all experiments, except the mixed events (Section 3.3.3), where τ = 1E3. We found setting ϵ = 0 during inference gave the fastest and most accurate estimate of S, at the expense of not allowing for direct propagation of uncertainty. While uncertainty estimation is not the focus of this paper, we do provide some examples of setting ϵ non-zero in Appendix A.7. Pseudo-code for the full inference scheme is given in Appendix Algorithm 1.We can use the trained model to obtain an individual-level likelihood distribution over stages, i.e., the likelihood at each state k given by Equation 3, where stage k = n corresponds to the first n events occurred and the remaining N -n events not occurred. Here we simply take the maximum likelihood stage for each individual, but alternative summary statistics could be used. 
baselines: 
 We consider two baselines; i) the original EBM [1] ; and the Alzheimer's Disease Probabilistic Cascades (ALPACA) model [33] , both of which use maximum likelihood to estimate the ordered sequence, s. The EBM learns s using gradient descent and Markov Chain Monte Carlo (MCMC) sampling. The ALPACA model instead defines s as the central permutation of a Mallows model [57] , with a density over permutations given by p(s) ∼ exp(-λd(s, s 0 )), where λ scales the spread around the central ordering, s 0 , and d(s, s 0 ) = Σ N i=1 |s(i) -s 0 (i)| is the distance between permutations. The ALPACA model learns s using expectation-maximisation (EM) and Gibbs sampling. We use the default parameters; for the EBM, 10 3 gradient descent iterations with 10 initial seeds, and 10 6 MCMC samples; for ALPACA, 10 EM iterations and 100 Gibbs samples.To enable comparison between our model and the baselines, we simulate data generated by an ordered sequence, s, according to the limit version of Equation 3 (see Appendix A.3). In brief, s is randomly initialised and individuals are assigned a stage with uniform probability, reflecting that individuals can be observed at any disease stage. Individuals are assigned as either controls or patients according to an arbitrary threshold on the disease stage (here we choose the lowest 20% stages as controls). Feature data for each individual are then generated from the Gaussian models of normal and abnormal feature values, depending on their stage, with zero means for the control distributions, random uniform means for the case distributions, and variable standard deviations for both controls and cases set to achieve a desired level of noise. Exact parameter values and code to generate synthetic data is given in the GitHub repository. Here we repeat performance evaluation over 10 simulated datasets for each experiment to support statistical significance tests.Figure 2 shows the runtime for the vEBM and baselines for three experiments (I = 100, J = 10; I = 1000, J = 100; I = 2000, J = 200). We do not consider J > 200 here, as the baselines become intractable, but Figure 2 clearly illustrates the unique computational ability of the vEBM to work with much larger J, as we demonstrate throughout this section. The vEBM is a factor of 1000 times faster for the J = 200 experiment; this factor would only increase for larger models. Figure 3 shows the effect of increasing aleatoric (measurement) noise levels on inference accuracy, as measured by the Kendall's tau [58] between the true and inferred sequences. The vEBM outperforms or is comparable to the baselines in all datasets and noise settings, except for I = 100, J = 1000 and σ = 0.1; this is expected, because at low noise and smaller numbers of features the EBM's MCMC sampling should find the global minimum, while the vEBM will always have some uncertainty due to its variational approximation. Statistical significance was obtained at p < 0.001 using unpaired t-tests (note that only one datapoint is shown for ALPACA at J = 100, and none at J = 200, due to computational intractability). We highlight that the metric is sensitive to any departure from the correct ordering, even by a single sequence position; accordingly the visual correlation between the true and inferred sequence remains high, e.g., even at the highest noise (σ = 1, corresponding to a 1:1 signal:noise ratio), as the relationship is still approximately diagonal. Additional examples in other datasets and when setting the Gumbel noise, ϵ, non-zero are given in Appendix A.6, A.7.We use pre-processed tensor-based morphometry (TBM) data from the Alzheimer's Disease Neuroimaging Initiative (ADNI) study, a longitudinal observational study of AD. TBM data are derived from structural magnetic resonance imaging (MRI) data and represent voxel-level maps of intensity gradients with respect to a reference healthy brain template, providing a standardised measure of voxel-level volume loss (or gain) between individuals. The TBM dataset we use here is comprised of cross-sectional TBM maps from 816 individuals (299 controls, 399 mild cognitive impairment, 188 AD) [59] . In Section 3.3.3 we also use three cognitive test scores -Mini-Mental State Examination (MMSE); Clinical Rating Dementia scale Sum of Boxes (CDRSB); Rey Auditory Verbal Learning Test (RAVLT). Both datasets are available to download for users with an ADNI account (https://adni.loni.usc.edu/data-samples/access-data/, data collections: "TBM Jacobian Maps MDT-SC"; "Tadpole Challenge").We apply the vEBM to TBM data from ADNI to reveal the first pixel-level sequence of disease events in AD (Figure 4 ). We do not include the baselines here due to computational intractability (as demonstrated in Section 3.2). The pattern of volume change represented by the vEBM sequence recapitulates known large-scale changes due to AD; initial change in the ventricles, followed by other sub-cortical changes, then changes across the cortex [60] . Moreover, the vEBM finds a detailed pattern of grey and white matter changes throughout the sequence, providing new small-scale insights into AD aetiology that have not previously been possible, which we explore further in the next section. To evaluate our ADNI pixel-level model with respect to previous analyses that have used segmented regional brain volumes, we map the vEBM pixel-level events post hoc to pixel-level labels obtained from the FreeSurfer segmentation of the reference template (Figure 5 ). Note that the regions shown are a subset of the total regions available from the FreeSurfer segmentation tool, which were chosen according to those that were sufficiently represented in terms of number of pixels in the 2D slice that was used to train the model (N > 10). Also note that the number of points on each trajectory / line corresponds to the number of pixels available in each region; e.g., there are few pixels in "Putamen" (N = 11), an intermediary number in "Thalamus-Proper" (N = 43), and many in "Cerebral-Cortex" (N = 299), corresponding to their relative sizes (areas) in the 2D slice.Our findings are in broad agreement with previous results; sub-cortical changes (Thalamus-Proper, Putamen, Hippocampus) are earliest, followed by cortical (Cerebral-Cortex) and white matter (Cerebral-White-Matter), and finally ventricular change (Lateral-Ventricle, VentralDC). However, our model provides much more fine-grained insights than, e.g., [20] , we now obtain continuous trajectories of change, which capture interesting non-linearities, e.g., in the Thalamus-Proper, Brain-Stem, and Lateral-Ventricle; this contrasts with the more linear changes in the Hippocampus, Cerebral-Cortex, and Cerebral-White-Matter.Figure 5 : Trajectories of regional brain volumes in our ADNI cohort, obtained by mapping the vEBM pixel-level events to pixel-level labels obtained from the FreeSurfer segmentation of the reference template. 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, as defined by the vEBM event sequence.The vEBM is not limited to modelling only image-based features, which we demonstrate by including three cognitive test score features (MMSE, CDRSB, RAVLT) and re-training the model. Figure 6 represents the spatio-(pseudo)-temporal pixel event topology obtained by the vEBM as a 2D histogram, and shows the position of the cognitive events by vertical lines. We calculate the spatial distribution of pixel events according to their Euclidean distance from the centre of the image. The colour denotes the number of pixel events in each histogram bin, e.g., in the first bin of events (the first column), we can see the density of pixel events occurring as a function of the distance from the centre. The pixel event topology shows the earliest events near the centre of the brain, as expected [60] , before spreading out across the brain; these events are interleaved with cognitive events, which occur across the latter two thirds of the progression. This interleaving suggests that the vEBM could be used to provide fine-grained staging in between cognitive events, e.g., for stratification in clinical trials. Interestingly, the pixel event topology is asymmetric about the central axis of the brain in the early stages, suggesting that the vEBM can identify subgroups of individuals who display asymmetric progression, which has previously been reported in small groups of people with AD, e.g., [61] .Figure 6 : Mixed pixel and cognitive feature event topology in AD found by the vEBM. The vertical axis shows the distance from the centre of the image to the pixel event, and the horizontal axis shows the event ordering obtained by the vEBM. Note that the cognitive events, denoted by vertical lines, are assigned an arbitrary distance of zero.We use pre-processed optical coherence tomography (OCT) data from the Duke University (DU) Ophthalmology 2013 dataset, a cross-sectional study of AMD [62] . The OCT data represent thickness maps for retinal pigment epithelium and drusen complex (RPEDC), a marker of AMD progression. The OCT dataset is comprised of 384 individuals (115 controls, 269 AMD), and is publicly available to download: https://duke.app.box.com/s/l80j6ziooeyy1eeo7edy0il32zbyyzbg. To select only pixels with disease signal, we remove pixels with an effect size less than 4 on a pixel-wise t-test between the control and AMD groups.We apply the vEBM to OCT data from the DU cohort to reveal the first pixel-level sequence of disease events in AMD (Figure 7 ). The density of RPEDC spread around the centre of the eye reflects previous observations [62] , and the vEBM provides a much finer-detailed progression pattern.Figure 7 : Pixel-level disease progression sequence in AMD obtained by the vEBM. White pixels correspond to events that have occurred by the corresponding point of the sequence. We have selected 10 sequence positions at uniform steps of 50 across the total of 537 in the full sequence, with the top left figure corresponding to position 80 and the bottom right to position 530. Images were made from the vEBM output using 3D Slicer (https://www.slicer.org/).Figure 8 shows the stage distribution for individuals in the ADNI and DU cohorts using the vEBM trained on mixed and pixel-only data, respectively. We find a fine-grained distribution of individuallevel stages that reflects the clinical labels, demonstrating the utility of the vEBM for stratification tasks, e.g., to select cohorts for clinical trials [63] . Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
discussion: 
 We introduced the vEBM, a novel optimal transport formulation for discrete disease progression models that scales to hitherto impossible numbers of events. It allows pixel-level visualisation of the order of pathology appearance in chronic disease, as demonstrated for the first time in AD and AMD.Here we did not fully explore model uncertainty, which as a Bayesian model, the vEBM can estimate directly; we reserve this for future work. We acknowledge that a model ordering several thousand events is not fully identifiable with only a few hundred snapshots, but the orderings we obtain are still highly meaningful, as demonstrated by our results in AD and AMD. Moreover, our optimal transport formulation naturally lends itself to feature-sparsification [64] , allowing redundant events to be grouped together. While the vEBM can in principal be directly applied to raw image data, pre-processing of images to a common reference frame (i.e., image registration) is necessary to facilitate comparison between individuals; however pre-processing is necessary for any data type to be in a common reference frame (or scale). We do not explicitly account for feature-wise covariance in the model, e.g., in image-based data we would expect a high degree of collinearity between neighbouring pixels; this could be addressed by including an additional term that imposes local structure, e.g., a Markov random field [35] . Finally, we note that the main limitation on the current formulation is not computational tractability but computer memory due to dense matrix operations, which could be alleviated using, e.g., sparse matrix representation.The vEBM enables disease progression modelling at scale in multiple areas of medical imaging, not only the modalities demonstrated in this paper; such as other MRI modalities e.g., diffusion weighted imaging, microstructure modelling, connectivity; other imaging modalities, e.g., positron emission tomography, computed tomography, X-rays, ultrasound; and non-radiological imaging modalities, e.g., microscopy. Furthermore, the vEBM can run quickly on a relatively low-spec computer without the need for GPU infrastructure, making it accessible to research labs -and potentially clinics -that have limited resources, while further minimising its carbon impact by reducing compute time. In addition, it provides a new, more powerful, model for each component of mixture subtype models, e.g., [5] , which currently uses a variant of the basic EBM. Such models are highly influential in stratifying patients into disease subgroups for more precise clinical trials and treatment deployment.The analyses presented here were performed either on a laptop PC with a single AMD Ryzen 7 PRO 6860Z CPU with 32GB RAM, or a desktop PC with a single AMD Ryzen Threadripper PRO 5975WX CPU with 270GB RAM (we note that we also provide a GPU implementation of the code, but we do not use it here to allow fair comparison with the baselines). Model training for the synthetic data analysis took 48 hours (wall-clock time) in total for all experiments. Model training for the AD analysis took 5 minutes. Model training for the AMD analysis took 1 minute. No pre-training of the model was performed. Approximately 48 hours CPU compute, and less than one hour on a GPU, was performed for code testing and preliminary experiments that are not included in this paper. Python code to reproduce all results presented here is available from the first author's GitHub repository 1 .Starting from the joint probability (Equation 1), we use the assumption of independence between measured features j = {1, 2, 3, ..., J} to write the model likelihood,EQUATION)Using the chain rule, the joint probability distribution over Y and k can be factorised as,EQUATIONFor the likelihood model (Equation 10), here we choose a two-component Gaussian mixture model (though as noted in the main text, any probabilistic model can be chosen),J j=1 P (Y i,j |k i , θ s(j) , S) = ki j=1P (Y i,j |k i , θ p s(j) , S)EQUATION)Substituting Equation 12into Equation 11 and marginalising over k i , we have,EQUATION)Here we chose the prior over latent stages, parametrised by hidden variable π, to be uniform and constant, P (k i |S; π) ∼ Unif(0, k). Finally, assuming independence between measurements from different individuals i, we can write the following expression for the total likelihood,EQUATION)Bayes' theorem can now be used to obtain the posterior over S.We use the limit relation in Equation 7to reparametrise the model likelihood (Equation 3) in terms of a discrete permutation, s,EQUATION) S k i π Y i,j θ jFigure 9 : Graphical model of the variational event-based model (vEBM). Hidden variables are denoted by circles and observations by squares. S: event permutation matrix; π: initial probability vector (fixed to uniform distribution); k i : disease state for individual i; θ j : distribution parameters for biomarker j; Y i,j : observed data.To see this, we take the first likelihood term in the RHS of Equation 3 and write it in matrix form,EQUATIONwhere we have used the limit relation from Equation 7between lines three and four. The same steps can be applied to the second likelihood term on the RHS of Equation 3, P (Y i,j |S, θ c s(j) ), to obtain the full likelihood reparametrisation in s.For completeness, we restate the KL divergence term derived by [47] ,EQUATIONwhere S 1 = Σ i,j x i,j and S 2 = Σ i,j exp(-x i,j τ prior /τ ). Full a full derivation see [47] , Supplementary Methods B.3.Pseudo-code for the vEBM inference scheme is given in Algorithm 1. Tables 1,2,3 show evaluation metrics for varying combinations of the temperature hyperparameter, τ , and its prior, τ prior . Table 4 shows the same metrics for constant τ, τ prior and extreme values for the number of Sinkhorn-Knopp iterations, n s . Synthetic data with noise, σ = 0.5, were used. Justification: We explicitly discuss limitations of our method in the Section 4.1.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [Yes] Justification: The software presented in this paper could be misused in a medical setting, which is not recommended as the software has not been designed to be a medical grade device. To protect against this we will provide usage guidelines on the GitHub repository for the code. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We credit all creators of the supporting software used in this paper, specifically: Python, 3D Slicer. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.https://github.com/pawij/vebm 
 
 Appendix:  


Decision: Thanks to the reviewers and the authors for engaging in discussion, and the authors for adequately addressing reviewer concerns. The reviewers all lean toward acceptance, so I will favor acceptance as well.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: My questions were combined with weaknesses.
limitations: limitations: The authors adequately addressed the limitations.
paper_summary: 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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Yes, the authors adequately addressed the limitations and, if applicable, potential negative societal impact of their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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?
limitations: limitations: The limitation has been discussed at end of the paper.

Comments: title: Thanks to authors' efforts
Response to rebuttal
Thanks for reminder
Response to rebuttal
Additional question
comment: Thanks to the authors for the reply. Understood and best wishes.
I want to thank you authors for their rebuttal. I decided to increase the score to Accept.
Thanks for the responses. My concerns are mostly addressed and looking forward to your follow-up works.
Hello reviewer 5r2K,

Thanks for already engaging in discussion with the authors.

I just wanted to check whether the most recent author response has adequately addressed your concerns. Separately, please indicate whether you're sticking with your original score, or if this discussion has led to any change in score on your end. Note that the author/reviewer discussion period ends very soon (Aug 13, 11:59pm AoE).

Thanks,
Your AC
Hello reviewer SCze,

Does the author respond adequately address your concerns?

Note that the author/reviewer discussion period ends very soon (Aug 13, 11:59pm AoE).

Thanks,
Your AC
No problem! Thanks for the positive response.

In response to your points above:

1. Comparing the distributions of stages between clinical labels using a suitable distance metric is a nice idea - we'll add that to the text.
2. Ok we understand the criticism - indeed the figure does not look very professional! We will try to improve the style to make it more visually appealing. Regarding the plots, we need to point out that the colour choice was made to conform to accessibility requirements (the default colours that we use were designed by the authors of matplotlib to be accessible for people who are colour-blind). However we agree that the error bars are not very clear and will make them larger.
Yes your understanding is exactly correct.

Ok we see your point - comparing to a pixel-level, i.e., correlated feature, ground truth trajectory would indeed be instructive, and given the time, we would have set up a suitable simulation to test this. As it was, we only had time to validate the method using ground truth trajectories of uncorrelated features (Sections 3.2.2 and Supplementary A.7). However, as you can see from our results, our method performs very well at recovering the ground truth trajectory for various dataset properties and model parameter settings, and generally outperforms the two baselines. As you suggest, implementing pixel-level simulations are a priority for our future work.

On a related note - it would be very difficult to validate the results from the real data analyses in Alzheimer's disease and AMD, due to the standard problem in medical imaging of having no ground truth. The closest one could get is to have an additional dataset of post mortem histology images, matched to the in vivo medical images that we used here, in order to validate the spatial distribution of abnormality predicted by the model. Indeed there are approaches developed by other researchers to register medical images with histology, but they are still experimental and would require substantial time and cross-collaborative effort to implement. That being said, it would be a very interesting direction to take our model, as it provides the first pixel-level predictions of its type, which would support more direct comparison between MRI and histology.
Thanks for the authors for reminding that. I meant to update the rating from 5 to 6. Misclick from phone. Corrected.
Thanks for the clarifications.

Before we respond to points 1 and 2, we'd like to ask why you decided to downgrade your review from a 5 to a 4? Your comment seems to be positive so we're unsure as to what we did to cause the score to be downgraded - some detail would be very helpful, so we can try to respond as best as possible.
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.
Thanks for the clarifications. To double confirm my understanding, event number is basically timesteps, and fraction is the number of abnormal pixels over the total number of pixels in that region. Presumably, the disease progresses it should become higher. That is fair.

However, my comment about the quantitative measurement against the ground truth trajectory refers to something else. Since the algorithm is predicting pixel level label, at a specific event/timestep,  Intersection over Union should be computed to measure how realistic the progression simulation is, given a ground truth trajectory.
Thanks for the positive response. We will include the clarifications you highlight in the manuscript.
Thanks for your responses and clarifications, which are very helpful and address my concerns. I do think that including some of these details / clarifications in the text itself would be helpful. Most important are the clarifications you provide in your responses to Weaknesses point 2 and Questions points 1-2. I have increased my score.
Thanks for the question. The "Fraction occurred" label corresponds to the fraction of pixels in each region that have become abnormal, as defined by the vEBM event sequence. The fraction of pixel-events occurred increases with the event number because the event sequence represents a monotonic accumulation in pixel-level abnormality. For example, if we look at the "Cerebral-Cortex" line, it shows the fraction of pixel-events within that region that have occurred as a function of their position in the event sequence, e.g., at "Event" = 600, which is nearing the half way point in the event sequence, approximately 60% of the pixel-events in "Cerebral-Cortex" have occurred.
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?