Learnable Tree Filter for Structure-preserving Feature Transform
Lin Song Yanwei Li Zeming Li Gang Yu Hongbin Sun Jian Sun Nanning Zheng


Robust Factor Analysis Parameter EstimationThis work was supported by the Hong Kong RGC 16208917 research grant.
    Peter M. Oppeneer
    December 30, 2023
================================================================================================================




Learning discriminative global features plays a vital role in semantic segmentation. And most of the existing methods adopt stacks of local convolutions or non-local blocks to capture long-range context. However, due to the absence of spatial structure preservation, these operators ignore the object details when enlarging receptive fields. In this paper, we propose the learnable tree filter to form a generic tree filtering module that leverages the structural property of minimal spanning tree to model long-range dependencies while preserving the details. Furthermore, we propose a highly efficient linear-time algorithm to reduce resource consumption. Thus, the designed modules can be plugged into existing deep neural networks conveniently. To this end, tree filtering modules are embedded to formulate a unified framework for semantic segmentation. We conduct extensive ablation studies to elaborate on the effectiveness and efficiency of the proposed method. Specifically, it attains better performance with much less overhead compared with the classic PSP block and Non-local operation under the same backbone. Our approach is proved to achieve consistent improvements on several benchmarks without bells-and-whistles. Code and models are available at https://github.com/StevenGrove/TreeFilter-Torchhttps://github.com/StevenGrove/TreeFilter-Torch.






§ INTRODUCTION

Scene perception, based on semantic segmentation, is a fundamental yet challenging topic in the vision field. The goal is to assign each pixel in the image with one of several predefined categories. With the developments of convolutional neural networks (CNN), it has achieved promising results using improved feature representations. Recently, numerous approaches have been proposed to capture larger receptive fields for global context aggregation <cit.>, which can be divided into local and non-local solutions according to their pipelines.

Traditional local approaches enlarge receptive fields by stacking conventional convolutions <cit.> or their variants (e.g., atrous convolutions <cit.>). Moreover, the distribution of impact within a receptive field in deep stacks of convolutions converges to Gaussian <cit.>, without detailed structure preservation  (the pertinent details, which are proved to be effective in feature representation <cit.>). Considering the limitation of local operations, several non-local solutions have been proposed to model long-range feature dependencies directly, such as convolutional methods (e.g., non-local operations <cit.>, PSP <cit.> and ASPP modules <cit.>) and graph-based neural networks <cit.>. However, due to the absence of structure-preserving property, which considers both spatial distance and feature dissimilarity, the object details are still neglected. Going one step further, the abovementioned operations can be viewed as coarse feature aggregation methods, which means they fail to explicitly preserve the original structures when capturing long-range context cues.



In this work, we aim to fix this issue by introducing a novel network component that enables efficient structure-preserving feature transform, called learnable tree filter. Motivated by traditional tree filter <cit.>, a widely used image denoising operator, we utilize tree-structured graphs to model long-range dependencies while preserving the object structure. To this end, we first build the low-level guided minimum spanning trees (MST), as illustrated in Fig. <ref>. Then the distance between vertices in MST are calculated based on the high-level semantics, which can be optimized in backpropagation. For example, the dissimilarity w_k,m between vertex k and m in Fig. <ref> is calculated from semantic-rich feature embeddings. Thus, combined with the structural property of MST, the spatial distance and feature dissimilarity have been modeled into tree-structured graph simultaneously (e.g., the distance between vertex k and its spatially adjacent one n has been enlarged in Fig. <ref>, for that more edges with dissimilarities are calculated when approaching n). To enable the potential for practical application, we further propose an efficient algorithm which reduces the 𝒪(N^2) complexity of brute force implementation to linear-time consumption. Consequently, different from conditional random fields (CRF) <cit.>, the formulated modules can be embedded into several neural network layers for differentiable optimization.

In principle, the proposed tree filtering module is fundamentally different from most CNN based methods. The approach exploits a new dimension: tree-structure graph is utilized for structure-preserving feature transform, bring detailed object structure as well as long-range dependencies. With the designed efficient implementation, the proposed approach can be applied for multi-scale feature aggregation with much less resource consumption. Moreover, extensive ablation studies have been conducted to elaborate on its superiority in both performance and efficiency even comparing with PSP block <cit.> and Non-local block <cit.>. Experiments on two well-known datasets (PASCAL VOC 2012 <cit.> and Cityscapes <cit.>) also prove the effectiveness of the proposed method.











§ LEARNABLE TREE FILTER

