Title: Unscrambling disease progression at scale: fast inference of event permutations with optimal transport
Venue: NeurIPS 2024
Accepted: yes
Keywords: optimal transport

variational inference

latent variable model

disease progression
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.

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

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.

1 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-9] or continuous [10-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-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-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-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.

1.1 Contributions

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.

2 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

C(Φ;0) = Eq(ZIy)[logPe(Y|Z)] - KL(q(Z|Y)IIP(Z))

P(Y|0P)

Eventpermutation

Snapshots, Y

matrix, S

i=1

whilenotconvergeddo

j=1

i=1

S

j=1

i=2

i=3

j=2

i=I

j=1

···

j=1..·j=4

.··

j=4

j=4

j=3

P(Y|0%)

i=1

UpdateS←VL(Φ);

Event sequence,s

j=4

end

s={4,1,2,3}

i=I

j=1..·j=4

A

B

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.

2.1 Model of disease progression

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),

Each 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 that

However 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),

Here 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 ) .

2.2 Optimal transport for permutations

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,

In the context of permutation matrices as couplings, S belongs to the Birkhoff polytope,

The 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,

Here 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 the 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,

Here 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 7 to 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.

2.3 Variational permutation inference

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.i.d. Gumbel noise,

The 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,

The Kullback-Leibler (KL) divergence term on the RHS of Equation 9 is 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.

2.4 Inference scheme

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 τ = 1 E 3 . 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.

2.5 Probabilistic staging

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.

3 Experiments

3.1 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 MCMCsamples; for ALPACA, 10 EM iterations and 100 Gibbs samples.

3.2 Synthetic data

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.

3.2.1 Faster inference

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.

3.2.2 Improved accuracy and robustness to noise

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

vEBM (ours)

104

EBM

ALPACA

time

103

Wall-clockt

102

101

100

50

100

150

200

Numberoffeatures

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.

3.3 Alzheimer's disease data

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

Noise = 0.1

Noise = 0.5

Noise = 1.0

1.0

Kendall's tau

0.8

Kendall'stau

0.8

Kendall's tau

0.3

vEBM (ours)

EBM

0.6

ALPACA

0.6

0.2

0.4

VEBM (ours)

0.4

+

VEBM(ours)

0.1

-

EBM

-T

EBM

0.2

ALPACA

0.2

ALPACA

0.0

50

100

150

200

50

100

150

200

50

100

150

200

95

Kendall's tau=0.93

95

Kendall's tau = 0.91

112

Kendall's tau = 0.28

VEBM

VEBM

VEBM

124

True

124

True

3

True

79

79

111

169

169

94

Feature

2

Feature

2

Feature

179

4

4

127

40

40

90

180

180

76

175

175

50

1

1

102

020 406080100120140160180

020 406080100120140160180

020 406080100120140160180

Event

Event

Event

( https://adni.loni.usc.edu/data-samples/access-data/ , data collections: 'TBM Jacobian Maps MDT-SC'; 'Tadpole Challenge').

3.3.1 Pixel-level disease progression events in AD

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 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.

3.3.2 Segmentation-based interpretation of pixel-level events

To evaluate our ADNI pixel-level model with respect to previous analyses that have used segmented regional brain areas, 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.

1.0

Brain-Stem

Cerebral-White-Matter

Cerebral-Cortex

Lateral-Ventricle

Thalamus-Proper

0.8

Putamen

Hippocampus

Fraction occurred

VentraIDC

0.6

0.4

0.2

0.0

0

200

400

600

800

1000

1200

1400

Event

3.3.3 Mixed feature disease progression events in AD

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].

RAVLT

CDRSB

MMSE

30

12

20

Distance (pixels)

OT

10

0

-10

-20

0

200

400

600

800

1000

1200

Event

3.4 Age-related macular degeneration data

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.

3.4.1 Pixel-level disease progression events in AMD

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.

3.5 Prediction of AD and AMD stage

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].

0.012

CN

0.035

CN

0.010

MCI

AMD

AD

0.030

0.008

0.025

Density

Density

0.006

0.020

D

0.015

0.004

0.010

0.002

0.005

0.000

0.000

0

200

400

600

800

1000

1200

1400

0

100

200

300

400

500

Stage

Stage

4 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.

4.1 Limitations

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.

4.2 Broader impacts

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.

Acknowledgments and Disclosure of Funding

