Paper: Demonstrating the capacity of a Path-Based variational inference formulation for robust hidden Markov modelling of complex and noisy binary trees (ICLR-2024)

Reviewer comment (Presentation):
I found the manuscript to be a relatively hard read. Methodological and implementation details are often not explained or reasoned about and there are many of them, moreover, figures and tables are too full of details. I would suggest this is a presentation problem, not all numbers and methodological details are equally important, and I feel like the authors could have made more of an effort in summarizing and selecting the important aspects. I believe this could make the work more interesting.

Author response:
**[R2-1] I found the manuscript to be a relatively hard read. Methodological and implementation details are often not explained or reasoned about and there are many of them**

We appreciate this feedback and aimed to address your concerns as follows:

- **Methodological Explanations:**  We have added details on the HMM model and spatial constraints design to Section 1 (Introduction); and Section 2.1.2 (Soft spatial constraints on termination of synthetic trees). 
- **Methodological Reasoning:** During development, we relied heavily on domain knowledge from our expert clinician co-author, who has > 10 years of experience in airway morphology, as well as medical literature, to ensure that our generative process includes all the variables and conditioning needed to yield high-fidelity synthetic airway trees. 
- **Implementation reasoning:**
    - **Generative model implementation:** We relied on domain-knowledge (from expert clinician and medical literature) to design a generative model suitable for human airway trees. 
    - **Validation experiments:** Benchmarking experiments showcase key challenges that exist in binary tree-structured data observed in 3-D medical images e.g., CT modalities. We handle real-case scenarios with noisy observations.

- **Implementation details:** We designed our implementation to integrate existing specialized open-source libraries.

Reviewer comment (Presentation):
Almost all tables and figures are not readable except in very enlarged versions. Could they perhaps be changed and/or resized to be readable at reasonable magnification/print size?

Author response:
**[R2-2] …, moreover, figures and tables are too full of details. I would suggest this is a presentation problem, not all numbers and methodological details are equally important, and I feel like the authors could have made more of an effort in summarizing and selecting the important aspects. I believe this could make the work more interesting. Almost all tables and figures are not readable except in very enlarged versions. Could they perhaps be changed and/or resized to be readable at reasonable magnification/print size?**

- We appreciate this important constructive comment. We have taken extensive measures to improve data presentation and summarization throughout the paper and ensure that all tables and figures are now readable at print size, especially in Section 3 (Results and Discussion).
-  We have added a section in Sup. D (Additional experimental results) for further details supporting the main paper.

Reviewer comment (Experiments):
There is no direct comparison with other competing approaches or baselines.

Author response:
**[R1-1] There is no direct comparison with other competing approaches or baselines.**
- The closest baseline for tree-HMMs modeling is cited in our paper as (Mohammadi et al., 2022) [1]. We tested this approach during our initial development but it did not work properly on our trees and the size of our dataset for the following reasons: CPU-implementation, design choices too domain specific (branching rules controlled by a genetically encoded subroutine). 
- We aim to simplify a specific modeling problem and integrate its implementation with open-source specialized libraries. Showing that this simplified modeling formulation provides correct and robust fits is our main goal. 
- In this context, we used synthetic benchmarking, a validation technique applied in studies such as that of Mohammadi et al. (2022) [1].
- We develop this method in an academic setting, with limited computational resources, which limits our capacity to test complex models.  

[1] Mohammadi F, Visagan S, Gross SM, Karginov L, Lagarde JC, Heiser LM, Meyer AS. A lineage tree-based hidden Markov model quantifies cellular heterogeneity and plasticity. Communications Biology. 2022;5(1):1258.

Reviewer comment (Novelty):
The technical novelty is limited.

Author response:
**[R1-2] The technical novelty is limited.**

We appreciate the reviewer's assessment regarding the novelty of our proposed method. We thank you for the opportunity to clarify the technical novelty of our work. 
We argue that the technical innovation of our contribution lies primarily in the following three aspects: 

**(1) Technical novelty**: Our proposed “paths-of-bifurcations” approximation is a novel approach, which raised several technical challenges: 
- Formulate the cost function on variables identified as needed and sufficient for an approximate variational inference to be able to recover ground truth parameters without accounting for the full tree structure; 
- Evaluation of quality of fits on real noisy binary tree-structured data observed in 3-D medical images; 
- Incorporate spatial constraints in the inference of HMM parameters; 
- Control optimization complexity with a limited number of model parameters;

**(2) Increasing accessibility of complex methods**: We provide our code-implementation compatible with state-of-the-art sophisticated open-source packages, thus facilitating broader adoption of generative modeling for binary tree architectures. 

**(3) Robust validation on realistic synthetic data**: We include code to generate realistic airway trees based on previously-validated procedural models for readers to reproduce our results and further develop their own synthetic benchmarks (eg. to extend to botanics applications).