To preserve object structures when capturing long-range dependencies, we formulate the proposed learnable tree filter into a generic feature extractor, called tree filtering module. Thus, it can be easily embedded in deep neural networks for end-to-end optimization. In this section, we firstly introduce the learnable tree filtering operator. And then the efficient implementation is presented for practical applications. The constructed framework for semantic segmentation is elaborated at last.



 §.§ Formulation
 
First, we represent the low-level feature as a undirected graph G = (V, E), with the dissimilarity weight ω for edges. The vertices V are the semantic features, and the interconnections of them can be denoted as E. Low-level stage feature map, which contains abundant object details, is adopted as the guidance for 4-connected planar graph construction, as illustrated in Fig.<ref>. Thus, a spanning tree can be generated from the graph by performing a pruning algorithm to remove the edges with the substantial dissimilarity. From this perspective, the graph G turns out to be the minimum spanning tree (MST) whose sum of dissimilarity weights is minimum out of all spanning trees. The property of MST ensures preferential interaction among similar vertices. Motivated by traditional tree filter <cit.>, a generic tree filtering module in the deep neural network can be formulated as:


    y_i=1/z_i∑ _∀ j∈Ω S(E_i,j)f(x_j).


Where i and j indicate the index of vertices, Ω denotes the set of all vertices in the tree G, x represents the input encoded features and y means the output signal sharing the same shape with x. E_i,j is a hyperedge which contains a set of vertices traced from vertex i to j in G. The similarity function S projects the features of the hyperedge into a positive scalar value, as described in Eq. <ref>. The unary function f(·) represents the feature embedding transformation. z_i is the summation of similarity S(E_i,j) alone with j to normalize the response.



    S(E_i,j)=exp(-D(i, j)),  where D(i, j)=D(j, i)=∑ _(k,m)∈E_i,jω _k,m.








According to the formula in Eq. <ref>, the tree filtering operation can be considered as one kind of weighted-average filter. The variable ω_k,m indicates the dissimilarity between adjacent vertices (k and m) that can be computed by a pairwise function (here we adopt Euclidean distance). The distance D between two vertices (i and j) is defined as summation of dissimilarity ω_k,m along the path in hyperedge E_i,j. Note that D degenerates into spatial distance when ω is set to a constant matrix. Since ω actually measures pairwise distance in the embedded space, the aggregation along the pre-generated tree considers spatial distance and feature difference simultaneously.


    y_i=1/z_i∑ _∀ j∈Ω f(x_j)∏ _(k,m)∈E_i,jexp(-ω _k,m),  where z_i=∑ _∀ j∈Ω∏ _(k,m)∈E_i,jexp(-ω _k,m).







The tree filtering module can be reformulated to Eq. <ref>. Obviously, the input feature x_j and dissimilarity ω_k,m take responsibility for the output response y_i. Therefore, the derivative of output with respect to input variables can be derived as Eq. <ref> and Eq. <ref>. V_m^i in Eq. <ref> is defined with the children of vertex m in the tree whose root node is the vertex i.


    ∂y_i/∂x_j=S(E_i,j)/z_i∂ f(x_j)/∂x_j,



    ∂y_i/∂ω _k,m=S(E_i,k)/z_i∂ S(E_k,m)/∂ω_k,m(∑ _j∈V_m^i S(E_m,j)f(x_j)-y_i z_m).









In this way, the proposed tree filtering operator can be formulated as a differentiable module, which can be optimized by the backpropagation algorithm in an end-to-end manner.



 §.§ Efficient Implementation

Let N denotes the number of vertices in the tree G. The tree filtering module needs to be accumulated N times for each output vertex. For each channel, the computational complexity of brute force implementation is 𝒪(N^2) that is prohibitive for practical applications. Definition of the tree determines the absence of loop among the connections of vertices. According to this property, a well-designed dynamic programming algorithm can be used to speed up the optimization and inference process. 

We introduce two sequential passes, namely aggregation and propagation, which are performed by traversing the tree structure. Let one vertex to be the root node. In the aggregation pass, the process is traced from the leaf nodes to the root node in the tree. For a vertex, its features do not update until all the children have been visited. In the propagation pass, the features will propagate from the updated vertex to their children recursively. 


    Aggr(ξ)_i   =ξ_i+∑ _par(j)=i S(E_i,j)Aggr(ξ)_j.
    
    
    
    
    
    Prop(ξ)_i   ={[                                                   Aggr(ξ)_r                                                                                                                   i = r; S(E_par(i),i)Prop(ξ)_par(i) + (1-S^2(E_i,par(i)))Aggr(ξ)_i)                                                                                                                   i ≠ r ].



