Convolutional Random Walk Networks for Semantic Image Segmentation

Introduction

Fully convolutional networks (FCNs) were first introduced in where they were shown to yield significant improvements in semantic image segmentation. Adopting the FCN approach, many subsequent methods have achieved even better performance . However, traditional FCN-based methods tend to suffer from several limitations. Large receptive fields in the convolutional layers and the presence of pooling layers lead to low spatial resolution in the deepest FCN layers. As a result, their predicted segments tend to be blobby and lack fine object boundary details. We report in Fig. 1 some examples illustrating typical poor localization of objects in the outputs of FCNs. Recently, Chen at al. addressed this issue by applying a Dense-CRF post-processing step on top of coarse FCN segmentations. However, such approaches often fail to accurately capture semantic relationships between objects and lead to spatially fragmented segmentations (see an example in the last column of Fig. 1).

To address these problems several recent methods integrated CRFs or MRFs directly into the FCN framework . However, these new models typically involve (1) a large number of parameters, (2) complex loss functions requiring specialized model training or (3) recurrent layers, which make training and testing more complicated. We summarize the most prominent of these approaches and their model complexities in Table 1.

We note that we do not claim that using complex loss functions always makes the model overly complex and too difficult to use. If a complex loss is integrated into an FCN framework such that the FCN can still be trained in a standard fashion, and produce better results than using standard losses, then such a model is beneficial. However, in the context of prior segmentation methods , such complex losses often require: 1) modifying the network structure (casting CNN into an RNN) , or 2) using a complicated multi-stage learning scheme, where different layers are optimized during a different training stage . Due to such complex training procedures, which are adapted for specific tasks and datasets, these models can be quite difficult to adapt for new tasks and datasets, which is disadvantageous.

Inspired by random walk methods , in this work we introduce a simple, yet effective alternative to traditional FCNs: a Convolutional Random Walk Network (RWN) that combines the strengths of FCNs and random walk methods. Our model addresses the issues of (1) poor localization around the boundaries suffered by FCNs and (2) spatially disjoint segments produced by dense CRFs. Additionally, unlike recent semantic segmentation approaches , our RWN does so without significantly increasing the complexity of the model.

Our proposed RWN jointly optimizes (1) pixelwise affinity and (2) semantic segmentation learning objectives that are linked via a novel random walk layer, which enforces spatial consistency in the deepest layers of the network. The random walk layer is implemented via matrix multiplication. As a result, RWN seamlessly integrates both affinity and segmentation branches, and can be jointly trained end-to-end via standard back-propagation with minimal modifications to the existing FCN framework. Additionally, our implementation of RWN requires only $131$ additional parameters. Thus, the effective complexity of our model is the same as the complexity of traditional FCNs (see Table 1). We compare our approach to several variants of the DeepLab semantic segmentation system , and show that our proposed RWN consistently produces better performance over these baselines for the tasks of semantic segmentation and scene labeling.

Related Work

The recent introduction of fully convolutional networks (FCNs) has led to remarkable advances in semantic segmentation. However, due to the large receptive fields and many pooling layers, segments predicted by FCNs tend to be blobby and lack fine object boundary details. Recently there have been several attempts to address these problems. These approaches can be divided into several groups.

The work in used FCN predictions as unary potentials in a separate globalization model that refines the segments using similarity cues based on regions or boundaries. One disadvantage of these methods is that the learning of the unary potentials and the training of the globalization model are completely disjoint. As a result, these methods often fail to capture semantic relationships between objects, which produces segmentation results that are spatially disjoint (see the right side of Fig. 1).

To address these issues, several recent methods have proposed to integrate a CRF or a MRF into the network, thus enabling end-to-end training of the joint model. However, the merging of these two models leads to a dramatic increase in complexity and number of parameters. For instance, the method in , requires to cast the original FCN into a Recurrent Neural Network (RNN), which renders the model much bigger in size (see Table 1). A recent method jointly predicts boundaries and segmentations, and then combines them using a recurrent layer, which also requires complex modifications to the existing FCN framework.