Reviewer comment (Writing):
The method could be more clearly described in the main paper, and variables/distributions could be more clearly defined.

Author response:
**[R1-3] The method could be more clearly described in the main paper, and variables/distributions could be more clearly defined.**

- The forms of HMM emission distributions, transition matrices, and start probabilities from Equations 1,2 are now detailed in Table 3 (Sup. B). We have also rephrased Section 2.1.1 (Generation of synthetic trees with hidden Markov modeling) for clarity.

Reviewer comment (Novelty):
"Fast and scalable inference of a generative process governing such tree-structures is key for classifying trees into categories." Are generative process modeling the prevalent approach to classification, or are there other (discriminative) approaches?

Author response:
**[R1-4] "Fast and scalable inference of a generative process governing such tree-structures is key for classifying trees into categories." Are generative process modeling the prevalent approach to classification, or are there other (discriminative) approaches?**

- For 3-D binary trees, discriminative approaches are prevalent. We have now emphasized this point in Section 1 (Introduction). 
- Our approach is the first to fit a generative model to a large, real dataset of 3-D human airway trees. We cite three other tree-HMM inference methods. However, none of these provide fast scalable open-source inference pipelines nor designed to be applicable beyond their specific use-case domain.

Reviewer comment (Writing):
"...  and split trees into ‘independent’ paths of bifurcations ..." why is independent in quotations?

Author response:
**[R1-5] "... and split trees into 'independent' paths of bifurcations ..." why is independent in quotations?**
- In our "paths-of-bifurcations" approximation, we split each single binary tree into multiple paths. The term “independent” refers to the fact that while our approximation treats each path as independent, some paths are not truly independent as they come from the same tree.
- We validate this assumption via our synthetic benchmark experiments, in particular, in the low probability rules (LPR) experiments LPR-A, LPR-B (Section 3.3), where we test whether we can recover rare state rules in deep branch generations despite using overlapping paths that artificially over-sample non-rare early branching patterns.

Reviewer comment (Writing):
Would it be possible in the introduction to set up a more clear problem formulation so that it is evident which types of data problem the paper addresses? For example, it is not clear whether the trees that are modelled have a graph-based or geometric representation, such as 3D coordinate-based or anatomical mesh representation.

Author response:
**[R1-6] Would it be possible in the introduction to set up a more clear problem formulation so that is evident which types of data problem the paper addresses? For example, it is not clear whether the trees that are modelled have a graph-based or geometric representation, such as 3d coordinat-based or anatomical mesh representation.**

- Our inference is designed for graph-based tree representations. We have added this distinction to Section 1 (Introduction).
- For our use-case on human airway trees, we included in the paper the steps required to go from geometric representations (segmentation masks) to graph representations (nodes and edges) in Sup. E.2.

Reviewer comment (Theory):
If I understand correctly, emission probabilities for branching angles and roll are Gaussian. Is anything done to ensure that the first angle "is the smallest angle by definition" and that angles are within a suitable range?

Author response:
**[R1-7] Is anything done to ensure that the first angle "is the smallest angle by definition"**
- We confirm that we enforce that the first angle "is the smallest angle by definition"  to ensure that we do not infer redundant branching rules (e.g., where we would infer two distinct branching rules, while only different in the order of alpha1 and alpha2) when we observe an asymmetric bifurcation in our data (almost always).
- We have clarified this point in Section 2.1.1 (Generation of synthetic trees with hidden Markov modeling)

Reviewer comment (Writing):
Would it be possible to provide some more details (possibly in the supplement) regarding the definitions of each term in the joint distribution (Eq. 1), such as a list with clear definitions of all variables and distributions.

Author response:
**[R1-8] …and that angles are within a suitable range?**
- During optimization angles are coded in radians ($-\pi$ to $\pi$), with $\alpha_1 < 0, \alpha_2 > 0$. But, we display angle values (in Tables 1-2, and Figures 2,3,6) and MAE (mean absolute error, Section 3) in degrees as this is more interpretable.
- We did not constrain the range of values on any angle variables.
- We have clarified these points in Section 2.3.2 (Training setup).

Reviewer comment (Experiments):
Have any other methods been used to model e.g. the human airways data? How does the proposed method perform in comparison?

Author response:
**[R1-10] Which variational distributions ('guide' in Pyro terminology) are used, and how important/sensitive is the choice?**
- We used the AutoDelta guide in Pyro and have added this information in Section 2.3.2 (Training setup). 
- The complete model and inference codes are available from our GitHub link.

Reviewer comment (Writing):
Minor issues: There are a few typos etc. Examples from abstract and introduction: "...prevalent multiple across fields..." ? Closing parenthesis doubled "))" in a few citations. "... and infer states..." should be 'and infers states'. "... human airway trees segmented on CT images." unclear.