The sequential passes can be formulated into the recursive operators for the input ξ: the aggregation pass and the propagation pass are respectively illustrated in Eq. <ref> and Eq. <ref>, where par(i) indicates the parent of vertex i in the tree whose root is vertex r. Prop(ξ)_r is initialized from the updated value Aggr(ξ)_r of root vertex r. 



14.0cm
    




As shown in the algorithm <ref>, we propose a linear-time algorithm for the tree filtering module, whose proofs are provided in the appendix. In the preparation stage, we uniformly sample a vertex as the root and perform breadth-first sorting (BFS) algorithm to obtain the topological order of tree G. The BFS algorithm can be accelerated by the parallel version on GPUs and ensure the efficiency of the following operations. 

To compute the normalization coefficient, we construct an all-ones matrix as J. Since the embedded feature f(x) is independent of the matrix J, the forward computation can be factorized into two dynamic programming processes. Furthermore, we propose an efficient implementation for the backward process. To reduce the unnecessary intermediate process, we combine the gradient of module and loss function. Note that the output y and normalization coefficient z have been already computed in the inference phase. Thus the key of the backward process is to compute the intermediate variables ψ and ν. The computation of these variables can be accelerated by the proposed linear time algorithm. And we adopt another iteration process for the gradient of pairwise distance.  

Computational complexity. Since the number of batch and channel is much smaller than the vertices in the input feature, we only consider the influence of the vertices. For each channel, the computational complexity of all the processes in algorithm <ref>, including the construction process of MST and the computation of pairwise distance, is 𝒪(N), which is linearly dependent on the number of vertices. It is necessary to point out that MST can be built in linear time using Contractive Boruvka algorithm if given a planar graph, as is designed in this paper. Note that the batches and channels are independent of each other. For the practical implementation on GPUs, we can naturally perform the algorithm for batches and channels in parallel. Also, we adopt an effective scheduling scheme to compute vertices of the same depth on the tree parallelly. Consequently, the proposed algorithm reduces the computational complexity and time consumption dramatically.




 §.§ Network Architecture for Segmentation





Based on the efficient implementation algorithm, the proposed tree filtering module can be easily embedded into deep neural networks for resource-friendly feature aggregation. To illustrate the effectiveness of the proposed module, here we employ ResNet <cit.> as our encoder to build up a unified network. The encoded features from ResNet are usually computed with output stride 32. To remedy for the resolution damage, we design a naive decoder module following previous works <cit.>. In details, the features in the decoder are gradually upsampled by a factor of 2 and summed with corresponding low-level features in the encoder, similar to that in FPN <cit.>. After that, the bottom-up embedding functions in the decoder are replaced by tree filter modules for multi-scale feature transform, as intuitively illustrated in Fig. <ref>.

To be more specific, given a low-level feature map M_l in top-down pathway, which riches in instance details, a 4-connected planar graph can be constructed easily with the guidance of M_l. Then, the edges with substantial dissimilarity are removed to formulate MST using the Boruvka <cit.> algorithm. High level semantic cues contained in X_l are extracted using a simplified embedding function (Conv 1×1). To measure the pairwise dissimilarity ω in Eq. <ref> (ω_l in Fig. <ref>), the widely used Euclidean distance <cit.> is adopted. Furthermore, different groups of tree weights ω_l are generated to capture component dependent features, which will be analyzed in Sec <ref>. To highlight the effectiveness of the proposed method, the feature transformation f(·) is simplified to identity mapping, where f(X_l)=X_l. Thus, the learnable tree filter can be formulated by the algorithm elaborated on Sec. <ref>. Finally, the low-level feature map M_l is fused with the operation output y_l using pixel-wise summation. For multi-stage feature aggregation, the building blocks (green nodes in Fig. <ref>) are applied to different resolutions (Stage 1 to 3 in Fig. <ref>). An extra global average pooling operation is added to capture global context and construct another tree filtering module (Stage 4). The promotion brought by the extra components will be detailed discussed in ablation studies.



§ EXPERIMENTS

In this section, we firstly describe the implementation details. Then the proposed approach will be decomposed step-by-step to reveal the effect of each component. Comparisons with several state-of-the-art benchmarks on PASCAL VOC 2012 <cit.> and Cityscapes <cit.> are presented at last.



 §.§ Implementation Details