PAW would like to acknowledge useful discussions with the Sussex PAL Book Club and logistical support from Siobhán and Lorcán. DCA would like to acknowledge support from Wellcome Trust grant 221915 and the NIHR UCLH Biomedical Research Centre. Both authors acknowledge support from the Wellcome Leap 1kD project. Neither author have competing financial interests to declare.

References

H. M. Fonteijn, M. Modat, M. J. Clarkson, et al. An event-based model for disease progression and its application in familial alzheimer's disease and huntington's disease. NeuroImage , 60: 1880-1889, 2012. doi: https://doi.org/10.1016/j.neuroimage.2012.01.062.

Y. Y. Liu, S. Li, F. Li, et al. Efficient learning of continuous-time hidden markov models for disease progression. Advances in Neural Information Processing Systems , 28, 2015.

A. M. Alaa, S. Hu, and M. van der Schaar. Learning from clinical judgments: Semi-markovmodulated marked hawkes processes for risk prognosis. International Conference on Machine Learning , 19(4):1-62, 2017.

A. M. Alaa and M. van der Schaar. A hidden absorbing semi-markov model for informatively censored temporal data: Learning and inference. Journal of Machine Learning Research , 70: 60-69, 2018. doi: http://dx.doi.org/10.5555/3305381.3305388.

A. L. Young, R. V. Marinescu, N. P. Oxtoby, et al. Uncovering the heterogeneity and temporal complexity of neurodegenerative diseases with subtype and stage inference. Nature Communications , 9, 2018. doi: https://doi.org/10.1038/s41467-018-05892-0.

A. M. Alaa and M. van der Schaar. Attentive state-space modeling of disease progression. Advances in Neural Information Processing Systems , 32, 2019.

Z. Sun, S. Ghosh, Y. Li, et al. A probabilistic disease progression modeling approach and its application to integrated huntington's disease observational data. JAMIA Open , 2(1):123-130, 2019. doi: https://doi.org/10.1093/jamiaopen/ooy060.

J. P. Williams, C. B. Storlie, T. M. Therneau, et al. A bayesian approach to multistate hidden markov models: Application to dementia progression. Journal of the American Statistical Association , 115:16-31, 2020. doi: https://doi.org/10.1080/01621459.2019.1594831.

C. Hadjichrysanthou, S. Evans, S. Bajaj, et al. The dynamics of biomarkers across the clinical spectrum of alzheimer's disease. Alzheimers Res Ther , 12, 2020. doi: https://doi.org/10.1186/ s13195-020-00636-z.

M. C Donohue, H. Jacqmin-Gadda, M. Le Goff, et al. Estimating long-term multivariate progression from short-term data. Alzheimer's & Dementia , 10:S400-410, 2014. doi: https: //doi.org/10.1016/j.jalz.2013.10.003.

J. B. Schiratti, S. Allassonnière, O. Colliot, et al. A bayesian mixed-effects model to learn trajectories of changes from repeated manifold-valued observations. J Machine Learning Research , 18:1-33, 2017.

N. P. Oxtoby, A. L. Young, D. M. Cash, et al. Data-driven models of dominantly-inherited alzheimer's disease progression. Brain , 141:1529-1544, 2018. doi: https://doi.org/10.1093/ brain/awy050.

M. Lorenzi, M. Filippone, G. B. Frisoni, et al. Probabilistic disease progression modeling to characterize diagnostic uncertainty: Application to staging and prediction in alzheimer's disease. NeuroImage , 190:56-68, 2019. doi: https://doi.org/10.1016/j.neuroimage.2017.08.059.

D. Li, S. Iddi, P. S. Aisen, et al. The relative efficiency of time-to-progression and continuous measures of cognition in presymptomatic alzheimer's disease. Alzheimer's and Dementia: Translational Research and Clinical Interventions , 5:308-318, 2019. doi: https://doi.org/10. 1016/j.trci.2019.04.004.

M. Bilgel and B. M. Jedynak. Predicting time to dementia using a quantitative template of disease progression. Alzheimers Dement , 11:205-215, 2019. doi: https://doi.org/10.1016/j. dadm.2019.01.005.

A. O'Connor, P. S. J Weston, I. M. Pavisic, et al. Quantitative detection and staging of presymptomatic cognitive decline in familial alzheimer's disease: a retrospective cohort analysis. Alzheimers Res Ther , 12, 2020. doi: https://doi.org/10.1186/s13195-020-00695-2.

