Learning Affinity via Spatial Propagation Networks
Sifei Liu, Shalini De Mello, Jinwei Gu, Guangyu Zhong, Ming-Hsuan Yang, Jan Kautz
Introduction
An affinity matrix is a generic matrix that determines how close, or similar, two points are in a space. In computer vision tasks, it is a weighted graph that regards each pixel as a node, and connects each pair of pixels by an edge . The weight on that edge should reflect the pairwise similarity with respect to different tasks. For example, for low-level vision tasks such as image filtering, the affinity values should reveal the low-level coherence of color and texture ; for mid to high-level vision tasks such as image matting and segmentation , the affinity measure should reveal the semantic-level pairwise similarities. Most techniques explicitly or implicitly assume a measurement or a similarity structure over the space of configurations. The success of such algorithms depends heavily on the assumptions made to construct these affinity matrices, which are generally not treated as part of the learning problem.
In this paper, we show that the problem of learning the affinity matrix can be equivalently expressed as learning a group of small row/column-wise, spatially varying linear transformation matrices. Since a linear transformation can be easily implemented as a differentiable module in a deep neural network, the transformation matrix can be learned in a purely data-driven manner as opposed to being constructed by hand. Specifically, we adopt an independent deep CNN to output all entities of the matrix with the input of the original RGB images, such that the affinity is learned from a deep model conditioned on the specific inputs. We show that using a three-way connection, instead of the fully connection, is sufficient for learning a dense affinity matrix and requires much fewer output channels of a deep CNN. Therefore, instead of using designed features and kernel tricks, our network outputs all entities of the affinity matrix in a data-driven manner.
The advantages of learning affinity matrix in a data-driven manner are multifold. First, the assumption of a similarity matrix based on distance metric in certain space (e.g., RGB or Euclidean ) may not describe the pairwise relations in mid-to-high-level feature space. To apply such designed pairwise kernels to tasks such as semantic segmentation, multiple iterations is required for a satisfactory performance. In contrast, the proposed method learns and outputs all entities of an affinity matrix under direct supervision of ultimate loss functions, where no iteration, specific design or assumption of kernel function is needed. Second, we can learn the high-level affinity measure initializing from hierarchical deep features from the VGG and ResNet where conventional metrics and kernels may not be applied. Due to the above properties, especially the first one, the framework is far more efficient than the related graphical models, such as Dense CRF.
Our proposed architecture, namely spatial propagation network (SPN), contains a deep CNN that learns the affinity entities and a spatial linear propagation module. Images or general 2D matrix are fed into the module, and propagated under the guidance of the learned affinity. All modules are differentiable and jointly trained using stochastic gradient descent (SGD) method. The spatial linear propagation module is computationally efficient for inference due to the linear time complexity of the recurrent architecture.
Related Work
Numerous methods explicitly design affinity matrices for image filtering , colorization , matting and image segmentation based on the physical nature of the problem. Other methods, such as total variation (TV) and learning to diffuse improve the modeling of pairwise relationships by utilizing different objectives, or incorporating more priors into diffusion partial differential equations (PDEs). However, due to the lack of an effective learning strategy, it is still challenging to produce a learning based affinity for complex visual analysis problems. Recently, Maire et al. trained a deep ConvNet to directly predict the entities of an affinity matrix, which demonstrated good performance on image segmentation. However, since the affinity is followed by a solver of spectral embedding as an independent part, it is not directly supervised for the classification/prediction task. Bertasius et al. introduced a random walk network that optimizes the objectives of pixel-wise affinity for semantic segmentation. Differently, the affinity matrix is additionally supervised by a ground-truth sparse pixel similarities, which limits the potential connections between pixels.
On the other hand, many graphical model-based methods have successfully improved the performance of image segmentation. In the deep learning framework, the conditional random fields (CRFs) with efficient mean field inference is frequently used to model the pairwise relations in the semantic labeling space. Some of them use it as a post-processing module , while others integrate it as a jointly-trained part . While both dense CRFs and the proposed method describe the densely connected pairwise relationships, dense CRFs relies on designed kernels, while ours utilizes propagation structure with directly learned pairwise links, which is much less explored before, but more suitable as an embedded deep-learning module. Since in this paper, SPN is trained as a universal segmentation refinement module, we specifically compare it with one of the methods that relies on the dense CRF as a post-processing strategy. Our architecture is also related to the multi-dimensional RNN or LSTM . However, both the standard RNN and LSTM contain multiple non-linear units and thus do not fit into the proposed affinity framework.
Proposed Approach
In this work, we construct a spatial propagation network (SPN) that can transform a two-dimensional (2D) map (e.g., the result of coarse image segmentation) to a new one with desired properties (e.g., a new segmentation map with significantly refined details). With spatially varying parameters that supports the propagation process, we show theoretically in Section 3.1 that this module is equivalent to the standard anisotropic diffusion process . As proved, the transformation of maps is controlled by a Laplacian matrix that is constituted by the parameters of the spatial propagation module. Since the propagation module is differentiable, those parameters can be learned by any type of neural network (e.g., a typical deep CNN) that is connected to this module through joint training. We introduce the propagation network in Section 3.2, and specifically analyze the properties of different types of connections within the framework for learning the affinity matrix.
We apply a linear transformation by means of the spatial propagation network, where a matrix is scanned row/column-wise in four fixed directions: left-to-right, top-to-bottom, and verse-vise. This strategy is used widely in . We take the left-to-right direction as an example for the following discussion. Other directions are processed independently in the same manner.
We denote and as two 2D maps of size , with exactly the same dimensions as the matrix before and after spatial propagation, where and , respectively, represent their columns with elements each. We linearly propagate information from left-to-right between adjacent columns using an linear transform matrix as:
where is the identity matrix, the initial condition , and is a diagonal matrix, where the element is the sum of all the elements of the row of :
As shown, the matrix , where , is updated in a column-wise manner recursively. For each column, is a linear, weighted combination of the previous column , and the corresponding column in .
When the recursive scanning is finished, the updated 2D matrix can be expressed with an expanded formulation of Eq. (1):
where is a lower triangular, transformation matrix, which relates and . and are vectorized versions of and , respectively, with the dimension of . Specifically, they are created by concatenating and along the same, single dimension, i.e., and . All the parameters are sub-matrices, where .
In the following section, we validate that Eq. (3) can be expressed as a spatial anisotropic diffusion process, with the corresponding propagation affinity matrix constituted by all .
The summation of elements in each row of equals to one.
Since G contains sub-matrices, each representing the transformation between the corresponding columns of and , we denote all the weights used to compute as the block-row . On setting , the constituent sub-matrix of is:
To prove that the summation of any row in equals to one, we instead prove that for , each row of has the summation of one.
Denoting as an vector, we need to prove that . Equivalently , because is a lower triangular matrix. In the following part, we first prove that when , we have by mathematical induction .
When , , which satisfies the assertion.
Assume there is a , such that , we must prove the formula is true for .
According to the formulation of the diagonal matrix in Eq. (2) we have . Therefore, the assertion is satisfied. When , we have:
which yields the equivalence of Theorem 1. ∎
We define the evolution of a 2D matrix as a time sequence , where is the initial state. When the transformation between any two adjacent states follows Eq. (3), the sequence is a diffusion process expressed with a partial differential equation (PDE):
where is the Laplacian matrix, is the degree matrix composed of in Eq. (2), and is the affinity matrix composed by the off-diagonal elements of .
We substitute the and as two consecutive matrices and in (3). According to Theorem 1, we ensure that the sum of each row is that can formulate a standard Laplacian matrix. Since has the diagonal sub-matrix , we can rewrite (3) as:
where , is an diagonal matrix containing all the and is the off-diagonal part of . It then yields , a discrete formulation of (7) with the time discretization interval as one. ∎
Theorem 2 shows the essential property of the row/column-wise linear propagation in Eq. (1): it is a standard diffusion process where defines the spatial propagation and , the affinity matrix, describes the similarities between any two points. Therefore, learning the image affinity matrix in Eq. (8) is equivalent to learning a group of transformation matrices in Eq. (1).
In the following section, we show how to build the spatial propagation (1) as a differentiable module that can be inserted into a standard feed-forward neural network, so that the affinity matrix can be learned in a data-driven manner.
2 Learning Data-Driven Affinity
Since the spatial propagation in Eq.(1) is differentiable, the transformation matrix can be easily configured as a row/column-wise fully-connected layer. However, we note that since the affinity matrix indicates the pairwise similarities of a specific input, it should also be conditioned on the content of this input (i.e., different input images should have different affinity matrices). Instead of setting the matrices as fixed parameters of the module, we design them as the outputs of a deep CNN, which can be directly conditioned on an input image.
One simple way is to set the output of the deep CNN to use the same size as the input matrix. When the input has channels (e.g., an RGB image has ), the output needs channels (there are connections from the previous row/column per pixel per channel, and with four different directions). Obviously, this is too many (e.g., an feature map needs an output of ) to be implemented in a real-world system. Instead of using full connections between the adjacent rows/columns, we show that certain local connections, corresponding to a sparse row/column-wise transform matrix, can also formulate a densely connected affinity. Specifically, we introduce the (a) one-way connection and the (b) three-way connection as two different ways to implement Eq. (1).
The one-way connection enables every pixel to connect to only one pixel from the previous row/column (see Figure 1(a)). It is equivalent to an one dimensional (1D), linear recurrent propagation that scans each row/column independently as an 1D sequence. Following Eq. (1), we denote and as the pixels in the column, where the left-to-right propagation for one-way connection is:
where is a scaler weight indicating the propagation strength between the pixels at and . Equivalently, in Eq. (1) is a diagonal matrix, with the elements constituted by .
The one-way connection is a direct extension of sequential recurrent propagation . The exact formulation of Eq. (9) has been used previously for semantic segmentation and for learning low-level vision filters . In , Chen et al.explain it by domain transform, where in semantic segmentation, corresponds to the object edges. Liu et al. explain it by arbitrary-order recursive filters, where corresponds to more general image properties (e.g., low-level image/color edges, missing pixels, etc.). Both of these can be explained as the same linear propagation framework of Eq. (1) with one-way connection.
As introduced in Theorem 2, the affinity matrix with linear propagation is composed of the off-diagonal elements of in Eq. (3). The one-way connection formulates a spares affinity matrix, since each sub-matrix of has nonzero elements only along its diagonal, and the multiplication of several individual diagonal matrics will also results in a diagonal matrix. On the other hand, the three-way connection, also with a sparse , can form a relatively dense with the multiplication of several different tridiagonal matrices. It means pixels can be densely and globally associated, by simply increasing the number of connections of each pixel during spatial propagation from one to three. As shown in Figures 2(a) and 2(b), the propagation of one-way connections is restricted to a single row, while the three-way connections can expand the region to a triangular 2D plane with respect to each direction. The summarization of the four directions result in dense connections of all pixels to each other (see Figure. 2(b)).
Model stability is of critical importance for designing linear systems. In the context of spatial propagation (Eq. 1), it refers to restricting the responses or errors that flow in the module from going to infinity, and preventing the network from encountering the vanishing of gradients in the backpropagation process . Specifically, the norm of the temporal Jacobian should be equal to or less than one. In our case, it is equivalent to regularizing each transformation matrix with its norm satisfying
where denotes the largest singularity value of . This condition, provides a sufficient condition for stability. In the supplementary material, we show the requirements with respect to the elements in .
See supplementary material. Let be the eigenvalue of matrix and be the largest one. According to Gershgorin´s Theorem , where every eigenvalue of a square matrix satisfies:
then . According to the triangle inequality, and since , we have
which satisfies the model stability condition. ∎
Theorem 3 shows that the stability of a linear propagation model can be maintained by regularizing the all weights of each pixel in the hidden layer , with the summation of their absolute values less than one. For the one-way connection, Chen et al. limited each scalar output to be within . Liu et al. extended the range to , where the negative weights showed preferable effects for learning image enhancers. It indicates that the affinity matrix is not necessarily restricted to be positive/semi-positive definite (e.g., the setting is also applied in .) For the three-way connection, we simply regularize the three weights (the output of a deep CNN) according to Theorem 3 without restriction to be any positive/semi-positive definite.
Implementation
We describe the implementation of the three-way connection-based network. We specify two separate branches: (a) a deep CNN, namely the guidance network that outputs all elements of the transformation matrix, and (b) a linear propagation module that outputs the propagation result (see Figure 2). The structure of a guidance network can be any regular deep CNN, which is designed for the task at hand. Examples of this network are described in Section 5. The propagation module receives an input map and output a refined or transformed result. It also takes the weights learned by the deep CNN guidance network as the second input.
The guidance network takes, as input, any 2D matrix that can help with learning the affinity matrix (e.g., typically an RGB image). It outputs all the weights that constitute the transformation matrix . The linear propagation module takes, as inputs, a 2D map that needs to be propagated (e.g., a coarse segmentation mask), and the weights generated by the guidance network. Suppose that we have a map of size that is input into the propagation module, the guidance network needs to output a weight map with the dimensions of , i.e., each pixel in the input map is paired with scalar weights per direction, and directions in total. The propagation module contains independent hidden layers for the different directions, where each layer combines the input map with its respective weight map using Eq. (10). All submodules are differentiable and jointly trained using stochastic gradient descent (SGD). We use node-wise max-pooling, similarly to , to integrate the hidden layers and to obtain the final propagation result.
Experimental Results
The SPN can be trained jointly with any segmentation model by being inserted on top of the last layer that outputs probability maps, or trained separately as a segmentation refinement model. In this paper we choose the second option. Given a coarse image segmentation mask as the input to the spatial propagation module, we show that the SPN can produce higher-quality masks with significantly refined details at object boundaries (see Figure 2). Many models generate low-resolution segmentation masks with coarse boundary shapes to seek a balance between computational efficiency and semantic accuracy. In specific, producing an original high-resolution segmentation mask usually requires the network to neither reduce the size of the input nor that of the output. In such settings, configuring a network with both sufficient capacity and a global receptive field is usually impractical due to the huge model size. Such problem is often solved in some sacrifice of the output resolution. The majority of work choose to firstly produce an output probability map with smaller resolution, and then refine the result using either post-processing or jointly trained modules . It is a non-trivial task for producing high-quality segmentation results. In this work, we train only one SPN model for each task, and treat it as an universal refinement tool for different public available segmentation models.
We carry out the refinement of segmentation on two tasks: (a) generating high-resolution segmentation results on the HELEN face parsing dataset ; and (b) refining generic object segmentation on top of a pretrained model (e.g., VGG and ResNet based models . For the HELEN dataset, we directly use low-resolution RGB face images to train a baseline parser, which successfully catches the global semantic information. The SPN is then trained on top of the coarse segmentation to generate high-resolution output. For the Pascal VOC dataset, we train the SPN on top of the coarse segmentation results generated by the FCN-8s , and directly generalize it to any other pretrained model. We implement the network with a modified CAFFE . The SPN is parallelized during propagating each row-column to the next one with CUDA. We used SGD optimizer, and set the base learning rate to 0.0001. In general, we train the HELEN and VOC segmentation tasks for about and epochs, respectively. The inference time (we do not use cuDNN) of SPN on HELEN and Pascal VOC is about 7ms and 84ms for an image of resolution, respectively. In comparison, the dense CRF costs about 1s , 3.2s and 4.4s with different versions of publicly available implementations (CPU only). We note that the majority time is spend on the guidance network, which can be accelerated by utilizing various existing network compressing strategies, applying smaller models, or sharing weights with the segmentation model if they are jointly trained. During inference, a single SPN hidden layer takes 1.3ms with the same computational settings.
For the HELEN dataset, we train the SPN with smaller patches cropped from the original high-resolution images, their corresponding coarse segmentation maps produced by our baseline parser, and with the corresponding high-resolution ground-truth segmentation masks for supervision. All coarse segmentation maps are obtained by applying the baseline (for HELEN) or pre-trained (for Pascal VOC) image segmentation CNNs on their standard training splits . Since the baseline HELEN parser produces low-resolution segmentation results, we upsample them using a bi-linear filter to be of the same size as the desired higher output resolution. For the Pascal VOC dataset, we use the original output image segmentation probability maps produced by the pre-trained CNN models as its input. These CNN models contain upsampling layers, that typically upsample the internal feature representations by (e.g., in ) and produce output segmentation masks that are of the same size as that of the input images. We set the SPN as a patch refinement model on top of the coarse map with basic semantic information. We fix the size of our input patches to , use the softmax loss, and use the SGD solver for all the experiments. During training, the patches are sampled from image regions that contain more than one ground-truth segmentation label (e.g., a patch with all pixels labeled as “background” will not be sampled). During testing for the VOC dataset, we restrict the classes in the refined results to be contained within the corresponding coarse input.
We combine the guidance network and the spatial propagation module similarly to . We use two propagation units (e.g., the bottom part in Figure. 2 is one propagation unit) with cascaded connections to achieve better results. Differently, we feed in the integrated hidden map of the first unit to the second unit, instead of cascading each direction separately and integrate them at the end of the second unit. We use two more convolutional layers with channels before and after the propagation units to transfer the input map to an intermediate feature map, to make it compatible with the node-wise max-pooling. In addition, we maintain a smaller size of the propagation layer to make the model more efficient w.r.t computational speed and memory. This is carried out by bi-linearly downsampling/upsampling after the two convolutional layers, so that the hidden maps of propagation module is with a smaller dimension of . Note that to compare the one-way with the three-way connection, we use exactly the same structure except the propagation units. We do not apply any configuration used by or .
The HELEN dataset provides high-resolution photography-style face images ( in total), with high-quality manually labeled facial components including eyes, eyebrows, nose, lips, and jawline, which makes the high-resolution segmentation tasks applicable. All prevIoUs work utilize low-resolution parsing output as their final results for evaluation. Although many achieve preferable performance, their results cannot be directly adopted by high-quality facial image editing applications. We use the setting that splits 100 samples for test following . We still take the hair region as one category, but do not evaluate it for fair comparisons with the state-of-the work . We use similarity transformation according to the results of 5-keypoint detection to align all face images to the center. Keeping the original resolution, we then crop or pad them to the size of .
We first train a baseline CNN with a symmetric downsample/upsample structure. The input image is downsampled from the original version. The downsampling part of the network is equipped with five consecutive conv+relu+max-pooling (with stride of 2) layers. Starting from , each one has double the number of channels, resulting in a feature maps at the bottleneck. In order to use the information at different levels of image resolution, we add skipped-links by summing features maps of the same dimensions from the corresponding upsample and dowsample layers. The upsample part has symmetric configurations, except that the max-pooling is replaced with bilinear upsampling, and the last sub-module has 11 channels for the 11 classes. We apply the multi-objective loss as to improve the accuracy along the boundaries. We note that the symmetric structure is powerful, since the results we obtained for the baseline CNN are comparable (see Table. 1) to that of , who apply a much larger model (38 MB vs. 12 MB) in comparison. We then train a SPN on top of the baseline CNN results, with patches of input RGB image and coarse segmentations masks sampled from the preprocessed high-resolution image. For the guidance network, we use the same structure as that of the baseline segmentation network, except that its upsampling part ends at a resolution of , and its output layer has channels. In addition, we train another face parsing CNN with sized inputs and outputs (CNN-Highres) for better comparison. It has three more sub-modules at each end of the baseline network, where all are configured with channels to process higher resolution images.
We show quantitative and qualitative results in Table. 1 and 3 respectively. We compared the one/three way connection SPNs with the baseline, the CNN-Highres and the most relevant state-of-the-art technique for face parsing . Note that the results of baseline and The original output (also for evaluation) size it . are bi-linearly upsampled to before evaluation. Overall, both SPNs outperform the other techniques with a significant margin of over 6 intersection-over-union (IoU) points, especially for the smaller facial components (e.g., eyes and lips) where with smaller resolution images, the segmentation network performs poorly. We note that the one-way connection-based SPN is quite successful on relatively simple tasks such as the HELEN dataset, but fails for more complex tasks, as revealed by the results of Pascal VOC dataset in the following section.
The PASCAL VOC 2012 segmentation benchmark involves 20 foreground object classes and one background class. The original dataset contains training, validation and testing images, with pixel-level annotations. The performance is mainly measured in terms of pixel IoU averaged across the 21 classes. We train our SPNs on the train split with the coarse segmentation results produced by the FCN-8s model . The model is fine-tuned on a pre-trained VGG-16 network, where different levels of features are upsampled and concatenated to obtain the final, low-resolution segmentation results ( smaller than the original image size). The guidance network of the SPN also fine-tunes the VGG-16 structure from the beginning till the pool5 layer as the downsampling part. Similar to the settings for the HELEN dataset, the upsampling part has a symmetric structure with skipped links until the feature dimensions of . The spatial propagation module has the same configuration as that of the SPN that we employed for the HELEN dataset. The model is applied on the coarse segmentation maps of the validation and test splits generated by any image segmentation algorithm without fine-tuning. We test the refinement SPN on three base models: (a) FCN-8s , (b) the atrous spatial pyramid pooling (ASPP-L) network fine-tuned with VGG-16, denoted as Deeplab VGG, and (c) the ASPP-L: a multi-scale network fine-tuned with ResNet-101 (pre-trained on the COCO dataset), denoted as Deeplab ResNet-101. Among them, (b) and (c) are the two basic models from , which are then refined with dense CRF conditioned on the original image.
Table 2 shows that through the three-way SPN, the accuarcy of segmentation is significantly improved over the coarse segmentation results for all the three baseline models.
It has strong capability of generalization and can successfully refine any coarse maps from different pre-trained models by a large margin. Different with the Helen dataset, the one-way SPN fails to refine the segmentation, which is probably due to its limited capability of learning preferable affinity with a sparse form, especially when the data distribution gets more complex. Table 4 shows that by replacing the dense CRF module with the same refinement model, the performance is boosted by a large margin, without fine-tune. One the test split, the DeepNet ResNet-101 based SPN achieves the mean IoU of , while the dense CRF gets . The three-way SPN produces fine visual results, as shown in the red bounding box of Figure 4. By comparing the probability maps (column 3 versus 7), SPN exhibits fundamental improvement in object details, boundaries, and semantic integrity.
In addition, we show in table 3 that the same refinement model can also be generalize to dilated convolution based networks . It significantly improves the quantitative performance on top of the “Front end” base model, as well as adding a multi-scale refinement module, denoted as “+Context”. Specifically, the SPN improves the base model with much larger margin compared to the context aggregation module (see “+3 way” vs “+Context” in table 3).
Conclusion
We propose spatial propagation networks for learning the affinity matrix for vision tasks. The spatial propagation network is a generic framework that can be applied to numerous tasks, and in this work we demonstrate the effectiveness in semantic ßobject segmentation. Experiments on the HELEN face parsing and PASCAL VOC object semantic segmentation tasks show that the spatial propagation network is general, effective and efficient for generating high-quality segmentation results.