Original: L EARNING P ROBABILISTIC T OPOLOGICAL R EPRESENTATIONS U SING D ISCRETE M ORSE T HEORY
Revised: L EARNING P ROBABILISTIC T OPOLOGICAL R EPRESENTATIONS U SING D ISCRETE M ORSE T HEORY

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Accurate delineation of fine-scale structures is a very important yet challenging problem.
Revised: Accurate delineation of fine-scale structures is a very important yet challenging problem.

Original: Existing methods use topological information as an additional training loss, but are ultimately making pixel-wise predictions.
Revised: Existing methods use topological information as an additional training loss, but are ultimately making pixel-wise predictions.

Original: In this paper, we propose the first deep learning based method to learn topological/structural  representations.
Revised: In this paper, we propose a novel deep learning based method to learn topological/structural 1 representations.

Original: We use discrete Morse theory and persistent homology to construct an one-parameter family of structures as the topological/structural representation space.
Revised: We use discrete Morse theory and persistent homology to construct an one-parameter family of structures as the topological/structural representation space.

Original: Furthermore, we learn a probabilistic model that can perform inference tasks in such a topological/structural representation space.
Revised: Furthermore, we learn a probabilistic model that can perform inference tasks in such a topological/structural representation space.

Original: Our method generates true structures rather than pixel-maps, leading to better topological integrity in automatic segmentation tasks.
Revised: Our method generates true structures rather than pixel-maps, leading to better topological integrity in automatic segmentation tasks.

Original: It also facilitates semi-automatic interactive annotation/proofreading via the sampling of structures and structure-aware uncertainty
Revised: It also facilitates semi-automatic interactive annotation/proofreading via the sampling of structures and structure-aware uncertainty.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Accurate segmentation of fine-scale structures, e.g., vessels, neurons and membranes is crucial for downstream analysis.
Revised: Accurate segmentation of fine-scale structures, e.g., vessels, neurons and membranes is crucial for downstream analysis.

Original: In recent years, topology-inspired losses have been proposed to improve structural accuracy (Hu et al., 2019; 2021; Shit et al., 2021; Mosinska et al., 2018; Clough et al., 2020).
Revised: In recent years, topology-inspired losses have been proposed to improve structural accuracy (Hu et al., 2019; 2021; Shit et al., 2021; Mosinska et al., 2018; Clough et al., 2020).

Original: These losses identify topologically critical locations at which a segmentation network is error-prone, and force the network to improve its prediction at these critical locations.
Revised: These losses identify topologically critical locations at which a segmentation network is error-prone, and force the network to improve its prediction at these critical locations.

Original: However, these loss-based methods are still not ideal.
Revised: However, these loss-based methods are still not ideal.

Original: They are based on a standard segmentation network, and thus only learn pixel-wise feature representations .
Revised: They are based on a standard segmentation network, and thus only learn pixel-wise feature representations .

Original: This causes several issues.
Revised: This causes several issues.

Original: First, a standard segmentation network makes pixel-wise predictions.
Revised: First, a standard segmentation network makes pixel-wise predictions.

Original: Thus, at the inference stage, topological errors, e.g. broken connections, can still happen, even though they may be mitigated by the topologyinspired losses.
Revised: Thus, at the inference stage, topological errors, e.g. broken connections, can still happen, even though they may be mitigated by the topologyinspired losses.

Original: Another issue is in uncertainty estimation, i.e., estimating how certain a segmentation network is at different locations.
Revised: Another issue is in uncertainty estimation, i.e., estimating how certain a segmentation network is at different locations.

Original: Uncertainty maps can direct the focus of human annotators for efficient proofreading.
Revised: Uncertainty maps can direct the focus of human annotators for efficient proofreading.

Original: However, for fine-scale structures, existing pixel-wise uncertainty map is not effective.
Revised: However, for fine-scale structures, existing pixel-wise uncertainty maps are not effective.

Original: As shown in Fig. 1(d), every pixel adjacent to a vessel branch is highly uncertain, in spite of whether the branch is salient or not.
Revised: As shown in Fig. 1(d), every pixel adjacent to a vessel branch is highly uncertain, in spite of whether the branch is salient or not.

Original: What is more desirable is a structural uncertainty map that can highlight uncertain branches (e.g., Fig. 1(f)).
Revised: What is more desirable is a structural uncertainty map that can highlight uncertain branches (e.g., Fig. 1(f)).

Original: To fundamentally address these issues, we propose to directly model and reason about the structures.
Revised: To fundamentally address these issues, we propose to directly model and reason about the structures.

Original: In this paper, we propose the first deep learning based method that directly learns the topological/structural1 representation of images .
Revised: In this paper, we propose a novel deep learning based method that directly learns the topological/structural representation of images .

Original: To move from pixel space to structure space, we apply the classic discrete Morse theory (Milnor, 1963; Forman, 2002) to decompose an image into a Morse complex, consisting of structural elements like branches, patches, etc.
Revised: To move from pixel space to structure space, we apply the classic discrete Morse theory (Milnor, 1963; Forman, 2002) to decompose an image into a Morse complex, consisting of structural elements like branches, patches, etc.

Original: These structural elements are the hypothetical structures one can infer from the input image.
Revised: These structural elements are the hypothetical structures one can infer from the input image.

Original: Their combinations constitute a space of structures arising from the input image.
Revised: Their combinations constitute a space of structures arising from the input image.

Original: See Fig. 2(c) for an illustration.
Revised: See Fig. 2(c) for an illustration.

Original: For further reasoning with structures, we propose to learn a probabilistic model over the structure space.
Revised: For further reasoning with structures, we propose to learn a probabilistic model over the structure space.

Original: The challenge is that the space consists of exponentially many branches and is thus of very high dimension.
Revised: The challenge is that the space consists of exponentially many branches and is thus of very high dimension.

Original: To reduce the learning burden, we introduce the theory of persistent homology (Sousbie, 2011; Delgado-Friedrichs et al., 2015; Wang et al., 2015) for structure pruning.
Revised: To reduce the learning burden, we introduce the theory of persistent homology (Sousbie, 2011; Delgado-Friedrichs et al., 2015; Wang et al., 2015) for structure pruning.

Original: Each branch has its own persistence measuring its relative saliency.
Revised: Each branch has its own persistence measuring its relative saliency.

Original: By continuously thresholding the complete
Revised: By continuously thresholding the complete

Original: Morse complex in terms of persistence, we obtain a sequence of Morse complexes parameterized by the persistence threshold, ϵ .
Revised: Morse complex in terms of persistence, we obtain a sequence of Morse complexes parameterized by the persistence threshold, ϵ .

Original: See Fig. 2(d).
Revised: See Fig. 2(d).

Original: By learning a Gaussian over ϵ , we learn a parametric probabilistic model over these structures.
Revised: By learning a Gaussian over ϵ , we learn a parametric probabilistic model over these structures.

Original: [Figure]
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: This parametric probabilistic model over structure space allows us to make direct structural predictions via sampling (Fig. 2(e)) and to estimate the empirical structure-level uncertainty via sampling (Fig. 2(g)).
Revised: This parametric probabilistic model over structure space allows us to make direct structural predictions via sampling (Fig. 2(e)) and to estimate the empirical structure-level uncertainty via sampling (Fig. 2(g)).

Original: The benefit is two-fold: First, direct prediction of structures will ensure the model outputs always have structural integrity, even at the inference stage.
Revised: The benefit is two-fold: First, direct prediction of structures will ensure the model outputs always have structural integrity, even at the inference stage.

Original: This is illustrated in Fig. 1(e).
Revised: This is illustrated in Fig. 1(e).

Original: Samples from the probabilistic model are all feasible structural hypotheses based on the input image, with certain variations at uncertain locations.
Revised: Samples from the probabilistic model are all feasible structural hypotheses based on the input image, with certain variations at uncertain locations.

Original: This is in contrast to state-of-the-art methods using pixel-wise representations (Fig. 1(c)-(d)).
Revised: This is in contrast to state-of-the-art methods using pixel-wise representations (Fig. 1(c)-(d)).

Original: Note the original output structure (Fig. 2(e), also called skeleton) is only 1-pixel wide and may not serve as a good segmentation output.
Revised: Note the original output structure (Fig. 2(e), also called skeleton) is only 1-pixel wide and may not serve as a good segmentation output.

Original: In the inference stage, we use a post-processing step to grow the structures without changing topology as the final segmentation prediction (Fig. 2(f)).
Revised: In the inference stage, we use a post-processing step to grow the structures without changing topology as the final segmentation prediction (Fig. 2(f)).

Original: More details are provided in Sec. 3.2 and Fig.
Revised: More details are provided in Sec. 3.2 and Fig.

Original: Second, the probabilistic structural model can be seamlessly incorporated into semi-automatic interactive annotation/proofreading workflows to facilitate large scale annotation of these complex structures.
Revised: Second, the probabilistic structural model can be seamlessly incorporated into semi-automatic interactive annotation/proofreading workflows to facilitate large scale annotation of these complex structures (see Fig. 5).

Original: This is especially important in the biomedical domain where fine-scale structures are notoriously difficult to annotate, due to the complex 2D/3D morphology and low contrast near extremely thin structures.
Revised: This is especially important in the biomedical domain where fine-scale structures are notoriously difficult to annotate, due to the complex 2D/3D morphology and low contrast near extremely thin structures.

Original: Our probabilistic model makes it possible to identify uncertain structures for efficient interactive annotation/proofreading.
Revised: Our probabilistic model makes it possible to identify uncertain structures for efficient interactive annotation/proofreading.

Original: Note that the structure space is crucial for uncertainty reasoning.
Revised: Note that the structure space is crucial for uncertainty reasoning.

Original: As shown in Fig. 1(f) and Fig. 2(g), our structural uncertainty map highlights uncertain branches for efficient proofreading.
Revised: As shown in Fig. 1(f) and Fig. 2(g), our structural uncertainty map highlights uncertain branches for efficient proofreading.

Original: On the contrary, traditional pixel-wise uncertainty map (Fig. 1(d)) is not helpful at all; it highlights all pixels on the boundary of a branch.
Revised: On the contrary, traditional pixel-wise uncertainty map (Fig. 1(d)) is not helpful at all; it highlights all pixels on the boundary of a branch.

Original: The main contributions of this paper are:
Revised: The main contributions of this paper are:

Original: We propose the first deep learning method which learns structural representations, based on discrete Morse theory and persistent homology.
Revised: We propose a novel deep learning method which learns structural representations, based on discrete Morse theory and persistent homology.

Original: We learn a probabilistic model over the structure space, which facilitates different tasks such as topology-aware segmentation, uncertainty estimation and interactive proofreading.
Revised: We learn a probabilistic model over the structure space, which facilitates different tasks such as topology-aware segmentation, uncertainty estimation and interactive proofreading.

Original: We validate our method on various datasets with rich and complex structures .
Revised: We validate our method on various datasets with rich and complex structures .