The work in proposes to use local convolutional layers, which leads to a significantly larger number of parameters. Similarly, the method in proposes to model unary and pairwise potentials by separate multi-scale branches. This leads to a network with at least twice as many parameters as the traditional FCN and a much more complex multi-stage training procedure.

In addition to the above methods, it is worth mentioning deconvolutional networks , which use deconvolution and unpooling layers to recover fine object details from the coarse FCN predictions. However, in order to effectively recover fine details one must employ nearly as many deconvolutional layers as the number of convolutional layers, which yields a large growth in number of parameters (see Table 1).

Unlike these prior methods, our implementation of RWN needs only $131$ additional parameters over the base FCN. These additional parameters represent only $0.0008\%$ of the total number of parameters in the network. In addition, our RWN uses standard convolution and matrix multiplication. Thus, it does not need to incorporate complicated loss functions or new complex layers . Finally, unlike the methods in that predict and refine the segmentations disjointly, our RWN model jointly optimizes pixel affinity and semantic segmentation in an end-to-end fashion. Our experiments show that this leads to spatially smoother segmentations.

Background

Random Graph Walks. Random walks are one of the most widely known and used methods in graph theory . Most notably, the concept of random walks led to the development of PageRank and Personalized PageRank , which are widely used for many applications. Let $G=(V,E)$ denote an undirected graph with a set of vertices $V$ and a set of edges $E$. Then a random walk in such graph can be characterized by the transition probabilities between its vertices. Let $W$ be a symmetric $n\times n$ affinity matrix, where $n$ denotes the number of nodes in the graph and where $W_{ij}\in$ denotes how similar the nodes $i$ and $j$ are. In the context of a semantic segmentation problem, each pixel in the image can be viewed as a separate node in the graph, where the similarity between two nodes can be evaluated according to some metric (e.g. color or texture similarity etc). Then let $D$ indicate a diagonal $n\times n$ matrix, which stores the degree values for each node: $D_{ii}=\sum_{j=1}^{n}W_{ij}$ for all $j$ except $i=j$. Then, we can express our random walk transition matrix as $A=D^{-1}W$.

Given this setup, we want to model how the information in the graph spreads if we start at a particular node, and perform a random walk in this graph. Let $y_{t}$ be a $n\times 1$ vector denoting a node distribution at time $t$. In the context of the PageRank algorithm, $y_{t}$ may indicate the rank estimates associated with each of the $n$ Web pages at time $t$. Then, according to the random walk theory, we can spread the rank information in the graph by performing a one-step random walk. This process can be expressed as $y_{t+1}=Ay_{t}$, where $y_{t+1}$ denotes a newly obtained rank distribution after one random walk step, the matrix $A$ contains the random walk transition probabilities, and $y_{t}$ is the rank distribution at time step $t$. Thus, we can observe that the information among the nodes can be diffused, by simply multiplying the random walk transition probability matrix $A$, with the rank distribution $y_{t}$ at a particular time $t$. This process can be repeated multiple times, until convergence is reached. For a more detailed survey please see .

Difference from MRF/CRF Approaches. CRFs and MRFs have been widely used in structured prediction problems . Recently, CRFs and MRFs have also been integrated into the fully convolutional network framework for semantic segmentation . We want to stress that while the goals of CRF/MRF and random walk methods are the same (i.e. to globally propagate information in the graph structures), the mechanism to achieve this objective is very different in these two approaches. While MRFs and CRFs typically employ graphs with a fixed grid structure (e.g., one where each node is connected to its four closest neighbors), random walk methods are much more flexible and can implement any arbitrary graph structure via the affinity matrix specification. Thus, since our proposed RWN is based on random walks, it can employ any arbitrary graph structure, which can be beneficial as different problems may require different graph structures.

Additionally, to globally propagate information among the nodes, MRFs and CRFs need to employ approximate inference techniques, because exact inference tends to be intractable in graphs with a grid structure. Integrating such approximate inference techniques into the steps of FCN training and prediction can be challenging and may require lots of domain-specific modifications. In comparison, random walk methods globally propagate information among the nodes via a simple matrix multiplication. Not only is the matrix multiplication efficient and exact, but it is also easy to integrate into the traditional FCN framework for both training and prediction schemes. Additionally, due to the use of standard convolution and matrix multiplication operations, our RWN can be trivially trained via standard back-propagation in an end-to-end fashion.