I. Koval, A. Bône, M. Louis, et al. Ad course map charts alzheimer's disease progression. Sci Rep , 11:8020, 2021. doi: https://doi.org/10.1038/s41598-021-87434-1.

A. M. Staffaroni, M. Quintana, B. Wendelberger, et al. Temporal order of clinical and biomarker changes in familial frontotemporal dementia. Nature Medicine , 28:2194-2206, 2022. doi: https://doi.org/10.1038/s41591-022-01942-9.

A. L. Young, N. P. Oxtoby, S. Garbarino, et al. Data-driven modelling of neurodegenerative disease progression: thinking outside the black box. Nature Reviews Neuroscience , 25:111-130, 2024. doi: https://doi.org/10.1038/s41583-023-00779-6.

A. L. Young, N. P. Oxtoby, P. Daga, et al. A data-driven model of biomarker changes in sporadic alzheimer's disease. Brain , 137:2564-2577, 2014. doi: https://doi.org/10.1093/brain/awu176.

N. C. Firth, C. M. Startin, R. Hithersay, et al. Aging related cognitive changes associated with alzheimer's disease in down syndrome. Ann Clin Transl Neurol , 5:1665-1677, 2018. doi: https://doi.org/10.1002/acn3.571.

J. W. Vogel, A. L. Young, N. P. Oxtoby, et al. Four distinct trajectories of tau deposition identified in alzheimer's disease. Nature Medicine , 27:871-881, 2021. doi: https://doi.org/10. 1038/s41591-021-01309-6.

L. M. Byrne, F. B. Rodrigues, K. Blennow, et al. Neurofilament light protein in blood as a potential biomarker of neurodegeneration in huntington's disease: a retrospective cohort analysis. Lancet Neurol , 16:601-609, 2017. doi: https://doi.org/10.1016/S1474-4422(17)30124-2.

L. M. Byrne, F. B. Rodrigues, E. B. Johnson, et al. Evaluation of mutant huntingtin and neurofilament proteins as potential markers in huntington's disease. Sci Transl Med , 10: eaat7108, 2018. doi: https://doi.org/10.1126/scitranslmed.aat7108.

P. A. Wijeratne, A. L. Young, N. P. Oxtoby, et al. An image-based model of brain volume biomarker changes in hungtington's disease. Ann Clin Transl Neurol , 5:570-582, 2018. doi: https://doi.org/10.1002/acn3.558.

P. A. Wijeratne, S. Garbarino, S. Gregory, et al. Revealing the timeline of structural mri changes in premanifest to manifest huntington disease. Neurol Genetics , 7, 2021. doi: https: //doi.org/10.1212/NXG.0000000000000617.

A. Eshaghi, R. V. Marinescu, A. L. Young, et al. Progression of regional grey matter atrophy in multiple sclerosis. Brain , 141:1665-1677, 2018. doi: https://doi.org/10.1093/brain/awy088.

A. Eshaghi, , A. L. Young, P. A. Wijeratne, et al. Identifying multiple sclerosis subtypes using unsupervised machine learning and mri data. Nat Comms , 12, 2021. doi: https://doi.org/10. 1038/s41467-021-22265-2.

N. P. Oxtoby, L. A. Leyland, L. M. Aksman, et al. Sequence of clinical and neurodegeneration events in parkinson's disease progression. Brain , 144:975-988, 2021. doi: https://doi.org/10. 1093/brain/awaa461.

R. Pascuzzo, N. P. Oxtoby, A. L. Young, et al. Prion propagation estimated from brain diffusion mri is subtype dependent in sporadic creutzfeldt-jakob disease. Acta Neuropathologica , 140: 169-181, 2020. doi: https://doi.org/10.1007/s00401-020-02168-0.

M. C. Gabel, R. J. Broad, A. L. Young, et al. Evolution of white matter damage in amyotrophic lateral sclerosis. Ann Clin Transl Neurol. , 7:722-732, 2020. doi: https://doi.org/10.1002/acn3. 51035.

A. L. Young, F. J. S. Bragman, B. Rangelov, et al. Disease progression modeling in chronic obstructive pulmonary disease. Am J Respir Crit Care Med , 201:294-302, 2020. doi: https: //doi.org/10.1164/rccm.201908-1600OC.

J. Huang and D. C. Alexander. Probabilistic event cascades for alzheimer's disease. Advances in Neural Information Processing Systems , 25, 2012.