Original: It outperforms state-of-the-art methods in both deterministic and probabilistic categories.
Revised: It outperforms state-of-the-art methods in both deterministic and probabilistic categories.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Structure/Topology-aware deep image segmentation.
Revised: Structure/Topology-aware deep image segmentation.

Original: A number of recent works have tried to segment with correct topology with additional topology-aware losses (Mosinska et al., 2018;
Revised: A number of recent works have tried to segment with correct topology with additional topology-aware losses (Mosinska et al., 2018;

Original: Hu et al., 2019; Clough et al., 2020; Hu et al., 2021; Shit et al., 2021).
Revised: Hu et al., 2019; Clough et al., 2020; Hu et al., 2021; Shit et al., 2021).

Original: Specifically, UNet-VGG (Mosinska et al., 2018) detects linear structures with pretrained filters.
Revised: Specifically, UNet-VGG (Mosinska et al., 2018) detects linear structures with pretrained filters.

Original: clDice (Shit et al., 2021) introduces an additional Dice loss for extracted skeleton structures.
Revised: clDice (Shit et al., 2021) introduces an additional Dice loss for extracted skeleton structures.

Original: TopoLoss (Hu et al., 2019; Clough et al., 2020) learns to segment with correct topology explicitly with a differentiable loss by leveraging the concept of persistent homology.
Revised: TopoLoss (Hu et al., 2019; Clough et al., 2020) learns to segment with correct topology explicitly with a differentiable loss by leveraging the concept of persistent homology.

Original: Similarly, DMT-loss (Hu et al., 2021) tries to identify the topological critical structures via discrete Morse theory, and then force the network to make correct pixel-wise prediction on these structures.
Revised: Similarly, DMT-loss (Hu et al., 2021) tries to identify the topological critical structures via discrete Morse theory, and then force the network to make correct pixel-wise prediction on these structures.

Original: All these losses, although aiming at topological integrity, still cannot change the pixel-wise prediction nature of the backbone network.
Revised: All these losses, although aiming at topological integrity, still cannot change the pixel-wise prediction nature of the backbone network.

Original: To the best of our knowledge, no existing methods generate structural representations/predictions like our proposed method.
Revised: To the best of our knowledge, no existing methods generate structural representations/predictions like our proposed method.

Original: Additionally, the discrete Morse complex has been used for image analysis, but only as a preprocessing step (Delgado-Friedrichs et al., 2015; Robins et al., 2011; Wang et al., 2015; Dey et al., 2019) or as a conditional input of a neural network (Banerjee et al., 2020).
Revised: Additionally, the discrete Morse complex has been used for image analysis, but only as a preprocessing step (Delgado-Friedrichs et al., 2015; Robins et al., 2011; Wang et al., 2015; Dey et al., 2019) or as a conditional input of a neural network (Banerjee et al., 2020).

Original: Segmentation uncertainty.
Revised: Segmentation uncertainty.

Original: Uncertainty estimation has been the focus of research in recent years (Graves, 2011; Gal & Ghahramani, 2016; Lakshminarayanan et al., 2017; Moon et al., 2020).
Revised: Uncertainty estimation has been the focus of research in recent years (Graves, 2011; Gal & Ghahramani, 2016; Lakshminarayanan et al., 2017; Moon et al., 2020).

Original: However, most existing work focuses on the classification problem.
Revised: However, most existing work focuses on the classification problem.

Original: In terms of image segmentation, the research is still relatively limited.
Revised: In terms of image segmentation, the research is still relatively limited.

Original: Some existing methods directly apply classification uncertainty to individual pixels, e.g., dropout (Kendall et al., 2015; Kendall & Gal, 2017).
Revised: Some existing methods directly apply classification uncertainty to individual pixels, e.g., dropout (Kendall et al., 2015; Kendall & Gal, 2017).

Original: This, however, is not taking into consideration the image structures.
Revised: This, however, is not taking into consideration the image structures.

Original: Several methods estimate the uncertainty by generating an ensemble of segmentation (Lakshminarayanan et al., 2017) or using multi-heads (Rupprecht et al., 2017; Ilg et al., 2018).
Revised: Several methods estimate the uncertainty by generating an ensemble of segmentation (Lakshminarayanan et al., 2017) or using multi-heads (Rupprecht et al., 2017; Ilg et al., 2018).

Original: Notably, Probabilistic-UNet (Kohl et al., 2018) learns a distribution over the latent space and then samples over the latent space to produce segmentation samples.
Revised: Notably, Probabilistic-UNet (Kohl et al., 2018) learns a distribution over the latent space and then samples over the latent space to produce segmentation samples.

Original: When it comesdown to uncertainty, however, these methods can still only generate a pixel-wise uncertainty map, using the frequency of appearance of each pixel in the sample segmentations.
Revised: When it comes to uncertainty, however, these methods can still only generate a pixel-wise uncertainty map, using the frequency of appearance of each pixel in the sample segmentations.

Original: These methods are fundamentally different from ours, which makes predictions on the structures.
Revised: These methods are fundamentally different from ours, which makes predictions on the structures.

Original:  M ETHOD
Revised:  M ETHOD

Original: Our key innovation is to restructure the output of a neural network so that it is indeed making predictions over a space of structures.
Revised: Our key innovation is to restructure the output of a neural network so that it is indeed making predictions over a space of structures.

Original: This is achieved through insights into the topological structures of an image and the usage of several important tools in topological data analysis.
Revised: This is achieved through insights into the topological structures of an image and the usage of several important tools in topological data analysis.

Original: To move from pixel space to structure space, we apply discrete Morse theory to decompose an image into a Morse complex, consisting of structures like branches, patches, etc.
Revised: To move from pixel space to structure space, we apply discrete Morse theory to decompose an image into a Morse complex, consisting of structures like branches, patches, etc.

Original: For simplification, we will use the term "branch" to denote a single piece of Morse structure.
Revised: For simplification, we will use the term "branch" to denote a single piece of Morse structure.

Original: These Morse branches are the hypothetical structures one can infer from the input image.
Revised: These Morse branches are the hypothetical structures one can infer from the input image.

Original: This decomposition is based on a likelihood function produced by a pixel-wise segmentation network trained in parallel.
Revised: This decomposition is based on a likelihood function produced by a pixel-wise segmentation network trained in parallel.

Original: Thus it is of a good quality, i.e., the structures are close enough to the true structures.
Revised: Thus it is of a good quality, i.e., the structures are close enough to the true structures.

Original: Any binary labeling of these Morse branches is a legitimate segmentation; we call it a structural segmentation .
Revised: Any binary labeling of these Morse branches is a legitimate segmentation; we call it a structural segmentation .

Original: But for full-scope reasoning of the structure space, instead of classifying these branches one-by-one, we would like to have the full inference, i.e., predicting a probability distribution for each branch.
Revised: But for full-scope reasoning of the structure space, instead of classifying these branches one-by-one, we would like to have the full inference, i.e., predicting a probability distribution for each branch.

Original: To further reduce the degrees of freedom to make the inference easier, we apply persistent homology to filter these branches with regard to their saliency.
Revised: To further reduce the degrees of freedom to make the inference easier, we apply persistent homology to filter these branches with regard to their saliency.

Original: This gives us a linear size family of structural segmentations, parameterized by a threshold ϵ .
Revised: This gives us a linear size family of structural segmentations, parameterized by a threshold ϵ .

Original: Finally, we learn an 1D Gaussian distribution for the ϵ as our probabilistic model.
Revised: Finally, we learn a 1D Gaussian distribution for the ϵ as our probabilistic model.

Original: This gives us the opportunity not only to sample segmentations, but also to provide a probability for each branch, which can be useful in downstream tasks including proofreading.
Revised: This gives us the opportunity not only to sample segmentations, but also to provide a probability for each branch, which can be useful in downstream tasks including proofreading.

Original: In Sec. 3.1, we introduce the discrete Morse theory and how to construct the space of Morse structures.
Revised: In Sec. 3.1, we introduce the discrete Morse theory and how to construct the space of Morse structures.

Original: We also explain how to use persistent homology to reduce the search space of reasoning into an 1-parameter family.
Revised: We also explain how to use persistent homology to reduce the search space of reasoning into a 1-parameter family.

Original: In Sec. 3.2, we will provide details on how our deep neural network is constructed to learn the probabilistic model over the structure space, as illustrated in Fig. 4.
Revised: In Sec. 3.2, we will provide details on how our deep neural network is constructed to learn the probabilistic model over the structure space, as illustrated in Fig. 4.

Original: 3. C ONSTRUCTING THE S TRUCTURE S PACE
Revised: 3. C ONSTRUCTING THE S TRUCTURE S PACE

Original: In this section, we explain how to construct a structural representation space using discrete Morse theory.
Revised: In this section, we explain how to construct a structural representation space using discrete Morse theory.

Original: The resulting structural representation space will be used to build a probabilistic model.
Revised: The resulting structural representation space will be used to build a probabilistic model.

Original: We will then discuss how to reduce the structure space into a 1-parameter family of structural segmentations, using persistent homology.
Revised: We will then discuss how to reduce the structure space into a 1-parameter family of structural segmentations, using persistent homology.

Original: We assume a 2D input image, although the method naturally extend to 3D images.
Revised: We assume a 2D input image, although the method naturally extend to 3D images.

Original: Given a reasonably clean input (e.g., the likelihood map of a deep neural network, Fig. 2(b)), we treat the 2D likelihood as a terrain function, and Morse theory (Milnor, 1963) can help to capture the structures regardless of weak/blurred conditions.
Revised: Given a reasonably clean input (e.g., the likelihood map of a deep neural network, Fig. 2(b)), we treat the 2D likelihood as a terrain function, and Morse theory (Milnor, 1963) can help to capture the structures regardless of weak/blurred conditions.

Original: See Fig.
Revised: See Fig.

Original: 3 for an illustration.
Revised: 3 for an illustration.

Original: The weak part of a line in the continuous map can be viewed as the local dip in the mountain ridge of the terrain.
Revised: The weak part of a line in the continuous map can be viewed as the local dip in the mountain ridge of the terrain.

Original: In the language of Morse theory, the lowest point of this dip is a saddle point ( S in Fig. 3(b)), and the mountain ridges which are connected to the saddle point ( M 1 S and M 2 S ) are called the stable manifolds of the saddle point.
Revised: In the language of Morse theory, the lowest point of this dip is a saddle point ( S in Fig. 3(b)), and the mountain ridges which are connected to the saddle point ( M 1 S and M 2 S ) are called the stable manifolds of the saddle point.

Original: We mainly focus on 2D images in this paper, although extending to 3D images is natural.
Revised: We mainly focus on 2D images in this paper.

Original: We consider two dimensional continuous function f : R 2 → R .
Revised: We consider two dimensional continuous function f : R 2 → R .

Original: For a point x ∈ R 2 , the gradient can be computed as ∇ f ( x ) =
Revised: For a point x ∈ R 2 , the gradient can be computed as ∇ f ( x ) =

