Learning Riemannian metric for disease progression modeling

Abstract: Linear mixed-effect models provide a natural baseline for estimating disease progression using longitudinal data. They provide interpretable models at the cost of modeling assumptions on the progression profiles and their variability across subjects. A significant improvement is to embed the data in a Riemannian manifold and learn patient-specific trajectories distributed around a central geodesic. A few interpretable parameters characterize subject trajectories at the cost of a prior choice of the metric, which determines the shape of the trajectories. We extend this approach by learning the metric from the data allowing more flexibility while keeping the interpretability. Specifically, we learn the metric as the push-forward of the Euclidean metric by a diffeomorphism. This diffeomorphism is estimated iteratively as the composition of radial basis functions belonging to a reproducible kernel Hilbert space. The metric update allows us to improve the forecasting of imaging and clinical biomarkers in the Alzheimer’s Disease Neuroimaging Initiative (ADNI) cohort. Our results compare favorably to the 56 methods benchmarked in the TADPOLE challenge.


Reviews and author responses:

This paper proposes a mixed effects model for modeling disease trajectories. The method builds on prior work that considers geometric models (which represent trajectories as geodesics on prespecified Riemannian manifolds), by allowing for the Riemannian metric to be learned from data. The proposed method combines desirable aspects of different disease trajectory modeling approaches: interpretable individual-specific and population specific parameters governing space and time variability and a metric that can be learned from data. On synthetic experiments the authors study learning properties of the proposed method. Experiments on an Alzheimer's dataset challenge show that the method has strong and competitive performances (in addition to desirable interpretable properties from the mixed effects modeling structure).  While I have experience with time-series modeling of disease trajectories, unfortunately, I must admit that I am not familiar with the literature on geometric trajectory models. As such, I cannot speak much to the novelty or technical soundness of the proposed method.

That said, the problem of disease progression modeling is important as motivated by the authors in the introduction, and these methods are very useful for, e.g., brain imaging data. They do a good job of identifying desirable features of disease progression models and then walk through the development of their modeling approach. The empirical results on the TADPOLE challenge appear to be good evidence of the strong performance of the proposed approach relative to existing alternatives.

I have a few suggestions to make the paper more accessible to a broader audience:
- The introduction reviews various strategies for disease progression modeling, but it does not clearly draw attention to the problems that this paper is intended to address until Line 67. I think streamlining the introduction to be more focused on the challenges addressed by and the benefits of the proposed approach will help readers more quickly understand what they will find in this paper.
- I think this paper would benefit from a dedicated and broad related work section. The introduction provides a somewhat narrow (and hasty) review of some methodology for disease trajectory modeling. A dedicated related work section would help a broader set of readers understand where the proposed method (and related works like [6,9]) fits in the landscape of work on disease trajectory modeling.

For example, I am familiar with work like the following on modeling longitudinal disease trajectory data:
- Schulam, P., & Arora, R. (2016). Disease Trajectory Maps. Advances in Neural Information Processing Systems, 29, 4709-4717.
- Schulam, P., Wigley, F., & Saria, S. (2015, February). Clustering longitudinal clinical marker trajectories from electronic health data: Applications to phenotyping and endotype discovery. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 29, No. 1).
- Proust-Lima, C., Séne, M., Taylor, J. M., & Jacqmin-Gadda, H. (2014). Joint latent class models for longitudinal and time-to-event data: A review. Statistical methods in medical research, 23(1), 74-90.
- Lasko, T. A., Denny, J. C., & Levy, M. A. (2013). Computational phenotype discovery using unsupervised feature learning over noisy, sparse, and irregular clinical data. PloS one, 8(6), e66341.


UPDATE:
Thanks the authors for their response. I will be maintaining my score. I think this paper will be greatly improved by broadening the perspective in the introduction, including a dedicated related work section, and discussing some of the other competitor methods that were used in the TADPOLE challenge to provide more context for readers. Reviewer w1hQ has suggestions that I also believe would improve the paper. Yes.

