Learning Probabilistic Topological Representations Using Discrete Morse Theory
Authors: Xiaoling Hu
Dimitris Samaras
Chao Chen
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Topological Representation
Discrete Morse Theory
Persistent Homology

Abstract: Accurate delineation of fine-scale structures is a very important yet challenging problem. Existing methods use topological information as an additional training loss, but are ultimately making pixel-wise predictions. In this paper, we propose a novel deep learning based method to learn topological/structural. We use discrete Morse theory and persistent homology to construct a one-parameter family of structures as the topological/structural representation space. Furthermore, we learn a probabilistic model that can perform inference tasks in such a topological/structural representation space. Our method generates true structures rather than pixel-maps, leading to better topological integrity in automatic segmentation tasks. It also facilitates semi-automatic interactive annotation/proofreading via the sampling of structures and structure-aware uncertainty.

Published as a conference paper at ICLR 2023
LEARNING PROBABILISTIC TOPOLOGICAL REPRESEN- TATIONS USING DISCRETE MORSE THEORY
Xiaoling Hu
Dimitris Samaras Stony Brook University
Chao Chen Stony Brook University
Stony Brook University
ABSTRACT
Accurate delineation of fine-scale structures is a very important yet challenging problem. Existing methods use topological information as an additional training loss, but are ultimately making pixel-wise predictions. In this paper, we propose a novel deep learning based method to learn topological/strrcturaal representations. We use discrete Morse theory and persistent homology to construct one-parameter family of structures as the topological/strrctun we learn probabilistic model that can perform inference tasks in such a topolog- ical/structural representation space. Our method generates true structures rather than pixel-maps, leading to better topological integrity in automatic segmentation tasks. It also facilitates semi-automatic interactive amnotation/proofradding via the sampling of structures and structure- aware uncertainty.
representation space. Furthermore,
INTRODUCTION
Accurate segmentation of fine-scale structures, e.g., vessels, neurons and membranes is crucial for downstream analysis. 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). 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.
However, these loss-based methods are still not ideal. They are based on a standard segmentation network, and thus only learn pixel-wise feature representations. This causes several issues. First, standard segmentation network makes pixel-wise predictions. Thus, at the inference stage, topological errors, e.g. broken connections, can still happen, even though they may be mitigated by the topology- inspired losses. Another issue is in uncertainty estimation, i.e., estimating how certain a segmentation network is at different locations. Uncertainty maps can direct the focus of human annotators for efficient proofreading. However, for fine-scale structures, existing pixel-wise uncertainty maps are not effective. As shown in Fig. d)), every pixel adjacent to a vessel branch is highly uncertain, in spite of whether the branch is salient or not. What is more desirable is a structural uncertainty map that can highlight uncertain branches (e.g., Fig. (f).
To fundamentally address these issues, we propose to directly model and reason about the structures. In this paper, we propose a novel deep learning based method that directly learns the topologi- cal/structural representation of images. To move from pixel space to structure space, we apply the classic discrete Morse theory (Milnor 1963 Forman 2002) to decompose an image into Morse complex, consisting of structural elements like branches, patches, etc. These structural elements are the hypothetical structures one can infer from the input image. Their combinations constitute a space of structures arising from the input image. See Fig. c) for an illustration.
For further reasoning with structures, we propose to learn a probabilistic model over the structure space. The challenge is that the space consists of exponentially many branches and is thus of very high dimension, To reduce the learning burden, we introduce the theory of persistent homology (Sousbie 2011 Felgado-Friedrichs
al. 2015 Wang al., 2015) for structure pruning. Each branch has its own persistence measuring its relative saliency. By continuously thresholding the complete Morse complex in terms of persistence, we obtain sequence of Morse complexes parameterized by the persistence threshold, e. See Fig. dd). By learning a Gaussian over we learn a parametric probabilistic model over these structures.
"Email: Xiaoling Hu (xiaolhu@cs.stongge wwt..essttobooosccciii cdu).
We will he using the terms topology/topological and structure/strretuuaal interchangeably in this paper.

---