Original: [ ∂f∂x 1 , ∂f∂x 2 ] T .
Revised: [ ∂f∂x 1 , ∂f∂x 2 ] T .

Original: We call a point x = ( x 1 , x 2 ) critical if ∇ f ( x ) = 0 .
Revised: We call a point x = ( x 1 , x 2 ) critical if ∇ f ( x ) = 0 .

Original: For a Morse function defined on R 2 , a critical point could be a minimum, a saddle or a maximum.
Revised: For a Morse function defined on R 2 , a critical point could be a minimum, a saddle or a maximum.

Original: Consider a continuous line (the red rectangle region in Fig. 3(a)) in a 2D likelihood map.
Revised: Consider a continuous line (the red rectangle region in Fig. 3(a)) in a 2D likelihood map.

Original: Imagine if we put a ball on one point of the line, then −∇ f ( x ) indicates the direction which the ball will flow down.
Revised: Imagine if we put a ball on one point of the line, then −∇ f ( x ) indicates the direction which the ball will flow down.

Original: By definition, the ball will eventually flow to the critical points where ∇ f ( x ) = 0 .
Revised: By definition, the ball will eventually flow to the critical points where ∇ f ( x ) = 0 .

Original: The collection of points whose ball eventually flows to p ( ∇ f ( p ) = 0 ) is defined as the stable manifold (denoted as S ( p ) ) of point p .
Revised: The collection of points whose ball eventually flows to p ( ∇ f ( p ) = 0 ) is defined as the stable manifold (denoted as S ( p ) ) of point p .

Original: Intuitively, for a 2D function f , the stable manifold S ( p ) of a minimum p is the entire valley of p (similar to the watershed algorithm); similarly, the stable manifold S ( q ) of a saddle point q consists of the whole ridge line which connects two local maxima and goes through the saddle point.
Revised: Intuitively, for a 2D function f , the stable manifold S ( p ) of a minimum p is the entire valley of p (similar to the watershed algorithm); similarly, the stable manifold S ( q ) of a saddle point q consists of the whole ridge line which connects two local maxima and goes through the saddle point.

Original: See Fig. 3(b) as an illustration. 
Revised: See Fig. 3(b) as an illustration. [Equation]

Original: [Figure]
Revised: [Figure]

Original: For a link-like structure, the stable manifold of a saddle contains the topological structures (usually curvilinear) of the continuous likelihood map predicted by deep neural networks, and they are exactly what we want to recover from noisy images.
Revised: For a link-like structure, the stable manifold of a saddle contains the topo logical structures (usually curvilinear) of the continuous likelihood map predicted by deep neural networks, and they are exactly what we want to recover from noisy images.

Original: In practice, we adopt the discrete version of Morse theory for images.
Revised: In practice, we adopt the discrete version of Morse theory for images.

Original: Discrete Morse theory.
Revised: Discrete Morse theory.

Original: Take a 2D image as a 2-dimensional cubical complex.
Revised: Take a 2D image as a 2-dimensional cubical complex (Wagner et al., 2012; Kaczynski et al., 2004).

Original: A 2-dimensional cubical complex then contains 0 -, 1 -, and 2dimensional cells, which correspond to vertices (pixels), edges and squares, respectively.
Revised: A 2-dimensional cubical complex then contains 0 -, 1 -, and 2 -dimensional cells, which correspond to vertices (pixels), edges and squares, respectively.

Original: In the setting of discrete Morse theory (DMT) (Forman, 1998; 2002), a pair of adjacent cells, termed as discrete gradient vectors, compose the gradient vector.
Revised: In the setting of discrete Morse theory (DMT) (Forman, 1998; 2002), a pair of adjacent cells, termed as discrete gradient vectors, compose the gradient vector.

Original: Critical points ( ∇ f ( x ) = 0 ) are those critical cells which are not in any discrete gradient vectors.
Revised: Critical points ( ∇ f ( x ) = 0 ) are those critical cells which are not in any discrete gradient vectors.

Original: In the 2D domain, a minimum, a saddle and a maximum correspond to a critical vertex, a critical edge and a critical square respectively.
Revised: In the 2D domain, a minimum, a saddle and a maximum correspond to a critical vertex, a critical edge and a critical square respectively.

Original: An 1-stable manifold (the stable manifold of a saddle point) in 2D corresponds to a V-path , i.e., connecting two local maxima and a saddle.
Revised: A 1-stable manifold (the stable manifold of a saddle point) in 2D corresponds to a V-path , i.e., connecting two local maxima and a saddle.

Original: See Fig. 3(b).
Revised: See Fig. 3(b).

Original: And the Morse complex generated by the DMT algorithm is illustrated in Fig. 3(c).
Revised: And the Morse complex generated by the DMT algorithm is illustrated in Fig. 3(c).

Original: Constructing the full structure space.
Revised: Constructing the full structure space.

Original: In this way, by using discrete Morse theory, for a likelihood map from the deep neural network, we can extract all the stable manifolds of the saddles, whose compositions constitute the full structure space.
Revised: In this way, by using discrete Morse theory, for a likelihood map from the deep neural network, we can extract all the stable manifolds of the saddles, whose compositions constitute the full structure space.

Original: Formally, we call any combinations of these stable manifolds a structure .
Revised: Formally, we call any combinations of these stable manifolds a structure .

Original: Fig. 2(c) illustrates 5 different structures.
Revised: Fig. 2(c) illustrates 5 different structures.

Original: This structure space, however, is of exponential size.
Revised: This structure space, however, is of exponential size.

Original: Assume there are N pieces 
Revised: Assume there are N pieces of stable manifolds/branches for a given likelihood map.

Original: Any combinations of these stable manifolds/branches will be a potential structure.
Revised: Any combinations of these stable manifolds/branches will be a potential structure.

Original: We will have 2 N possible structures in total.
Revised: We will have 2 N possible structures in total.

Original: This can be computationally prohibitive to construct and to model.
Revised: This can be computationally prohibitive to construct and to model.

Original: We need a principled way to reduce the structural search space.
Revised: We need a principled way to reduce the structural search space.

Original: Reducing the structural search space with persistent homology.
Revised: Reducing the structural search space with persistent homology.

Original: We propose to use persistent homology (Sousbie, 2011; Delgado-Friedrichs et al., 2015; Wang et al., 2015) to reduce the structural search space.
Revised: We propose to use persistent homology (Sousbie, 2011; Delgado-Friedrichs et al., 2015; Wang et al., 2015) to reduce the structural search space.

Original: Persistent homology is an important tool for topological data analysis (Edelsbrunner &
Revised: Persistent homology is an important tool for topological data analysis (Edelsbrunner &

Original: Harer, 2010; Edelsbrunner et al., 2000).
Revised: Harer, 2022; Edelsbrunner et al., 2000).

Original: Intuitively, we grow a Morse complex by gradually including[Figure] more and more discrete elements (called cells) from empty.
Revised: Intuitively, we grow a Morse complex by gradually including more and more discrete elements (called cells) from empty.

Original: A branch of the Morse complex is a specialtype of cell.
Revised: A branch of the Morse complex is a special

Original: Other types include vertices, patches, etc. Cells will be continuously added to the complex.
Revised: [Figure]

Original: New branches will be born and existing branches will die.
Revised: New branches will be born and existing branches will die.

Original: The persistence algorithm (Edelsbrunner et al., 2000) pairs up all these critical cells as birth and death pairs.
Revised: The persistence algorithm (Edelsbrunner et al., 2000) pairs up all these critical cells as birth and death pairs.

Original: The difference of their function values is essentially the life time of the specific topological structure/branch, which is called the persistence .
Revised: The difference of their function values is essentially the life time of the specific topological structure/branch, which is called the persistence .

Original: The importance of a branch is associated with its persistence.
Revised: The importance of a branch is associated with its persistence.

Original: Intuitively, the longer the persistence of a specific branch is, the more important the branch is.
Revised: Intuitively, the longer the persistence of a specific branch is, the more important the branch is.

Original: Recall that our original construction of the structure space considers all possible combination of branches, and thus can have exponentially many combinations.
Revised: Recall that our original construction of the structure space considers all possible combinations of branches, and thus can have exponentially many combinations.

Original: Instead, we propose to only select branches with high persistence as important ones.
Revised: Instead, we propose to only select branches with high persistence as important ones.

Original: By doing this, we will be able to prune the less important/noisy branches very efficiently, and recover the branches with true signals.
Revised: By doing this, we will be able to prune the less important/noisy branches very efficiently, and recover the branches with true signals.

Original: Specifically, the structure pruning is done via Morse cancellation (more details are included in Appendix A.3) operation.
Revised: Specifically, the structure pruning is done via Morse cancellation (more details are included in Appendix A.3) operation.

Original: The persistence thresholding provides us the opportunity to obtain a structure space of linear size .
Revised: The persistence thresholding provides us the opportunity to obtain a structure space of linear size .

Original: We start with the complete Morse complex, and continuously increase the threshold ϵ .
Revised: We start with the complete Morse complex, and continuously increase the threshold ϵ .

Original: At each threshold, we obtain a structure by filtering with ϵ and only keeping the branches whose persistence is above ϵ .
Revised: At each threshold, we obtain a structure by filtering with ϵ and only keeping the branches whose persistence is above ϵ .

Original: This gives a sequence of structures parametrized by ϵ .
Revised: This gives a sequence of structures parametrized by ϵ .

Original: As shown in Fig. 2(d), the family of structures represents different structural densities.
Revised: As shown in Fig. 2(d), the family of structures represents different structural densities.

Original: The one-parameter space allows us to easily learn a probabilistic model and carry out various inference tasks such as segmentation, sampling, uncertainty estimation and interactive proofreading.
Revised: The one-parameter space allows us to easily learn a probabilistic model and carry out various inference tasks such as segmentation, sampling, uncertainty estimation and interactive proofreading.

Original: Specifically, we will learn a Gaussian distribution over the persistence threshold ϵ , ϵ ∼ N ( µ, σ ) .
Revised: Specifically, we will learn a Gaussian distribution over the persistence threshold ϵ , ϵ ∼ N ( µ, σ ) .

Original: Denote the persistence of a branch b as ϵ b .
Revised: Denote the persistence of a branch b as ϵ b .

Original: Any branch b belongs to the structure map M
Revised: Any branch b belongs to the structure map M

Original: (we also call the structure map M a structural segmentation) as long as its persistence is smaller or equal to themaximal persistence threshold of M , i.e., b ∈ M if and only if ϵ b ≤ ϵ M , where ϵ M is used to generate M ( ϵ M ≥ max b ∈ M ϵ b ).
Revised: (we also call the structure map M a structural segmentation) as long as its persistence is higher or equal to the persistence threshold of M , i.e., b ∈ M if and only if ϵ b ≥ ϵ M , where ϵ M is used to generate M .

Original: More details will be provided in Sec. 3.2.
Revised: More details will be provided in Sec. 3.2.

Original: Approximation of Morse structures for volume data.
Revised: Approximation of Morse structures for volume data.