Convolutional Random Walk Networks

In this work, our goal is to integrate a random walk process into the FCN architecture to encourage coherent semantic segmentation among pixels that are similar to each other. Such a process introduces an explicit grouping mechanism, which should be beneficial in addressing the issues of (1) poor localization around the boundaries, and (2) spatially fragmented segmentations.

A schematic illustration of our proposed RWN architecture is presented in Fig. 2. Our RWN is a network composed of two branches: (1) one branch that predicts semantic segmentation potentials, and (2) another branch devoted to predicting pixel-level affinities. These two branches are merged via a novel random walk layer that encourages spatially coherent semantic segmentation. The entire RWN can be jointly optimized end-to-end. We now describe each of the components of the RWN architecture in more detail.

For the semantic segmentation branch, we present results for several variants of the DeepLab segmentation systems, including DeepLab-LargeFOV , DeepLab-attention , and DeepLab-v2, which is one of the top performing segmentation systems. DeepLab-largeFOV is a fully convolutional adaptation of the VGG architecture, which contains $16$ convolutional layers. DeepLab-attention , is a multi-scale VGG based network, for which each multi-scale branch focuses on a specific part of the image. Finally, DeepLab-v2 is a multi-scale network based on the residual network implementation. We note that even though we use a DeepLab architecture in our experiments, other architectures such as and many others could be integrated into our framework.

2 Pixel-Level Affinity Branch

To learn the pairwise pixel-level affinities, we employ a separate affinity learning branch with its own learning objective (See Fig. 2). The affinity branch is connected with the input $n\times n\times 3$ RGB image, and low-level $conv1\_1$ and $conv1\_2$ layers. The feature maps corresponding to these layers are $n\times n$ in width and height but they have a different number of channels ($3,64,$ and $64$ respectively). Let $k$ be the total number of affinity learning parameters (in our case $k=3+64+64=131$). Then, let $F$ be a sparse $n^{2}\times n^{2}\times k$ matrix that stores $L1$ distances between each pixel and all of its neighbors within a radius $R$, according to each channel. Note that the distances are not summed up across the $k$ channels, but instead they are computed and stored separately for each channel. The resulting matrix $F$ is then used as an input to the affinity branch, as shown in Figure 2.

The affinity branch consists of a $1\times 1\times k$ convolutional layer and an exponential layer. The output of the exponential layer is then attached to the Euclidean loss layer and is optimized to predict the ground truth pixel affinities, which are obtained from the original semantic segmentation annotations. Specifically, we set the ground truth affinity between two pixels to $1$ if the pixels share the same semantic label and have distance less than $R$ from each other. Note that $F$, which is used as an input to the affinity branch, is a sparse matrix, as only a small fraction of all the entries in $F$ are populated with non-zero values. The rest of the entries are ignored during the computation.

Also note that we only use features from RGB, $conv1\_1$ and $conv1\_2$ layers, because they are not affected by pooling, and thus, preserve the original spatial resolution. We also experimented with using features from deeper FCN layers such as fc6, and fc7. However, we observed that features from deeper layers are highly correlated to the predicted semantic segmentation unary potentials, which causes redundancy and little improvement in the segmentation performance. We also experimented with using more than one convolutional layer in the affinity learning branch, but observed that additional layers provide negligible improvements in accuracy.

3 Random Walk Layer

To integrate the semantic segmentation potentials and our learned pixel-level affinities, we introduce a novel random walk layer, which propagates the semantic segmentation information based on the learned affinities. The random walk layer is connected to the two bottom layers: (1) the fc8 layer containing the semantic segmentation potentials, and (2) the affinity layer that outputs a sparse $n^{2}\times n^{2}$ random walk transition matrix $A$. Then, let $f$ denote the activation values from the fc8 layer, reshaped to the dimensions of $n^{2}\times m$, where $n^{2}$ refers to the number of pixels, and $m$ is the number of object classes in the dataset. A single random walk layer implements one step of the random walk process, which can be performed as $\hat{y}=Af$, where $\hat{y}$ indicates the diffused segmentation predictions, and $A$ denotes the random walk transition matrix.

