Unscrambling disease progression at scale: fast inference of event permutations with optimal transport
Authors: Peter A. Wijeratne
Daniel C. Alexander
Venue: NeurIPS 2024 poster (2024)
Keywords: optimal transport
variational inference
latent variable model
disease progression
Primary area: probabilistic_methods
TL;DR: We introduce a new method that substantially speeds up inference of discrete disease progression models, allowing them to scale to large feature sets.

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
Peter Wijeratne
Daniel C. Alexander
Sussex Al Centre
Hawkes Institute
Department of Informatics University of Sussex Brighton, BN1 9RH United Kingdom
Department of Computer Science University College London London, WCIE 6BT United Kingdom d. allexaderruuc.cco ac.
p. p.wijeratneesusseex ac. uk
Abstract
Disease progression models infer group-level temporal trajectories of change patients' features chronic degenerative condition plays out. They provide unique insight into disease biology and staging systems with individual-level clini- latent permutation
cal utility. Discrete models consider disease progression of events, where each event corresponds to a feature becoming measurably ab- normal. However, permutation inference using traditional maximum likelihood approaches becomes prohibitive due to combinatoric explosion, severely limiting model dimensionality and utility. Here leverage ideas from optimal transport 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 factor of 1000 times faster inference than the current state of the art and, correspondingly, supports models with several orders of mag- nitude more features than the current state of the art can consider. Experiments demonstrate 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 s discase 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 low compute, interpretable and applicable to any progressive condition and data modality, giving it broad potential clinical utility.
Introduction
The main aim of disease progression modelling is to learn 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 treat the problem dynamically, using either discrete [[49] continuous [0-18] models with latent variables to describe the hidden disease stage or time. An abundance of
such models have been proposed (see [19] for comprehensive review) and have found extensive success in providing unique interpretability and utility across a wide range progressive diseases, including Alzheimer's disease (AD) [I] 20 12 22). Huntington's 's disease 223.26], multiple
sclerosis 27 (88) Parkinson's 's disease 29 prion disease 30), amyotrophic lateral sclerosis [1],
and chronic obstructive pulmonary disorder 32
However all previous approaches make compromise: they are either i) interpretable in feature
space but sacrifice computational tractability (33 34
13
35 6
36-39; or ii) are made
computationally tractable by encoding to latent space but sacrifice direct interpretability [40] 41). Models of type (i) often require preprocessing dimensionality reduction to extract modest number
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