Original: In the 2D setting, the stable manifold of saddles is composed of curvilinear structures, and the captured Morse structures will essentially contain the non-boundary edges , which fits well with vessel data.
Revised: In the 2D setting, the stable manifold of saddles is composed of curvilinear structures, and the captured Morse structures will essentially contain the non-boundary edges , which fits well with vessel data.

Original: However, the output structures should always be boundary edges for volume data, which can’t be dealt with the original discrete
Revised: However, the output structures should always be boundary edges for volume data, which cannot be dealt with the original discrete

Original: Morse theory.
Revised: Morse theory.

Original: Consequently, we approximate the Morse structures of 2D volume data with the boundaries of the stable manifolds of local minima.
Revised: Consequently, we approximate the Morse structures of 2D volume data with the boundaries of the stable manifolds of local minima.

Original: As mentioned above, the stable manifold of a local minimum p in the 2D setting corresponds to the whole valley, and the boundaries of these valleys construct the approximation of the Morse structures for volume data.
Revised: As mentioned above, the stable manifold of a local minimum p in the 2D setting corresponds to the whole valley, and the boundaries of these valleys construct the approximation of the Morse structures for volume data.

Original: Similar to the original discrete
Revised: Similar to the original discrete

Original: Morse theory, we also introduce a persistence threshold parameter ϵ and use persistent homology to prune the less important branches.
Revised: Morse theory, we also introduce a persistence threshold parameter ϵ and use persistent homology to prune the less important branches.

Original: The details of the proposed persistent-homology filtered topology watershed algorithm are illustrated in Appendix A.4.
Revised: The details of the proposed persistent-homology filtered topology watershed algorithm are illustrated in Appendix A.4.

Original: 3. N EURAL N ETWORK A RCHITECTURE
Revised: 3. N EURAL N ETWORK A RCHITECTURE

Original: In this section, we introduce our neural network that learns a probabilistic model over the structural representation to obtain structural segmentations.
Revised: In this section, we introduce our neural network that learns a probabilistic model over the structural representation to obtain structural segmentations.

Original: See Fig.
Revised: See Fig.

Original: 4 for an illustration of the overall pipeline.
Revised: 4 for an illustration of the overall pipeline.

Original: Since the structural reasoning needs a sufficiently clean input to construct discrete Morse complexes, our method first obtains such a likelihood map by training a segmentation branch which is supervised by the standard segmentation loss,  cross-entropy loss.
Revised: Since the structural reasoning needs a sufficiently clean input to construct discrete Morse complexes, our method first obtains such a likelihood map by training a segmentation branch which is supervised by the standard segmentation loss, binary cross-entropy loss.

Original: Formally, L seg =
Revised: Formally, L seg =

Original: L bce ( Y, S ( X ; ω seg )) , in which  S ( X ; ω seg ) is the output likelihood map, ω seg is the segmentation branch’s weight.
Revised: L bce ( Y, S ( X ; ω seg )) , in which X is the input image, ω seg is the segmentation branch’s weight, S ( X ; ω seg ) is the output likelihood map, and Y is the ground truth.

Original: The output likelihood map, S ( X ; ω seg ) , is used as the input for the discrete Morse theory algorithm (DMT), which generates a discrete Morse complex consisting of all possible Morse branches from the likelihood map.
Revised: The output likelihood map, S ( X ; ω seg ) , is used as the input for the discrete Morse theory algorithm (DMT), which generates a discrete Morse complex consisting of all possible Morse branches from the likelihood map.

Original: Thresholding these branches using persistent homology with different ϵ ’s will produce different structures.
Revised: Thresholding these branches using persistent homology with different ϵ values will produce different structures.

Original: We refer to the DMT computation and the persistent homology thresholding operation as f DMT and f PH .
Revised: We refer to the DMT computation and the persistent homology thresholding operation as f DMT and f PH .

Original: So given a likelihood map S ( X ; ω seg ) and a threshold ϵ , we can generate a structure (which we call a skeleton): S skeleton ( ϵ ) =
Revised: So given a likelihood map S ( X ; ω seg ) and a threshold ϵ , we can generate a structure (which we call a skeleton): S skeleton ( ϵ ) =

Original: f PH ( f DMT ( S ( X ; ω seg )); ϵ ) .
Revised: f PH ( f DMT ( S ( X ; ω seg )); ϵ ) .

Original: Next, we discuss how to learn the probabilistic model.
Revised: Next, we discuss how to learn the probabilistic model.

Original: Recall  we want to learn a Gaussian distribution over the persistent homology threshold, ϵ ∼ N ( µ, σ ) .
Revised: Recall that we want to learn a Gaussian distribution over the persistent homology threshold, ϵ ∼ N ( µ, σ ) .

Original: The parameters µ and σ are learned by a neural network called the posterior network .
Revised: The parameters µ and σ are learned by a neural network called the posterior network .

Original: The network uses the input image X and the corresponding ground truth mask Y as input, and outputs the parameters µ ( X, Y ; ω post ) and σ ( X, Y ; ω post ) .
Revised: The network uses the input image X and the corresponding ground truth mask Y as input, and outputs the parameters µ ( X, Y ; ω post ) and σ ( X, Y ; ω post ) .

Original: ω post is the parameter of the network.
Revised: ω post is the parameter of the network.

Original: During training, at each iteration, we draw a sample ϵ from the distribution ( ϵ ∼ N ( µ, σ ) ).
Revised: During training, at each iteration, we draw a sample ϵ from the distribution ( ϵ ∼ N ( µ, σ ) ).

Original: Using the sample ϵ , together with the likelihood map, we can generate the corresponding sampled structure,
Revised: Using the sample ϵ , together with the likelihood map, we can generate the corresponding sampled structure,

Original: S skeleton ( ϵ ) .
Revised: S skeleton ( ϵ ) .

Original: This skeleton will be compared with the ground truth for supervision.
Revised: This skeleton will be compared with the ground truth for supervision.

Original: To compare a sampled skeleton, S skeleton ( ϵ ) , with ground truth Y , we use the skeleton to mask both Y and the likelihood map S ( X ; ω seg ) , and then compare the skeleton-masked ground truth and the likelihood using  cross-entropy loss: L bce ( Y ◦ S skeleton ( ϵ ) , S ( X ; ω seg ) ◦ S skeleton ( ϵ )) .
Revised: To compare a sampled skeleton, S skeleton ( ϵ ) , with ground truth Y , we use the skeleton to mask both Y and the likelihood map S ( X ; ω seg ) , and then compare the skeleton-masked ground truth and the likelihood using binary cross-entropy loss: L bce ( Y ◦ S skeleton ( ϵ ) , S ( X ; ω seg ) ◦ S skeleton ( ϵ )) , in which ◦ denotes the Hadamard product.

Original: To learn the distribution, we use the expected loss:
Revised: To learn the distribution, we use the expected loss:

Original: [Equation]
Revised: [Equation]

Original: The losscan be backpropagated through the posterior network through reparameterization technique  (Kingma & Welling, 2013).
Revised: We backpropagate this loss through the posterior network using the reparameterization technique in (Kingma & Welling, 2013).

Original: More details are provided in Appendix A.5.
Revised: More details are provided in Appendix A.5.

Original: Note that this loss will also provide supervision to the segmentation network through the likelihood map.
Revised: Note that this loss will also provide supervision to the segmentation network through the likelihood map.

Original: Learning a prior network from the posterior network.
Revised: Learning a prior network from the posterior network.

Original: Although our posterior network can learn the distribution well, it relies on the ground truth mask Y as input.
Revised: Although our posterior network can learn the distribution well, it relies on the ground truth mask Y as input.

Original: This is not available at inference stage.
Revised: This is not available at inference stage.

Original: To address this issue, inspired by Probabilistic-UNet (Kohl et al., 2018), we use another network to learn the distribution of ϵ with only the image X as input.
Revised: To address this issue, inspired by VAE (Kingma & Welling, 2013; Kohl et al., 2018), we use another network to learn the distribution of ϵ with only the image X as input.

Original: We call this network the prior net .
Revised: We call this network the prior net .

Original: We denote by P the distribution using parameters predicted by the prior network, and denote by Q the distribution predicted by the posterior network.
Revised: We denote by P the distribution using parameters predicted by the prior network, and denote by Q the distribution predicted by the posterior network.

Original: Duringthe training, we want to use the prior net to mimic the posterior net; and then in the inference stage, wecan use the prior net to obtain a reliable distribution over ϵ with only the image X .
Revised: During training, we want to force the prior net to mimic the posterior net; and then in the inference stage, we use the prior net to obtain a reliable distribution over ϵ with only the image X .

Original: Thus, we incorporate the Kullback-Leibler divergence of these two distributions, D KL ( Q || P ) =
Revised: Thus, we incorporate the Kullback-Leibler divergence of these two distributions, D KL ( Q || P ) =

Original: E ϵ ∼ Q (log QP ) , which measures  thedifferences of prior distribution P ( N ( µ prior , σ prior ) ) and the posterior distribution Q ( N ( µ post , σ post ) ).
Revised: E ϵ ∼ Q (log QP ) , which measures how close the prior distribution P ( N ( µ prior , σ prior ) ) is to the posterior distribution Q ( N ( µ post , σ post ) ).

Original: Training the neural network.
Revised: Training the neural network.

Original: The final loss is composed by the standard segmentation loss, the skeleton loss L skeleton , and the KL divergence loss, with two hyperparameters α and β to balance the three terms,
Revised: The final loss is composed by the standard segmentation loss, the skeleton loss L skeleton , and the KL divergence loss, with two hyperparameters α and β to balance the three terms,

Original: [Equation]
Revised: [Equation]

Original: The network is trained to jointly optimize the segmentation branch and the probabilistic branch (containing both prior and posterior nets) simultaneously.
Revised: The network is trained to jointly optimize the segmentation branch and the probabilistic branch (containing both prior and posterior nets) simultaneously.

Original: During the training stage, the KL divergence loss ( D KL ) pushes the prior distribution towards the posterior distribution.
Revised: During the training stage, the KL divergence loss ( D KL ) pushes the prior distribution towards the posterior distribution.

Original: The training scheme is also illustrated in Fig.
Revised: The training scheme is also illustrated in Fig.

Original: Inference stage: generating structure-preserving segmentation maps.
Revised: Inference stage: generating structure-preserving segmentation maps.

Original: In the inference stage, given an input image, we are able to produce unlimited number of plausible structure-preserving skeletons via sampling.
Revised: In the inference stage, given an input image, we are able to produce unlimited number of plausible structure-preserving skeletons via sampling.

Original: We use a postprocessing step to grow the 1-pixel wide structures/skeletons without changing its topology as the final structural segmentation.
Revised: We use a postprocessing step to grow the 1-pixel wide structures/skeletons without changing their topology as the final structural segmentation.

Original: Specifically, the skeletons are overlaid on the binarized initial segmentation map (Fig. 5(c)), and only the connected components which exist in the skeletons are kept as the final segmentation maps (Fig. 5(e)).
Revised: Specifically, the skeletons are overlaid on the binarized initial segmentation map (Fig. 5(c)), and only the connected components which exist in the skeletons are kept as the final segmentation maps (Fig. 5(e)).