The random walk layer is then attached to the softmax loss layer, and is optimized to predict ground truth semantic segmentations. One of the advantages of our proposed random walk layer is that it is implemented as a matrix multiplication, which makes it possible to back-propagate the gradients to both (1) the affinity branch and (2) the segmentation branch. Let the softmax-loss gradient be an $n^{2}\times m$ matrix $\frac{\partial{L}}{\partial{\hat{y}}}$, where $n^{2}$ is the number of pixels in the fc8 layer, and $m$ is the number of predicted object classes. Then the gradients, which are back-propagated to the semantic segmentation branch are computed as $\frac{\partial{L}}{\partial{f}}=A^{T}\frac{\partial{L}}{\partial{\hat{y}}}$, where $A^{T}$ is the transposed random walk transition matrix. Also, the gradients, that are back-propagated to the affinity branch are computed as $\frac{\partial{L}}{\partial{A}}=\frac{\partial{L}}{\partial{\hat{y}}}f^{T}$, where $f^{T}$ is a $m\times n^{2}$ matrix that contains transposed activation values from the fc8 layer of the segmentation branch. We note that $\frac{\partial{L}}{\partial{A}}$ is a sparse $n^{2}\times n^{2}$ matrix, which means that the above matrix multiplication only considers the pixel pairs that correspond to the non-zero entries in the random walk transition matrix $A$.

4 Random Walk Prediction at Testing

In the previous subsection, we mentioned that the prediction in the random walk layer can be done via a simple matrix multiplication operation $\hat{y}=Af$, where $A$ denotes the random walk transition matrix, and $f$ depicts the activation values from the fc8 layer. Typically, we want to apply multiple random walk steps until convergence is reached. However, we also do not want to deviate too much from our initial segmentation predictions, in case the random walk transition matrix is not perfectly accurate, which is a reasonable expectation. Thus, our prediction scheme needs to balance two effects: (1) propagating the segmentation information across the nodes using a random walk transition matrix, and (2) not deviating too much from the initial segmentation.

This tradeoff between the two quantities is very similar to the idea behind MRF and CRF models, which try to minimize an energy formed by unary and pairwise terms. However, as discussed earlier, MRF and CRF methods tend to use 1) grid structure graphs and 2) various approximate inference techniques to propagate segmentation information globally. In comparison, our random walk approach is advantageous because it can use 1) any arbitrary graph structure and 2) an exact matrix multiplication operation to achieve the same goal.

Let us first denote our segmentation prediction after $t+1$ random walk steps as $\hat{y}^{t+1}$. Then our general prediction scheme can be written as:

where $\alpha$ denotes a parameter $\mathopen{[}0,1\mathclose{]}$ that controls the tradeoff between (1) diffusing segmentation information along the connections of a random walk transition matrix and (2) not deviating too much from initial segmentation values (i.e. the outputs of the last layer of the FCN). Let us now initialize $\hat{y}^{0}$ to contain the output values from the fc8 layer, which we denoted with $f$. Then we can write our prediction equation by substituting the recurrent expressions:

Now, because we want to apply our random walk procedure until convergence we set $t=\infty$. Then, because our random walk transition matrix is stochastic we know that $\lim_{t\to\infty}(\alpha A)^{t+1}=0$. Furthermore, we can write $S_{t}=\sum_{i=0}^{t}A^{i}=I+A+A^{2}+\ldots+A^{t}$, where $I$ is an identity matrix, and where $S_{t}$ denotes a partial sum of random walk transitions until iteration $t$. We can then write $S_{t}-AS_{t}=I-A^{t+1}$, which implies:

From our previous derivation, we already know that $\lim_{t\to\infty}(A)^{t+1}=0$, which implies that

Thus, our final prediction equation, which corresponds to applying repeated random walk steps until convergence, can be written as

In practice, the random walk transition matrix $A$ is pretty large, and inverting it is impractical. To deal with this problem, we shrink the matrix $(I-\alpha A)$ using a simple and efficient technique presented in , and then invert it to compute the final segmentation. In the experimental section, we show that such a prediction scheme produces solid results and is still pretty efficient ($\approx$ 1 second per image). We also note that we use this prediction scheme only during testing. During training we employ a scheme that uses a single random walk step (but with a much larger radius), which is faster. We explain this procedure in the next subsection.