Published as a conference paper at ICLR 2023
(a) Image
GT
Prob.UNet
Pixel Uncer.
Ours
Stru. Uncer.
Figure I: Illusrationn of structural segmentation and structure-Ievel uncertainty. Compared with Probabilistic-UNe (Kohl al. 2018) (Fig. (d)),)), the proposed method able to generate
structure- strueture-preserving segmentation map (Fig. e)), and structure-level uncertainty (Fig. ()).
(b) Likelihood
(c) Space
(d) 1-d Distribution
Sample
DMT
Reducing Search Space
Final Seg
(g) Uncertainty Mup
(a) Input
Maltipla 4
Figure 2: The probabilistic topological/structural representation. (a) is sample input, (b) is the predicted likelihood map from the deep neural network, (c) is the whole structure space obtained by rumning discrete Morse theory algorithm the likelihood map. (d) the 1-d structural family parametrized by the persistence threshold c, as well as a Gaussian distribution over e, (e) a sampled skeleton, (f) the final structural segmentation map generated using the skeleton sample, and (g) the uncertainty map generated by multiple segmentations.
This parametric probabilistic model over structure space allows us to make direct structural predictions via sampling (Fig. [e))), and to estimate the empirical structure-level uncertainty via sampling (Fig. [gg) The benefit is two-fold: First, direct prediction of structures will ensure the model outputs always have structural integrity, Samples from the probabilistic model are all feasible structural hypotheses based on the input image, with certain variations uncertain locations. This is in contrast to state-of-the-art methods using pixel-wise representations (Fig. [c))- Note the original output structure (Fig. 2]e), also called skeleton) only l-pixel wide and may not serve as good segmentation output. In the inference
the inference stage. This is illustrated in Fig. Ie).
stage, we use post-pne post-processing step to grow the structures without changing topology as the final segmentation prediction (Fig. (()). More details are provided in Sec. 3.2 and Fig.
Second, the probabilistic structural model can be seamlessly incorporated into semi-autoomatic interactive amnotation/proofreading workflows to facilitate large scale annotation of these complex structures (see Fig. (5). 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. Our probabilistic model makes it possible to identify uncertain structures for efficient interactive ammotation/proofreaa uncertainty reasoning. As shown in Fig. f and Fig. [g. structural uncertainty map highlights uncertain branches for efficient proofreading. On the contrary. traditional pixel-wise uncertainty map (Fig. D d)) not helpful all; highlights all pixels the boundary of a branch.
Note that the structure space crucial for
The main contributions this paper are:
1. We propose novel deep learning method which learns structural representations, based on discrete Morse theory and persistent homology.
2. We learn probabilistic model over the structure space, which facilitates different tasks such as topology-aware segmentation, uncertainty estimation and interactive proofreading.
3. We validate method on various datasets with rich and complex structures. It outperforms ti-the-art art methods in both deterministic and probabilistic categories.
state-of sice--ofthe-aate
2

---

Published as a conference paper at ICLR 2023
2 RELATED WORK
Structure/ Structure/Topology--ee Topology- deep image segmentation. A number of recent works have tried to
segment with correct topology with additional topology-aware losses (Mosinska et al. 2018 Hu et al.
2019 Clough al. 2020; Hu et al. 2021, Shit et al. 2021). Specifically, UNet- V GG (Mosinska et al. 2018) detects linear structures with pretrained filters. clDice (Shit et al. 2021) introduces an
additional Dice loss for extracted skeleton structures. 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. Similarly, DMT- (Hu et al., (02T)) 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. All these losses, although aiming at topological integrity, still cannot change the pixel-wise prediction nature of the backbone network. To the best of our knowledge, no existing methods generate structural representations//reciiiioss like our proposed method.
Additionally, the discrete Morse complex has been used for image analysis. but only as preprocessing step Belgado--Friedricle et al. 2015 Robins et al., 201 I Wang et al. 2015 Dey et al. 2019) or as a
conditional input of neural network (Banerjee et al.) (020).
Segmentation uncertainty. Uncertainty estimation has been the focus of research in recent years (Graves) 2011] Gal & Ghahramani| 2016 [Lakshminarayanan et al., 2017 Moon et al. 2020).
However, most existing work focuses on the classification problem. In terms of image segmentation, the research is still relatively limited. Some existing methods directly apply classification uncertainty to individual pixels, e.g., dropout (Kendall et al., 2015; Kendall & Gal (017). This, however, is not taking into consideration the image structures. Several methods estimate the uncertainty by generating an ensemble of segmentation (Lakshminarayanaa et al. 2017) or using multi-heads (Rupprecht et al. 2017. Mlg et al., 2018). Notably, Probabilistic- UNer (Kohl et al. 2018) learns a distribution over the latent space and then samples over the latent space to produce segmentation samples. 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. These methods are fundamentally different from ours, which makes predictions on the structures.
METHOD
Our key innovation is to restructure the output of a neural network so that it is indeed making predictions over space of structures. This is achieved through insights into the topological structures of an image and the usage of several important tools in topological data analysis.
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. For simplification, we will use the term "branch" to denote single piece of Morse structure. These Morse branches the hypothetical structures one can infer from the input image. This decomposition is based on a likelihood function produced by pixel-wise segmentation network trained in parallel. Thus it is of good quality, i.e., the structures are close enough to the true structures.
Any binary labeling of these Morse branches is legitimate segmentation; we call it structural segmentation. 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. To further reduce the degrees of freedom to make the inference easier, we apply persistent homology filter these branches with regard their saliency. This gives us a linear size family of structural segmentations, parameterized by a threshold e. Finally, we learn ID Gaussian distribution for the as our probabilistic model. This gives us the opportunity not only to sample segmentations, but also to provide probability for each branch, which can be useful in downstream tasks including proofreading. In Sec. we introduce the discrete Morse theory and how to construct the space of Morse structures. We also explain how to use persistent homology to reduce the search space of reasoning into a 1-parameter family. In Sec. 3.2, we will provide details on how our deep neural network constructed to learn the probabilistic model over the structure space, illustrated in Fig.
3.1 CONSTRUCTING THE STRUCTURE SPACE
In this section, explain how to construct structural representation space using discrete Morse theory. The resulting structural representation space will be used to build a probabilistic model. We will then discuss how reduce the structure space into 1-parameter family of structural
3

---

Published as a conference paper at ICLR 2023
segmentations, using persistent homology. We assume 2D imput image, although the method naturally extends to 3D images.
Given a reasonably clean imput (e.g., the likelihood map of deep neural network, Fig. 2(b),, treat the 2D likelihood as a terrain function, and Morse theory (Milnor 1963) can help to capture the structures regardless of weak/blurred conditions. See Fig. for illustration. The weak part of a line in the continuous map can be viewed as the local dip in the mountain ridge of the terrain. In the language of Morse theory, the lowest point of this dip is saddle point (S in Fig. 3))), and the mountain ridges which are connected to the saddle point (M,S and M2S) are called the stable manifolds of the saddle point.
We mainly focus on 2D images in this paper. We consider two dimensional continuous function 2f
f: IR2 + R. For a point x E IR², the gradient can be computed as V fx()) = Dan o We call point (x1, m22) critical if V = O. For Morse function defined on R22 critical point could be a minimum, a saddle or a maximum.
Consider continuous line (the red rectangle region in Fig. B()) in 2D likelihood map. Imagine if we put a ball on one point of the line, then -V /V[(e (x: indicates the direction which the ball will flow down. By definition, the ball will eventually flow the critical points where Vf(x)= - O. The collection of points whose ball eventually flows to f ) = 0) is defined as the stable manifold (denoted S()) of point p. Intuitively, for 2D function f. the stable manifold (p) of minimum p is the entire valley of (similar to the watershed algorithm); similarly, the stable manifold Sgg) of a saddle point & consists of the whole ridge line which connects two local maxima and goes through the saddle point. See Fig. (b) an illustration.
For a link-like structure, the stable manifold of saddle contains the topo- logical structures (usually curvilinear) of the continuous likelihood map pre- dicted by deep neural networks, and they are exactly what want to re- cover from noisy images. In practice, we adopt the discrete version of Morse theory for images.
M
(a) map
(b)
(c) Morse complex
Figure 3: (a) shows sample likelihood map from the deep neural network, and (b) is the terrain view of the red patch in (a) and illustrates the stable manifold of saddle point in 2D) case for a line-like structure. (c) is the 2D Morse complex generated by DMT from (a).
Discrete Morse theory. Take 2D image as a 2-dimensional cubical com- plex (Wagner et al. 20122] Kaczynski et al. 2004). A 2-dimensional cubical complex then contains 0-, 1-, and 2-dimensional cells, which correspond to vertices (pixels), edges and squares, respectively. In the setting of discrete Morse theory (DMT) (Forman 11988] 2002). a pair of adjacent cells, termed as discrete gradient vectors, compose the gradient vector. Critical points (v)(z)= 0) are those critical cells which are not in any discrete gradient vectors. In the 2D domain, a minimum, a saddle and a maximum correspond to a critical vertex, a critical edge and a critical square respectively. 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. See Fig. (().. And the Morse complex generated by the DMT algorithm is illustrated in Fig. g).
Constructing the full structure space. 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. Formally, we call any combinations of these stable manifolds structure. Fig. c) illustrates different structures. This structure space, however, is of exponential size. Assume there are N pieces of stable manifolds/branches for a given likelihood map. Any combinations of these stable manifolds/bran
will be a potential structure. We will have 2 possible structures in total. This can be computationally prohibitive to construct and to model. We need a principled way to reduce the structural search space.
Reducing the structural search space with persistent homology. We propose to use persistent homology (Sousbie) 2011: Delgado-Friedrrihs et al.) 2015 Wang et al. 2015) to reduce the structural search space. Persistent homology is important tool for topological data analysis (Edelsbrunner & Harer 2022 [Edelsbrunner et al. 2000). Intuitively, we grow a Morse complex by gradually including more and more discrete elements (called cells) from empty. A branch of the Morse complex is special
4

---

Published as a conference paper at ICLR 2023
Likelihood map o(Xg
Postrier Nei
S(X: (lngg)
Skevtto
Dyt
Sample
PH
Skelcton S
baateton
Pror Tiee
Poweroor Tiis
Segnestation
Bransh
Litelitood
Grunnal Inchy
Ground truth
Y
(a)
Overview of the proposed framework
(b) Details of Lakeletonn
Figure 4: The overall workflow of the training stage.
The red arrows indicate supervision.
type cell. Other types include vertices, patches, etc. Cells will be continuously added to the complex. New branches will be born and existing branches will die. The persistence algorithm (Edelsbrunner et al. 2000) pairs up all these critical cells as birth and death pairs. The difference of their function values essentially the life time of the specific topological structure/branch, which called the persistence. The importance of a branch is associated with its persistence. Intuitively, the longer the persistence of a specific branch is, the more important the branch
Recall that our original construction of the structure space considers all possible combinations of branches, and thus can have exponentially many combinations. Instead, we propose to only select branches with high persistence as important ones. By doing this, we will be able to prune the less important/noisy branches very efficiently, and recover the branches with true signals. Specifically, the structure pruning is done via Morse cancellation (more details are included in Appendix (A.3) operation. The persistence thresholding provides us the opportunity to obtain a structure space of linear size. We start with the complete Morse complex, and continuously increase the threshold c. At each threshold, we obtain a structure by filtering with c and only keeping the branches whose persistence above e. This gives sequence of structures parametrized by e. As shown in Fig. dd) the family of structures represents different structural densities.
The one-parameter space allows inference tasks such as segmentation, sampling, uncertainty estimation and interactive proofreading. Specifically, will learn a Gaussian distribution over the persistence threshold c, e N N(,, o.). Denote the persistence of branch b as 6-- Any branch b belongs to the structure map M (we also call the structure map M structural segmentation) as long as its persistence is higher or equal to the persistence threshold of M, i.e., b E M if and only if c 2 CM. where c is used to generate M. More details will be provided Sec. 3.2
easily learn probabilistic model and carry various
Approximation of Morse structures for volume data. 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. However, the output structures should always be boundary edges for volume data, which cannot be dealt with the original discrete Morse theory. Consequently, boundaries of the stable manifolds of local minima. As mentioned above, the stable manifold of 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. Similar to the original discrete Morse theory, we also introduce persistence threshold parameter and use persistent homology prune the less important branches. The details of the proposed persistent persistent-homolpge homology filtered topology watershed algorithm illustrated in Appendix A.4
approximate the Morse structures of 2D volume data with the
3.2 NEURAL NETWORK ARCHITECTURE
In this section, we introduce neural network that learns probabilistic model over the structural representation to obtain structural segmentations. See Fig. 4 for an illustration the overall pipeline.
Since the structural reasoning needs sufficiently clean input to construct discrete Morse complexes, our method first obtains such likelihood map by training a segmentation branch which is supervised by the standard segmentation loss, binary cross-entropy loss. Formally, Leeg Lboe (Y, S(X; (XWwey)). which X is the input image, Wseg is the segmentation branch's 's weight, S(X; wang) is the output likelihood map, and Y is the ground truth.
5

---

Published as a conference paper at ICLR 2023
The output likelihood map, S(X; is used 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. Thresholding these branches using persistent homology with different values will produce different structures. We refer to the DMT computation and the persistent homology thresholding operation SDMT and frH- So given a likelihood map 3(X:wse)) and threshold 6, we can generate a structure (which we call a skeleton): Sakeleton (E) = (PH(fDMT(SS) (X; (Se))) Next, we discuss how to learn the probabilistic model. Recall that we want to learn Gaussian distribution over the persistent homology threshold, ~ (H,, The parameters learned by neural network called the posterior network. The network uses the imput image and the corresponding ground truth mask Y input, and outputs the parameters μ² (X, Y; W and or (X, Y; Wpost). Wpost is the parameter of the network.
and a are
During training, at each iteration, we draw sample from the distribution (c ~ N(,, r)). Using the sample 6, together with the likelihood map, we can generate the corresponding sampled structure, Sskeleton (e). This skeleton will be compared with the ground truth for supervision. To compare sampled skeleton, Sskeleton (t), with ground truth Y. we use the skeleton to mask both and the likelihood map S(X; Wseg), ). and then compare the seeleton-masked ground truth and the likelihood using binary cross-entropy loss: Loce (Y SYooS (c), (Weg)) Sekkleton (c)), in which o denotes the Hadamard product.
To learn the distribution, we use the expected loss:
Lakeleton Ee_,(( Lacee (Y Sskeieton (c), S(X; wsegg) Sskeleton(())
We backpropagate this loss through the posterior network using the reparameterization technique in (Kingma & Welling 2013). More details are provided in Appendix A.3 Note that this loss will also provide supervision to the segmentation network through the likelihood map.
Learning prior network from the posterior network. Although our posterior network can learn the distribution well, it relies on the ground truth mask Y as input. This is not available inference stage. To address this issue, inspired by VAE (Kingma & Welling 2013 Kohl et al. 2018), we use another network to learn the distribution of c with only the image as input. We call this network the prior net. We denote by P the distribution using parameters predicted by the prior network, and denote by Q the distribution predicted by the posterior network.
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. Thus, we incorporate the Kullback-Leibe divergence of these two distributions, D DKr((e (K(()) = I (log Q), which measures how close the prior distribution P (N (Mprior, Oriio) )) is to the posterior distribution Q (N (ppost: ppst))..
Training the neural network. The final loss is composed by the standard segmentation loss, the skeleton loss Lakeleton, and the KL divergence loss, with two hyperparameters a and 33 to balance the three terms,
(2) The network is trained to jointly optimize the segmentation branch and the probabilistic branch (containing both prior and posterior nets) simultancously. During the training stage, the KL divergence loss (DxL) pushes the prior distribution towards the posterior distribution. The training scheme is also illustrated in Fig. 4
I(X,Y) = Leeg + a.Iskeletton + (QI((Q)P)
generating drructure-preserving segmentation maps. In the inference stage, given an input image, we are able to produce unlimited number of plausible structure- prreturvi-peeeeiingg skeletons via sampling. We use a postprocessing step to grow the 1-pixel wide structures/skelltons without changing their topology as the final structural segmentation. Specifically, the skeletons are overlaid on the binarized initial segmentation map (Fig-Ste)) and only the connected components which exist in the skeletons are kept the final segmentation maps (Fig. ((e)). In this way, each plausible skeleton (Fig. d)) generates one final segmentation map (Fig. ((e) and it has exact the same topology the corresponding skeleton. The pipeline of the procedure is illustrated in Fig.
Inference stage:
Uncertainty of structures. Given a learned prior distribution, P, over the family of structural segmentations, we can naturally calculate the probability of each Morse structure/branch. Recall a branch has its persistence é,. And the prior probability of a structural segmentation map M is P(Eas), in which eM is used to generate M. Also any branch b whose persistence is higher

---

Published as a conference paper at ICLR 2023
Mulliple ilerations
Prior Dis.
Skeleton
Prior
Sample
PH
(a) Inpot
DMT
Segmentation
User
(b) Likelihood
Initial
(c) Final
Figure 5: The inference and interactive annotation//roofrradiin pipeline.
or equal to the persistence threshold of M belongs to M, i.c., M if and only if Eb 2 EM- Therefore, the probability of branch b being in segmentation map M such EM P follows Berouuli distribution with the probability Pr(b) being the cumulative distribution function (CDF) of P. CDFp(G)) P(c S c.). This can be directly calculated at inference, and the absolute difference of the CDF from 0.5 is the confidence (which equals I l-uncertainty) of the Morse structure/branch b. We add illustration of uncertainty estimation for branches in Fig. (Sec. o2 Appendix).
EXPERIMENTS
Our method directly makes prediction and inference on structures rather than on pixels. This can significantly benefit downstream tasks. 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 direct prediction can ensure topological integrity even better than previous topology-awad losses. This is not surprising as our prediction is on structures. On memi-automatic proofreading task, we show our structure-level uncertainty can assist human annotators obtain satisfying segmentation annotations a much more efficient manner than previous methods.
pnootaatio/proofraddin tasks. On automatic image segmentation, we show that
4.1 AUTOMATIC TOPOLOGY-AWARE IMAGE SEGMENTATION
Datasets. We use three datasets to validate the efficacy of the proposed method: ISBI13 (Arganda Carreras al. 2013) (volume), CREMI (volume), and DRIVE (Staal et 2004) (vessel). More
details are included Appendix A.6
Evaluation metries. We use four different evaluation metrics: Dice score, ARI, VOI, Betti num- ber error. Dice is popular pixel-wise segmentation metric, and the other three are structure/topology- aware segmentation metrics. More details are included in Appendix A.7
Baselines. We compare the proposed method with two kinds of baselines: 1) Standard segmentation
baselines: DIVE (Fakhry et al. 2016), UNet (Ronneberger et al. 2015), UNet-VGG (Mosinska et al. 2018), TopoLoss (Hu al. 2019) and DMT (Hu et al. (021) 1.21).20 2) Probabilistic-basee segmentation methods: Dropout UNet (endall et al. 2015) and Probabilistic-UNe (Kohl et al. 2018). More details about these baselines included in Appendix A.8
Quantitative and qualitative results. Table D shows the quantitative results comparing to several baselines. Note that for deterministic methods, the numbers computed directly based 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). We use (-test (95% confidence interval) to determine the statistical significance and highlight the significant better results. From the table, we can observe that the proposed method achieves significant better performances in terms of topology-awar metrics (ARI, VOI and Betti Error).
Fig. shows qualitative results. Comparing with DMT (Hu et al. 2021). ) our method is able to produce set of true structure prwpurrrereeeriigggg segmentation maps, as illustrated in Fig. ((e-g). Note that compared with the existing topology-awas -wwar segmentation methods, our method more capable of recovering the weak connections by using Morse skeletons as hints. More qualitative results are included Appendix A.2
7

---

Published as a conference paper at ICLR 2023
Table I: Quantitative results for different models on three different biomedical datasets.
Method
Dice
ARI
VOI
Betti Error
ISBI13 (Volume)
DIVE
0,9658 0.0020
0.6923 = 0.0134
2.790 + 0.025
3.875 0.326
UNet
0,9649 = 0.0057
0.7031 = 0.0256
2.583 + 0.078
3.463 0.435
UNet- VGG
0,9623 = 0.0047
0.7483 = 0.0367
1.534 + 0.063
2.952 + 0.379
TopoLoss
0.9689 = 0.0026
0.8064
0.0112
1.436 0.008
1.253 0.172
DMT
0.9712 0.0047
0.8289
0.0189
1.176 + 0.052
1.102 0.203
Dropout UNet
0.9591 = 0.0031
0.7127 = 0.0181
2.483 0.046
3.189 = 0.371
Prob.-UNet
0.9618 + 0.0019
0.7091 - 0.0201
2.319 1 0.041
3.019 = 0.233
Ours
0.9637 0.0032
0.8417 0.0114
1.013 + 0.081
0.972 0.141
CREMI (Volume)
DIVE
0,9542 = 0.0037
0.6532 = 0.0247
2.513 0.047
4.378
0.152
UNet
0.9523 + 0.0049
0.6723 + 0.0312
2.346 + 0.105
3.016 + 0.253
UNet-VGG
0.9489 + 0.0053
0.7853 + 0.0281
1.623 + 0.083
1.973 0.310
Toso
0.9596 0.0029
0.8083 0.0104
1.462 0.028
1.113 + 0.224
DMT
0.9653 0.0019
0.8203 = 0.0147
1.089 + 0.061
0,982 0.179
Dropout UNet
0.9518 + 0.00T8
0.68T4 = 0.0202
2.195 0.087
3.190
+ 0.198
Prob--UNet
0.9531 + 0.0022
0.6961 = 0.0115
1.901 + 0.107
2.931
+ 0.177
Ours
0.9681 0.0016
0.8475 = 0.0043
0.935 1 0.069
0.919
0.059
DRIVE (Vessel)
DIVE
0.7543 + 0.0008
0.8407 = 0.0257
1.936 + 0.127
3.276 0.642
UNet
0.7491 = 0.0027
0.8343 = 0,0413
1.975 + 0.046
3.643 + 0.536
UNet- VGG
0.7218 = 0.0013
0.8870 = 0.0386
1.167 + 0.026
2.784 = 0.293
TopoLoss
0.7621 0.0036
0.9024 = 0.0113
1.083 + 0.006
1.076 0.265
DMT
0.7733 0.0039
0.9077 = 0.0021
0.876 + 0.038
0.873 0.402
Dropout UNet
0.7410 = 0.0019
0.8331 = 0.0152
2.013 1 0.072
3.121 = 0.334
Prob--UNet
0.7429 + 0.0020
0.8401
0.1881
1.873 1 0.081
3.080 0.206
Ours
0.7814 0.0026
0.0019 0.9109
0.804 + 0.047
0.767
0.098
(a) Image
(b) GT
(c) LH
(d) DMT
(e) Sample1
(f) Sample2
(g) Sample3
Figure 6: Qualitative results of our method compared to DMT-loss (Hu et al. 2021). From left to right: (a) image. (b) ground truth, (c) continuous likelihood map and (d) thresholded binary mask for DMT (Hu et al. 2021), and (e-g) three sampled segmentation maps generated by our method.
VOI v5
Ablation study of loss weights. We observe that the perfor- mances of our method are quite robust to the loss weights and (3. 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 (3) on DRIVE dataset. The results reported in Fig. 7 When 3 = 10, the model achieves slightly better performance in terms of VOI (0.804 + 0.047, the smaller the better) than other choices. Note that, for all the experiments, 1.
Figure Ablation study for B.
8

---

Published as a conference paper at ICLR 2023
(a) Image
(b) GT
(c) Samplel
(d) Sample2
(e) Sample3
(f) Uncertainty
Figure 8: An illustration of structure-level uncertainty.
Illustration of the structure-level uncertainty. We also explore the structure-level uncertainty with the proposed method here. We show three sampled masks (Fig. ((c-)) in the inference stage for given image (Fig. [ a)), and the structure-Ievel uncertainty map (Fig. (()). Note that in practice, for simplification, different from Sec. [.2 we empirically generate uncerfainty map by taking variance across all the samples (the number is 10 in our case). Different from pixel-wiss wiie uncertainty, each small branch has the same uncertainty value with our method. By looking at the original image, we find that the uncertainties are usually caused by the weak signals of the original image. The weak signals of the original image make the model difficult to predict these locations correctly and confidently, especially in structure wise. Actually this also makes sense real cases. Different from natural images, even experts can not always reach a consensus for biomedical image annotation (Armato III et al. 2011 Clark et al. 2013). More details of structure-leve uncertainty are included in Appendix A.I
The ad advantage of the joint training and optimization. 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. More discussions provided in Appendix A.9
4.2 SEMI-AUTOMATIC EFFICIENT ANNOTATION/PROOFEEDDDI WITH USER INTERACTION
Proofreading a struggling while essential step for the annotation of images with fine-scale structures. We develop semi- autimaut interactive innotation/proorraddin pipeline based on the proposed method. As the proposed method is able to generate structural segmentations (Fig. (c), the anno- tators can efficiently proofread one image with rich structures by removing the unnecessaryleedraiing the missing structures with the help of structure-level uncertainty map (Fig. 8 f)) The whole inference and structure-awant interactive image annotation pipeline is illustrated in Fig.
VS Click
We conduct empirical experiments to demonstrate the efficiency of proofreading by using the proposed method. We randomly select a few samples from ISBI13 dataset and simulate the user interaction process. For both the proposed and baseline methods, we get started from the final segmentations and correct one mis- classified branch each time. For deterministic method (DMTT, the user draws one false-negative or erases one false-positive branch for each click. For pixel-wise probabilistic method (Prob. UNet), the user does the same while taking the uncertainty map as guidance. For the proposed method, the user checks each branch based on the descending order of structure-level uncer- tainty. VOI is used to evaluate the performances.
2.2
1.8
Ours
DMT
1.4
Prob.-UNet
1.2
1.0
0.8
10
of Click
Fig. shows the comparative results of semi-automatii proof- Figure 9: Interaction simulation.
reading with user interactions. By always checking branches with highest uncertainty, the proposed method clearly achieves better results and improves the results much faster than baseline methods. Our developed pipeline achieves higher efficiency because of following two perspectives: 1) the generated structural segmentations are essentially method provides the structure-level uncertainty map to guide the human proofreading.
botlog//strucuur-preeevini 2) the proposed
5 CONCLUSION
Instead of making pixel-wise predictions, we propose to learn structural representation with a proba- bilistic model obtain structural segmentations. Specifically, we construct the structure space by leveraging classical discrete Morse theory. We then build probabilistic model to learn a distribution over structures. In the inference stage, explore the structure-level uncertainty, which is beneficial for interactive proofreading. Extensive experiments have been conducted to demonstrate the efficacy of the proposed method.
able generate a set of structural segmentations and

---

Published as a conference paper at ICLR 2023
Reproducibility Statement: We provide the necessary experimental details in Sec. More specif- ically, the details of the data are provided in Appendix [A.6] The details of baseline methods are described in Appendix A.8 Appendix conttans the evaluations metrics used in this paper. The used computation resources are specified in Appendix A.13
The authors thank anonymous reviewers for their constructive feedback. This material based upon work supported by the Defense Advanced Research Projects Agency (DARPA) under Agreement No. HR001 1-22-9-0077. Also, this research Xiaoling Hu and Chao Chen was partially supported by NSF CCF-214490) The research of Dimitris Samaras was partially supported by NSF IIS-2212046, IIS-2123920.
Acknowledgement:
REFERENCES
I Arganda-Carraa Carreras, HS Seung. Vishwanathan, and D Berger. 3d segmentation of neurites in EM images challenge-ISBI 2013, 2013.
Samuel G Armato III, Geoffrey McLennan, Luc Bidaut, Michael F McNitt-Gray, Charles R Meyer, Anthony P Reeves, Binsheng Zhao, Denise Aberle, Claudia Henschke, Eric A Hoffman, et al. The lung image database consortium (lide) and image database resource initiative (idri): completed reference database of lung nodules on et scans. Medical physics, 2011.
Samik Banerjee, Lucas Magee, Dingkang Wang, Xu Li, Bing- Huo, Jaikishan Jayakumar, Katherine Matho, Meng-Kuan Lin, Keerthi Ram, Mohanasankar Sivaprakasam, et al. Semantic segmentation of microscopic neuroanatomical data by combining topological priors with encoder- decoder deep networks. Nature Machine Intelligence, 2020.
Kenneth Clark, Bruce Vendt, Kirk Smith, John Freymann, Justin Kirby, Paul Koppel, Stephen Moore, Stanley Phillips, David Maffitr, Michael Pringle, et al. The cancer imaging archive (tcia): maintaining and operating a public information repository. Journal of digital imaging, 2013.
J Clough, N Byrne, Oksuz, VA Zimmer, JA Schnabel, and A King. A topological loss function for deep-learning based image segmentation using persistent homology. TPAMI, 2020.
O. Delgado-Friedrricss, Robins, and A. Sheppard. Skeletonization and partitioning of digital images using discrete Morse theory. TPAMI, 2015.
T. Dey, J. Wang, and Y. Wang. Road network reconstruction from satellite images with machine learning supported by topological methods. In Proc. 27th ACM SIGSPATIAL Inil. Conf. Adv. Geographic Information Systems (GIS), 2019.
Tamal K. Dey, Jiayuan Wang, and Yusu Wang. Graph reconstruction by discrete Morse theory. In Proc. 34th Insl. Sympos. Comput. Geom. (SoCG), 2018.
Herbert Edelsbrunner and John Harer. Computational topology: an introduction. Mathematical Society, 2022.
American
Herbert Edelsbrunner, David Letscher, and Afra Zomorodian. Topological persistence and simplifica- tion. In FOCS, 2000.
Ahmed Fakhry, Hanchuan Peng, and Shuiwang Ji. Deep models for brain em image segmentation: novel insights and improved performance. Bioinformatics, 2016.
R. Forman. Morse theory for cell complexes. Advances in mathematics, 1998.
Robin Forman. A user's guide to discrete Morse theory. Sém. Lothar. Combin, 2002.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In ICML, 2016.
Alex Graves. Practical variational inference for neural networks. In NeurlPS, 2011.
Xiaoling Hu, Fuxin Li, Dimitris Samaras, and Chao Chen. tation. In NewrIPS, 2019.
deep image segmen-
10

---

Published as a conference paper at ICLR 2023
Xiaoling Hu, Yusu Wang, Li Fuxin, Dimitris Samaras, and Chao Chen. Topology-aware segmentation using discrete Morse theory. In ICLR, 2021.
Eddy llg, Ozgun Cicek, Silvio Galesso, Aaron Klein, Osama Makansi, Frank Hutter, and Thomas Brox. Uncertainty estimates and multi-hypotheses networks for optical flow. ECCV, 2018.
Viren Jain, Benjamin Bollmann, Mark Richardson, Daniel R Berger, Moritz N Helmstaedter, Kevin L Briggman, Winfried Denk, Jared B Bowden, John M Mendenhall, Wickliffe Abraham, et al. Boundary learning by optimization with topological constraints. In CVPR, 2010.
Tomasz Kaczynski, Konstantin Michael Mischaikow, and Marian Mrozek. Computational homology. Springer, 2004.
Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? In NeurlPS, 2017.
Alex Kendall, Vijay Badrinarayanan, and Roberto Cipolla. Bayesian segnet: Model uncertainty in deep convolutional encoder-decoder architectures for scene understanding. arXiv preprint arXiv: 026680, 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint 1rXiv:1312.6114, 2013.
Simon Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph Ledsam, Klaus Maier-Hein, SM Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. A probabilistic u-net for segmentation of ambiguous images. In NeurIPS, 2018.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. NewrIPS, 2017.
John Milnor. Morse theory(AA--5), Princeton university press, 1963.
Jooyoung Moon, Jihyo Kim, Younghak Shin, and Sangheum Hwang. Confidence-aware learning for deep neural networks. In ICML, 2020.
Agata Mosinska, Pablo Marquez-Neila, Mateusz Koziński, and Pascal Fua. Beyond the pixel-wise loss for topology-aware delineation. In CVPR, 2018.
William M Rand. Objective criteria for the evaluation of clustering methods. Journal of the American Statistical association, 1971.
Vanessa Robins, Peter John Wood, and Adrian P Sheppard. Theory and algorithms for constructing discrete Morse complexes from grayscale digital images. TPAMI, 2011.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In MICCAI, 2015.
Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In ICCV, 2017.
Suprosanna Shit, Johannes C Paetzold. Anjany Sekuboyina, Ivan Ezhov, Alexander Unger, Andrey Zhylka, Josien PW Pluim, Uirich Bauer, and Bjoern H Menze, cIDice-a novel loss function for tubular structure segmentation. In CVPR, 2021.
Thierry Sousbie. The persistent cosmic web and its filamentary structure-i. theory and implementation. Monthly Notices of the Royal Astronomical Society, 2011.
Joes Staal, Michael D Abràmoff, Meindert Niemeijer, Max A Viergever, and Bram Van Ginneken. Risge-based vessel segmentation color images of the retina. TMI, 2004.
Hubert Wagner, Chao Chen, and Erald Vuçini. Efficient computation of persistent homology for cubical data. In Topological methods in data analysis and visualization II. Springer, 2012.

---

Published as a conference paper at ICLR 2023
S. Wang, Y. Wang, and Y. Li. Efficient map reconstruction and augmentation via topologica methods. In Proc. 23rd ACM SIGSPATIAL, 2015.
Afra Zomorodian and Gunnar Carlsson. Computing persistent homology. Discrete & Computational Geometry, 2005.
12

---

Published as a conference paper at ICLR 2023
A
APPENDIX
Appendix
illustrates structure-level uncertainty compared with traditional pixel-wise uncertainty.
Appendix
shows more qualitative results.
Appendix
provides the details of Morse cancellation.
Appendix
illstrates the details of persistent-homology filtered topology watershed algorithm.
Appendix
describes the reparameterization technique.
Appendix
A.6
provides the details the datasets used in this paper.
Appendix
illustrates the details of the metrics used in this paper.
Appendix
8.8
describes the details the baselines used in this paper.
Appendix
discusses the advantage of joint training and optimization.
Appendix
provides few instance priors.
Appendix
shows the comparison against noisy likelihood map.
Appendix
12
llsstrtts the uncertainty estimation.
Appendix
A.J. 13 provides the computational resources for all the conducted experiments.
A.1 ILLUSTRATION OF THE STRUCTURE-LEVEL UNCERTAINTY
Fig. ssowss the comparison of traditional pixel-wise uncertainty and the proposed structure-level uncertainty. Specifically, Fig. [e) is sampled segmentation result by Prob.-UNet (Kohl al. 2018), and Fig. [ d) is the pixel uncertainty map from Prob.-UNet (Kohl et al. 2018). Different from traditional pixel-wise uncertainty, proposed structure-leve uncertainty (Fig If) can focus the structures.
We also overlay the structure-level uncertainty (Fig. I()) on the original image (Fig. Ta)), which is shown in Fig. TO, 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.
(a) Image
(b) Overlay
Figure 10: Overlaying the structure-level uncertainty on the original image.
A.2 QUALITATIVE RESULTS
Fig. II shows more qualitative results. From Fig. [ we can observe that the proposed method can generate both diverse and preserin segmentation maps.
13

---

Published as a conference paper at ICLR 2023
(a) Image
(b) GT (c) LH
(d) DMT
(e) Samplel
(f) Sample2
(g) Sample3
Figure 11: Qualitative results of the proposed method compared DMT-loss (Hu et al. 2021). From
left to right: (a) sample image, (b) ground truth, (c) continuous likelihood map and (d) thresholded
binary mask for DMT (Hu et al. 202T), and (e-g) three sampled segmentation maps generated by our method.
A.3 ADDITIONAL. DETAILS ON THE METHOD
A.3.1 DISCRETE MORSE THEORY
A 2D image is viewed 2-dimensional cubical complex, consisting of 0-, 1-, and 2-dimensional cells corresponding vertices, edges and squares its building blocks.
Discrete Morse theory (DMT) (Forman 1998 2002) is combinatorial version of Morse theory for general cell complexes. We will briefly introduce some relevant concepts for the present paper, and describe it in the setting of cubical complexes for images.
Discrete gradient vector (also called a V-pair for simplicity): Let K be cubical complex. Given p-cell T. we denote by pair (T,o) where a T.
if a is 1)-dimensional face for Discrete gradient vector is
Given collection of V-pairs M(K) over the cubical complex K. A sequence of cells T .+11 " -op .o..... off = .0411 where the superscript p in a² stands for the dimension of cell form V-path if (Ta, ox) E M(K) for for any E k] and a; < Ti-1 for any i [1, k 1]. V-path æ is acyclic (To, ok+1) & M(K). This collection of V-pairs M(K) form a discrete gradient vector fi (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. Given a discrete gradient vector field M(K) simplex K is critical if it not any V-pair in M(K).
Even though a discrete gradient vector (a V-pair), say (T,O) )) is combinatorial pair instead of real vector, still indicates "flow" from 7 to its face a. A V-path thus corresponds to flow path (integral line) in the smooth setting. However, to make a collection of V-pairs valid analog of gradient field, (cond-i) says that each simplex there should only be one "flow" direction; while
We will not introduce the concept of discrete Morse function, the discrete gradient vector field is sufficient define all relevation notations.
14

---

Published as a conference paper at ICLR 2023
(cond-ii) is necessarily flow lines traced by gradient can only go down in function values and thus never come back (thus acyclic).
A critical simplex has "vanishing gradient" it is not involved in any V-pair in M(K) (i.e, there is no flow at this simplex). Given 2D cubical complex K a a discrete gradient vector field M(K), we can view critical 0-, 1- and 2-cells minima, saddle points, and maxima, respectively.
Hence, a I-stable manifold in 2D will correspond to a V-path connecting a critical square (a maximum) and critical edge (a saddle).
Morse cancellation. 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. Fortunately, the discrete Morse theory provides an elegant and purely combinatorial way to cancel pairs of critical simplices (and thus reduce their stable manifolds). In particular, given M(K), a pair of critical simplices (5(++1)7") is cancellable if there is unique V-path
d du. Yı. 51.-... . 5g. Ys+1 = from 8 to 7. The Morse cancellation operation simple reverse all V-pairs along this path by removing all V-pairs along these path, and adding (ō.-1, Ti) to M(K) for any i [1, + 1]. It is easy to check that after the cancellation neither nor y is critical.
A.3.2 PERSISTENCE PRUNING
By setting p(or) for each cell be the maximum p-value of each vertex in a, we are able to extend this vertex-valuee function to function p K + R.
Then how obtain discrete gradient vector field from such function p 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. Let > 0 be threshold for simplification. We then perform persistence algorithm (Edelsbrunner et al. 2000) induced by the super-level set filtration of p and pair up all cells in K, denoted by (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 & grow the complex K by starting from the empty set and gradually
Carlsson| 2005). Imagine
include more and more cells in decreasing o values, which is the so-called super-level set filtration of K induced by p. Through this course, new topological feature a simplex σ, and sometiems a feature can be destroyed upon adding a simplex т. Persistence algorithm (Edelsbrunner et al. (2000) will pair up simplices; that is, its output simplices (,( [(o, })}. where each pair captures the birth and death of topological features during this evolution. The persistence
be created upon adding
set of pairs of
pair, say p = (a, r). is defined as pers(p) = p(o) - p(T). measuring how long the topological feature captured by lives in term of In this case, we also write pers(o) = pers(r)= pers(p) the persistence of a simplex (say a or T) 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 (a,T) E (K) in increasing order their persistence if (i) its persistence pers(ō, 7) < (i.e, this pair has low persistence and thus not important); and (ii) this pair (6, 7) is cancellable.
Let (K) be the resulting discrete gradient field after simplifying all low- 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 (e) be the resulting collection of 1-stable manifolds. In particular, see illustration of V-path (highlighted in black) corresponding to 1-stable manifold of the green saddle in Fig. 3(b).
A.4 APPROXIMATION FOR VOLUME DATA
As illustrated in the main text, propose messitent-homollpe filtered topology watershed algorithm to obtain the approximation of Morse structures for volume data. The details are illustrated in Alg. 0
15

---

Published as a conference paper at ICLR 2023
Algorithm 1: Persistent- filtered Topology Watershed Algorithm
Input: a grid image, and a threshold 0 Output: Morse structures for volume data
Definition: G (V, F)) denote a graph; f(()) is the intensity value of node vs lower_star(v)
{(u,0) v) E (l) < f(v)}; cc(v) is the connected component id of node
1: PD =0: Build the proximity graph (4-connectivity) for 2D grid image;
2: U = V sorted according f(v); T a sub-graph, which includes all the nodes and edges whose
value < I.
3: for U do
4: 1. /=f(v), = + {u)
5:
for (u, in (eveersttarv))) do
5:
Assert T
6:
if Tc((u) cc(v) then
6:
Edge_tag(14 u, v) = loop edge
6:
Continue
7:
else
7:
Edge_tage v) tree edge
7:
younger cc = arg max madd) f(w)
7:
older_cc cc arg min wooow)))) fww) pers = =tf((younge__cc)
7:
8:
if pers >= 8 then
8:
Edge Edge_lag(u, v) watershed edge Continue
8:
9:
end for in younger_e do
10:
ce(w) = older_cc end for
11:
PD PD + (f(younger_cce),) t)
12: end if
13: end for
14:
end for
15: return Membrane_ verteeeeeeeeeeeeeeeeeeeee U vertices of watershed_eddee edge__sce set
A.5 REPARAMETERIZATION TECHNIQUE
We adopt the reparameterization technique of VAE to make the network differentiable and be able to backpropagate.
The posterior net randomly draw samples from posterior distribution implement the posterior net as neural network, will need to backpropagate through random sampling. The issue is that backpropagation cannot flow through random node; to overcome this obstacle, we adopt the reparameterization technique proposed in (Kingma Welling 2013).
N ([ppss, post). To
Assuming the posterior is normally distributed, we can approximate it with another normal distribution. We approximate with normally distribution 7 (Z (O,
C N(H,o),
c= 14 oZ.
Now instead of saying that is sampled from Q(X, (poeet) we can say e is a function that takes parameter (Z.(p, σ)) and these μr, come from deep neural network. Therefore all we need partial derivatives w.r.t. μl, a and irrelevant for taking derivatives for backpropagation.
A.6 DATASETS
Both volume and vessel datasets are used to validate the efficacy the proposed method, and the details of the datasets are as follows:
16

---

Published as a conference paper at ICLR 2023
1. ISBI13 (volume): ISBI13 (Arganda- Aigganda-Carreras et al. 2013 is a EM dataset, containing 100 images with resolution 1024x1024.
2. CREMI (volume): CREMI is another EM dataset, containing 125 images, and each of them has resolution of 1250x1250
3. DRIVE (vessel): DRIVE (Staal et al. 2004) is a retinal vessel dataset with 40 images. The resolution for each image is 584x565.
We use a 3-fold cross-validation for all the methods to report the numbers over the validation set.
A.7 EVALUATION METRICS
The details of the metrics used in this paper are listed as follows:
DICE: DICE score is usually measure the volumetric overlap between the predicted and ground truth masks.
2. Adapted Rand Index (ARI): ART is the maximal F-score of the foreground-restriited Rand index (Rand 1971). measure of similarity between two clusters.
3. VOI Jain al. 2010): VOI is measure of the distance between two clusterings.
4. Betti Error (Hu et al. 2019): Betti Error measures the topology difference between the predicted and the ground truth mask. We randomly sample patches over the predicted segmentation and compute the average absolute enror between their Betti numbers and the corresponding ground truth patches.
A.8 BASELINES
compare the proposed method with two kinds of baselines: 1) Standard segmentation baselines:
1. DIVE (Fakhry et al. 2016) is originally designed for EM data segmentation.
2. UNet (Ronneberger et al. 2015) achieves good performances in different contexts with encoder-decoder scheme and skip connections.
(Mosinska al. 2018 proposes a vopology-aware loss based on the detected 3. linear structures with pretrained filters.
4. TopoLoss (Hu et al.) 2019) identifies the topological critical points with persistence homol- ogy and derives a novel topological loss.
5. DMT (Hu et al. 2021 identifies the whole topological structures and introduces additional penalty on the whole structures instead of isolated critical points.
2) Probabilistic-based segmentation methods:
1. Dropout UNet (Kendall et al.) 2015) dropouts the three inner-most encoder and decoder blocks with probability of 0.5 during both the training and inference.
2. Probabilistic-UN. UNet (Kohl et al. 2018) introduces probabilistic segmentation method by combining UNet with a VAE.
For all methods, generate binary segmentations by thresholding predicted likelihood maps at 0.5.
A.9 THE ADVANTAGE OF JOINT TRAINING AND OPTIMIZATION
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.
this way, we can still obtain structure- prenttrr--reeevinggg segmentation maps, but there are two main issues: I) if the segmentation network itself is structure-agnostin 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. The proposed joint training strategy overcomes both these issues. First, during the training, we incorporate the structure-aware loss (Isskoleton) Consequently, the trained segmentation branch itself
17

---

Published as a conference paper at ICLR 2023
Probability density function
64
455
-0
8
Probability density function
Probability function
(a) Image
(b) GT
(c) LH
(d) Gaussian
Figure 12: A few instance priors for specific images: (a) image, (b) ground truth, (c) continuous likelihood map and (d) Gaussian priors for specific images.
atrcttur-awaae essentially. On the other hand, with the prior and posterior nets, we are able to learn a reliable distribution of the persistence threshold (c) given an image in the inference stage. Sampling over the distribution makes possible to generate satisfactory prructur--preeervii maps within few trials (the inference will not take long), which is more much efficient.
segmentation
We conduct an empirical experiment to demonstrate the advantage of the joint training and optimiza- tion. For the postprocessing, given the predicted likelihood maps from the standard segmentation networks, we randomly choose five persistence thresholds and generate the segmentation masks sepa- rately, and select the most reasonable one as the final segmentation mask report the performances. The results Tab. demonstrate the advantage of our joint training and optimization strategy.
Table 2: Quantitative results for comparison of postprocessing on DRIVE dataset.
Method
Dice
ARI
VOL
Betti Error
Postprocessing
0.7653 0.0052
0.8841 1 0.0046
1.165 : 0.086
1.249 1 0.388
Ours 0.7814 0.0026
0.9109 + 0.0019
0.804 + 0.047
0.767 + 0,098
A.10
INSTANCE PRIORS FOR SPECIFIC IMAGES
Here we provide few instance priors for specific images (Fig. 12
A.11
ILLUSTRATION OF NOISY LIKELIHOOD MAP
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. In other words, the predicted likelihood maps
18

---

Published as a conference paper at ICLR 2023
(a) Image (b) GT
(c) Noisy likelihood (d) Clean likelihood
Figure 13: A comparison of noisy likelihood map: (a) image, (b) ground truth, (c) noisy likelihood
map and (d) clean likelihood map.
P( s F2)
Eg = 0.89
82 0.63
0.25
= 0.1 = 0.25 == 0.42
8 = 0.63
0.89
Figure 14: The learned Gaussian distribution of We use E1, 62 and e3 to denote the persistence of three different branches shown in green, blue and red, respectively. The shaded regions denotes the probability of branch2 belongs to the final segmentation map M.
from the segmentation network are good enough (see Fig. ( and Fig. (Tc)) to distinguish the skeletons/structuress of the original images, with issues only happening at locations with weak signals.
Here we provide a comparison of noisy and clean likelihood maps (Fig. [3). the model trained to be converged, usually we are able obtain the clean likelihood map (Fig 13 d)), and then could apply discrete Morse theory on to construct the structural space. the other hand,
model
not converged, the obtained likelihood map will be noisy (Fig. 13 c)), and we will not be able to recover the true signals the original image.
A.12 ILLUSTRATION OF UNCERTAINTY ESTIMATION (CDF)
As we mentioned in Sec. 3.2 the probability branch b being in segmentation map M such as ~ P follows a Bernoulli distribution with the probability Pr(b) being the cumulative distribution function (CDF) of P, CDFr((s) = P(s S ct). Fig. 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.
A.13 COMPUTATIONAL RESOURCES
All the experiments are performed on RTX A5000 GPU (24G Memory), and AMD EPYC 7542 32-Core Processor.
19


Reviews:
[{"review_id": "yFXXjmaNlgJ", "reviewer": "Reviewer_PPHt", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a new method for learning topological representations of 2D and 3D images. In contrast to existing work, the paper proposes a method that leverages the intrinsic structure of the \"representation space\", making it possible to derive (a) topological representations in a probabilistic fashion, while (b) also enabling and facilitating the analysis of structural uncertainty. The new method is then used to to primarily address segmentation task, exhibiting high-quality performance in different data sets.\n", "summary_of_the_review": "This paper was a joy to read; it provides a substantial contribution and, for the most part, presents all results and concepts in a very clear fashion. The clarity issues that I listed can be easily addressed by some restructuring and rewriting within the revision cycle. I am thus happy to provisionally endorse this paper for publication.", "strength_and_weaknesses": "This paper presents a strong high-quality contribution for obtaining topological representations of 2D and 3D images, which in turn can be used to perform downstream tasks, such as segmentation. The main strengths of the paper lie in its **firm theoretical foundation** and its **thorough experimental setup**. There are only minor issues to be rectified in the present write-up, mostly concerning the clarity of the approach: to make the paper fully accessible to an audience that may not be entirely familiar with all the concepts, some additional modifications of the text (see below) are required. Moreover, certain claims in the paper can be slightly rephrased. That being said, the paper at hand has no substantial weaknesses that would prevent publication. On aspect that slightly marred the review for me was that no code is being provided; this does not detract from reproducibility, though. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "uZFbt-H-JUX", "reviewer": "Reviewer_4b8b", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a novel technique for segmentation, especially for data with fine-scale structure. The authors questioned the conventional segmentation technique, which evaluates the segmentation on a pixel-by-pixel basis and overlooks the perspective of topological structure, and proposed a method to solve this problem. The proposed method proposes a Deep Learning structure that introduces a method for handling topological data, and compares its effectiveness with a typical baseline method.", "summary_of_the_review": "It is interesting to point out the problems with conventional segmentation methods and to propose methods to solve them. On the other hand, it is difficult to judge the value of the proposed method because it is unclear why the issue is important and there are many parts of the text that are unclear, such as the proposed method. At the very least, the sentence structure should be reviewed.\n\n**Conclusion following discussion with the authors**\n\nThe problem of sentence structure, which hinders understanding of the proposed method, has been greatly improved. There is still room for improvement of the method, but it does provide some progress on the segmentation problem. I recognize that this is a positive result, and I will raise my score.", "strength_and_weaknesses": "Strength\n\n- We raised the issue of conventional segmentation technology and created a method to solve it.\n- Experiments show that the proposed method significantly outperforms the conventional method\n\nWeakness\n\n- Before the substantive discussion, one major concern is that this paper is a difficult text to understand in terms of the proposed methodology. One of the reasons for this is the use of uncommon terms and symbols without explanation, although the terminology may be common knowledge among the authors. It is difficult to present all of them, but the following are representative ones.\n\t- In P6L1, Y remains undefined. A short time later, Y is shown to be ground truth, but it is extremely difficult to read. Also, not a big deal, but $L_{bce}$ is marked as cross-entropy, but I don't know what bce stands for.\n\t- In equation (1), $Y\\circ S$ is written, but I don't know what it means.\n\t- The terminology regarding TDA, persistent homology, seems to be different from the common one. A persistent diagram in persisitent homology is one that extracts changes in the features of a figure that vary with filtration by a certain parameter. The proposed PD seems to refer to a method to create a figure (a single figure) by connecting point clouds when the parameter is a certain value. The proposed method is a major obstacle to understanding the correction by the simplex construction and its filtration, which seems to play an important role.\n- This method seems to clarify the segment boundaries and compare them with the ground truth while adjusting the degree of clarification. I don't understand why ground truth Y and $S_{skelton}$ are combined in (1) with the above understanding, partly because the meaning of $Y\\circ S$ is not clear. Also, in fig.1, Y and $S_{skelton}$ still do not seem to be combined, which is inconsistent with equation (1). This part should clarify definitions and implications.\n- I assume that by comparing skelton and ground truth with well-defined boundaries, we are trying to make the topological information, i.e., the information of segmented chunks, similar, but why use cross-entropy for this purpose? I think you are looking at pixel-by-pixel information after all. I think there is a way to get more chunky information on a Morse theory basis. A simple idea is the Weierstrass distance of a persistent diagram (in the general sense, not in the sense of this paper) using Morse theory-based sub-level filtration, for example.\n- The results shown in the experiment are an evaluation metric for topological information. This is not surprising, since the proposed method takes topology information into account, whereas conventional methods do not. If these indices are more valid than conventional indices for segmentation evaluation, they should be discussed. Alternatively, a comparison should be made with the evaluation indices provided by conventional methods.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "3lBR5B376lT", "reviewer": "Reviewer_Jetp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The focus of the paper is on extraction of structures like vessels in 2D, 3D image datasets.  The central idea in the paper is a deep learning algorithm that can deal with topological representations.  Topological structure variability is modeled by a one-parameter family facilitated by discrete Morse theory and persistent homology.  Unlike other schemes that deal with pixelwise probabilities for describing segmentation variability, this paper learns a probabilistic model that performs inference directly in topological (structural) representation space.  Results are shown to illustrate that the new method outperforms other methods that rely on pixelwise probabilities. Moreover, the utility of the uncertainty estimates on the topological structures for ease of editing is illustrated.", "summary_of_the_review": "Overall I find the paper quite interesting and feel it has value. The main concern I have is on the justification of the empirical claims and in having more details on the inner working of the devised pipeline (please see weakness section).", "strength_and_weaknesses": "Key strengths:\nThe deep architectural formulation and integration of ideas from morse theory and persistence homology is interesting.\n\nWeaknesses:\nThe topic of vessel segmentation or topological structure estimation from 2D/3D data is something that has been studied in the literature for several decades.  While the paper claims that the new method outperforms other baselines, it appears that the new technique's performs is very close in performance to that of DMT method.  The impact of the connectivity improvements on end application goals must be described in order to really gauge in what concrete scenarios the DMT method was outperformed.    I feel that explicit knowledge of the noise characteristics of the sensing modality are critical for adaptively rank ordering possible branch hypotheses (for a given input image instance) in low signal to noise conditions where the improvements are being sought.  While I understand the rationale for the prior network and the KL divergence loss between the prior and posterior network parameter estimates,  it will be helpful to see in the appendix details of how well the instance specific priors look for specific images.  \n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%