Original: In this way, each plausible skeleton (Fig. 5(d)) generates one final segmentation map (Fig. 5(e)) and it has exact the same topology as the corresponding skeleton.
Revised: In this way, each plausible skeleton (Fig. 5(d)) generates one final segmentation map (Fig. 5(e)) and it has exact the same topology as the corresponding skeleton.

Original: The pipeline of the procedure is illustrated in Fig.
Revised: The pipeline of the procedure is illustrated in Fig.

Original: Uncertainty of structures.
Revised: Uncertainty of structures.

Original: Given a learned prior distribution, P , over the family of structural segmentations, we can naturally calculate the probability of each Morse structure/branch.Recall [Figure]
Revised: Given a learned prior distribution, P , over the family of structural segmentations, we can naturally calculate the probability of each Morse structure/branch.

Original:  a branch b has its persistence ϵ b .
Revised: Recall a branch b has its persistence ϵ b .

Original: And the prior probability of a structural segmentation map M is P ( ϵ M ) , in which ϵ M is used to generate M .
Revised: And the prior probability of a structural segmentation map M is P ( ϵ M ) , in which ϵ M is used to generate M .

Original: Also any branch b whose persistence is smaller or equal to themaximal persistence threshold of M belongs to M , i.e., b ∈ M if and only if ϵ b≤ ϵ M . Thus we have ϵ M ≥max b ∈ M ϵ b .
Revised: Also any branch b whose persistence is higher [Figure] or equal to the persistence threshold of M belongs to M , i.e., b ∈ M if and only if ϵ b ≥ ϵ M .

Original: Therefore, the probability of a branch b being in a segmentation map M such as ϵ M ∼ P follows a Bernoulli distribution with the probability Pr ( b ) being the cumulative distribution function (CDF) ofP , CDF P ( ϵ b ) =
Revised: Therefore, the probability of a branch b being in a segmentation map M such as ϵ M ∼ P follows a Bernoulli distribution with the probability Pr ( b ) being the cumulative distribution function (CDF) of

Original: P  ( ϵ ≤ ϵ b ) .
Revised: P , CDF P ( ϵ b ) = P ( ϵ ≤ ϵ b ) .

Original: This can be directly calculated at inference, and the absolute difference of the CDF from 0.5 is the confidence ( which equals 1-uncertainty ) of the Morse structure/branch b . 
Revised: This can be directly calculated at inference, and the absolute difference of the CDF from 0.5 is the confidence ( which equals 1-uncertainty ) of the Morse structure/branch b . We add an illustration of uncertainty estimation for branches in Fig. 14 (Sec. A.12 of Appendix).

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our method directly makes prediction and inference on structures rather than on pixels.
Revised: Our method directly makes prediction and inference on structures rather than on pixels.

Original: This can significantly benefit downstream tasks.
Revised: This can significantly benefit downstream tasks.

Original: While probabilities of structures can certainly be used for further analysis of the structural system, in this paper we focus on both automatic image segmentation and semi-automatic annotation/proofreading tasks.
Revised: While probabilities of structures can certainly be used for further analysis of the structural system, in this paper we focus on both automatic image segmentation and semi-automatic annotation/proofreading tasks.

Original: On automatic image segmentation, we show that direct prediction can ensure topological integrity even better than previous topology-aware losses.
Revised: On automatic image segmentation, we show that direct prediction can ensure topological integrity even better than previous topology-aware losses.

Original: This is not surprising as our prediction is on structures.
Revised: This is not surprising as our prediction is on structures.

Original: On semi-automatic proofreading task, we show our structure-level uncertainty can assist human annotators to obtain satisfying segmentation annotations in a much more efficient manner than previous methods.
Revised: On semi-automatic proofreading task, we show our structure-level uncertainty can assist human annotators to obtain satisfying segmentation annotations in a much more efficient manner than previous methods.

Original: 4. A UTOMATIC T OPOLOGY -A WARE I MAGE S EGMENTATION
Revised: 4. A UTOMATIC T OPOLOGY -A WARE I MAGE S EGMENTATION

Original: Datasets.
Revised: Datasets.

Original: We use three datasets to validate the efficacy of the proposed method: ISBI13 (ArgandaCarreras et al., 2013) (volume), CREMI (volume), and DRIVE (Staal et al., 2004)
Revised: We use three datasets to validate the efficacy of the proposed method: ISBI13 (ArgandaCarreras et al., 2013) (volume), CREMI (volume), and DRIVE (Staal et al., 2004)

Original: (vessel).
Revised: (vessel).

Original: More details are included in Appendix A.6.
Revised: More details are included in Appendix A.6.

Original: Evaluation metrics.
Revised: Evaluation metrics.

Original: We use four different evaluation metrics: Dice score , ARI , VOI , and Betti number error .
Revised: We use four different evaluation metrics: Dice score , ARI , VOI , and Betti number error .

Original: Dice is a popular pixel-wise segmentation metric, and the other three are structure/topologyaware segmentation metrics.
Revised: Dice is a popular pixel-wise segmentation metric, and the other three are structure/topologyaware segmentation metrics.

Original: More details are included in Appendix A.7.
Revised: More details are included in Appendix A.7.

Original: Baselines.
Revised: Baselines.

Original: We compare the proposed method with two kinds of baselines: 1) Standard segmentation baselines:
Revised: We compare the proposed method with two kinds of baselines: 1) Standard segmentation baselines:

Original: DIVE (Fakhry et al., 2016), UNet (Ronneberger et al., 2015), UNet-VGG (Mosinska et al., 2018), TopoLoss (Hu et al., 2019) and DMT (Hu et al., 2021).
Revised: DIVE (Fakhry et al., 2016), UNet (Ronneberger et al., 2015), UNet-VGG (Mosinska et al., 2018), TopoLoss (Hu et al., 2019) and DMT (Hu et al., 2021).

Original: 2) Probabilistic-based segmentation methods: Dropout UNet (Kendall et al., 2015) and Probabilistic-UNet (Kohl et al., 2018).
Revised: 2) Probabilistic-based segmentation methods: Dropout UNet (Kendall et al., 2015) and Probabilistic-UNet (Kohl et al., 2018).

Original: More details about these baselines are included in Appendix A.8.
Revised: More details about these baselines are included in Appendix A.8.

Original: Quantitative and qualitative results.
Revised: Quantitative and qualitative results.

Original: Table 1 shows the quantitative results comparing to several baselines.
Revised: Table 1 shows the quantitative results comparing to several baselines.

Original: Note that for deterministic methods, the numbers are computed directly based on the outputs; while for probabilistic methods, we generate five segmentation masks and report the averaged numbers over the five segmentation masks for each image (for both the baselines and the proposed method).
Revised: Note that for deterministic methods, the numbers are computed directly based on the outputs; while for probabilistic methods, we generate five segmentation masks and report the averaged numbers over the five segmentation masks for each image (for both the baselines and the proposed method).

Original: We use t-test (95% confidence interval) to determine the statistical significance and highlight the significant better results.
Revised: We use t-test (95% confidence interval) to determine the statistical significance and highlight the significant better results.

Original: From the table, we can observe that the proposed method achieves significant better performances in terms of topology-aware metrics (ARI, VOI and Betti Error).
Revised: From the table, we can observe that the proposed method achieves significant better performances in terms of topology-aware metrics (ARI, VOI and Betti Error).

Original: Fig.
Revised: Fig.

Original: 6 shows qualitative results.
Revised: 6 shows qualitative results.

Original: Comparing with DMT (Hu et al., 2021), our method is able to produce a set of true structure-preserving segmentation maps, as illustrated in Fig.
Revised: Comparing with DMT (Hu et al., 2021), our method is able to produce a set of true structure-preserving segmentation maps, as illustrated in Fig.

Original: 6(e-g).
Revised: 6(e-g).

Original: Note that compared with the existing topology-aware segmentation methods, our method is more capable of recovering the weak connections by using Morse skeletons as hints.
Revised: Note that compared with the existing topology-aware segmentation methods, our method is more capable of recovering the weak connections by using Morse skeletons as hints.

Original: More qualitative results are included in Appendix A.2.
Revised: More qualitative results are included in Appendix A.2.

Original: [Table]
Revised: [Table]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Ablation study of loss weights.
Revised: Ablation study of loss weights.

Original: We observe that the performances of our method are quite robust to the loss weights α and β .
Revised: We observe that the performances of our method are quite robust to the loss weights α and β .

Original: As the learned distribution over the persistence threshold might affect the final performances, we conduct an ablation study in terms of the weight of KL divergence loss ( β ) on DRIVE dataset.
Revised: As the learned distribution over the persistence threshold might affect the final performances, we conduct an ablation study in terms of the weight of KL divergence loss ( β ) on DRIVE dataset.

Original: The results are reported in
Revised: The results are reported in

Original: Fig.
Revised: Fig.

Original: When β = 10 , the model achieves slightly better performance in terms of VOI (0.804 ± 0.047, the smaller the better) than other choices.
Revised: When β = 10 , the model achieves slightly better performance in terms of VOI (0.804 ± 0.047, the smaller the better) than other choices.

Original: Note that, for all the experiments, we set α = 1 .
Revised: Note that, for all the experiments, we set α = 1 .

Original: [Figure]
Revised: [Figure]

Original: Illustration of the structure-level uncertainty.
Revised: Illustration of the structure-level uncertainty.

Original: We also explore the structure-level uncertainty with the proposed method here.
Revised: We also explore the structure-level uncertainty with the proposed method here.

Original: We show three sampled masks (Fig.
Revised: We show three sampled masks (Fig.

Original: 8(c-e)) in the inference stage for a given image (Fig. 8(a)), and the structure-level uncertainty map (Fig. 8(f)).
Revised: 8(c-e)) in the inference stage for a given image (Fig. 8(a)), and the structure-level uncertainty map (Fig. 8(f)).

Original: Note that in practice, for simplification, different from Sec.
Revised: Note that in practice, for simplification, different from Sec.

Original: 3.2, we empirically generate uncertainty map by taking variance across all the samples (the number is 10 in our case).
Revised: 3.2, we empirically generate uncertainty map by taking variance across all the samples (the number is 10 in our case).

Original: Different from pixel-wise uncertainty, each small branch has the same uncertainty value with our method.
Revised: Different from pixel-wise uncertainty, each small branch has the same uncertainty value with our method.

Original: By looking at the original image, we find that the uncertainties are usually caused by the weak signals of the original image.
Revised: By looking at the original image, we find that the uncertainties are usually caused by the weak signals of the original image.

Original: The weak signals of the original image make the model difficult to predict these locations correctly and confidently, especially in structure wise.
Revised: The weak signals of the original image make the model difficult to predict these locations correctly and confidently, especially in structure wise.

Original: Actually this also makes sense in real cases.
Revised: Actually this also makes sense in real cases.

