Topology-Aware Segmentation Using Discrete Morse Theory

1 INTRODUCTION . Segmenting objects while preserving their global structure is a challenging yet important problem . Various methods have been proposed to encourage neural networks to preserve fine details of objects ( Long et al. , 2015 ; He et al. , 2017 ; Chen et al. , 2014 ; 2018 ; 2017 ) . Despite their high per-pixel accuracy , most of them are still prone to structural errors , such as missing small object instances , breaking thin connections , and leaving holes in membranes . These structural errors can significantly damage downstream analysis . For example , in the segmentation of biomedical structures such as membranes and vessels , small pixel errors at a junction will induce significant structure error , leading to catastrophic functional mistakes . See Fig . 1 for an illustration . Topology is a very global characterization that needs a lot of observations to learn . Any training set is insufficient in teaching the network to correctly reason about topology , especially near challenging spots , e.g. , blurred membrane locations or weak vessel connections . A neural network tends to learn from the clean-cut cases and converge quickly . Meanwhile , topologically-challenging locations remain mis-classified , causing structural/topological errors . We note that this issue can not be alleviated even with more annotated ( yet still unbalanced ) images . We propose a novel approach that identifies critical topological structures during training and teaches a neural network to learn from these structures . Our method can produce segmentations with correct topology , i.e. , having the same Betti number ( i.e. , number of connected components and handles/tunnels ) as the ground truth . Underlying our method is the classic Morse theory ( Milnor , 1963 ) , which captures singularities of the gradient vector field of the likelihood function . Intuitively speaking , we treat the likelihood as a terrain function and Morse theory helps us capture terrain structures such as ridges and valleys . See Fig . 1 for an illustration . These structures , composed of 1D and 2D manifold pieces , reveal the topological information captured by the ( potentially noisy ) likelihood function . We consider these Morse structures as topologically critical ; they encompass all potential skeletons of the object . We propose a new loss that identifies these structures and enforce higher penalty along them . This way , we effectively address the sampling bias issue and ensure that the networks predict correctly near these topologically difficult locations . Since the Morse structures are identified based ∗Email : Xiaoling Hu ( xiaolhu @ cs.stonybrook.edu ) . on the ( potentially noisy ) likelihood function , they can be both false negatives ( a structure can be a true structure but was missed in the segmentation ) and false positives ( a hallucination of the model and should be removed ) . Our loss ensures that both kinds of structural mistakes are corrected . Several technical challenges need to be addressed . First , classical Morse theory was defined for smooth functions on continuous domains . Computing the Morse structures can be expensive and numerically unstable . Furthermore , the entire set of Morse structures may include an excessive amount of structures , a large portion of which can be noisy , irrelevant ones . To address these challenges , we use the discrete version of Morse theory by Forman ( 1998 ; 2002 ) . For efficiency purposes , we also use an approximation algorithm to compute 2D Morse structures with almost linear time . The idea is to compute zero dimensional Morse structures of the dual image , which boils down to a minimum spanning tree computation . Finally , we use the theory of persistent homology ( Edelsbrunner et al. , 2000 ; Edelsbrunner & Harer , 2010 ) to prune spurious Morse structures that are not relevant . Our discrete-Morse-theory based loss , called the DMT-loss , can be evaluated efficiently and can effectively train the neural network to achieve high performance in both topological accuracy and per-pixel accuracy . Our method outperforms state-of-the-art methods in multiple topology-relevant metrics ( e.g. , ARI and VOI ) on various 2D and 3D benchmarks . It has superior performance in the Betti number error , which is an exact measurement of the topological fidelity of the segmentation . Related work . Closely related to our method are recent works on persistent-homology-based losses ( Hu et al. , 2019 ; Clough et al. , 2019 ) . These methods identify a set of critical points of the likelihood function , e.g. , saddles and extrema , as topologically critical locations for the neural network to memorize . However , only identifying a sparse set of critical points at every epoch is inefficient in terms of training . Instead , our method identifies a much bigger set of critical locations at each epoch , i.e. , 1D or 2D Morse skeletons ( curves and patches ) . This is beneficial in both training efficiency and model performance . Extending the critical location sets from points to 1D curves and 2D patches makes it much more efficient in training . Compared with TopoLoss ( Hu et al. , 2019 ) , we observe a 3-time speedup in practice . Furthermore , by focusing on more critical locations early , our method is more likely to escape poor local minima of the loss landscape . Thus it achieves better topological accuracy than TopoLoss . The shorter training time may also contribute to better stability of the SGD algorithm , and thus better test accuracy ( Hardt et al. , 2016 ) . Another topology-aware loss ( Mosinska et al. , 2018 ) uses pretrained filters to detect broken connections . However , this method can not be generalized to unobserved geometry and higher dimensional topology ( loops and voids ) . We also refer to other existing work on topological features and their applications ( Adams et al. , 2017 ; Reininghaus et al. , 2015 ; Kusano et al. , 2016 ; Carriere et al. , 2017 ; Ni et al. , 2017 ; Wu et al. , 2017 ) . Deep neural networks have also been proposed to learn from topological features directly extracted from data ( Hofer et al. , 2017 ; Carrière et al. , 2019 ) . Persistent-homology-inspired objective functions have been proposed for graphics ( Poulenard et al. , 2018 ) machine learning ( Chen et al. , 2019 ; Hofer et al. , 2019 ) . Discrete Morse theory has been used to identify skeleton structures from images ; e.g. , ( Delgado-Friedrichs et al. , 2015 ; Robins et al. , 2011 ; Wang et al. , 2015 ) . The resulting 1D Morse structure has been used to enhance neural network architecture : e.g. , in ( Dey et al. , 2019 ) it is used to both pre- and post-process images , while in ( Banerjee et al. , 2020 ) , the 1D Morse skeleton is used as a topological prior ( part of input ) for an encoder-decoder deep network for semantic segmentation of microscopic neuroanatomical data . Our work , in contrast , uses Morse structures ( beyond 1D ) of the output to strengthen the global structural signal more explicitly in an end-to-end training of a network . Many methods have leveraged the power of deep neural networks ( Ronneberger et al. , 2015 ; Long et al. , 2015 ; Badrinarayanan et al. , 2017 ; Ding et al. , 2019 ; Kervadec et al. , 2019 ; Karimi & Salcudean , 2019 ; Mosinska et al. , 2018 ) for fine-scale structure segmentation . One may also enforce connectivity constraints when postprocessing the likelihood map ( Han et al. , 2003 ; Le Guyader & Vese , 2008 ; Sundaramoorthi & Yezzi , 2007 ; Ségonne , 2008 ; Wu et al. , 2017 ; Gao et al. , 2013 ; Vicente et al. , 2008 ; Nowozin & Lampert , 2009 ; Zeng et al. , 2008 ; Chen et al. , 2011 ; Andres et al. , 2011 ; Stuhmer et al. , 2013 ; Oswald et al. , 2014 ; Estrada et al. , 2014 ) . However , when the deep neural network itself is topology-agnostic , the likelihood map may be fundamentally flawed and can not be salvaged topologically . Specific to neuron image segmentation , some methods ( Funke et al. , 2017 ; Turaga et al. , 2009 ; Januszewski et al. , 2018 ; Uzunbas et al. , 2016 ; Ye et al. , 2019 ) directly find neuron regions instead of their boundary/membranes . These methods can not be generalized to other types of data such as satellite images , retinal images , vessel images , etc . 2 METHOD . We propose a novel loss to train a topology-aware network end-to-end . It uses global structures captured by discrete Morse theory ( DMT ) to discover critical topological structures . In particular , through the language of 1- and 2-stable manifolds , DMT helps identify 1D skeletons or 2D sheets ( separating 3D regions ) that may be critical for structural accuracy . These Morse structures are used to define a DMT-loss that is essentially the cross-entropy loss constrained to these topologically critical structures . As the training continues , the neural network learns to better predict around these critical structures , and eventually achieves better topological accuracy . Please refer to Fig . 2 for an overview of our method . 2.1 MORSE THEORY . Morse theory ( Milnor , 1963 ) identifies topologically critical structures from a likelihood map ( Fig . 3 ( a ) ) . In particular , it views the likelihood as a terrain function ( Fig . 3 ( b ) ) and extracts its landscape features such as mountain ridges and their high-dimensional counterparts . The broken connection in the likelihood map corresponds to a local dip in the mountain ridge of the terrain in Fig . 3 ( b ) and Fig . 3 ( c ) . The bottom of this dip is captured by a so-called saddle point ( S in Fig . 3 ( c ) ) of the likelihood map . The mountain ridge connected to this bottom point captures the main part of missing pixels . Such “ mountain ridges ” can be captured by the so-called stable manifold w.r.t . the saddle point using the language of Morse theory . By finding the saddle points and the stable manifold of the saddle points on the likelihood map , we can ensure the model learns to “ correctly ” handle pixels near these structures . We note that an analogous scenario can also happen with such 1D signals ( such as blood vessels ) as well as 2D signals ( such as membranes of cells ) in 3D images – they can also be captured by saddles ( of different indices ) and their stable manifolds . In this paper , we focus on the application of segmenting 2D and 3D images . Specifically , suppose we have a smooth function f : Rd → R to model the likelihood ( density ) map . Given any point x ∈ Rd , the negative gradient −∇f ( x ) = − [ ∂f∂x1 , ∂f ∂x2 , . . . , ∂f∂xd ] T indicates the steepest descending direction of f . A point x = ( x1 , x2 , . . . , xk ) is critical if the function gradient at this point vanishes ( i.e. , ∇f ( x ) = 0 ) . For a well-behaved function ( more formally , called Morse function ) defined on Rd , a critical point could be a minimum , a maximum , or d− 1 types of saddle points . See Fig . 3 ( c ) for an example . For d = 2 , there is only one saddle point type . For d = 3 , there are two saddle point types , referred to as index-1 and index-2 saddles . Formally , when taking the eigen-values of the Hessian matrix at a critical point , its index is equal to the number of negative eigenvalues . Intuitively , imagine we put a drop of water on the graph of f ( i.e , the terrain in Fig . 3 ( b ) ) at the lift of x onto this terrain , then −∇f ( x ) indicates the direction along which the water will flow down . If we track the trajectory of this water drop as it flows down , this gives rise to a so-called integral line ( a flow line ) . Such flow lines can only start and end at critical points1 , where the gradient vanishes . The stable manifold S ( p ) of a critical point p is defined as the collection of points whose flow line ends at p. For a 2D function f : R2 → R , for a saddle q , its stable manifold S ( q ) starts from local maxima ( mountain peaks in the terrain ) and ends at q , tracing out the mountain ridges separating different valleys ( Fig . 3 ( c ) ) . The stable manifold S ( p ) of a minimum p , on the other hand , corresponds to the entire valley around this minimum p. See the valley point V and its corresponding stable field in Fig . 3 ( c ) . For a 3D function f : R3 → R , the stable manifold w.r.t . an index-2 saddle connects mountain peaks to saddles , tracing 1D mountain ridges as in the case of a 2D function . The stable manifold w.r.t . an index-1 saddle q consists of flow lines starting at index-2 saddles and ending at q . Their union , called the 2-stable manifold of f , consists of a collection of 2-manifold pieces . These stable manifolds indicate important topological structures ( graph-like or sheet-like ) based on the likelihood of the current neural network . Using these structures , we will propose a novel loss ( Sec . 2.3 ) to improve the topological awareness of the model . In practice , for images , we will leverage the discrete version of Morse theory for both numerical stability and easier simplification . Discrete Morse theory . Due to space limitations , we briefly explain discrete Morse theory , leaving technical details to Appendix A.2.1 . We view a d-dimensional image , d = 2 or 3 , as a d-dimensional cubical complex , meaning it consists of 0- , 1- , 2- and 3-dimensional cells corresponding to vertices ( pixels/voxels ) , edges , squares , and cubes as its building blocks . Discrete Morse theory ( DMT ) , originally introduced in ( Forman , 1998 ; 2002 ) , is a combinatorial version of Morse theory for general cell complexes . In this setting , the analog of a gradient vector is a pair of adjacent cells , called discrete gradient vectors . The analog of an integral line ( flow line ) is a sequence of such cell-pairs ( discrete gradient vectors ) , forming a so-called V-path . Critical points correspond to critical cells which do not participate in any discrete gradient vectors . A minimum , an index-1 saddle , an index-2 saddle and a maximum for a 3D domain intuitively correspond to a critical vertex , a critical edge , a critical square and a critical cube respectively . A 1-stable manifold in 2D will correspond to a V-path , i.e. , a sequence of cells , connecting a critical square ( a maximum ) and a critical edge ( a saddle ) . See Fig . 3 ( c ) for an illustration . In 3D , it will be a V-path connecting a critical cube and a critical square . 1More precisely , flow lines only tend to critical points in the limit and never reach them .

Summary: The paper focus on the segmentation of images in which the correct topology of the segmentation an important role plays. The key idea of the paper is the use of Discrete Morse Theory to identify those areas of the segmentation/likelihood map that are important for ensuring the correct topology of the segmentation. This is achieved by modulating the normal binary cross entropy loss function so that areas of Morse structures are focused on during the learning.