This paper presents a method for disease trajectory modeling that embeds patient trajectories in a Riemannian manifold to allow for more flexible models than simpler linear mixed effects models.   The primary methodological contribution appears to be a technique for learning the Riemannian metric from data, rather than using a fixed prior choice of a metric. This technique is applied to real data from patients with Alzheimer’s.
  Originality:  This paper appears to make a relatively straightforward extension of existing work in disease trajectory modeling by learning the Riemannian metric directly, rather than using a fixed metric. Although previous work is well cited,  it was hard to determine whether this idea to learn the metric from data had been previously attempted, and whether that was the key piece of methodological novelty.  I am not an expert in this subarea and did not track down all of the references cited, but some of the references cited in the introduction imply that this had already been done before, so the specific technical contribution remains unclear to me.

Quality:  The submission appears technically sound and the methods appear appropriate for the task.  However, there is limited empirical evaluation of the methods, and no comparison with simpler baselines such as linear mixed effects models to motivate the improvements that these nonlinear models offer. The synthetic experiments in section 3.1 were also hard to follow, and it was unclear what the important takeaways were. 

Clarity: The paper was hard to follow, and heavily relies on references for some of the details and background.  It was hard to read as a standalone paper.  I found the methods section overly mathy and lacking in basic intuition as to what was going on, and it assumed a lot of prior knowledge about Riemannian geometry and related works. 

Significance:  It’s hard to tell if the results are impactful, and in general some additional motivation on why this particular class of methods is a good idea was lacking.  It’s unclear from reading this paper in what situations both this specific method and this broader subarea of Riemannian-manifold-embedded nonlinear models for disease trajectories are helpful, and what their failure modes are.  Although the proposed GB improves over the existing DCM, a simpler baseline like a linear mixed effects model would be helpful to get a sense for how much these more complex methods improve over something far simpler.  Although predictive performance is generally improved for GB vs DCM it’s hard to tell how to generalize this information to other problems and contexts.  