Following traditional protocols <cit.>, ResNet <cit.> is adopted as our backbone for the following experiments. Specifically, we employ the “poly” schedule with an initial learning rate 0.004 and power 0.9. The networks are optimized for 40K iterations using mini-batch stochastic gradient descent (SGD) with a weight decay of 1e-4 and a momentum of 0.9. We construct each mini-batch for training from 32 random crops (512×512 for PASCAL VOC 2012 <cit.> and 800×800 for Cityscapes <cit.>) after randomly flipping and scaling each image by 0.5 to 2.0×.


































 §.§ Ablation Studies
 
To elaborate on the effectiveness of the proposed approach, we conduct extensive ablation studies. First, we give detailed structure-preserving relation analysis as well as the visualization, as presented in Fig. <ref>. Next, the equipped stage and group number of the tree filtering module is explored. Different building blocks are compared to illustrate the effectiveness and efficiency of the tree filtering module.

Structure-preserving relations. As intuitively presented in the first row of Fig. <ref>, given different positions, the corresponding instance details are fully activated with the high response, which means that the proposed module has learned object structures and the long-range intra-class dependencies. Specifically, the object details (e.g., boundaries of the train rather than the coarse regions in Non-local <cit.> blocks) have been highlighted in the affinity maps. Qualitative results are also given to illustrate the preserved structural details, as presented in Fig. <ref>.

Which stage to equip the module? Tab. <ref> presents the results when applying the tree filtering module to different stages (group number is fixed to 1). The convolutional operations are replaced by the building block (green nodes in Fig. <ref>) to form the equipped stage. As can be seen from Tab. <ref>, the network performance consistently improves with more tree filtering modules equipped. This can also be concluded from the qualitative results (the second row in Fig. <ref>), where the higher stage (Stage 3) contains more semantic cues and lower stages (Stage 1 and 2) focus more on complementary details. The maximum gap between the multi-stage equipped network and the raw one even up to 4.5% based on ResNet-50 backbone. Even with the powerful ResNet-101, the proposed approach still attains a 2.1% absolute gain, reaching 77.1% mIoU on PASCAL VOC 2012 val set.

Different group numbers. Different group settings are used to generate weights for the single-stage tree filtering module when fixing the channel number to 256. As shown in Tab. <ref>, the network reaches top performance (Group Num=16) when the group number approaching the category number (21 in PASCAL VOC 2012), and additional groups afford no more contribution. We guess the reason is that different kinds of tree weights are needed to deal with similar but different components, as shown in the third row of Fig. <ref>.

Different building blocks. We further compare the proposed tree filtering module with classic context modeling blocks (e.g., PSP block <cit.> and Non-local block <cit.>) and prove the superiority both in accuracy and efficiency. As illustrated in Tab. <ref>, the proposed module (TF) achieves better performance than others (PSP and NL) with much less resource consumption. What's more, the tree filtering module brings consistent improvements in different backbones (5.2% for ResNet-50 with stride 32 and 2.6% for ResNet-101) with additional 0.7M parameters and 1.3G FLOPs overheads. Due to the structural-preserving property, the proposed module achieves additional 1.1% improvement over the PSP block with neglectable consumption, as presented in Tab. <ref>. And the extra Non-local block contributes no more gain over the tree filtering module, which could be attributed to the already modeled feature dependencies in the tree filter.

Extra components. Same with other works <cit.>, we adopt some simple components for further improvements, including an extra global average pooling operation and additional ResBlocks in the decoder. In details, the global average pooling block combined with the Stage 4 module (see Fig. <ref>) is added after the backbone to capture global context, and the “Conv1×1” operations in the decoder (refer to the detailed diagram in Fig. <ref>) are replaced by ResBlocks (with “Conv3x3”). As presented in Tab. <ref>, the proposed method achieves consistent improvements and attains 79.4% on PASCAL VOC 2012 val set without applying data augmentation strategies.









































 §.§ Experiments on Cityscapes

To further illustrate the effectiveness of the proposed method, we evaluate the Cityscapes <cit.> dataset. Our experiments involve the 2975, 500, 1525 images in train, val, and test set, respectively. With multi-scale and flipping strategy, the proposed method achieves 80.8% mIoU on Cityscapes test set when trained on fine annotation data only. Compared with previous leading algorithms, our method achieves superior performance using ResNet-101 without bells-and-whistles.



 §.§ Experiments on PASCAL VOC