X. Wang, D. Sontag, and F. Wang. Unsupervised learning of disease progression models. Proceedings of the 20th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining , 2014. doi: http://dx.doi.org/10.1145/2623330.2623754.

R. V. Marinescu, A. Eshaghi, M. Lorenzi, et al. Dive: A spatiotemporal progression model of brain pathology in neurodegenerative disorders. NeuroImage , 192:166-177, 2019. doi: https://doi.org/10.1016/j.neuroimage.2019.02.053.

A. L. Young, J. W. Vogel, L. Aksman, et al. Ordinal sustain: Subtype and stage inference for clinical scores, visual ratings, and other ordinal data. Frontiers in Artificial Intelligence , 4, 2021. doi: https://doi.org/10.3389/frai.2021.613261.

P. A. Wijeratne and D. C. Alexander. Learning transition times in event sequences: The temporal event-based model of disease progression. IPMI LNCS , 12729:583-595, 2021. doi: https://doi.org/10.1007/978-3-030-78191-0\_45.

P. A. Wijeratne, A. Eshaghi, W. J. Scotton, et al. The temporal event-based model: Learning event timelines in progressive diseases. Imaging Neuroscience , 1:1-19, 2023. doi: https: //doi.org/10.1162/imag_a_00010.

A. L. Young, L. Aksman, D. C. Alexander, et al. Subtype and stage inference with timescales. Lecture Notes in Computer Science , 13939:15-26, 2023. doi: https://doi.org/10.1007/ 978-3-031-34048-2_2.

G. Martí-Juan, M. Lorenzi, and G. Piella. Mc-rvae: Multi-channel recurrent variational autoencoder for multimodal alzheimer's disease progression modelling. NeuroImage , 268, 2023. doi: https://doi.org/10.1016/j.neuroimage.2023.119892.

A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. International Conference on Learning Representations , 2023. doi: https://doi.org/10.48550/ arXiv.2312.00752.

G. Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning , 11:355-607, 2019. doi: https://doi.org/10.48550/arXiv.1803.00567.

C. S. Parker, N. P. Oxtoby, A. L. Young, et al. Parsimonious ebm: generalising the eventbased model of disease progression for simultaneous events. bioRxiv , 2024. doi: https: //doi.org/10.1101/2022.07.10.499471.

R. G. Krishnan, U. Shaltit, and D. Sontag. Structured inference networks for nonlinear state space models. AAAI , 28:2101-2109, 2017.

C. Villani. Optimal transport, old and new. Springer Verlag , 338, 2009.

M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transportation distances. Advances in Neural Information Processing Systems , 26:2292-2300, 2013. doi: https://doi.org/ 10.48550/arXiv.1306.0895.

G. Mena, D. Belanger, S. Linderman, et al. Learning latent permutations with gumbel-sinkhorn networks. International Conference on Learning Representations , 2018. doi: https://doi.org/10. 48550/arXiv.1802.08665.

G. Birkhoff. Tres observaciones sobre el algebra lineal. Univ. Nac. Tucumán , pages 147-151, 1946.

R. Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matrices. Annals of mathematical statistics , 35:876-879, 1964.

R. Sinkhorn and P. Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics , 21:343-348, 1967.

C. J. Maddison, A. Mnih1, and Y. W. Teh. The concrete distributions: a continuous relaxation of discrete random variables. International Conference on Learning Representations , 2016. doi: https://doi.org/10.48550/arXiv.1611.00712.

H. W. Kuhn. The hungarian method for the assignment problem. Navel Research Logistics , 2: 83-97, 1955.

J. Munkres. Algorithms for the assignment and transportation problems. Journal of the society for industrial and applied mathematics , 5:32-38, 1957.

D. M. Blei, A. Kucukelbir, and J. D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association , 112:859-877, 2017. doi: https://doi.org/10. 1080/01621459.2017.1285773.

D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv , 2013. doi: https: //doi.org/10.48550/arXiv.1312.6114.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations , 2014. doi: https://doi.org/10.48550/arXiv.1412.6980.

C. L. Mallows. Non-null ranking moidels. Biometrika , 44:114-130, 1957.

M. Kendall. A new measure of rank correlation. Biometrika , 30:81-89, 1938. doi: https: //doi.org/10.1093/biomet/30.1-2.81.