of interpretable features from high dimensional data, e.g., deriving features of anatomical regions from medical images, because computation time scales The preprocessing introduces uncertainty and is often computationally burdensome in itself.
with the number of features.
Here we introduce the variational event-based model (yEBM), 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, 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 their 'optimal' location in continuous latent permutation, unlocking benefits from recent advancements in the field of computational optimal transport 422 Our approach generalises discrete generative models
of disease progression, e.g.. [T 20 S 36 39 43 which it obtains as limit; and it directly infers
continuous probability over events, while the others require costly sampling methods. Crucially, 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 the sequence-
models proposed by [. These models underpin both a range of direct applications 20 5 23
24
12
21
25
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 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 most. In contrast, our new formulation enables this type of model 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., [4T]. However, with other deep these approaches require vast amounts of data with multiple time points 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.
state-space models, c.g., [44] 6
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 new generative latent variable model of disease progression, the variational event- 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 learn the vEBM efficiently from high dimensional data.
We synthetic data demonstrate that the vEBM achieves factor of 1000 faster inference than baselines, provides hetter inference accuracy, and is robust to noise.
We use the vEBM with data from Alzheimer's disease (AD) and age-related macular degeneration (AMD) 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 generative latent variable model of disease progression terms of latent permutation matrix of events (Section 2. Our key methodological contribution is reformulating the generative model in the context of optimal transport; we introduce the relevant mathematical tools do this 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

---

((: 0)
g (2)r) (21r) (log)(loo[[((
(L
P(Y
Event permutation matrix, S
Snapshots, Y
while not converged
1=1
1=2
1=1 I
1-1
4
)=3
P
1... k
Update veC(o);
Event sequence,
j=4
end
$=,4,1,2,3]
i=
=4
A
B
C
Figure I: Schematic of the variational event-based model for toy 4-feature dataset. A. The dataset contains snapshots from I individuals, with j. k = 2, 4} features and latent events; the features can be of any type and can be incomplete. B. Before inference, probabilistic models of normality and abnormality are fit to the dataset, giving the likelihood look-up tables P(Y ((PF) (Section 2.1): these are fixed throughout inference, as denoted by the inner box outside the training loop. To infer the permutation matrix S (Section 2.2). the ELBO (Section 2.3 is optimised and S is updated each iteration using the Sinkhorn- Kunkhom-KKoop algorithm (Section 2.4 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 event probabilities, which can be sampled to obtain uncertainty.
model in variational inference setting (Section 2.3), devise a suitable inference scheme (Section and provide a method for probabilistic individual-leve staging using the trained model (Section Figure provides schematic overview of the vEBM.
2.4h.
2.5
2.1
Model of disease progression
Consider generative latent variable model with observed data Y and latent variables Z {S, 0}, where E R latent permutation matrix of N events; E Z is the latent state of an individual; and 0 (0(1), ((2), 0(3), .... O(J)} are additional model parameters. We write the joint probability as a hierarchical Bayesian model using the chain rule (see Appendix Figure for the graphical model).
P(Z,Y) = P(Z) P(Z))
P(YIZ) = P(S) P(O) P(KSS) P(Y k.
Each element in the permutation matrix, S, defines an 'event', which corresponds becoming measurably abnormal with respect to a reference distribution. Following [I] 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, flj, standard deviation, o,. and mixture weights, Wj. so that
feature
P(Y,10999 ~ N oo... w.f.).
(2)
However any probabilistic characterisation that defines reference group anchor the progression is permissible. To enable 5 to be inferred using data from different individuals at single time-point ("snapshots"), make key assumptions: i) monotonic progression of events the group level; and i8) consistent event permutation across the whole population. We note that assumption ii) could be relaxed permit multiple
permutations (i.e., clusters) within the same population. Then for individuals = {1, 3,.., ..., }} with observed features = {1, 2, 3, .., J.. the model likelihood can be written (see Appendix A.2 for a full derivation).
P(Y|S;
I
Σ
P(k S) II P(Y. k, OP (
II
P(Y,,18, kn (o
(3)
J=1
j=k++1
3

---

Here E (erm(N) is discrete permutation of N events, corresponding to the hard permutation obtained from S (see the next section); and {x(s) = 6 .(j) U (G)) 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 B can be set equal and factorised, i... 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,) (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 the disease event sequence, 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 and S in model.
Optimal transport aims to identify the mass-cong 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 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 inferring latent permutations [47]. Here we are interested in learning a latent permutation matrix, S, with corresponding discrete permutation, 8, such that,
[1, i
v(i.)) E Z4,, St, 0, otherwise.
(4)
In the context of permutation matrices as couplings, S belongs to the Birkhoff polytope,
BN {S R7*" S.y 2 0. SS =1, S.,=11
St,
(5)
The Birkhoff-von Neumann theorem states that By 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[jdemand efficient algorithms to solve for S. Following [46), we use the Sinkhorn- algorithm with an entropy regularisation term, H(S) Ei, Si, (og(S)))))) as approximation to solving the optimal transport problem,
K(X/T) argmax(S, X)F + (S)..
(6)
SEBx
Here K(-) is the Sinkhorn-Knopp operator, which maps the positive orthant on to By by iteratively normalising rows and columns [49 [5]] X is the unnormalised assignment probability (transportation cost) matrix; and 7 is temperature parameter, analogous to the -eppraature-deeddentt softmax function for discrete categories [ST]. In our context, X corresponds the event likelihood distribu- tions given by E as such, we are looking to find the permutation matrix, S, that transports event probabilities
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 result from 47 who showed that M(), 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 Balkhorn-Knope operator. Cx(0))
M(X) = :
lim K(X/T).
(x(N)

---

Here en are basis (one-hot) vectors of size N with a value of in the n-th position and 0 everywhere else. In practice we compute the hard permutation matrix M (X) from S using the Hungarian 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 7 must be chosen to balance between the limit of a hard permutation, where the gradients are discontinuous (and hence non-differentiable) and uniform soft permutation, where the gradients are flat (and hence non-informative). A parametric analysis of T, Tpriors and the number of Sinkhorn-Knopp iterations, Ris, is presented in Appendix A.8
algorithm 52 53
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),, with matrix, e, of i.i.d. Gumbel noise,
G(X,7) K((X +
The Gumbel-Sinkhorn distribution effectively implements the reparameterisation trick 555 for permutations, and in the limit T + 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, Tprior), and for the posterior, T; with parameters Φ. We seek to optimise the corresponding ELBO,
lYg 2 [(0; 0) = En (z (oggPoY
(ZI)))
= log
= 0, T)
The Kullback-Leibe (KL) divergence term on the RHS of Equation 9 is intractable, but can be rewritten as (L c)/rllc/Tpmor) by substituting (X + c)/T for and estimated using random sampling 47
+ For completeness, we restate the full term derived by [4] in Appendix A.4
2.4 Inference scheme
To optimise the ELBO we use Adam 56 with nept 200 iterations and learning rate of 0.1. Temperature hyperparameters T and Tprior were set to for all experiments, except the mixed events EE3. 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
(Section 3.3.3. where
2.5 Probabilistic staging
We can use the trained model to obtain an individual-level likelihood distribution over stages, i.c., the likelihood at each state k given by Equation 3 where stage k == n corresponds to the first 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 [I: and the Alzheimer's Disease Probabilistic Cascades (ALPACA) model [33], both of which use maximum likelihood estimate the ordered sequence, s. The EBM learns s using gradient descent and Markov Chain Monte Carlo (MCMC) sampling. The ALPACA model instead defines as the central permutation of a Mallows model [57]. with density over permutations given by p(s) exp(-Xd( (d((s, so)), where À scales the spread around the central ordering, 8p, and d(xx, ao) MN |s() so(i) is the distance between permutations. The ALPACA model learns using expecaaiion-maximisttioon (EM) and Gibbs sampling. We use the default parameters; for the EBM, 10 gradient descent iterations with 10 initial seeds, and 106
MCMC samples; for ALPACA, 10 EM iterations and 100 Gibbs samples.
5

---

3.2 Synthetic data
To enable comparison between our model and the baselines, we simulate data generated by an ordered sequence, 8, according to the limit version of Equation 3 (see Appendix A.3) In brief, is randomly initialised and individuals are assigned a stage with uniform probability, reflecting that individuals can be observed any disease stage. Individuals are assigned either controls patients according 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 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
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 clearry illustrates the unique computational ability of the vEBM to work with much larger J. we demonstrate throughout this section. The vEBM is factor of 1000 times faster for the 200 experiment: this factor would only increase for larger models.
Figure shows the runtime for the vEBM and baselines for three experiments
3.2.2 Improved accuracy and robustness to noise
Figure 3 shows the effect of increasing aleatoric (measurement) noise levels on inference accu- racy, as measured by the Kendall's tau (58) be- tween the true and inferred sequences. The vEBM outperforms or is comparable the base- lines in all datasets and noise settings, except for / = 100, J 1000 and is expected, because at low noise and smaller numbers of features the EBM's MCMC sam- pling 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 shown for ALPACA at J 100, and none at J = 200, due to computational intractabil- ity). We highlight that the metric is sensitive to any departure from the correct ordering, even by single sequence position; accordingly the visual correlation between the true and inferred sequence remains high, e.g., even the highest noise (a 1, corresponding to a 1:1 sigsa ratio), as the relationship is still approximately diagonal. Additional examples in other datasets and when setting the Gumbel noise, e, non-zero are given in Appendix
VEBM (ours)
10a
EBM
Wall-clork time (S)
ALPACA
10a
0.1; this
104
10
10a
50
100
150
200
Number of features
Figure 2: Speed of inference as a function of model size, for the vEBM and baselines. Note that there is no datapoint for ALPACA at J = 200 due to computational intractability.
3.3 Alzheimer's disease data
We use pre-processed tensor-based morphometry (TBM) data from the Alzheimer's s Disease Neu- roimaging 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 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 Exam- ination (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

---

VEBM tours]
0.8
Kendail's tail
0.8
Kendallss tau
0.6
Kerdalls
0.6
0.2
ALPACA
0.4
VEBM (ours)
0.4
VEBM (ours)
0.1
ALPACA
0.2
ALPACA
100
150
200
150
200
100
150
200
Number of features
Number of features
Number of features
tau
112
Kendall's
VEBM
VEBM
VEBM
124
True
124
Inue
5
True
79
79-
111
169
169
Teauure
127
40
40
90
180
175
1
1-
102
0 20 40 60
80 180120140180180
0 20
40
60
80 1001100100001110
20 40
60
80 100120140160
Event
Event
Event
Figure 3: Accuracy of inference as function of aleatoric noise, obtained by the vEBM from synthetic data with low, medium, and high noise levels (left: 0.1; middle: 0.5; right: 1). Top row: Kendall's tau distance between the inferred and true sequences as a function of model size (number of features). Standard errors of the mean shown, from 10 repeats per experiment. Bottom row: example positional variance diagrams. The vertical axis lists the sequence of events inferred by the vEBM with the earliest event (order position 1) at the top. The true sequence is overlaid as red squares. Datasets have = 2000 individuals and J = 200 features.
(https://adni. loni. usc www..uu...uueeeeee wttpa//aaccaccceesastteeee access data collections: "TBM Jaco- bian 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 computational intractability (as demonstrated 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 small-scale insights into AI)
and white matter changes throughout the sequence, providing aetiology that have not previously been possible, which we explore further in the next section.
Figure 4: Pixel-level disease progression sequence in AD obtained by the vEBM. White pixels correspond to events that have occurred by the corresponding point of the sequence. The figure shows 10 sequence positions at uniform steps of 100 across the total of 1344, with the top left figure corresponding to position 50 (the first 50 events have occurred) and the bottom right to position 950. Images were made from the vEBM output using 3D Slicer (https: //www. slicer org/)

---

3.3.2 Degmentation-basee 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 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 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 TThalamus-Propee (N = 43), and many in "Cerebral-Core (N 299), corresponding to their relative sizes (areas) the 2D slice.
Our findings are in broad agreement with previous results; sub-cortical changes (Thalamus- Putamen. Hippocampus) are earliest, followed by cortical (Corebral-Cortex) and white matter (Cerebral- (Crrbbra-White- White-Matter) and finally ventricular change
VentralDC). However, our model provides much more fine- areeggr insights than, e.g.. [20], we now obtain continuous trajectories of change, which capture interesting non-linearities, e.g., in the Thaalamus- Proper, Brain-Stem, and Lateral- Ventricle; this contrasts with the more linear changes in the Hippocampus, Cerebral-Cortex, and Cerebral- W Chite-Mateer
Srath Stom
Cemabral Mateer
Carabral-CCn Co-too
Ventricle
ha Prooer
Putamen en
H #paocampus VentrelDC
Fraittinn occured
64
00-
2000 1200
Evert
Figure 5: Trajectories of regional brain areas 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 area at the corresponding event number, as defined by the vEBM event sequence.
3.3.3 Mixed feature disease progression events in AD
The vEBM is not limited 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 histogram, and shows the position of the cognitive events by vertical lines. We calculate the spatial distribution of pixel events according 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 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, expected 600, 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
pixel event topology obtained by the vEBM as 2D

---

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., [6]]
RAVIT
CDRSB
MMSE
30
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 the pixel event, and the horizontal axis shows the event ordering obtained by the vEBM. Note that the cognitive events, de- noted by vertical lines, are assigned an arbitrary distance of zero.
20
10
Distarcc
-10
200
400
600
800
1000
1200
Event
3.4 Age-related macular degeneration data
We pre-processed optical coherence tomography (OCT) data from the Duke University (DU) Ophthalmology 2013 dataset, a cross-sectiona. study of AMD (62 The OCT data represent thickness maps for retinal pigment epithelium and drusen complex (RPEDC), 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/ com/s/180 B/180j62iooee0 only pixels with disease signal, we remove pixels with an effect size less than 4 on pixel-wise t-test between the control and AMD groups.
To select
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 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: //wwp: slicer. wwicer.org//
3.5 Prediction of AD and AMD stage
Figure 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 individual- level 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.001
MCI
AMD
0.003
Denscy
Densiyy
0.002
0.023
0.015
0.004
0.913
0.002
0.005
0.007
0.007
230
1200
Stage
Stage
Figure 8: Individual stages obtained by the vEBM in AD (left) and AMD (right). CN: control; MCI: mild cognitive impairment; AD: diagnosed AD; AMD: diagnosed AMD.
4 Discussion
We introduced the vEBM, novel optimal transport formulation for discrete disease progression models that scales to hitherto impossible numbers of events. It allows pixel-level visualisation the order of pathology appearance in chronic disease, demonstrated for the first time in AD and AMD.
4.1 Limitations
Here we did not fully explore model uncertainty, which Bayesian model, the vEBM can estimate directly; we reserve this for future work. We acknowledge that a model ordering several thousand events not fully identifiable with only 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 feature-spassification (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 a common reference frame (or scale). We do not explicitly account for feature-wise covariance in the model, e.g., in image- based data would expect high degree of collinearity between neighbouring pixels; this could be addressed by including an additional term that imposes local structure, c.g., Markov random field (35). Finally, 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 scale in multiple areas of medical imaging, not only the modalities demonstrated in this paper: such other MRI modalities e.g., diffusion weighted imaging, microstructure modelling, connectivity; other imaging modalities, c.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 research labs and potentially clinies that have limited resources, while further minimising its carbon impact by reducing compute time. In addition, it provides 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 Loreá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 declare.
10

---

References
[1] H. M. Fonteijn, M. Modat, M. Clarkson, et al. An event- model for disease progression and its application in familial alzheimer's ' disease and huntington's 's disease. Neurolmage, 60: 1880-1889, 2012. doi: http:///diiorg111110000 neuroimageee200000 2012.01.000
[2] 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.
[3] A. M. Alaa, S. Hu, and M. van der Schaar. Learning from clinical judgments: Semi-markows modulated marked hawkes processes for risk prognosis. International Conference on Machine Learning. 19(4):1-62. 2017.
[4] A. M. Alaa and M. van der Schaar. A hidden absorbing semi-markon model for informatively censored temporal data: Learning and inference. Journal of Machine Leurning Research, 70: 60-69, 2018. doi: http://dx.doo www.koooo/oo/oooooooo wg//1..55553305550 S13033000311.335338 1.3305388.
[5] A. L. Young, R. V. Marinescu, N. P. Oxtoby, et al. Uncovering the heterogeneity and tem- poral complexity of neurodegeneraaive diseases with subtype and stage inference. Nature Communications, 9, 2018. doi: https://di..or//0/ http://0oo..r/11000000/4/1111110 1038//414677-088-----1 05892-0.
[6] A. M. Alaa and M. van der Schaar. Attentive state-space modeling of disease progression. Advances in Neural Information Processing Systems, 32, 2019.
[7] Z. Sun, S. Ghosh, Y. Li. et al. A probabilistic discase progression modeling approach and its application to integrated huntington's 's disease observational data. JAMIA Open, 2(1):123-130, 2019. doi: https://doi.org orgg 1tt//1.....9///// 3/jamaaappe/oooy060
[8] Williams, C. B. Storlie, T. M. Therneau, et bayesian approach to multistate hidden markov models: Application to dementia progression. Journal of the American Statistical Association, 115:16-3,, 2020. doi: https://doiorr//10 10800006114999209
[9] C. Hadjichrysanthou, S. Evans, Bajaj, et al. The dynamics of biomarkers across the clinical spectrum of alzheimer's disease. Alzheimers Res Ther, 12, 2020. doi: https://doi.org/100 10. 186/ s 15195-020-0036--2
[10] M. Donohue, H. Jacqmin-Gadda, M. Le Goff, et al. Estimating long-term multivariate progression from short-term data. Alzheimer's & Dementia, 10:5400-410, 2014. doi: https: //doi.org/10 wtt.oorr//.10// wttg//0..016////. jal/// 201...0.1.0..
[11] 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.
[12] N. P. Oxtoby, A. L. Young, D. M. Cash, et al. Data-driven models of dominandly-inherieed alzheimer's s disease progression. Brain, 141:1529 brain/awy050.
1145-29-444,, 1544, 2018. doi: https://doi.org/10. 1093/
[13] M. Lorenzi, M. Filippone, G. B. Frisoni, al. Probabilistic disease progression modeling to characterize diagnostic uncertainty: Application to staging and prediction in alzheimer's disease. Neurolmage, 190:56-66 68, 2019. doi:
[14] D. Li, S. Iddi, S. Aisen, et al. The relative efficiency of toproggopooo and continuous measures of cognition in presymptomatic alzheimer's disease. Alzheimer's and Dementia: Translational Research and Clinical Interventions, (016//.trci.20199 19.00491494.......
2019. doi: https://doi.org//0.
[15] M. Bilgel and B. M. Jedynak. Predicting time to dementia using quantitative template of disease progression. Alzheimers Dement, 11:205-215. 2019. doi: https://doi.rrg/0.110116 (itmm2209900.0000
[16] 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: http:://oooooo/11111116/11111000
11

---

[17] I. Koval, A. Bóne, M. Louis, et al. Ad course map charts alzheimer's 's disease progression. Sci Rep, 11:8020, 2021. doi: htpps//tb..rr/10000384411155..9 -21-8884-21--14111
[18] A. M. Staffaroni, M. Quintana, B. Wendelberger, et al. Temporal order of clinical and biomarker changes in familial frontotemporal dementia. Nature Medicine, 28:2194-22066 2022. doi: http://doi.org//00 100. 1038/441511-022-0194222
[19] 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/
[20] A. Young, N. P. Oxtoby, P. Daga, et al. A data-driven model of biomarker changes in sporadic alzheimer's 's disease. Brain, 137:2564-255 2577, 2014. doi: https://doi.org/ 10. 1093/brain/awul/ 176.
[21] N. C. Firth, C. M. Startin, R. Hithersay, et al. Aging related cognitive changes associated with alzheimer's 's disease down syndrome. Ann Clin Transl Neurol, 5:1665-1677, 2018. doi: https://doi.org/I0 http://oo...rr/11..0002/////... 1002/acn3.571.
[22] W. Vogel, L.. Young, N. Oxtoby, et al. Four distinct trajectories of deposition identified in alzheimer's s disease. Nature Medicine, 27:871-881, 2021. doi: https://doi.org//0. 098/s41591-021-013000
[23] L. M. Byrne, F. B. Rodrigues, K. Blennow, et al. Neurofilament light protein in blood as potential biomarker of neurodegeneration in huntington's 's disease: a retrospective cohort analysis. Lancer Neurol, 16:601-609, 2017. doi: https://doi.org/ 1tt//1...016.111 1016/S16/S 474-4442474432331122
[24] L. M. Byme, F. B. Rodrigues, E. B. Johnson, et al. Evaluation of mutant huntingtin and neurofilament proteins as potential markers huntington's 's disease. Sci Transl Med, 10: eaat7108. 2018. doi: https://doi.org//0. wtt..////.....//sitttataaaaa 126/scitranslmeed.at771
[25] A. Wijeratne, A. L. Young,
Oxtoby, et al. An image- masee model of brain volume biomarker changes in hungtington's 's disease. Ann Clin Transl Neurol, 3:570-582, 2018. doi: http://doi..rr//10 http://bo....r11.000000222..33.. 1002/acn3.558.
[26] P. A. Wijeratne, S. Garbarino, S. Gregory, et al. Revealing the timeline of structural mri changes in premanifest to manifest huntington disease. Neurol Genetics, 2021. doi: https: //doi.org/10 www.oor///.1....... 1...1...2222N 1000X000000000000000006777
[27] A. Eshaghi. R. V. Marinescu, A. L. Young, et al. Progression of regional grey matter atrophy in multiple sclerosis. Brain, 141 1665-6666167, 2018. doi: https://oo..rg//10.0999/// 1093/brain//wwy088..
[28] 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.rr//I 1038/s41467-001--22666
[29] N. P. Oxtoby, I., A. Leyland, I. M. Aksman, al. Sequence of clinical and neurodegeneration events in parkinson's 's disease progression. Brain, 1449975-888, 2021. doi: https://doi.org/10 1093/brain/awaae61..
[30] R. Pascuzzo, N. P. Oxtoby, A. L. Young, et al. Prion propagation estimated from brain diffusion mri is subtype dependent in sporadic creutzfeldt-jakob discase. Acta Neuropathologica, 140: 169-181, 2020, doi: https://doi.or// 10. 1000000000001--000022-6-----
[31] 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// 1002//cnn. 5 51035.
[32] A. I.. Young, F. J. Bragman, B. Rangelov, et al. Disease progression modeling in chronic obstructive pulmonary disease. Am Respir Crit Care Med, 201 201:294-300 2020. doi: https: //doi.org/10.11 2014/rcm.20190--1666
[33] Huang and Alexander. Probabilistic event cascades for alzheimer's disease. Advances in Neural Information Processing Systems, 25, 2012.
12

---

[34] X. Wang, D. Sontag, and F. Wang. Unsupervised learning of disease progression models. Proceedings of the 20th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, 2014. doi: http://dxx wwt.d....ggg1111155522203000001.
[35] R. V. Marinescu, A. Eshaghi, M. Lorenzi, et al. Dive: A spatiotemporal progression model of brain pathology in neurodegenerative disorders. Neurolmage, 192:166- 192:166-177, 177, 2019. doi: https://doi.org/100 http://dii..r/1111111/////////// www.oimagg..219900222.
[36] 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 wtt.d//i.rrr//133/9/f///ee w..38///fri.22222
[37] Wijeratne and D. C. Alexander. Learning transition times in event sequences: The temporal event- event-based model of disease progression. IPMI LNCS, 12729:583-599 2021. doi: https://doi.org/100 1000//990000077---1-
1007/978- 0000-7891--00 45.
[38] A. Wijeratne, A. Eshaghi, W. Scotton, et al. The temporal event-based model: Learning event timelines in progressive diseases. Imaging Neuroscience, 1:1-19, 2023, doi: hitps: //doi.org/10. 1122/m22maammmmm 0000010.
[39] 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..rr//10 10100/ 978-3-031-34048-221
[40] G. Marti-Juan, M. Lorenzi, and G. Piella. Mc-rvae: Multi-channel recurrent variational autoencoder for multimodal alzheimer's 's disease progression modelling. Neurolmage, 268, 2023. doi: https://do..or wrg/oioor///////// 1tgg/1..11166 wee.:ooo.e..uunoiaaeee..... neuroimage.2023. 198922.
[41] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. International Conference on Learning Representations, 2023. doi: https://doi.or//10..44 arXiv.2312.00752
[42] G. Peyré and M. Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11 11:5556077, 2019. doi: https://doi.opg http://dbo.or//10045555/////aa 00.48550/arXiv. 1803.vv.110000066.
[43] C. S. Parker, N. P. Oxtoby, A. L. Young, et al. Parsimonious ebm: generalising the event- based model of disease progression for simultaneous events. bioRxiv, 2024. doi: https: //bbi.or//10.111//02......
[44] R. G. Krishnan, U. Shaltit, and D. Sontag. Structured inference networks for nonlinear state space models. AAAI, 28:2101-2109, 2017.
[45] C. Villani. Optimal transport, old and new. Springer Verlag, 338, 2009.
[46] 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. 10.48550/arXiv. 20t.://..iv.131.89999909 1306.0895.
[47] G. Mena, D. Belanger, S. Linderman, et al. Learning latent permutations with gumbel-sinkhom networks. International Conference on Learning Representations, 2018. doi: https://doi.org// 48550/arXiv. 180.008665.
[48] G. Birkhoff. Tres observaciones sobre el algebra lineal. Univ. Nac. Tucumán, pages 147-151, 1946.
[49] R. Sinkhorn. A relationship between arbitrary positive matrices and doubly stochastic matrices. Annals of mathematical statistics, 35:876-888 -799, 1964.
[50] R. Sinkhorn and Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21:343-348, 1967.
[51] C. J. Maddison, A. Mnihl, and Y. W. Teh. The concrete distributions: a continuous relaxation of discrete random variables. International Conference on Learning Representations, 2016. doi: http:://kb..og//11888aaaavvvv...
13

---

[52] H. W. Kuhn. The hungarian method for the assignment problem. Navel Research Logistics, 2: 83-97. 1955.
[53] Munkres. Algorithms for the assignment and transportation problems. Journal of the society for industrial and applied mathematics, 5:32-38, 1957.
[54] 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//0. 088000661455..017..00
[55] D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv, 2013. doi: https: //doi.org/10.44. 48550//rrxvvv 13 14.
[56] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. International Conference 1412.6980.
on Learning Representations, 2014. doi:
[57] C. Mallows. Non-null ranking moidels. Biometrika, 44:114-130, 1957.
[58] M. Kendall. A new measure rank correlation. Biometrika, 30:81-89, 1938. doi: https: //doi.org/10. 1ooo.oog/10.103//ioome///// 10093/biomet33011-2.1.1.
[59] X. Hua, D. P. Hilbar, C. R. K. Ching, al. Unbiased tensor-based morphometry: Improved robustness and sample size estimates for alzheimer's 's disease clinical trials. Neuroimage, 66: 648-661, 2013, doi: https://doi.org// 10.. 1016/ nwuroimage.201 201.1mm.....22.0.00066
[60] B. Dubois, H. Hampel, H. H. Feldman, natural history, and diagnostic criteria. Alzheimers Dement, 12(3):292-323, 2016. doi: http: //dx.doi.org/10.
Preclinical alzheimer's disease: definition,
[61] A. Low, E. Mak, M. Malpetti, et Asymmetrical atrophy of thalamic subnuclei in alzheimer's discase and amyloid sooloid--ositvve mild cognitive impairment is associated with key clinical features. Alzheimer's 's & Dementia: Diagnosis, Assessment & Disease Monitoring, 11:690-699 2019. doi: https://doi.org http://dooorr//11.. 1tgg/10.100 wtwp///..dd...111.........
[62] S. Farsiu, S. J. Chiu, V. O'Connell, al. Quantitative classification of eyes with and without intermediate age-related macular degeneration using optical coherence tomography. Ophthalmology, 121:162-172, 2014. doi: htps://doi.orr//00 o0t66...pttha.2113..7...
[63] E. Dorsey, Venuto, Venkataraman, al. Novel methods and technologies for 21st- century clinical trials. JAMA Neurol, 72(5):582-55) -88,, 2015. doi: jumaneurol.2014.4555
1001/
[64] M. Cuturi, M. Klein, and P. Ablin. Monge, bregman and occam: Interpretable optimal transport in high- high-dimmnsioos with feature-sparse maps. arXiv, 2024. doi: https://doi.org/ 00.55/// 2302.04065.
14

---

A Appendix supplemental material
A.1 Compute information
The analyses presented here were performed either on laptop PC with a single AMD Ryzen PRO 6860Z CPU with 32GB RAM, desktop PC with 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 I 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 not included in this paper. Python code to reproduce all results presented here is available from the first author's GitHub repository
A.2 Derivation of vEBM
Starting from the joint probability (Equation I measured features j = (1,2,3, ..., J} to write the model likelihood,
we use the assumption of independence between
P(Y S,0) II P(Y,) O.n),
(10)
Using the chain rule, the joint probability distribution over Y and k can be factorised as,
P(Y,, (st)) f=1
II
(1I)
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),
J
Av
II
S) II P(Y,
0P 2o () II P(Y. ky G.))
(12)
j=
Substituting Equation
12 into Equation I1 and marginalising over hi we have,
J
P(Y|S,0) Σ P(k(S, x) II P(Y, OP u.) 8)
II (2)) (13) j=k,11
Here chose the prior over latent stages, parametrised by hidden variable to be uniform and constant, (kk, (S; T) Unif(0, k). Finally, assuming independence between measurements from different individuals i, we write the following expression for the total likelihood,
N Σ
R n P(KSS
P( = II
II 1k S) II
(14)
j=1
j=ke+1
Bayes' theorem can now be used to obtain the posterior over S.
A.3 Posterior limit
We the limit relation in Equation 7 to reparametrise the model likelihood (Equation 3 in terms of a discrete permutation, 8,
P) lim S(
(15)
15

---

$
Figure 9: Graphical model of the variational event-based model (vEBM). Hidden variables are denoted by circles and observations by squares. S: event permutation matrix; x: initial probability vector (fixed to uniform distribution); k; disease state for individual i; 0z: distribution parameters for biomarker j; Yı,; observed data.
0j
Yij
To see this, we take the first likelihood term in the RHS of Equation 3
and write it in matrix form,
T PPYo,ol (())
p(Yo.>l0* ((1))
p(Y (,))))))))
p(Yr,ol00 (())))
PVVi,,10M ((1))
p(Yo,ol0))
P(Y.JJ103) 103)
(e(Co)
(16)
p(Yr,o00))
p(V))119))
(a(J)
P(Y))103) M(X)T
P(XG))07) lim ((X/T)*. 7-+
where we have used the limit relation from Equation 7 between lines three and four. The same steps can be applied the second likelihood term on RHS of Equation P(Y,s (, to obtain the full likelihood reparametrisation in s.
A.4 Kullback-Leibler divergence
For completeness, restate the KI. divergence term derived by [47
KL(X + +e/ ( N2 (log(T//Tpin) 1 (Tprior/ 1)) Sp + S2 + Tprior/
(17)
where S1 Dig and S2 Eijexp(- 7pipi /T. Full 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 I0 and I show positional variance diagrams for datasets with I = 100, J 10, and 1000,
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, 6, set non-zero. The corresponding uncertainty obtained from 1000 random samples of the posterior is shown by greyscale shading the positional variance diagrams.
16

---

Algorithm 1: Pscudo-code of the variational event-based model (vEBM) inference scheme. For the experiments this paper use n2 20 Sinkhorn- Knopp iterations, unless otherwise stated.
Input :Y, T, "priors Mopt, Rs. 6, learning rate
Output:
// fit mixture models Compute P(Y - EM(Y, (P:));
// infer permutation matrix Initialise S: for Ropt do
Sample e; //// matrix Gumbel noise
for iterations do Update + ( + e/T):// Sinkhorn-Knopp algorithm end
Compute (Zlr))oooPo KL(g( /
// ELBO
Update VoL(o);
end
// compute sequence
Compute + S [0, 1, N
Kendall's
Kendal's
Kendall's tau 0.23
VEBM
Testure
Fruturre
3
Event
Event
Figure 10: Example positional variance diagrams obtained by the vEBM from synthetic data with low, medium, and high noise levels (left: = 0.1; middle: = 0.5; right: = 1). The vertical axis lists the sequence of events inferred by the vEBM with the earliest event (order position the top. The true sequence is overlaid as red squares, Datasets have = 100 individuals and J 10 features.
A.8 Hyperparameter study
Tables 123 show evaluation metrics for varying combinations of the temperature hyperparameter, and its prior, Tprior Table shows the same metrics for constant T, Tprior and extreme values for the number of Baikhorn-Knoph iterations, n4, Synthetic data with noise, 0.5, were used.
Table hyperparameter study, for Tprior 1; learning rate 0.1: n, 10; Rop 100.
0.1
10.0
J
100 10
1000 100
2000
200
100 10
1000 100
2000 200
Kendall's tau
0.69
0.01
0.05
1.0
0.81
0.54
Frac. correct
0.8
0.06
0.03
1.0
0.84
0.62
17

---

Kendell's tau 0.93
Kendallss tau = 0.7
Ke eddil's tou 0.0
VEBM
VEBM
VEB
35
11
47
05
49
Foature
79-
84
Figure 11: Example positional variance diagrams obtained by the vEBM from synthetic data with low, medium, and high noise levels (left: o = 0.1: middle: o 0.5; right: o = 1). The vertical axis lists the sequence of events inferred by the vEBM with the earliest event (order position 1) the top. The true sequence is overlaid as red squares. Datasets have = 1000 individuals and 100 features.
Kandall's
VEBM
VEBM
vEBM
True
True
Tue
Reature
3-
Event
Event
Event
Example positional variance diagrams obtained by the vEBM from synthetic data with low, medium, and high noise levels (left: 0.1: middle: 0.3; right: 1). The vertical axis lists the sequence of events inferred by the vEBM with the earliest event (order position 1) at the top. The matrix shows uncertainty in the ordering: dark squares the diagonal indicate high certainty of event position; lighter colors and off-diagonal squares indicate uncertainty in the event position. The true sequence is overlaid red squares. Datasets have = 100 individuals and J 10 features.
Figure 12:
Table 2: 7 hyperparameter study, for Tprior = 0.1; learning rate 0.1; nt = 10; Mop = 100.
0.1
10.0
I J
100 x 10
1000
100
2000
200
100 x 10
1000 100
2000 200
Kendall's tau
1.0
0.87
0.5
1.00
0.79
0.14
Frac. correct
1.0
0.94
0.54
1.00
0.8
0.24
Table 3: hyperparameter study, for Tprior
10.0; learning rate 0.1: n,
10; Repp 100.
T = 0.1
10.0
T J
100 x 10
1000
100
2000
200
100 x 10
1000 100
2000 200
Kendall's tau
0.07
0.1
0.01
0.69
-0.09
0.08
Frac. correct
0.1
0.0
0.01
0.8
0.06
0.04
18

---

Kendall's tau 0.72
Kendall's tau 0.66
Kendall's tau 0.05
83
70
VEEM
VEBM
VEEM
11
49
49-
Feature
34-
79
63
84
Figure 13: Example positional variance diagrams obtained by the vEBM from synthetic data with low, medium, and high noise levels (left: = 0.1; middle: a = 0.5; right: a = 1). The vertical axis lists the sequence of events inferred by the vEBM with the earliest event (order position 1) the top. The matrix shows uncertainty in the ordering: dark squares on the diagonal indicate high certainty of event position; lighter colors and off-diagonal squares indicate uncertainty in the event position. The true sequence is overlaid as red squares. Datasets have = 1000 individuals and J = 100 features.
tau
Kendallss tau
Kandall's tau
VEBM
VEBM True
YEBM
33
Tuu
65
115
Tue
Feature
2
Festure
2
Feature
111
32
120
148
108
175
192
130
1001201401880
280
14001201401880 160 180 Event
Event
Event
Figure 14: Example positional variance diagrams obtained by the vEBM from synthetic data with low, medium, and high noise levels (left: 0.1; middle: 0.5; right: 1). The vertical axis lists the sequence of events inferred by the vEBM with the earliest event (order position 1) at the top. The matrix shows uncertainty in the ordering: dark squares the diagonal indicate high certainty of event position; lighter colors and off- di-gdiagonal squares indicate uncertainty in the event position. The true sequence is overlaid red squares. Datasets have 2000 individuals and J 200 features.
Table 4: n, hyperparameter study, for T, Tprior 1.0; learning rate 0.1; nl.
10; Nopt 100.
100
Ix J
100 x 10
1000
100
2000
200
100 10
1000
100
2000 200
Kendall's
0.07
0.16
0.07
1.0
0.79
0.94
Frac. correct 0.4
0.19
0.12
1.0
0.83
0.95
19

---

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 that our method is least factor of 1000 times faster than the baselines, and in Figures 6 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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We explicitly discuss limitations of our method in the Section
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 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 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 few datasets with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers penalize honesty concerning limitations.
will be specifically instructed
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?
20

---

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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: Most of the results can be reproduced exactly using the code provided (upon publication) and either synthetic data 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 thttps://adhe 1oni. usc. aadu/data-saaae saaetta-saampeeeeeceese----s data/ Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, 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 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 novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions 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) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or 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 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 (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
21

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Complete details 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 NeurlPS code and data submission guidelines (https://niis.co cc/ public//uides//ooeeeeeee http://gudeeeee/eeeeemmaaaaaas.a
http:eeuubbsssiooooooocoo 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 (c.g., for a new open-sourge benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
ctipss.c//ublii//uuie guwdes/ CodeeSumiiiseeeeee onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, 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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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.
22

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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 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 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 figures symmetric error bars that would yield results that are out range (e.g. negative 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 compute required for each of the individual experimental runs as well 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 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 http::////earpppsscccppp//////aa
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations 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 work discussed in Section 4.2
Guidelines:
23

---

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 considcrations.
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 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 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 system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: The software presented in this paper could be misused in medical setting, which is not recommended as the software has not been designed to be 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 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 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 encourage authors to take this into account and make 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 used and, possible, include URL.
24

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.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, papptsssspppsiiiipp.... thcode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 encouraged to reach out to the asset's 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include 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 sereenshots, 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 research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, 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 equivalent approval/review based 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
25

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing 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 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 the institution conducting the review.
26


Reviews:
[{"review_id": "h2OOqYcw7v", "reviewer": "Reviewer_wy4q", "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.", "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.\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?", "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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Results clearly demonstrate that their method is scalable, and it's effective in simulation where the true data match their model. The writing style is clear, and the visualizations are excellent, particularly Figures 4 and 6. Experimental results are interesting and clearly presented.", "weaknesses": "- The related work section is very brief, which made me question whether it is comprehensive.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Weaknesses point 1 (“The related work section is very brief…”)\n\nThis section was deliberately kept brief to focus on the most relevant comparable methods; specifically those that can infer group-level sequences from cross-sectional data. While there are many models that use longitudinal data (see, e.g., reference [19] for a comprehensive review), they are not as relevant for direct comparison. We also note that we do discuss the relative pros and cons of alternative types of models (including longitudinal models) in the Introduction (L30-36). However we appreciate that this is not made completely clear in the Related Work section, so we will add a couple of sentences to clarify this point.\n\nWeaknesses point 2 (“I found the second half of section 2.2 very difficult to follow…”)\n\nThe transport cost matrix, X, defines the likelihood of an event occurring and the permutation matrix, S, defines the event ordering; we seek to find the optimal event permutation that maximises the overall model likelihood. As such we can think of the relationship between S and X as their (element-wise) product giving the likelihood of a given event permutation. 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. We will add clarification to the text, if the Reviewer believes these descriptions are helpful.\n\nWeaknesses point 3 (“I am not convinced this is solving a real problem...”)\n\nHere we use pixel-level regions just as a demonstration of the model’s ability to handle high dimensional data. However the model also enables, e.g., fine-grained regional brain volume analysis of the event ordering with 100s of regions, which current methods cannot cope with.\n\nThe class of models we focus on here are necessarily group level models rather than models of individuals, although they do capture variability over the group (but we do not leverage this information here). As we highlight in the Introduction (L25-29), these models have proved extremely important in understanding the temporal evolution of chronic disease and guiding treatment development. They are not just, or even primarily, a device for prediction at the individual level. This class of technique has proved very influential over the last 10 years; in the paper we reference 15 examples of applications in various diseases, and 4 examples of direct methodological developments on the original model, but there are many more applications and methods developments that have been published. The method we provide in our paper here will \"turbocharge\" all the current models of this class, including the highly popular SuStaIn model (reference [5]), and unlock pixel / voxel scale analyses that were previously intractable.\n\nQuestions point 1 (“I am confused about how the normal and abnormal distributions…”)\n\nThese distributions are defined by separating synthetically generated individuals into two groups, based on their randomly generated stage (the lowest 20% of stages are considered controls, see L194-195). Mixture models are then fitted to these groups, and are fixed throughout inference, as denoted in Figure 1. Equation 3 doesn’t explicitly distinguish between these fitted parameters (“\\theta”) and the learned parameter “S”, because in principle they could be learned jointly; however if the Reviewer thinks it is helpful, we can clarify the difference by using semicolon notation for \"\\theta\" in the likelihood equations.\n\nQuestions point 2 (“What is denoted by the color of the pixels in Figure 5?”)\n\nThe color 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. We will add a color bar and a sentence in the figure caption to clarify this.\n\nQuestions point 3 (“The authors mention that the method requires image registration…”)\n\nWe deliberately avoided going into detail regarding the registration, because a) we used a pre-processed data collection, produced using Tensor Based Morphometry, from the ADNI dataset, which has previously been described in detail (see reference [59]); and b) as we note in the Limitations (L292-295), some sort of pre-processing is always necessary for any method to enable comparison between individuals, so it isn’t a special limitation of our method. \n\nRegarding the point about aligning images from different stages of progression – if we understand correctly, this is actually the primary utility of disease progression modelling, which can be thought of as a type of temporal registration of images generated by latent stages to a common temporal reference frame. Our method captures this variability and learns this temporal reference frame in terms of a sequence of events. However if the Reviewer is referring to the spatial registration of images with respect to a common spatial reference template, the primary utility of Tensor Based Morphometry is to provide a method to map to a common reference template and calculate the voxel-level volumetric change with respect to this template, thus transforming all images from different stages to a common spatial reference frame.\n\nQuestions point 4 (“Do the authors envisage applications of this method outside medicine?”)\n\nYes, we envisage many applications in areas that involve learning progressive sequences of events, e.g., in environmental modelling, our method can infer trajectories of biodiversity loss “events”, from temporal eco-acoustic monitoring data. We plan to make the vEBM code available with the aim of opening up such applications by researchers in other fields.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tlfy1roUo1", "reviewer": "Reviewer_SCze", "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.", "questions": "My questions were combined with weaknesses.", "limitations": "The authors adequately addressed the limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. I am not sure about the results in Section 3.3.2. You see that CDR58, MMSE, and RAVLT are pretty far in the Event order. If you check the Temporal Event-Based Model (TEBM) (Wijeratne et al., 2021; Wijeratne et al., 2023), these cognitive test scores can be found earlier in the disease timeline. However, there are differences; they used T1 images, and in your case, you used TBM images. Before making new claims, it will be essential to consider experiments with a previously explored set of features. Importantly, the proposed solution is supposed to solve the scaling question. Hence, it should work with the previous set of features.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Weaknesses point 1 (“I am not sure about the results in Section 3.3.2…”)\n\nGenerally we expect cognitive test scores to appear later than structural changes in MRI - cognitive deficit is a consequence of loss of brain tissue. As stated in L257-260, the cognitive events occur across the latter 2/3rds of the event sequence; this is broadly consistent with the results in Wijeratne et al. 2023 and Wijeratne & Alexander 2021. \nThere are differences to the literature (e.g., RAVLT occurs late in our sequence, while in Wijeratne & Alexander 2021 it occurs mid-sequence). However, the cohorts used are different – we use a subset of ADNI individuals with TBM data, while e.g., Wijeratne & Alexander 2021 use a different subset, and Wijeratne et al. 2023 use the entire ADNI cohort. Inter-group heterogeneity in Alzheimer’s disease is well reported (e.g., reference [5]) and differences in the ordering of changes are expected between different subsets of the ADNI cohort. This would make for a very interesting future research direction for our model, to investigate image-based heterogeneity at the pixel-level - instead of just regional level that previous analyses have been limited to.\n\nWeaknesses point 2 (“The authors enable models that express progression at the pixel level compared to the previous region level…”)\n\nIncluding comparison to regional volumes is an interesting idea, and one which we would have included given the time. We have now included an additional analysis in the “Author Rebuttal” PDF, where we used the FreeSurfer segmentation tool to obtain pixel-level anatomical labels, which we mapped to our “pixel” events to enable interpretation via a subset of regional volumes  (Figure 1 in the uploaded PDF).\n\nIndeed what we see is broadly what we’d expect from 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 – 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. These are entirely new insights provided by the model that could be used to guide which regions are most useful as biomarkers at which stages in the disease; and also provide new insights into the underlying atrophy dynamics, which could inform mechanistic models of disease spread in the brain. We believe this result substantially improves the interpretation of our method and will add it to the paper.\n\nNote that to enable direct comparison with previous published results we would need to run the same segmentation tool that they used, e.g., the GIF segmentation tool used in Wijeratne et al. 2023 and Wijeratne & Alexander 2021. This would make an interesting additional analysis that we reserve for future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SHSiBtXIyA", "reviewer": "Reviewer_Vfdn", "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.", "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.\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?", "limitations": "Yes, the authors adequately addressed the limitations and, if applicable, potential negative societal impact of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. Correct me if I am wrong: the proposed method seems to be a method that models the disease progression on a population level. This might limit the method to population level studies for disease research purposes and render it unsuitable for predicting individual-level progressions which could facilitate personalized treatment plans.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Weaknesses point 1 (“Correct me if I am wrong…”) \n\nYes, that is exactly what our model does. It is less of a device for individual level predictions (although can be used that way; see Figure 7, where we use it to estimate individual-level stages along the group-level sequence) and more to elucidate group level patterns of change that aid disease understanding, biomarker discovery, structure heterogeneity, and inspire new treatment development. As we note in the Introduction (L25-29), this class of technique has proved very influential over the last 10 years; in the paper we reference 15 examples of applications in various diseases, and 4 examples of direct methodological developments on the original model, but there are many more applications and methods developments that have been published.\n\nWeaknesses point 2 (“While producing pixel-level disease progression sequences…”)\n\nYes that's true. A separate grouping procedure would be beneficial and something we conserve for future work here now we have the computational framework in place. As noted in Limitations (L295-298), a first approach would follow reference [35] and use a Markov random field to impose local structure. \n\nWeaknesses point 3 (“For the results in Figure 7…”)\n\nThe main point of these plots is to show the individual-level utility of the method, in terms of the distributions of stages for each group. But we are happy to discuss ideas to make this more quantitative; does the Reviewer have anything specific in mind?\n\nWeaknesses points 4 & 5 \n\nWe have deliberately made Figure 1 greyscale for printing purposes & accessibility, if that is what the Reviewer is suggesting regarding aesthetics. We will add a color bar to Figure 5.\n\nQuestions 1 (“May I ask the authors to explain…”)\n\nThe positional variance diagram shows the most likely ordering of events (horizontal axis), as characterised by changes in the features (vertical axis). The model also estimates the uncertainty in the positioning of the events, shown by the shaded heatmap in the Supplementary Figures 11-13. Typically these diagrams are used as the main visualisation of event-based model sequences, but due to the large numbers of features our method enables, they become sub-optimal for visualisation purposes. However the main purpose of including them in Figure 3 was to visually demonstrate the agreement between the vEBM inferred sequence and the true sequence, using synthetic data. In real data, e.g., the ADNI experiments, it makes more sense to visualise the results in pixel space, to preserve the spatial information more clearly.\n\nQuestions 2 (“Would the authors consider comparing with alternative deep-learning-based models…”)\n\nThese model classes make sense for making predictions, but they lack the interpretability of the EBM class of model, which is a key output. Additionally, they typically require longitudinal data, while the EBM class models operate using only cross-sectional data; thus limiting ODE-type models to using only cross-sectional data would be an unfair comparison.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "B3XM34uqkc", "reviewer": "Reviewer_5r2K", "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.", "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.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* Interesting experimental results showing wall-clock improvement on synthetic dataset\n* 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.\n* 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. \n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Weaknesses point 1 (“On the Pixel-level experiments…”)\n\nYes, it is a good idea to compare with regional-level results, and one we would have included given the time. We have now included an additional analysis in the “Author Rebuttal” PDF, where we used the FreeSurfer segmentation tool to obtain pixel-level anatomical labels, which we mapped to our “pixel” events to enable comparison with regional volumes (Figure 1 in the uploaded PDF).\n\nIndeed what we see is broadly what we’d expect from 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 – 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. These are entirely new insights provided by our model that could be used to guide which regions are most useful as biomarkers at which stages in the disease; and also provide new insights into the underlying atrophy dynamics, which could inform mechanistic models of disease spread in the brain. We believe this result substantially improves the interpretation of our method and will add it to the paper.\n\nNote that to enable direct comparison with previous published results we would need to run the same segmentation tool that they used, e.g., the GIF segmentation tool used in Wijeratne et al. 2023 and Wijeratne & Alexander 2021. This would make an interesting additional analysis that we reserve for future work.\n\nWeaknesses point 2 (“Overall, this work gives the audience…”)\n\nIt is a fair point that a key component of the methodological improvement is the application of the Sinkhorn algorithm. However, the reframing of latent variable disease progression models as an optimal transport problem is entirely novel, and has implications beyond just improving the speed of inference (which is what we focused on this paper). Indeed one of the other Reviewers (Vfdn) notes this as one of the paper's strengths: \"It is quite innovative to view the disease progression modeling task from the optimal transport perspective.\".\n\nThe key innovation from an optimisation perspective is moving from a discrete sequence of events to a continuous probability density matrix over permutations – it is this that allows us to unlock gradient-based optimisation via variational inference. Furthermore - and which we don’t explore in depth in the main paper but do provide some initial analysis of in the supplementary - this formulation gives us direct access to the distribution over event permutations (as opposed to having to sample it using MCMC). This allows for direct sampling of model uncertainty in our Bayesian formulation. \n\nBut on a more fundamental perspective, one can think of modelling disease as morphing healthy probability distributions into unhealthy ones, (e.g., the diffusion of pathogens spreading in the brain), which can be ideally solved using the tools of optimal transport. For example, one could imagine many types of state-space-based latent variable models that would work with our permutation-based optimal transport coupling, e.g., Hidden Markov models. We’re excited about exploring the many recent developments in optimal transport theory from this perspective in the future.\n\nWeaknesses point 3 (“Practicality of method in real life problems…”)\n\nWe need to point out that this criticism is not quite fair – while the simulations show 200 features (a decision we made to facilitate visualisation in the positional variance diagrams shown in Figure 3), the ADNI analysis used 1344 features (as noted in the Figure 4 caption). In principle we could include 10s of thousands of pixels / voxels, the main limitation being identifiability and computer memory, as noted in the Limitations section. However we appreciate that this detail is not easy to find in the text, so we have moved the information about the number of pixel features used in each analysis to the main text in Sections 3.3 and 3.4.\n\nQuestions (“In both figure 4, and figure 6, …”)\n\nYes this is not very clear – we will clarify by replacing these sentences with: “White pixels correspond to events that have occurred by the corresponding point of the sequence.”.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their constructive comments, which we have endeavored to address as faithfully as possible in our responses below.

As part of our response, please find attached a PDF containing a new result.

We look forward to continuing the constructive discussion!

Meta-review: 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.
Decision: Accept (poster)