Original: Different from natural images, even experts can not always reach a consensus for biomedical image annotation (Armato III et al., 2011;
Revised: Different from natural images, even experts can not always reach a consensus for biomedical image annotation (Armato III et al., 2011;

Original: Clark et al., 2013).
Revised: Clark et al., 2013).

Original: More details of structure-level uncertainty are included in Appendix A.1.
Revised: More details of structure-level uncertainty are included in Appendix A.1.

Original: The advantage of the joint training and optimization.
Revised: The advantage of the joint training and optimization.

Original: Another straightforward alternative of the proposed approach is to use the discrete Morse theory to postprocess the continuous likelihood map obtained from the standard segmentation networks.
Revised: Another straightforward alternative of the proposed approach is to use the discrete Morse theory to postprocess the continuous likelihood map obtained from the standard segmentation networks.

Original: More discussions are provided in Appendix A.9.
Revised: More discussions are provided in Appendix A.9.

Original: 4. S EMI -A UTOMATIC E FFICIENT A NNOTATION /P ROOFREADING WITH U SER I NTERACTION
Revised: 4. S EMI -A UTOMATIC E FFICIENT A NNOTATION /P ROOFREADING WITH U SER I NTERACTION

Original: Proofreading is a struggling while essential step for the annotation of images with fine-scale structures.
Revised: Proofreading is a struggling while essential step for the annotation of images with fine-scale structures.

Original: We develop a semi-automatic interactive annotation/proofreading pipeline based on the proposed method.
Revised: We develop a semi-automatic interactive annotation/proofreading pipeline based on the proposed method.

Original: As the proposed method is able to generate structural segmentations (Fig. 8(c)-(e)), the annotators can efficiently proofread one image with rich structures by removing the unnecessary/redrawing the missing structures with the help of structure-level uncertainty map (Fig. 8(f)).
Revised: As the proposed method is able to generate structural segmentations (Fig. 8(c)-(e)), the annotators can efficiently proofread one image with rich structures by removing the unnecessary/redrawing the missing structures with the help of structure-level uncertainty map (Fig. 8(f)).

Original: The whole inference and structure-aware interactive image annotation pipeline is illustrated in Fig.
Revised: The whole inference and structure-aware interactive image annotation pipeline is illustrated in Fig.

Original: [Figure]
Revised: [Figure]

Original: We conduct empirical experiments to demonstrate the efficiency of proofreading by using the proposed method.
Revised: We conduct empirical experiments to demonstrate the efficiency of proofreading by using the proposed method.

Original: We randomly select a few samples from ISBI13 dataset and simulate the user interaction process.
Revised: We randomly select a few samples from ISBI13 dataset and simulate the user interaction process.

Original: For both the proposed and baseline methods, we get started from the final segmentations and correct one misclassified branch each time.
Revised: For both the proposed and baseline methods, we get started from the final segmentations and correct one misclassified branch each time.

Original: For deterministic method (DMT), the user draws one false-negative or erases one false-positive branch for each click.
Revised: For deterministic method (DMT), the user draws one false-negative or erases one false-positive branch for each click.

Original: For pixel-wise probabilistic method (Prob.-
Revised: For pixel-wise probabilistic method (Prob.-

Original: UNet), the user does the same while taking the uncertainty map as guidance.
Revised: UNet), the user does the same while taking the uncertainty map as guidance.

Original: For the proposed method, the user checks each branch based on the descending order of structure-level uncertainty.
Revised: For the proposed method, the user checks each branch based on the descending order of structure-level uncertainty.

Original: VOI is used to evaluate the performances.
Revised: VOI is used to evaluate the performances.

Original: Fig.
Revised: Fig.

Original: 9 shows the comparative results of semi-automatic proofreading with user interactions.
Revised: 9 shows the comparative results of semi-automatic proofreading with user interactions.

Original: By always checking branches with highest uncertainty, the proposed method clearly achieves better results and improves the results much faster than baseline methods.
Revised: By always checking branches with highest uncertainty, the proposed method clearly achieves better results and improves the results much faster than baseline methods.

Original: Our developed pipeline achieves higher efficiency because of following two perspectives: 1) the generated structural segmentations are essentially topology/structure-preserving; 2) the proposed method provides the structure-level uncertainty map to guide the human proofreading.
Revised: Our developed pipeline achieves higher efficiency because of following two perspectives: 1) the generated structural segmentations are essentially topology/structure-preserving; 2) the proposed method provides the structure-level uncertainty map to guide the human proofreading.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Instead of making pixel-wise predictions, we propose to learn structural representation with a probabilistic model to obtain structural segmentations.
Revised: Instead of making pixel-wise predictions, we propose to learn structural representation with a probabilistic model to obtain structural segmentations.

Original: Specifically, we construct the structure space by leveraging classical discrete Morse theory.
Revised: Specifically, we construct the structure space by leveraging classical discrete Morse theory.

Original: We then build a probabilistic model to learn a distribution over structures.
Revised: We then build a probabilistic model to learn a distribution over structures.

Original: In the inference stage, we are able to generate a set of structural segmentations and explore the structure-level uncertainty, which is beneficial for interactive proofreading.
Revised: In the inference stage, we are able to generate a set of structural segmentations and explore the structure-level uncertainty, which is beneficial for interactive proofreading.

Original: Extensive experiments have been conducted to demonstrate the efficacy of the proposed method.
Revised: Extensive experiments have been conducted to demonstrate the efficacy of the proposed method.

Original: Reproducibility Statement: We provide the necessary experimental details in Sec. 4.
Revised: Reproducibility Statement: We provide the necessary experimental details in Sec. 4.

Original: More specifically, the details of the data are provided in Appendix A.6.
Revised: More specifically, the details of the data are provided in Appendix A.6.

Original: The details of baseline methods are described in Appendix A.8.
Revised: The details of baseline methods are described in Appendix A.8.

Original: Appendix A.7 contains the evaluations metrics used in this paper.
Revised: Appendix A.7 contains the evaluations metrics used in this paper.

Original: The used computation resources are specified in Appendix A.10.
Revised: The used computation resources are specified in Appendix A.13.

Original: A A PPENDIX
Revised: A A PPENDIX

Original: Appendix A.1 illustrates structure-level uncertainty compared with traditional pixel-wise uncertainty.
Revised: Appendix A.1 illustrates structure-level uncertainty compared with traditional pixel-wise uncertainty.

Original: Appendix A.2 shows more qualitative results.
Revised: Appendix A.2 shows more qualitative results.

Original: Appendix A.3 provides the details of Morse cancellation.
Revised: Appendix A.3 provides the details of Morse cancellation.

Original: Appendix A.4 illustrates the details of persistent-homology filtered topology watershed algorithm.
Revised: Appendix A.4 illustrates the details of persistent-homology filtered topology watershed algorithm.

Original: Appendix A.5 describes the reparameterization technique.
Revised: Appendix A.5 describes the reparameterization technique.

Original: Appendix A.6 provides the details of the datasets used in this paper.
Revised: Appendix A.6 provides the details of the datasets used in this paper.

Original: Appendix A.7 illustrates the details of the metrics used in this paper.
Revised: Appendix A.7 illustrates the details of the metrics used in this paper.

Original: Appendix A.8 describes the details of the baselines used in this paper.
Revised: Appendix A.8 describes the details of the baselines used in this paper.

Original: Appendix A.9 discusses the advantage of joint training and optimization.
Revised: Appendix A.9 discusses the advantage of joint training and optimization.

Original: Appendix A.10 provides the computational resources for all the conducted experiments.
Revised: Appendix A.13 provides the computational resources for all the conducted experiments.

Original: A. I LLUSTRATION OF THE STRUCTURE - LEVEL UNCERTAINTY
Revised: A. I LLUSTRATION OF THE STRUCTURE - LEVEL UNCERTAINTY

Original: Fig.
Revised: Fig.

Original: 1 shows the comparison of traditional pixel-wise uncertainty and the proposed structure-level uncertainty.
Revised: 1 shows the comparison of traditional pixel-wise uncertainty and the proposed structure-level uncertainty.

Original: Specifically, Fig.
Revised: Specifically, Fig.

Original: 1(c) is a sampled segmentation result by Prob.-UNet (Kohl et al., 2018), and Fig. 1(d) is the pixel uncertainty map from Prob.-UNet (Kohl et al., 2018).
Revised: 1(c) is a sampled segmentation result by Prob.-UNet (Kohl et al., 2018), and Fig. 1(d) is the pixel uncertainty map from Prob.-UNet (Kohl et al., 2018).

Original: Different from traditional pixel-wise uncertainty, our proposed structure-level uncertainty (Fig. 1(f)) can focus on the structures.
Revised: Different from traditional pixel-wise uncertainty, our proposed structure-level uncertainty (Fig. 1(f)) can focus on the structures.

Original: We also overlay the structure-level uncertainty (Fig. 1(f)) on the original image (Fig. 1(a)), which is shown in Fig.
Revised: We also overlay the structure-level uncertainty (Fig. 1(f)) on the original image (Fig. 1(a)), which is shown in Fig.

Original: By comparison with the original image, we can observe that the structure-level uncertainty is mainly caused by the weak signals in the original image.
Revised: By comparison with the original image, we can observe that the structure-level uncertainty is mainly caused by the weak signals in the original image.

Original: [Figure]
Revised: [Figure]

Original: A. Q UALITATIVE RESULTS
Revised: A. Q UALITATIVE RESULTS

Original: Fig.
Revised: Fig.

Original: 11 shows more qualitative results.
Revised: 11 shows more qualitative results.

Original: From Fig.
Revised: From Fig.

Original: 11, we can observe that the proposed method can generate both diverse and structure-preserving segmentation maps
Revised: 11, we can observe that the proposed method can generate both diverse and structure-preserving segmentation maps.

Original: . A.  M ORSE CANCELLATION
Revised: A. A DDITIONAL DETAILS ON THE METHOD A.3. D ISCRETE M ORSE T HEORY

Revised: A 2D image is viewed as a 2-dimensional cubical complex, consisting of 0 -, 1 -, and 2 -dimensional cells corresponding to vertices, edges and squares as its building blocks. [Figure] Discrete Morse theory (DMT) (Forman, 1998; 2002) is a combinatorial version of Morse theory for general cell complexes.

Revised: We will briefly introduce some relevant concepts for the present paper, and describe it in the setting of cubical complexes for images.

Revised: Discrete gradient vector (also called a V-pair for simplicity):

Revised: Let K be a cubical complex.

Revised: Given a p -cell τ , we denote by σ < τ if σ is a ( p − 1 )-dimensional face for τ .

Revised: Discrete gradient vector is a pair ( τ, σ ) where σ < τ .

Revised: Given a collection of V-pairs M ( K ) over the cubical complex K .

Revised: A sequence of cells π : τ p + , σ p 1 , τ p + , σ p 2 , · · · , σ pk , τ p +1 k , σ pk +1 , where the superscript p in α p stands for the dimension of cell α , form a V-path if ( τ i , σ i )