5 Implementation Details

We jointly train our RWN in an end-to-end fashion for $2000$ iterations, with a learning rate of $10^{-5}$, $0.9$ momentum, the weight decay of $5\cdot 10^{-5}$, and $15$ samples per batch. For the RWN model, we set the tradeoff parameter $\alpha$ to $0.01$. During testing we set the random walk connectivity radius $R=5$ and apply the random walk procedure until convergence. However, during training we set $R=40$, and apply a single random walk step. This training strategy works well because increasing the radius size eliminates the need for multiple random walk steps, which speeds up the training. However, using $R=5$ and applying an infinite number of random walk steps until convergence still yields slightly better results (see study in 5.4), so we use it during testing. For all of our experiments, we use a Caffe library . During training, we also employ data augmentation techniques such as cropping, and mirroring.

Experimental Results

In this section, we present our results for semantic segmentation on the SBD dataset, which contains objects and their per-pixel labels for $20$ Pascal VOC classes (excluding the background class). We also include scene labeling results on the commonly used Stanford background and Sift Flow datasets. We evaluate our segmentation results on these tasks using the standard metric of the intersection-over-union (IOU) averaged per pixels across all the classes from each dataset. We also include the class-agnostic overall pixel intersection-over-union score, which measures the per-pixel IOU across all classes.

We experiment with several variants of the DeepLab system as our main baselines throughout our experiments: DeepLab-LargeFOV , DeepLab-attention , and DeepLab-v2.

Our evaluations provide evidence for four conclusions:

In subsections 5.1, 5.2, we demonstrate that our proposed RWN outperforms DeepLab baselines for both semantic segmentation, and scene labeling tasks.

In subsection 5.1, we demonstrate that, compared to the dense CRF approaches, RWN predicts segmentations that are spatially smoother.

In Subsection 5.3, we show that our approach is more efficient than the denseCRF inference.

Finally, in Subsection, 5.4, we demonstrate that our random walk layer is beneficial and that our model is flexible to use different graph structures.

Standard Evaluation. In Table 2, we present semantic segmentation results on the Pascal SBD dataset , which contains $8055$ training and $2857$ testing images. These results indicate that RWN consistently outperforms all three of the DeepLab baselines. In Figure 3, we also compare qualitative segmentation results of a DeepLab-v2 network and our RWN model. We note that the RWN segmentations contain fewer false positive predictions and are also spatially smoother across the object regions.

Furthermore, in Table 3, we present experiments where we compare RWN with models using dense CRFs to post-process the predictions of DeepLab systems. We also include DeepLab-DT , which uses domain-transfer filtering to refine the segmentations inside an FCN. Based on these results, we observe that, despite not using any post-processing, our RWN produces results similar to or even better than the DeepLab models employing post-processing. These results indicate that RWN can be used as a globalization mechanism to ensure spatial coherence in semantic segmentation predictions. In Figure 4 we present qualitative results where we compare the final segmentation predictions of RWN and the DeepLab-v2-CRF system. Based on these qualitative results, we observe that RWN captures more accurately the fine details of the objects, such as the bike wheels, or plane wings. The DeepLab-v2-CRF system misses some of these object parts.

Localization Around the Boundaries. Earlier we claimed that due to the use of large receptive fields and many pooling layers, FCNs tend to produce blobby segmentations that lack fine object boundary details. We want to show that our RWN produces more accurate segmentations around object boundaries the traditional FCNs. Thus, adopting the practice from , we evaluate segmentation accuracy around object boundaries. We do so by counting the relative number of misclassified pixels within a narrow band (“trimap”) surrounding the ground truth object boundaries. We present these results in Figure 5. The results show that RWN achieves higher segmentation accuracy than the DeepLab (DL) system for all trimap widths considered in this test.