X. Hua, D. P. Hilbar, C. R. K. Ching, et al. Unbiased tensor-based morphometry: Improved robustness and sample size estimates for alzheimer's disease clinical trials. Neuroimage , 66: 648-661, 2013. doi: https://doi.org/10.1016/j.neuroimage.2012.10.086.

B. Dubois, H. Hampel, H. H. Feldman, et al. Preclinical alzheimer's disease: definition, natural history, and diagnostic criteria. Alzheimers Dement , 12(3):292-323, 2016. doi: http: //dx.doi.org/10.1016/j.jalz.2016.02.002.

A. Low, E. Mak, M. Malpetti, et al. Asymmetrical atrophy of thalamic subnuclei in alzheimer's disease and amyloid-positive mild cognitive impairment is associated with key clinical features. Alzheimer's & Dementia: Diagnosis, Assessment & Disease Monitoring , 11:690-699, 2019. doi: https://doi.org/10.1016/j.dadm.2019.08.001.

S. Farsiu, S. J. Chiu, R. V. O'Connell, et al. Quantitative classification of eyes with and without intermediate age-related macular degeneration using optical coherence tomography. Ophthalmology , 121:162-172, 2014. doi: https://doi.org/10.1016/j.ophtha.2013.07.013.

E. R. Dorsey, C. Venuto, V. Venkataraman, et al. Novel methods and technologies for 21stcentury clinical trials. JAMA Neurol , 72(5):582-588, 2015. doi: http://dx.doi.org/10.1001/ jamaneurol.2014.4524.

M. Cuturi, M. Klein, and P. Ablin. Monge, bregman and occam: Interpretable optimal transport in high-dimensions with feature-sparse maps. arXiv , 2024. doi: https://doi.org/10.48550/arXiv. 2302.04065.

A Appendix / supplemental material

A.1 Compute information

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 .

A.2 Derivation of vEBM

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,

Using the chain rule, the joint probability distribution over Y and k can be factorised as,

For 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),

Substituting Equation 12 into Equation 11 and marginalising over k i , we have,

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,

Bayes' theorem can now be used to obtain the posterior over S .

A.3 Posterior limit

We use the limit relation in Equation 7 to reparametrise the model likelihood (Equation 3) in terms of a discrete permutation, s ,

S

k i

π

Yi,j

θ

j

To see this, we take the first likelihood term in the RHS of Equation 3 and write it in matrix form,

where we have used the limit relation from Equation 7 between 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 .

A.4 Kullback-Leibler divergence

For completeness, we restate the KL divergence term derived by [47],

where 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.

A.5 Inference scheme

Pseudo-code for the vEBM inference scheme is given in Algorithm 1.

A.6 Positional variance diagrams

Figures 10 and 11 show positional variance diagrams for datasets with I = 100 , J = 10 , and I = 1000 , J = 100 , individuals and features, respectively.

A.7 Uncertainty estimation

Figures 12, 13, 14 show equivalent positional variance diagrams to Figures 3, 10, 11, but with the Gumbel noise term, ϵ , set to non-zero. The corresponding uncertainty obtained from 1000 random samples of the posterior is shown by greyscale shading on the positional variance diagrams.

Kendall's tau =1.0

Kendall's tau =1.0

Kendall's tau =0.73

8

VEBM

8

VEBM

VEBM

1

True

1

口

True

1

口

True

5

5

5

0

0

0

Feature

7

Feature

7

Feature

8

2

2

2

9

9

9

4

4

4

3

3

6

6

6

3

0

1

2

3

45

6

5789

0

1

2

3

45

6789

0

1

2

3

45

6

7

8

9

Event

Event

Event

A.8 Hyperparameter study

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.

83

Kendall's tau = 0.93

83

Kendall's tau = 0.7

70

Kendall's tau = 0.0

18

口

VEBM

True

18

口

VEBM

True

30

口

VEBM

True

55

55

55

11

47

85

Feature

49

Feature

49

Feature

34

24

24

3

89

89

99

79

79

41

84

84

9

91

91

23

0

1020304050607080

90

0

10203040506070

80

90

0

10

203040506070

80

90

Event

Event

Event

Kendall's tau = 0.42

Kendall's tau = 1.0

Kendall's tau=0.73

8

VEBM

8

VEBM

7

二

VEBM

1

True

1

True

1

True

5

5

5

7

to

0

Feature

0

Feature

7

Feature

8

2

2

2

9

9

9

4

4

4

3

3

6

6

6

3

0

1

2

3

45

6

5789

0

1