Revised: ∈ M ( K ) for for any i ∈ [1 , k ] and σ i <

Revised: τ i − 1 for any i ∈ [1 , k + 1] .

Revised: A V-path π is acyclic if ( τ 0 , σ k +1 ) / ∈ M ( K ) .

Revised: This collection of V-pairs M ( K ) form a discrete gradient vector field 2 if (cond-i) each cell in M ( K ) can only appear in at most one pair in M ( K ) ; and (cond-ii) all V-paths in M ( K ) are acyclic.

Revised: Given a discrete gradient vector field M ( K ) , a simplex σ ∈ K is critical if it is not in any V-pair in M ( K ) .

Revised: Even though a discrete gradient vector (a V-pair), say ( τ, σ ) is a combinatorial pair instead of a real vector, it still indicates a “flow" from τ to its face σ .

Revised: A V-path thus corresponds to a flow path (integral line) in the smooth setting.

Revised: However, to make a collection of V-pairs a valid analog of gradient field, (cond-i) says that at each simplex there should only be one “flow" direction; while (cond-ii) is necessarily as flow lines traced by gradient can only go down in function values and thus never come back (thus acyclic).

Revised: A critical simplex has “vanishing gradient" as it is not involved in any V-pair in M ( K )

Revised: (i.e, there is no flow at this simplex).

Revised: Given a 2D cubical complex K a a discrete gradient vector field M ( K ) , we can view critical 0 -, 1 - and 2 -cells as minima, saddle points, and maxima, respectively.

Revised: Hence, a 1-stable manifold in 2D will correspond to a V-path connecting a critical square (a maximum) and a critical edge (a saddle).

Revised: Morse cancellation.

Original: As the predicted likelihood map is noisy, the extracted discrete gradient field M ( K ) couldalso be noisy.
Revised: A given discrete gradient field M ( K ) could be noisy, e.g, there are shallow valleys where the mountain ridge around it should be ignored.

Original: Fortunately, the discrete Morse theory provides an elegant  way to cancelcritical simplicial pairs and ignore the the unimportant Morse branchs.
Revised: Fortunately, the discrete Morse theory provides an elegant and purely combinatorial way to cancel pairs of critical simplices (and thus reduce their stable manifolds).

Original: Particularly, if there is a unique V-path π
Revised: In particular, given M ( K ) , a pair of critical simplices ⟨ δ ( p +1) ,γ p ⟩ is cancellable if there is a unique V-path π

Original: = δ = δ 0 , γ 1 , δ 1 , . .
Revised: = δ = δ 0 , γ 1 , δ 1 , . .

Original: , δ s , γ s +1 =
Revised: , δ s , γ s +1 =

Original: γ from δ to γ , then the pair of critical simplices ⟨ δ  ( p +1) ,  γp ⟩ is cancellable .
Revised: γ from δ to γ . The Morse cancellation operation simple reverse all V-pairs along this path by removing all V-pairs along these path, and adding ( δ i − 1 , γ i ) to M ( K ) for any i ∈ [1 , s + 1] . It is easy to check that after the cancellation neither δ nor γ is critical.

Original: By removing all V-pairs along these path, and adding ( δ i − 1 , γ i )  to M ( K ) for any
Revised: By setting ρ ( σ ) for each cell to be the maximum ρ -value of each vertex in σ , we are able to extend this vertex-valued function ρ to a function ρ : K → R . Then how to obtain a discrete gradient vector field from such function ρ : K → R ? Following the approach developed in (Wang et al., 2015; Dey et al., 2018), we initialize a trivial discrete gradient vector field where all cells are initially critical.

Original: [Figure]
Revised: Let ϵ > 0 be a threshold for simplification. We then perform persistence algorithm (Edelsbrunner et al., 2000) induced by the super-level set filtration of ρ and pair up all cells in K , denoted by P ρ ( K ) . Persistent homology is one of the most important development in the field of topological data analysis in the past two decades (Edelsbrunner & Harer, 2022; Edelsbrunner et al., 2000; Zomorodian & Carlsson, 2005).

Revised: Imagine we grow the complex K by starting from the empty set and gradually include more and more cells in decreasing ρ values, which is the so-called super-level set filtration of K induced by ρ .

Revised: Through this course, new topological feature can be created upon adding a simplex σ , and sometiems a feature can be destroyed upon adding a simplex τ .

Original: i ∈ [1 , s + 1] ,  the Morse cancellation operation reverses all V-pairs along this path.
Revised: Persistence algorithm (Edelsbrunner et al., 2000) will pair up simplices; that is, its output is a set of pairs of simplices P ρ ( K ) = { ( σ, τ ) } , where each pair captures the birth and death of topological features during this evolution. The persistence of a pair, say p = ( σ, τ ) , is defined as pers( p ) =

Original:  In this way, neither δ nor γ is critical after the cancellation operation and we  can prune/remove the corresponding stable manifold/branch.
Revised: ρ ( σ ) − ρ ( τ ) , measuring how long the topological feature captured by p lives in term of ρ . In this case, we also write pers( σ ) = pers( τ ) = pers( p ) – the persistence of a simplex (say σ or τ ) can be viewed as the importance of this simplex. With this intuition of the persistence pairings, we next perform Morse-cancellation operation to all pairs of these cells (

Revised: σ, τ ) ∈ P ρ ( K ) in increasing order their persistence if (i) its persistence pers( δ, γ ) <

Revised: ϵ (i.e, this pair has low persistence and thus not important); and (ii) this pair ( δ, γ ) is cancellable.

Original: More details can be found in (Hu et al., 2021).
Revised: Let M ϵ ( K ) be the resulting discrete gradient field after simplifying all low-persistence critical simplices. We then construct the 1-stable manifolds for the remaining (high persistence, and thus important) saddles (critical 1-cells) from M ϵ ( K ) . Let S 1 ( ϵ ) be the resulting collection of 1-stable manifolds. In particular, see an illustration of a V-path (highlighted in black) corresponding to a 1-stable manifold of the green saddle in Fig. 3(b).

Original: A. A PPROXIMATION FOR VOLUME DATA
Revised: A. A PPROXIMATION FOR VOLUME DATA

Original: As illustrated in the main text, we propose a persistent-homology filtered topology watershed algorithm to obtain the approximation of Morse structures for volume data.
Revised: As illustrated in the main text, we propose a persistent-homology filtered topology watershed algorithm to obtain the approximation of Morse structures for volume data.

Original: The details are illustrated in Alg. 1.
Revised: The details are illustrated in Alg. 1.

Original: A. R EPARAMETERIZATION TECHNIQUE
Revised: A. R EPARAMETERIZATION TECHNIQUE

Original: We adopt the reparameterization technique of VAE to make the network differentiable and be able to backpropagate.
Revised: We adopt the reparameterization technique of VAE to make the network differentiable and be able to backpropagate.

Original: The posterior net randomly draw samples from posterior distribution ϵ ∼ N ( µ post , σ post ) .
Revised: The posterior net randomly draw samples from posterior distribution ϵ ∼ N ( µ post , σ post ) .

Original: To implement the posterior net as a neural network, we will need to backpropagate through randomsampling.
Revised: To implement the posterior net as a neural network, we will need to backpropagate through random

Original: The issue is that backpropagation cannot flow through random node; to overcome this obstacle, we adopt the reparameterization technique proposed in (Kingma & Welling, 2013).

Original: Assuming the posterior is normally distributed, we can approximate it with another normal distribution.

Original: We approximate ϵ with normally distribution Z ( Z ∼ N (0 , I ) ).

Original: [Equation]

Original: Now instead of saying that ϵ is sampled from Q ( X, Y ;

Original: ω post ) , we can say ϵ is a function that takes parameter ( Z ,( µ , σ ) ) and these µ , σ come from deep neural network.

Original: Therefore all we need is partial derivatives w.r.t.

Original: Algorithm
Revised: Algorithm

Original: Input: a grid 2D image, and a threshold θ
Revised: Input: a grid 2D image, and a threshold θ

Original: Output: Morse structures for volume data
Revised: Output: Morse structures for volume data

Original: Definition : G =
Revised: Definition : G =

Original: ( V, E ) denote a graph; f ( v ) is the intensity value of node v ; lower_star ( v ) =
Revised: ( V, E ) denote a graph; f ( v ) is the intensity value of node v ; lower_star ( v ) =

Original: { ( u, v ) ∈ E | f ( u ) <
Revised: { ( u, v ) ∈ E | f ( u ) <

Original: f ( v ) } ; cc ( v ) is the connected component id of node v .
Revised: f ( v ) } ; cc ( v ) is the connected component id of node v .

Original: 1: PD = ∅ ; Build the proximity graph (4-connectivity) for 2D grid image; 2: U =
Revised: 1: PD = ∅ ; Build the proximity graph (4-connectivity) for 2D grid image; 2: U =

Original: V sorted according f ( v ) ; T a sub-graph, which includes all the nodes and edges
Revised: V sorted according f ( v ) ; T a sub-graph, which includes all the nodes and edges

Original: whose value < t .
Revised: whose value < t .

Original: 3: for v in U do 4: t = f ( v ) , T = T + { v } 5: for ( u, v ) in lower_star( v ) do 5:
Revised: 3: for v in U do 4: t = f ( v ) , T = T + { v } 5: for ( u, v ) in lower_star( v ) do 5:

Original: Assert u ∈ T 6: if cc ( u ) =
Revised: Assert u ∈ T 6: if cc ( u ) =

Original: cc ( v ) then 6:
Revised: cc ( v ) then 6:

Original: Edge_tag ( u, v ) =
Revised: Edge_tag ( u, v ) =

Original: loop edge 6:
Revised: loop edge 6:

Original: Continue 7: else 7:
Revised: Continue 7: else 7:

Original: Edge_tag ( u, v ) =
Revised: Edge_tag ( u, v ) =

Original: tree edge 7: younger_cc = arg max w = cc ( u ) ,cc ( v ) f ( w ) 7: older_cc =
Revised: tree edge 7: younger_cc = arg max w = cc ( u ) ,cc ( v ) f ( w ) 7: older_cc =

Original: arg min w = cc ( u ) ,cc ( v ) f ( w ) 7: pers =
Revised: arg min w = cc ( u ) ,cc ( v ) f ( w ) 7: pers =

Original: t - f (younger_cc) 8: if pers >= θ then 8:
Revised: t - f (younger_cc) 8: if pers >= θ then 8:

Original: Edge_tag ( u, v ) =
Revised: Edge_tag ( u, v ) =

Original: watershed edge 8:
Revised: watershed edge 8:

Original: Continue 9: end if 10: for w in younger_cc do 10: cc ( w ) = older_cc 11: end for 11:
Revised: Continue 9: end if 10: for w in younger_cc do 10: cc ( w ) = older_cc 11: end for 11:

Original: PD = PD + ( f (younger_cc), t ) 12: end if 13: end for 14: end for 15: return Membrane_vertex_set = ∪ vertices of watershed_edge_set 
Revised: PD = PD + ( f (younger_cc), t ) 12: end if 13: end for 14: end for 15: return Membrane_vertex_set = ∪ vertices of watershed_edge_set sampling.