Spatial Smoothness. We also argued that applying the dense CRF as a post-processing technique often leads to spatially fragmented segmentations (see the right side of Fig. 1). How can we evaluate whether a given method produces spatially smooth or spatially fragmented segmentations? Intuitively, spatially fragmented segmentations will produce many false boundaries that do not correspond to actual object boundaries. Thus, to test the spatial smoothness of a given segmentation, we extract the boundaries from the segmentation and then compare these boundaries against ground truth object boundaries using the standard maximum F-score (MF) and average precision (AP) metrics, as done in the popular BSDS benchmark . We perform this experiment on the Pascal SBD dataset and present these results in Table 4. We can see that the boundaries extracted from the RWN segmentations yield better MF and AP results compared to the boundaries extracted from the different variants of the DeepLab-CRF system. Thus, these results suggest that RWN produces spatially smoother segmentations than DeepLab-CRF.

2 Scene Labeling

We also tested our RWN on the task of scene labeling using two popular datasets: Stanford Background and Sift Flow . Stanford Background is a relatively small dataset for scene labeling. It contains $715$ images, which we randomly split into $600$ training images and $115$ testing images. In contrast, the Sift Flow dataset contains $2489$ training examples and $201$ testing images. For all of our experiments, we use the DeepLab-largeFOV architecture since it is smaller and more efficient to train and test. To evaluate scene labeling results, we use the overall IOU evaluation metric which is a commonly used metric for this task. In Table 5, we present our scene labeling results on both of these datasets. Our results indicate that our RWN method outperforms the DeepLab baseline by $2.57\%$, and $2.54\%$ on these two datasets, respectively.

3 Runtime Comparisons

We also include the runtime comparison of our RWN approach versus the denseCRF inference. We note that using a single core of a $2.7$GHz Intel Core i7 processor, the denseCRF inference requires 3.301 seconds per image on average on a Pascal SBD dataset. In comparison, a single iteration of a random walk, which is simply a sparse matrix multiplication, takes 0.032 seconds on average on the same Pascal SBD dataset. A DeepLab_v2 post-processed with denseCRF achieves $81.9\%$ IOU score on this same Pascal SBD dataset. In comparison, RWN_v2 with a single random walk iteration and with R=40 (radius) achieves $82.2\%$ IOU, which is both more accurate and more than $100$ times more efficient than the denseCRF inference.

4 Ablation Experiments

Optimal Number of Random Walk Steps. In Figure 7, we illustrate how the IOU accuracy changes when we use a different number of random walk steps. We observe that the segmentation accuracy keeps increasing as we apply more random walk steps, and that it reaches its peak performance when the random walk process converges, which indicates the effectiveness of our random walk step procedure. In Figure 6, we also illustrate how the predicted object segmentation probabilities change as we apply more random walk steps. We observe that the object boundaries become much better localized as more iterations of random walk are applied.

Radius Size. To analyze the effect of a radius size in the RWN architecture, we test alternative versions of our model with different radii sizes. Our results indicate, that the RWN model produces similar results with different radii in the interval of $R>3$ and $R<20$ if the random walk step process is applied until convergence. We also note that, if we select $R=40$, and apply a random walk step only once, we can achieve the segmentation accuracy of $75.5\%$ and $77.6\%$ according to the two evaluation metrics, respectively. In comparison, choosing $R=5$ and applying random walk until convergence yields the accuracies of $75.8\%$ and $77.9\%$, which is slightly better. However, note that selecting $R=40$, and applying multiple random walk steps does not yield any improvement in segmentation accuracy. These experiments show the flexibility of our model compared to the MRF or CRF models, which typically use graphs with a fixed grid structure. Our model has the ability to use different graph structures depending on the problem.

Conclusion

In this work, we introduced Random Walk Networks (RWNs), and showed that, compared to traditional fully convolutional networks (FCNs), they produce improved accuracy for the same model complexity. Our RWN addresses the issues of 1) poor localization around the segmentation boundaries and 2) spatially disjoint segmentations. Additionally, our implementation of RWN uses only $131$ additional learnable parameters ($0.0008\%$ of the original number of the parameters in the network) and it can be easily integrated into the standard FCN learning framework for a joint end-to-end training. Finally, RWN provides a more efficient alternative to the denseCRF approaches.

Our future work includes experimenting with alternative RWN architectures and applying RWN to new domains such as language processing or speech recognition.

References