2

3

45

6789

0

1

2

3

4

5

6

7

8

9

Event

Event

Event

Kendall's tau = 0.72

Kendall's tau=0.66

Kendall's tau = 0.05

83

70

45

30

口

VEBM

True

30

口

True

VEBM

30

VEBM

True

55

88

55

11

11

85

Feature

49

Feature

49

Feature

34

24

24

3

56

56

54

79

79

58

9

9

84

87

82

23

0

1020304050607080

90

0

1020304050607080

90

0

10203040506070

80

90

Event

Event

Event

Kendall's tau = 0.39

Kendall's tau=0.36

Kendall's tau = 0.06

95

95

95

93

口

口

VEBM

True

65

VEBM

True

115

口

口

VEBM

True

79

79

24

169

97

12

eature

2

Feature

2

Feature

111

32

120

148

108

40

28

133

180

133

72

175

72

192

71

130

0

20

4060 80 100120140160180

0

20

0406080100120140160180

0

20

40

06080100120140160180

Event

Event

Event

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We clearly show in Figure 2 that our method is at least a factor of 1000 times faster than the baselines, and in Figures 4, 6, 7 the novel application of our method to Alzheimer's disease and age-related macular degeneration data.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Most of the results can be reproduced exactly using the code provided (upon publication) and either synthetic data or open access data (Duke University ophthalmology dataset). The experiments involving the ADNI dataset are reproducible but require that the user gains access to the ADNI dataset, which is granted upon request via their website ( https://adni.loni.usc.edu/data-samples/access-data/ ).

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Complete details are given for accessing the code (upon publication) and all datasets. Scripts will be provided in the code repository to reproduce a subset of experiments that use open access data.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Full information on all experimental details and datasets is provided in the main text, Sections 3.2, 3.3, 3.4. Additional information is included in the Appendix Section A.5.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We run multiple tests with synthetic data (10 per experiment) to support significance testing, using an unpaired t-test, as described in Section 3.2.

Guidelines:

The answer NA means that the paper does not include experiments.

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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We have provided compute information in Appendix Section A.1.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed and agree with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The broader impacts of our work are discussed in Section 4.2.

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).

11. 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.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We will release the code for the variational event-based model (vEBM) open source with an MIT license upon publication of the paper.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [No]

Justification: All data used in this paper were previously collected by the study coordinators (ADNI and Duke University), who provided the relevant information to participants in order to obtain their consent to enter the study [59, 62].

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [No]

Justification: All data used in this paper were previously collected by the study coordinators (ADNI and Duke University), who obtained informed consent from all participants [59, 62].

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.

Peer reviews:
comments: 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?

LIMITATIONS:
The limitation has been discussed at end of the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. It is quite innovative to view the disease progression modeling task from the optimal transport perspective.\n2. Figure 1 makes the paper slightly easier to read, as it outlines the proposed method with references to subsequent paper sections. \n3. 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.\n4. 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.\n\nWEAKNESSES:\n1. 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.\n2. 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.\n3. For the results in Figure 7, while it is visually informative, it will be great if the authors can incorporate quantitative metrics.\n4. Minor issue: For Figure 1, it would be great if the authors can improve the aesthetics. \n5. Minor issue: For Figure 5, it would be helpful to provide the color bar.\n\nQUESTIONS:\n1. 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.\n2. 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?\n\nLIMITATIONS:\nYes, the authors adequately addressed the limitations and, if applicable, potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. The experiments have been performed with multiple datasets: synthetic data, neuroimaging, and optical coherence tomography.\n2. 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.\n3. The methodology development was pretty clear and supported by relevant literature.\n4. Novel formulation of Event-Based Model via variational inference and optimal transport.\n\nWEAKNESSES:\n1. 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.\n2. 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. \n\nWijeratne, 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.\n\nWijeratne, Peter A., et al. \"The temporal event-based model: Learning event timelines in progressive diseases.\" Imaging Neuroscience 1 (2023): 1-19.\n\nQUESTIONS:\nMy questions were combined with weaknesses.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nResults 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.\n\nWEAKNESSES:\n- The related work section is very brief, which made me question whether it is comprehensive.\n- 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.\n- 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.\n\nQUESTIONS:\n- 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.\n- What is denoted by the color of the pixels in Figure 5?\n- 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?\n- Do the authors envision applications of this method outside of medicine?\n\nLIMITATIONS:\nSomewhat, 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}