We carry out experiments on PASCAL VOC 2012 <cit.> that contains 20 object categories and one background class. Following the procedure of <cit.>, we use the augmented data <cit.> with 10,582 images for training and raw train-val set for further fine-tuning. In inference stage, multi-scale and horizontally flipping strategy are adopted for data augmentation. As shown in Tab. <ref>, the proposed method achieves the state-of-the-art performance on PASCAL VOC 2012 <cit.> test set. In details, our approach reaches 84.2% mIoU without MS-COCO <cit.> pre-train when adopting vanilla ResNet-101 as our backbone. If MS-COCO <cit.> is added for pre-training, our approach achieves the leading performance with 86.3% mIoU. 



§ CONCLUSION

In this work, we propose the learnable tree filter for structure-preserving feature transform. Different from most existing methods, the proposed approach leverages tree-structured graph for long-range dependencies modeling while preserving detailed object structures. We formulate the tree filtering module and give an efficient implementation with linear-time source consumption. Extensive ablation studies have been conducted to elaborate on the effectiveness and efficiency of the proposed method, which is proved to bring consistent improvements on different backbones with little computational overhead. Experiments on PASCAL VOC 2012 and Cityscapes prove the superiority of the proposed approach on semantic segmentation. More potential domains with structure relations (e.g., detection and instance segmentation) remain to be explored in the future.



§ ACKNOWLEDGMENT

We would like to thank Lingxi Xie for his valuable suggestions. This research was supported by the National Key R&D Program of China (No. 2017YFA0700800).






unsrt









§ ALGORITHM PROOF

In this section, we present the detailed proofs for the formula in Algorithm 1. Note that the symbols follow the definition in the main paper.

Given one vertex as i, the Eq. <ref> proves the backward process of ∂ loss/∂x in Algorithm 1.

    ∂ loss/∂x_i   =∂ f(x_i)/∂x_i∑_j ∈Ω∂ loss/∂y_j∂y_j/∂x_i
       =∂ f(x_i)/∂x_i∑_j ∈ΩS(E_i, j)(∂ loss/∂y_j1/z_j)
       =∂ f(x_i)/∂x_iψ_i




Given one edge with a pair of connected vertices i and j, the Eq. <ref>-<ref> proves the backward process of ∂ loss/∂ω in Algorithm 1, where Ω_i indicates the set of children of vertex i in the tree whose root is j and Ω_j indicates the set of children of vertex j in the tree whose root is i.


    ∂ loss/∂ω_i, j   =∂ S(E_i, j)/∂ω_i, j∑∑_m ∈Ω∂ loss/∂y_m∂y_m/∂ S(E_i, j)
       =∂ S(E_i, j)/∂ω_i, j∑(∑_m ∈Ω1/z_m∂ loss/∂y_m∂ z_m y_m/∂ S(E_i, j) - ∑_m ∈Ω∂ loss/∂y_my_m/z_m∂ z_m/∂ S(E_i, j))
       =∂ S(E_i, j)/∂ω_i, j∑(γ_i^s - γ_i^z)


The computational complexity of the component γ_i^s can be reduce to linear by the dynamic programming procedure below.


    γ_i^s   =∑_m ∈Ω∂ loss/∂y_m1/z_m∂ z_m y_m/∂ S(E_i, j)
       =∑_m ∈Ω(ϕ/z)_m∑_k ∈Ω∂ S(E_m, k) f(x_k)/∂ S(E_i, j)
       =∑_m ∈Ω_i(ϕ/z)_m∑_k ∈Ω_jS(E_m, i)S(E_i, k)f(x_k) + ∑_m ∈Ω_j(ϕ/z)_m∑_k ∈Ω_iS(E_m, j)S(E_j,k)f(x_k)
       =∑_m ∈Ω_iS(E_m, i)(ϕ/z)_m∑_k ∈Ω_jS(E_i, k)f(x_k) + ∑_m ∈Ω_jS(E_m, j)(ϕ/z)_m∑_k ∈Ω_iS(E_j,k)f(x_k)
       =ψ̂_i·ρ̂_j + ψ̂_j·ρ̂_i
       =ψ̂_i·(ρ_i - S(E_i, j)ρ̂_i)) + ρ̂_i·(ψ_i - S(E_j, i)ψ̂_i))
       =ψ̂_i·ρ_i + ψ_i·ρ̂_i - 2S(𝐄_i,j)ψ̂_i·ρ̂_i

The same procedure can be easily adapted to obtain the component γ_i^z.

    γ_i^z   =∑_m ∈Ω∂ loss/∂y_my_m/z_m∂ z_m/∂ S(E_i, j)
        =ν̂_i z_i + ν_i ẑ_i - 2S(E_i,j)ν̂_i ẑ_i