Author response:
**[R1-12] What would be the primary application of the proposed method?**
- We are developing this method in the context of unsupervised discovery of new sub-phenotypes of human airway trees in the general population and in specific diseases (e.g., Chronic obstructive pulmonary disease, COPD). This is for us our primary application. 
- Sub-phenotyping within large cohorts enables to increase the power of downstream analysis on sub-populations of patients to: (1) formulate hypothesis on disease mechanisms, (2) detect biomarkers in specific sub-populations associated with individual susceptibility to disease, (3) create high-fidelity synthetic digital twins (here synthetic airway tree structures) for multiple applications such as particle deposition studies (deposition of smoke particles / environmental pollutants at bifurcation structures in the lung), association with genomic markers or inclusion in clinical trials.

**[R1-13] Minor issues: There are a few typos etc. Examples from abstract and introduction: "...prevalent multiple across fields..."? Closing parenthesis doubled"))" in a few citations. "... and infer states..." should be 'and infers states' "... human airway trees segmented on CT images." Unclear**
- We appreciate your attention to details when reviewing our paper and we have carefully proof-read the paper for typos.

Reviewer comment (Theory):
While I agree in spirit with the mathematical technique, I think section 2.3 needs work in clarifying notation (perhaps a block diagram?), and has possible errors due to its compression into a single page of 2 equations.

Author response:
- Many thanks for your attention to detail, we have corrected a typo in Equation 1.
- The forms of HMM emission distributions, transition matrices, and start probabilities from Equations 1,2 have been added in Table 3 (Sup. B), along with line numbers corresponding to Algorithm 1 (Generative process) to increase clarity.

Reviewer comment (Theory):
It is not clear that air pathways actually have these complex state transitions? Could they instead be modeled using the (constant) branching process? Why are there hidden states, instead of constant branching factors or branching factors that are a function of depth? Basically, do the physical branching angles, rate of branching, etc., actually change?

Author response:
- Metzger et al., 2008 [1] hypothesizes and provides evidence to support lineage-based, genetically coded branching subroutines within the bronchial tree in animal models rather than a constant branching process. This formulation aligns with hidden Markov states rather than depth-based branching factors. 
- Sauret et al., 2002 [2] present the complex patterns of branching angles measured in two human bronchial trees. They report no significant difference in the plane rotation angle (= $\phi$, in our paper) with increasing branching generation. They highlight the wide distribution of all branching angles within branching generation=9 as an example. All the above contra-indicates a depth-based rule system for branching angles.
- Figure 15 in Sup. E of our paper displays airway tree binary masks from the ATM dataset (used in Section 3.5). One can see the complex variation and patterning between subjects.

Reviewer comment (Experiments):
CT has imaging resolution on the order of mm (1mm isotropic is a very high quality human CT, multi-mm thick slice or gapped slices are standard clinical CT resolution). Do we expect the actual alveolar trees to be visible? If not, would these resolutions lead to non-tree like observations? (i.e., multi-compartment mergers)

Author response:
- No, alveolar trees are not visible. We are working with clean airway binary segmentations with no alveolar trees,  as depicted in Figures 1c and 15.

Reviewer comment (Theory):
I'm fairly certain the ELBO should be a lower bound. Why can we have the log inside of the integral/summand in equation (1) without an inequality? Usually the latent variables are introduced at this stage, then either the spectral gap is derived, or Jensen's inequality used.

Author response:
- We agree with the properties of the ELBO that you described. However, Equation 1 is the log-joint and not the ELBO. In our work, the ELBO is derived automatically and numerically by Pyro.
- We have chosen to detail the log joint (rather than the ELBO) in Equation 1 because from there, any optimisation method (e.g., expectation-maximization, variational inference) can be used to infer latent variables. In our GitHub repository, we demonstrate inference in Pyro and package ‘hmmlearn’, a Scikit-Learn extension package for Hidden Markov Models.

Reviewer comment (Experiments):
One motivating factor for this work is rule heterogeneity. However, is this actually possible using these priors? I think this is what your experiments are trying to show in 3.2; I encourage the authors to work on this. Branching processes with differing rules based on interesting conditions can be very interesting in applied fields, but Figures 4 and 5 are incredibly small and unreadable, as are table 1 and 2, but even more importantly, the authors should use the synthetic data to prove that if such a rule shift (so, a regime change between lung tissue structures) existed, it could be learned by this HMM, and wouldn't have been learned by the fixed branching process.

Author response:
- Yes, rule heterogeneity is what we are testing in Section 3.2. We have clarified the text and figures to demonstrate these results more clearly. 
- Please refer to Table 2 for the types of synthetic trees we validated our inference method on including  trees with rule heterogeneity, low probability rules and noisy data quality. 
We present the capabilities of our approximate inference method to recover HMM rules under these conditions in Section 3. 

- We have taken extensive measures to improve data presentation and summarization throughout the paper and ensure that all tables and figures are now readable at print size, especially in Section 3 (Results and Discussion).