Revised: The issue is that backpropagation cannot flow through random node; to overcome this obstacle, we adopt the reparameterization technique proposed in (Kingma & Welling, 2013).

Revised: Assuming the posterior is normally distributed, we can approximate it with another normal distribution.

Revised: We approximate ϵ with normally distribution Z ( Z ∼ N (0 , I ) ).

Revised: [Equation]

Revised: Now instead of saying that ϵ is sampled from Q ( X, Y ;

Original:  A. D ATASETS
Revised: ω post ) , we can say ϵ is a function that takes parameter ( Z ,( µ , σ ) ) and these µ , σ come from deep neural network. Therefore all we need is partial derivatives w.r.t. µ , σ and Z is irrelevant for taking derivatives for backpropagation. A. D ATASETS

Original: Both volume and vessel datasets are used to validate the efficacy of the proposed method, and the details of the datasets are as follows:
Revised: Both volume and vessel datasets are used to validate the efficacy of the proposed method, and the details of the datasets are as follows:

Original: ISBI13 (volume): ISBI13 (Arganda-Carreras et al., 2013) is a EM dataset, containingimages with resolution of 1024x1024.
Revised: ISBI13 (volume): ISBI13 (Arganda-Carreras et al., 2013) is a EM dataset, containingimages with resolution of 1024x1024.

Original: CREMI (volume): CREMI is another EM dataset, containing 125 images, and each of them has a resolution of 1250x1250.
Revised: CREMI (volume): CREMI is another EM dataset, containing 125 images, and each of them has a resolution of 1250x1250.

Original: DRIVE (vessel): DRIVE (Staal et al., 2004) is a retinal vessel dataset with 40 images.
Revised: DRIVE (vessel): DRIVE (Staal et al., 2004) is a retinal vessel dataset with 40 images.

Original: The resolution for each image is 584x565.
Revised: The resolution for each image is 584x565.

Original: We use a 3-fold cross-validation for all the methods to report the numbers over the validation set.
Revised: We use a 3-fold cross-validation for all the methods to report the numbers over the validation set.

Original: A. E VALUATION METRICS
Revised: A. E VALUATION METRICS

Original: The details of the metrics used in this paper are listed as follows:
Revised: The details of the metrics used in this paper are listed as follows:

Original: DICE :
Revised: DICE :

Original: DICE score is usually to measure the volumetric overlap between the predicted and ground truth masks.
Revised: DICE score is usually to measure the volumetric overlap between the predicted and ground truth masks.

Original: Adapted Rand Index (ARI) : ARI is the maximal F-score of the foreground-restricted Rand index (Rand, 1971), a measure of similarity between two clusters.
Revised: Adapted Rand Index (ARI) : ARI is the maximal F-score of the foreground-restricted Rand index (Rand, 1971), a measure of similarity between two clusters.

Original: VOI (Jain et al., 2010): VOI is a measure of the distance between two clusterings.
Revised: VOI (Jain et al., 2010): VOI is a measure of the distance between two clusterings.

Original: Betti Error (Hu et al., 2019):
Revised: Betti Error (Hu et al., 2019):

Original: Betti Error measures the topology difference between the predicted and the ground truth mask.
Revised: Betti Error measures the topology difference between the predicted and the ground truth mask.

Original: We randomly sample patches over the predicted segmentation and compute the average absolute error between their Betti numbers and the corresponding ground truth patches.
Revised: We randomly sample patches over the predicted segmentation and compute the average absolute error between their Betti numbers and the corresponding ground truth patches.

Original: A. B ASELINES
Revised: A. B ASELINES

Original: We compare the proposed method with two kinds of baselines:  Standard segmentation baselines:
Revised: We compare the proposed method with two kinds of baselines: 1) Standard segmentation baselines:

Original: DIVE (Fakhry et al., 2016) is originally designed for EM data segmentation.
Revised: DIVE (Fakhry et al., 2016) is originally designed for EM data segmentation.

Original: UNet (Ronneberger et al., 2015) achieves good performances in different contexts with encoder-decoder scheme and skip connections.
Revised: UNet (Ronneberger et al., 2015) achieves good performances in different contexts with encoder-decoder scheme and skip connections.

Original: UNet-VGG (Mosinska et al., 2018) proposes a topology-aware loss based on the detected linear structures with pretrained filters.
Revised: UNet-VGG (Mosinska et al., 2018) proposes a topology-aware loss based on the detected linear structures with pretrained filters.

Original: TopoLoss (Hu et al., 2019) identifies the topological critical points with persistence homology and derives a novel topological loss.
Revised: TopoLoss (Hu et al., 2019) identifies the topological critical points with persistence homology and derives a novel topological loss.

Original: DMT (Hu et al., 2021) identifies the whole topological structures and introduces additional penalty on the whole structures instead of isolated critical points.
Revised: DMT (Hu et al., 2021) identifies the whole topological structures and introduces additional penalty on the whole structures instead of isolated critical points.

Original: 2) Probabilistic-based segmentation methods: 1.
Revised: 2) Probabilistic-based segmentation methods: 1.

Original: Dropout UNet (Kendall et al., 2015) dropouts the three inner-most encoder and decoder blocks with a probability of 0.5 during both the training and inference.
Revised: Dropout UNet (Kendall et al., 2015) dropouts the three inner-most encoder and decoder blocks with a probability of 0.5 during both the training and inference.

Original: Probabilistic-UNet (Kohl et al., 2018) introduces a probabilistic segmentation method by combining UNet with a VAE.
Revised: Probabilistic -UNet (Kohl et al., 2018) introduces a probabilistic segmentation method by combining UNet with a VAE.

Original: For all methods, we generate binary segmentations by thresholding predicted likelihood maps at 0.5.
Revised: For all methods, we generate binary segmentations by thresholding predicted likelihood maps at 0.5.

Original: A. T HE ADVANTAGE OF THE JOINT TRAINING AND OPTIMIZATION
Revised: A. T HE ADVANTAGE OF THE JOINT TRAINING AND OPTIMIZATION

Original: As mentioned in the main text, another straightforward alternative of the proposed approach is to use the discrete Morse theory to postprocess the continuous likelihood map obtained from the standard segmentation networks.
Revised: As mentioned in the main text, another straightforward alternative of the proposed approach is to use the discrete Morse theory to postprocess the continuous likelihood map obtained from the standard segmentation networks.

Original: In this way, we can still obtain structure-preserving segmentation maps, but there are two main issues: 1) if the segmentation network itself is structure-agnostic, we’ll not be able to generate satisfactory results even with the postprocessing, and 2) we have to manually choose the persistence threshold to prune the unnecessary branches for each image, which is cumbersome and unrealistic in practice.
Revised: In this way, we can still obtain structure-preserving segmentation maps, but there are two main issues: 1) if the segmentation network itself is structure-agnostic, we will not be able to generate satisfactory results even with the postprocessing, and 2) we have to manually choose the persistence threshold to prune the unnecessary branches for each image, which is cumbersome and unrealistic in practice.

Original: The proposed joint training strategy overcomes both these issues.
Revised: The proposed joint training strategy overcomes both these issues.

Original: First, during the training, we incorporate the structure-aware loss ( L skeleton ).
Revised: First, during the training, we incorporate the structure-aware loss ( L skeleton ).

Original: Consequently, the trained segmentation branch itself is structure-aware essentially.
Revised: Consequently, the trained segmentation branch itself is structure-aware essentially.

Original: On the other hand, with the prior and posterior nets, we are able to learn a reliable distribution of the persistence threshold ( ϵ ) given an image in the inference stage.
Revised: On the other hand, with the prior and posterior nets, we are able to learn a reliable distribution of the persistence threshold ( ϵ ) given an image in the inference stage.

Original: Sampling over the distribution makes it possible to generate satisfactory structure-preserving segmentation maps within a few trials (the inference won’t take long), which is more much efficient. 
Revised: Sampling over the distribution makes it possible to generate satisfactory structure-preserving segmentation maps within a few trials (the inference will not take long), which is more much efficient. [Figure]

Original: We conduct an empirical experiment to demonstrate the advantage of the joint training and optimization.
Revised: We conduct an empirical experiment to demonstrate the advantage of the joint training and optimization.

Original: For the postprocessing, given the predicted likelihood maps from the standard segmentation networks, we randomly choose five persistence thresholds and generate the segmentation masks separately, and select the most reasonable one as the final segmentation mask to report the performances.
Revised: For the postprocessing, given the predicted likelihood maps from the standard segmentation networks, we randomly choose five persistence thresholds and generate the segmentation masks separately, and select the most reasonable one as the final segmentation mask to report the performances.

Original: The results in Tab.
Revised: The results in Tab.

Original: 2 demonstrate the advantage of our joint training and optimization strategy.
Revised: 2 demonstrate the advantage of our joint training and optimization strategy.

Original: [Table]
Revised: [Table]

Original: A. C OMPUTATIONAL RESOURCES
Revised: A. I NSTANCE PRIORS FOR SPECIFIC IMAGES

Revised: Here we provide a few instance priors for specific images (Fig. 12).

Revised: A. I LLUSTRATION OF NOISY LIKELIHOOD MAP

Revised: As we know, for image segmentation tasks, with the power of deep neural networks, we are able to achieve quite high pixel-wise segmentation accuracy.

Revised: In other words, the predicted likelihood maps from the segmentation network are good enough (see Fig. 6(c) and Fig.

Revised: 11(c))

Revised: to distinguish the skeletons/structures of the original images, with issues only happening at locations with weak signals.

Revised: Here we provide a comparison of noisy and clean likelihood maps (Fig. 13).

Revised: If the model is trained to be converged, usually we are able to obtain the clean likelihood map (Fig. 13(d)), and

Revised: then we could

Revised: [Figure]

Revised: [Figure]

Revised: apply discrete Morse theory on it to construct the structural space.

Revised: On the other hand, if the model is not converged, the obtained likelihood map will be noisy (Fig. 13(c)), and we will not be able to recover the true signals of the original image.

Revised: A. I LLUSTRATION OF UNCERTAINTY ESTIMATION (CDF)

Revised: As we mentioned in Sec. 3.2, the probability of a branch b being in a segmentation map M such as ϵ M ∼ P follows a Bernoulli distribution with the probability Pr ( b ) being the cumulative distribution function (CDF) of P , CDF P ( ϵ b )

Revised: =

Revised: P ( ϵ ≤ ϵ b ) .

Revised: In Fig.

Revised: 14, we use dotted lines to denote the filtered out branches and the shaded regions denotes the probability of branch2 belongs to the final segmentation map M .

Revised: A. C OMPUTATIONAL RESOURCES

Original: All the experiments are performed on a RTX A5000 GPU (24G Memory), and AMD EPYC32-Core Processor.
Revised: All the experiments are performed on a RTX A5000 GPU (24G Memory), and AMD EPYC32-Core Processor.