Other comments:
-	Does this approach assume that all d variables are observed at each observation time? Can this handle settings where only a subset of variables are observed at each time point?
-	Some individual examples of fitted trajectories along with the observed data would be helpful to get a better sense of how well GB, and DCM, fit the observed data.  A lot can be lost in just showing overall MAEs
-	Did you do any subgroup analyses to try to better understand which types of patient trajectories are easier/harder to predict?  
-	How do these classes of approaches fare in the presence of clusters of trajectories, since they assume that all trajectories are subject-specific deviations from a global mean?  Eg if there are a small number of unknown disease subtypes, and each subtype exhibits a specific type of trajectory shape.
-	How many iterations does it take for convergence? What is the rough runtime per iteration? 
-	What is the main take-home message from the comparison with the TADPOLE challenge?  
-	How good are the reported MAEs? What level of error might be low enough in order to be clinically impactful, if you were to back-transform the variables back to their original scale?  Even if some of these mean differences are statistically significant, it is hard to contextualize how much better an error reduction of .002 is in practice
-	For the Imp(n=1) row of Table 1, the HIPP DCM column, is this a typo? Should that 0.31 really be 0.031?
-	Could use some copy-editing, as there are a fair number of typos throughout.
-	Not all of the notation is fully defined in the methods
-	There were several places where ] -1, 1 [ brackets were used, was this just a typo and these were supposed to be [-1,1]?


----------
Edit: after reading through the author response, I have bumped my score up to a 5.  More detailed comments below in thread.

 There was little discussion of limitations of the proposed approach.  A short section mentioning limitations in one place would be helpful.

Societal impact was only briefly mentioned in passing in the end of the conclusion around the need to pay attention to diversity and representativeness of the training data.  Were you able to do any subgroup analyses to see if your model performs roughly as well on different demographics?  If it does not, is there a chance that if your model were operationalized that it could exacerbate inequalities?  How might your model be useful in practice, if there is little we can do to slow or reverse progression of Alzheimer’s?  Could this lead to increased stress and anxiety for afflicted individuals?  


 Hi authors,

Apologies for the delay, but appreciate the detailed response and this addresses most of my concerns. As a result I am raising my score to 5.  

After re-reading your paper again a few times the contribution makes more sense to me, and I do like the underlying idea. At this point, my criticism is mainly around the presentation of your method and in better understanding its empirical strengths and weaknesses, since there is only a single baseline to directly compare against. Would be great to see some of the subsequent analyses / figures mentioned in your comment make it into the supplement, eg showing individual trajectories where GB fits especially well and especially poorly, and compare against DCM for those examples as well.  Perhaps a table or something teasing out the effect of time to prediction and number of past visits on performance would also be nice to see.  

If possible, direct comparison to a simpler baseline such as a standard linear mixed effects model would also be very instructive, as the TADPOLE results don't really get at that at all. Knowing that your method is performing generally as well as, or better than, a collection of previous models on a benchmark is one useful source of information.  But it tells us very little about how much improvement your method offers over some of the standard out-of-box methods that were used in the challenge, eg mixed models, random forests, RNNs. Do the TADPOLE results list the specific types of models used, and could some of those more granular results be shown as well if possible, instead of just the best overall performance? 


 Dear reviewers,

We are thankful for your work and kind comments on our manuscript following our submission for publication. We have highlighted several common criticisms from your reviews and gathered them in this answer. More specific answers will be directly available for each review.

### Quantitative comparison with state-of-the-art

For clarification purposes, note that the comparison with other disease progression methods as well as a variety of alternatives is done thanks to the TADPOLE challenge. Indeed, this challenge allows for a fair comparison of the prediction performances presented here against 56 alternative methods. These methods include penalized regression, linear mixed-effect models, recurrent neural networks, and multi-task learning. Ref. 24 in the manuscript is the main reference for analyzing the results. Since the challenge was closed, more algorithms have been presented for the same prediction task, using the same dataset, including  https://www.sciencedirect.com/science/article/pii/S1053811920306893 and https://link.springer.com/chapter/10.1007/978-3-030-32251-9_19 using deep RNNs, and   https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0211558 using random forests. 

The forecast of ventricle volumes (VENTS) and cognitive decline (ADAS-cog) presented in this paper have smaller mean absolute errors than all these competing methods.

We will better clarify the scope of this comparison in the final version of the paper.

### Qualitative comparison with state-of-the-art

Our method predicts the whole trajectory of change of the biomarkers. It contrasts therefore with black-box type machine learning methods that predict the value of the biomarkers after a fixed time interval (say after 1 year), a problem that can be formulated as a regression problem.

Recurrent neural networks have been proposed to predict a sequence of biomarker values. Such methods, however, need to cope with missing values, a limited and variable number of observations per patient, a variable time interval between observations, and a very variable pattern of change depending on disease stage and subject. 

Mixed-effect models make several assumptions about the pattern of progression. In particular, they assume that the progression falls within a parametric family of curves of small dimension, e.g. a linear curve or a logistic curve with unknown parameters. This is the case for instance of the DCM approach where the linear or logistic curve are seen as the solution of an Hamiltonian flow (a geodesic is a particular case of an Hamiltonian flow) with a given Hamiltonian function. Our main contribution is to learn the metric of the Riemannian manifold, and therefore the Hamiltonian, to better learn the dynamics of progression of the biomarker. In this respect, our contribution relates to the techniques of Hamiltonian and normalizing flows that are becoming popular in the machine learning community, see e.g. https://arxiv.org/abs/1909.13789 or https://arxiv.org/abs/1908.09257.

We will clarify the positioning of our approach with respect to the state of the art along these lines, and add a dedicated "related work" section in the introduction.

 We appreciate the efforts you put in our paper and thank you for your positive feedback.

The model presented in this work is an extension of the DCM, which is the backbone for papers [6] and [8] in the paper.

About the choice of the Riemannian framework as a non-linear modelization, we have several reasons. First, this construction is generic and applies for any data that can be represented in a manifold. This property allows the DCM technique (and hence our GB method) to be applied not only for biomarker data but also for structured data: images (https://www.frontiersin.org/articles/10.3389/fneur.2018.00235/full) and shapes (https://link.springer.com/article/10.1007/s11263-020-01343-w), thus avoiding the need to derive a new model for each data type. Second, the Riemannian metric offers a way to learn data structure in a theoretically well-grounded framework. The Riemannian structure offers specific tools, such as the possibility to define a spatiotemporal coordinate system in the observation space to position the progression of different subjects. In particular, it allows learning dynamical parameters (age at onset, acceleration factors) together with subject-specific progression parameters while ensuring the model to be identifiable (see https://jmlr.csail.mit.edu/papers/volume18/17-197/17-197.pdf for a detailed discussion). Eventually, a geodesic on a Riemannian manifold is a particular case of a Hamiltonian flow, where the Hamiltonian is given by the Riemannian metric. We propose here to learn the dynamics of the biomarker changes by learning the metric and hence the Hamiltonian of a dynamical system. Our work, therefore, relates to the Hamiltonian and normalizing flows that become popular in the machine learning community (see e.g. https://arxiv.org/abs/1909.13789 or https://arxiv.org/abs/1908.09257.) The counterpart of this more abstract construction is an overhead cost in understanding the method which might not be justified when applied on the simpler cases, such as the progression of a single one-dimensional biomarker.

 We appreciate the efforts you put in our paper and thank you for your positive feedback.

We would like to answer each of your points :
   - We will expand on the methods used in the TADPOLE challenge. Indeed, this challenge allows for a fair comparison of the prediction performances presented here against 56 alternative methods. These methods include penalized regression, linear mixed-effect models, recurrent neural networks, and multi-task learning. Ref. 24 in the manuscript is the main reference for analyzing the results. Since the challenge was closed, more algorithms have been presented for the same prediction task, using the same dataset, including  https://www.sciencedirect.com/science/article/pii/S1053811920306893 and https://link.springer.com/chapter/10.1007/978-3-030-32251-9_19 using deep RNNs, and  https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0211558 using random forests. We will clarify the positioning of our approach with respect to the state of the art along these lines, and add a dedicated "related work" section in the introduction where we will take your suggestions into account
   - We will make an effort to introduce methods and their specificities to highlight what the strong points are. Our method predicts the whole trajectory of change of the biomarkers. It contrasts therefore with black-box machine learning methods that predict the value of the biomarkers after a fixed time interval (say after 1 year), a problem that can be formulated as a regression problem.
Recurrent neural networks have been proposed to predict a sequence of biomarker values. Such methods, however, need to cope with missing values, a limited and variable number of observations per patient, a variable time interval between observations, and a very variable pattern of change depending on disease stage and subject. 
Mixed-effect models make several assumptions about the pattern of progression. In particular, they assume that the progression falls within a parametric family of curves of small dimension, e.g. a linear curve or a logistic curve with unknown parameters. Our contribution is exactly to extend this approach by broadening the class of functions.

 We would like to answer each of your points :

   - Originality: The reviewer is right in that our main contribution is learning a Riemannian metric, and hence the structure of the data, from the data itself. There were previous attempts to learn the structure of the data in a Riemannian setting, in particular the references 1 and 2 in the manuscript. Compared to these works, our contributions are threefold. 1) Theoretically, we provide a global Riemannian structure by learning a diffeomorphism, as opposed to the local Riemannian structure which was learned by a neural network in reference 1. The neural network does not provide guarantees about the bijectivity of the mapping. 2) Regarding the algorithmic perspective, we reduce and control the number of parameters. Compared to previous trials of metric learning where the metric was directly learned, thus learning a $d^2$ dimensional function (where $d$ is the dimension of the observations), we here restrict the framework to push-forward induced metrics, which amounts to learning a diffeomorphism which is only a $d$ dimensional function. This is much simpler and makes practical convergence faster and more consistent. 3) Regarding the assessment of the method, we measure the capacity of generalization compared to state-of-the-arts with the TADPOLE Challenge and highlight its specificity by comparing it with DCM.
    We will make these points clearer in the introduction.
   - Quality: For clarification purposes, note that the comparison with other disease progression methods as well as a variety of alternatives is done thanks to the TADPOLE challenge. Indeed, this challenge allows for a fair comparison of the prediction performances presented here against 56 alternative methods. These methods include penalized regression, linear mixed-effect models, recurrent neural networks, and multi-task learning. Ref. 24 in the manuscript is the main reference for analyzing the results. Since the challenge was closed, more algorithms have been presented for the same prediction task, using the same dataset, including  \url{https://www.sciencedirect.com/science/article/pii/S1053811920306893} and \url{https://link.springer.com/chapter/10.1007/978-3-030-32251-9_19} using deep RNNs, and   \url{https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0211558} using random forests. The forecast of ventricle volumes (VENTS) and cognitive decline (ADAS-cog) presented in this paper have smaller mean absolute errors than all these competing methods. We will better clarify the scope of this comparison in the final version of the paper.
About the section on Synthetic experiments: it was presented for illustrating the potential of our method and for giving guidelines for selecting the hyper-parameter sigma to avoid underfitting/over-fitting. We remark that we should take $\sigma$ in [0.1,0.5] and that $\sigma\approx 0.22$ is a good compromise when the parameters of data generation empirically match actual datasets of interest. 
For the experiments with real datasets, we did not perform cross-validation to avoid further time complexity. Instead, we used the guidelines of the Synthetic experiments by trying $\sigma=0.11,0.22,0.5$ and evaluating the results on a validation set.
   - Clarity: it is true that our contribution builds on previous works, and therefore may be difficult to understand in a standalone paper. Nevertheless, we provided mathematical background in the supplementary material.
   - Significance: Learning the structure of the data is a challenging task in Machine Learning and is tackled by various methods for instance variational autoencoders. Our construction applies to any data that can be represented in a manifold. This property allows the DCM technique (and hence our GB method) to be applied not only for biomarker data but also for structured data: images (https://www.frontiersin.org/articles/10.3389/fneur.2018.00235/full) and shapes (https://link.springer.com/article/10.1007/s11263-020-01343-w), thus avoiding the need to derive a new model for each data type. Second, the Riemannian metric offers a way to learn the data structure in a theoretically well-grounded framework. The Riemannian structure offers specific tools, such as the possibility to define a spatiotemporal coordinate system in the observation space to position the progression of different subjects. In particular, it allows learning dynamical parameters (age at onset, acceleration factors) together with subject-specific progression parameters while ensuring the model to be identifiable (see https://jmlr.csail.mit.edu/papers/volume18/17-197/17-197.pdf for a detailed discussion). Eventually, a geodesic on a Riemannian manifold is a particular case of a Hamiltonian flow, where the Hamiltonian is given by the Riemannian metric. We propose here to learn the dynamics of the biomarker changes by learning the metric and hence the Hamiltonian of a dynamical system. Our work, therefore, relates to the Hamiltonian and normalizing flows that become popular in the machine learning community (see e.g.  https://arxiv.org/abs/1909.13789 or https://arxiv.org/abs/1908.09257.) The counterpart of this more abstract construction is an overhead cost in understanding the method which might not be justified when applied on the simpler cases, such as the progression of a single one-dimensional biomarker.The method is a general model for non-linear mixed effect models, so the applications are not restricted to disease modeling.
    In the case of Alzheimer's disease and other neurodegenerative diseases, descriptive models with personalized trajectories are very useful because prediction allows for a better selection of patients into clinical trials. 
   - Other comments :
        - Our model is a generative statistical model, and as such does not need the imputation of missing data. The likelihood is only informed by available observations. See for instance the reference https://ieeexplore.ieee.org/document/8759198.
        - Being able to describe individualized trajectories is indeed one strength of the approach. We showed mean trajectories as a more representative curve for the disease. We will add individual examples in the supplementary material to illustrate individuals with good and bad fits
        - Subsequent analyses show that the main factors driving prediction errors are the time to prediction (the further in time the harder the prediction is) and the number of past visits (the more the better the individualized trajectory is)
        - Heterogeneity of data (clusters of population or disease subtypes) is indeed the main issue in medical cohorts. This problem can be handled in our Bayesian framework by adding a mixture component. We can refer to another work expanding the DCM model with a mixture to identify subtypes (see https://link.springer.com/chapter/10.1007/978-3-030-78191-0_44). The same idea can be applied along with the GB model to cope with subtypes, this is however outside the scope of this paper.    
       - The iteration time is described in the supplementary document, section B. The number of iterations is chosen beforehand but is usually overestimated in order to ensure convergence. The early stopping criterion is not easy to implement as the likelihood increases stochastically.
        - The take-home message with TADPOLE is that our method outperforms the best challenger in both outcomes (please note that the best challenger in ADAS is not the same challenger in VENTS) which is a strong argument in favor of validating our method. 
       - We refer the reviewer to the conclusions of the challenge for the quality of the MAEs obtained. The clinical impact is not only measured with the prediction capability but the qualitative description of the individualized trajectories is always helpful to the clinicians. Current research focuses on using such predictions for selecting patients for clinical trials. Improvement in MAE translates into sample size reduction, which has a great practical impact.
       - You are right, it is 0.031.
       - We indeed have some typos here and elsewhere, which were pinpointed by some reviewers. We will take care of it
       - We will make sure that all notations are clearly introduced
       - It should be $(-1, 1)$, the notation is not standard and we will correct it
   - Limitations: One of the main limitations of our approach comes from the fact that we need to learn the metric in addition to the parameters of the DCM model. The dimension is controlled by the number of control points, the kernel parameter, and the number of compositions. We need to tune these parameters to balance between overfitting and flexibility. The second limitation is that our method starts with a prior shape of the progression curve, for instance, a logistic or a linear function as in the DCM approach. It then bends the curve to better fit data. This bending occurs only in the age range where there are enough observations (usual for Bayesian methods) . Therefore, the forecast of biomarkers far in the future might not work well if the asymptotic shape of the curve greatly varies among subjects. This is the case for the volume of the ventricles in contrast to the ADAS-Cog assessment. We will add a small paragraph describing this
   - Social impact: The reviewer is right that the approach needs more careful evaluation before being used in clinical practice, such as the evaluation of biases in the prediction and clarification of the context of use. We will elaborate on these issues in the social impact section. These problems are usually handled in relation with clinicians and industrial partners following rigorous regulatory procedures.


 For clarification purposes, synthetic experiments were made for a qualitative understanding of the possibilities of the method and the TADPOLE challenge was the quantitative evaluation of it.

The comparison with other disease progression methods as well as a variety of alternatives is done thanks to the TADPOLE challenge. Indeed, this challenge allows for a fair comparison of the prediction performances presented here against 56 alternative methods. These methods include penalized regression, linear mixed-effect models, recurrent neural networks, and multi-task learning. Ref. 24 in the manuscript is the main reference for analyzing the results. Since the challenge was closed, more algorithms have been presented for the same prediction task, using the same dataset, including  \url{https://www.sciencedirect.com/science/article/pii/S1053811920306893} and \url{https://link.springer.com/chapter/10.1007/978-3-030-32251-9_19} using deep RNNs, and   \url{https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0211558} using random forests. 

The forecast of ventricle volumes (VENTS) and cognitive decline (ADAS-cog) presented in this paper have smaller mean absolute errors than all these competing methods.

We will better clarify the scope of this comparison in the final version of the paper.

The method is general thanks to the Riemannian framework and can thus be applied to any data lying on a manifold. This property allows the DCM technique (and hence our GB method) to be applied not only for biomarker data but also for structured data: images (https://www.frontiersin.org/articles/10.3389/fneur.2018.00235/full) and shapes (https://link.springer.com/article/10.1007/s11263-020-01343-w), thus avoiding the need to derive a new model for each data type.

This paper uses Riemannian metric learning within the context of mixed-effect models for modelling distributions of patient trajectories through a disease process.  Disease progression modelling is an active area with high relevance to NeurIPS and this is a novel take on its implementation.  Experiments demonstrate efficacy in simulation, illustrate the benefits of the trade-off between flexibility and interpretability the proposed approach provides, and show strongly competitive performance on a standard Alzheimer's disease data set (from the TADPOLE challenge) even against purely black-box data-driven models that offer much less in terms of understanding disease and its heterogeneity.  This is a well thought-out and novel formulation for disease progression modelling that is well motivated in terms of balancing flexibility and interpretability.  Experiments are clear and compelling and I believe these ideas will be of broad interest to the NeurIPS community, as well as computational neurology and disease progression modelling researchers.  The presentation bridges those fields nicely.  

I don't have much feedback to give on this paper, but would say that it would benefit from a clearer specification from the start of what benefits the authors expect from proposed formulation over existing implementations such as refs [6-8] in the manuscript.  In essence the ideas are similar - estimating trajectories as well as time shift and acceleration factors.  Is the Riemannian formulation primarily a mathematical nicety, or are there genuine performance issues with earlier implementations that the authors expect it to improve?  With more careful specification of these motivations and expectations, the experimental work could be more targeted to highlighting those hypothesised benefits.


 Fine

In this work, the authors proposed a mixed effect model in a Riemannian manifold to estimate disease progression. The proposed model is built upon previous work with an improvement to estimate a Riemannian metric from data. An alternating maximization algorithm is proposed to learn a sequence of diffeomorphisms, represented in the RKHS, to bend the initial metric toward optimal. The method was evaluated on synthetic data and a real-world dataset to demonstrate improvement upon previous work.   This work brings several ideas from previous studies together and combines them into an improved mixed effect model in a Riemannian manifold. The method was validated on synthetic data and a real dataset. The improvement was demonstrated on the ANDI dataset, but only one baseline was compared.  Evaluation on synthetic data could be further improved since only two qualitative examples were provided. The proposed method seems to be general and can be used in other applications, but no evidence was provided to justify that. The authors have adequately addressed the limitations

Meta-review: Thanks to the authors for their engaging submission on an important topic.  There was a lot of positive response in the reviews to this work, and all reviewers were positive about the underlying idea.  However, there are areas in which this submission could be strengthened. 

A major reviewer concern is understanding the strengths and weaknesses of the proposed method.  An explicitly stated motivating factor is that the parametric assumptions of standard mixed effects models are too restrictive, and that the flexibility afforded by learning the Reimannian manifold will allow models to better fit the data and form better trajectories. And intuitively this makes sense. But the empirical evidence of this is lacking.  As reviewer w1hQ points out, a standard linear mixed effects model would make for an informative baseline to empirically justify these claims.  As the abstract introduces this method as an improvement upon (and alternative to) a standard linear mixed effects model, having a direct empirical comparison would help identify what parts of the proposed method are most performant, and in what situations the proposed method may be inappropriate.  Further, a simulation study where aspects of the data generating process are known and manipulated could enhance such a direct comparison to a standard mixed effects model.

Decision: accept