Deeply Learning the Messages in Message Passing Inference

Introduction

Learning deep structured models has attracted considerable research attention recently. One popular approach to deep structured model is formulating conditional random fields (CRFs) using deep Convolutional Neural Networks (CNNs) for the potential functions. This combines the power of CNNs for feature representation learning and of the ability for CRFs to model complex relations. The typical approach for the joint learning of CRFs and CNNs , is to learn the CNN potential functions by optimizing the CRF objective, e.g., maximizing the log-likelihood. The CNN and CRF joint learning has shown impressive performance for semantic image segmentation.

For the joint learning of CNNs and CRFs, stochastic gradient descent (SGD) is typically applied for optimizing the conditional likelihood. This approach requires the marginal inference for calculating the gradient. For loopy graphs, marginal inference is generally expensive even when using approximate solutions. Given that learning the CNN potential functions typically requires a large number of gradient iterations, repeated marginal inference would make the training intractably slow. Applying an approximate training objective is a solution to avoid repeat inference; pseudo-likelihood learning and piecewise learning are examples of this kind of approach. In this work, we advocate a new direction for efficient deep structured model learning.

In conventional CRF approaches, the final prediction is the result of inference based on the learned potentials. However, our ultimate goal is the final prediction (not the potentials themselves), so we propose to directly optimize the inference procedure for the final prediction. Our focus here is on the extensively studied message passing based inference algorithms. As discussed in , we can directly learn message estimators to output the required messages in the inference procedure, rather than learning the potential functions as in conventional CRF learning approaches. With the learned message estimators, we then obtain the final prediction by performing message passing inference.

We explore a new direction for efficient deep structured learning. We propose to directly learn the messages in message passing inference as training deep CNNs in an end-to-end learning fashion. Message learning does not require any inference step for the gradient calculation, which allows efficient training. It can be cast into traditional classification problems.

The network output dimension for message estimation is the same as the number of classes ($K$), while the network output for general CNN potential functions in CRFs is $K^{a}$, which is exponential in the order ($a$) of the potentials (for example, $a=2$ for pairwise potentials, $a=3$ for triple-cliques, etc). Hence CNN based message learning has significantly fewer network parameters and thus is more scalable, especially in the cases of a large number of classes involved.

The number of iterations in message passing inference can be explicitly taken into consideration in the message learning procedure. In this paper, we are particularly interested in learning messages that are able to offer high-quality CRF prediction results with only one message passing iteration, making the message passing inference very fast.

We apply our method to semantic image segmentation on the PASCAL VOC 2012 dataset and achieve impressive performance.

Combining the strengths of CNNs and CRFs for segmentation has been explored in several recent methods. Some methods resort to a simple combination of CNN classifiers and CRFs without joint learning. DeepLab-CRF in first train fully CNN for pixel classification and applies a dense CRF method as a post-processing step. Later the method in extends DeepLab by jointly learning the dense CRFs and CNNs. RNN-CRF in also performs joint learning of CNNs and the dense CRFs. They implement the mean-field inference as Recurrent Neural Networks which facilitates the end-to-end learning. These methods usually use CNNs for modelling the unary potentials only. The work in trains CNNs to model both the unary and pairwise potentials in order to capture contextual information. Jointly learning CNNs and CRFs has also been explored for other applications like depth estimation . The work in explores joint training of Markov random fields and deep networks for predicting words from noisy images and image classification.

All these above-mentioned methods that combine CNNs and CRFs are based upon conventional CRF approaches. They aim to jointly learn or incorporate pre-trained CNN potential functions, and then perform inference/prediction using the potentials. In contrast, our method here directly learns CNN message estimators for the message passing inference, rather than learning the potentials.

The inference machine proposed in is relevant to our work in that it has discussed the idea of directly learning message estimators instead of learning potential functions for structured prediction. They train traditional logistic regressors with hand-crafted features as message estimators. Motivated by the tremendous success of CNNs, we propose to train deep CNNs based message estimators in an end-to-end learning style without using hand-crafted features. Unlike the approach in which aims to learn variable-to-factor message estimators, our proposed method aims to learn the factor-to-variable message estimators. Thus we are able to naturally formulate the variable marginals, which is the ultimate goal for CRF inference, as the training objective (see Sec. 3.3). The approach in jointly learns CNNs and CRFs for pose estimation, in which they learn the marginal likelihood of body parts but ignore the partition function in the likelihood. Message learning is not discussed in this work, and the exact relation between this pose estimation approach and message learning remains unclear.

Learning CRF with CNN potentials

Before describing our message learning method, we review the CRF-CNN joint learning approach and discuss limitations. An input image is denoted by ${\boldsymbol{x}}\in\mathcal{X}$ and the corresponding labeling mask is denoted by ${\boldsymbol{y}}\in\mathcal{Y}$. The energy function is denoted by $E({\boldsymbol{y}},{\boldsymbol{x}})$, which measures the score of the prediction ${\boldsymbol{y}}$ given the input image ${\boldsymbol{x}}$. We consider the following form of conditional likelihood:

Here $Z$ is the partition function. The CRF model is decomposed by a factor graph over a set of factors $\mathcal{F}$. Generally, the energy function is written as a sum of potential functions (factor functions):

Here $F$ indexes one factor in the factor graph; ${\boldsymbol{y}}_{F}$ denotes the variable nodes which are connected to the factor $F$; $E_{F}$ is the (log-) potential function (factor function). The potential function can be a unary, pairwise, or high-order potential function. The recent method in describes examples of constructing general CNN based unary and pairwise potentials.

Take semantic image segmentation as an example. To predict the pixel labels of a test image, we can find the mode of the joint label distribution by solving the maximum a posteriori (MAP) inference problem: ${\boldsymbol{y}}^{\star}=\operatorname*{argmax\,}_{{\boldsymbol{y}}}{\displaystyle P}({\boldsymbol{y}}|{\boldsymbol{x}}).$ We can also obtain the final prediction by calculating the label marginal distribution of each variable, which requires to solve a marginal inference problem:

Here ${\boldsymbol{y}}\backslash y_{p}$ indicates the output variables ${\boldsymbol{y}}$ excluding $y_{p}$. For a general CRF graph with cycles, the above inference problems is known to be NP-hard, thus approximate inference algorithms are applied. Message passing is a type of widely applied algorithms for approximate inference: loopy belief propagation (BP) , tree-reweighted message passing and mean-field approximation are examples of the message passing methods.

CRF-CNN joint learning is to learn CNN potential functions by optimizing the CRF objective, typically, the negative conditional log-likelihood, which is:

The energy function $E({\boldsymbol{y}},{\boldsymbol{x}})$ is constructed by CNNs, for which all the network parameters are denoted by ${\boldsymbol{\theta}}$. Adding regularization, minimizing negative log-likelihood for CRF learning is:

Here ${\boldsymbol{x}}^{(i)}$, ${\boldsymbol{y}}^{(i)}$ denote the $i$-th training image and its segmentation mask; $N$ is the number of training images; $\lambda$ is the weight decay parameter. We can apply stochastic gradient descent (SGD) to optimize the above problem for learning ${\boldsymbol{\theta}}$. The energy function $E({\boldsymbol{y}},{\boldsymbol{x}};{\boldsymbol{\theta}})$ is constructed from CNNs, and its gradient ${\nabla}_{\boldsymbol{\theta}}E({\boldsymbol{y}},{\boldsymbol{x}};{\boldsymbol{\theta}})$ can be easily computed by applying the chain rule as in conventional CNNs. However, the partition function $Z$ brings difficulties for optimization. Its gradient is:

Direct calculation of the above gradients is computationally infeasible for general CRF graphs. Usually it is necessary to perform approximate marginal inference to calculate the gradients at each SGD iteration . However, repeated marginal inference can be extremely expensive, as discussed in . CNN training usually requires a huge number of SGD iterations (hundreds of thousands, or even millions), hence this inference based learning approach is in general not scalable or even infeasible.

Learning CNN message estimators

In conventional CRF approaches, the potential functions are first learned, and then inference is performed based on the learned potential functions in order to generate the final prediction. In contrast, our approach directly optimizes the inference procedure for final prediction. We propose to learn CNN estimators to directly output the required intermediate values in an inference algorithm.

Here we focus on the message passing based inference algorithm which has been extensively studied and widely applied. In the CRF prediction procedure, the “message” vectors are recursively calculated based on the learned potentials. We propose to construct and learn CNNs to directly estimate these messages in the message passing procedure, rather than learning the potential functions. In particular, we directly learn factor-to-variable message estimators. Our message learning framework is general and can accommodate all message passing based algorithms such as loopy belief propagation (BP) , mean-field approximation and their variants. Here we discuss using loopy BP for calculating variable marginals. As shown by Yedidia et al. , loopy BP has a close relation with Bethe free energy approximation.

Here $\mathcal{F}_{p}$ is a set of factors connected to the variable $p$; $\mathcal{F}_{p}\backslash F$ is the set of factors $\mathcal{F}_{p}$ excluding the factor $F$. For loopy graph, the variable-to-factor message is normalized in each iteration:

The factor-to-variable message is computed as:

Here $\mathcal{N}_{F}$ is a set of variables connected to the factor $F$; $\mathcal{N}_{F}\backslash p$ is the set of variables $\mathcal{N}_{F}$ excluding the variable $p$. Once we get all the factor-to-variable messages of one variable node, we are able to calculate the marginal distribution (beliefs) of that variable:

in which $Z_{p}$ is a normalizer: $Z_{p}=\sum_{y_{p}}\operatorname*{exp\,}[\sum_{F\in\mathcal{F}_{p}}{\boldsymbol{\beta}}_{F\to p}(y_{p})].$

The calculation of factor-to-variable message ${\boldsymbol{\beta}}_{F\to p}$ depends on the variable-to-factor messages ${\boldsymbol{\beta}}_{p\to F}$. Substituting the definition of ${\boldsymbol{\beta}}_{p\to F}$ in (8), ${\boldsymbol{\beta}}_{F\to p}$ can be re-written as:

Here $q$ denotes the variable node which is connected to the node $p$ by the factor $F$ in the factor graph. We refer to the variable node $q$ as a neighboring node of $q$. $\mathcal{N}_{F}\backslash p$ is a set of variables connected to the factor $F$ excluding the node $p$. Clearly, for a pairwise factor which only connects to two variables, the set $\mathcal{N}_{F}\backslash p$ only contains one variable node. The above equations show that the factor-to-variable message ${\boldsymbol{\beta}}_{F\to p}$ depends on the potential $E_{F}$ and ${\boldsymbol{\beta}}_{F^{\prime}\to q}$. Here ${\boldsymbol{\beta}}_{F^{\prime}\to q}$ is the factor-to-variable message which is calculated from a neighboring node $q$ and a factor $F^{\prime}\neq F$.

Conventional CRF learning approaches learn the potential function then follow the above equations to compute the messages for calculating marginals. As discussed in , given that the goal is to estimate the marginals, it is not necessary to exactly follow the above equations, which involve learning potential functions, to calculate messages. We can directly learn message estimators, rather than indirectly learning the potential functions as in conventional methods.

Consider the calculation in (3.1). The message ${\boldsymbol{\beta}}_{F\to p}$ depends on the observation ${\boldsymbol{x}}_{pF}$ and the messages ${\boldsymbol{\beta}}_{F^{\prime}\to q}$. Here ${\boldsymbol{x}}_{pF}$ denotes the observations that correspond to the node $p$ and the factor $F$. We are able to formulate a factor-to-variable message estimator which takes ${\boldsymbol{x}}_{pF}$ and ${\boldsymbol{\beta}}_{F^{\prime}\to q}$ as inputs and outputs the message vector, and we directly learn such estimators. Since one message ${\boldsymbol{\beta}}_{F\to p}$ depends on a number of previous messages ${\boldsymbol{\beta}}_{F^{\prime}\to q}$, we can formulate a sequence of message estimators to model the dependence. Thus the output from a previous message estimator will be the input of the following message estimator.

There are two message passing strategies for loopy BP: synchronous and asynchronous passing. We here focus on the synchronous message passing, for which all messages are computed before passing them to the neighbors. The synchronous passing strategy results in much simpler message dependences than the asynchronous strategy, which simplifies the training procedure. We define one inference iteration as one pass of the graph with the synchronous passing strategy.

We propose to learn CNN based factor-to-variable message estimator. The message estimator models the interaction between neighboring variable nodes. We denote by $M$ a message estimator. The factor-to-variable message is calculated as:

We refer to $\mathbf{d}_{pF}$ as the dependent message feature vector which encodes all dependent messages from the neighboring nodes that are connected to the node $p$ by $F$. Note that the dependent messages are the output of message estimators at the previous inference iteration. In the case of running only one message passing iteration, there is no dependent messages for $M_{F}$, and thus we do not need to incorporate $\mathbf{d}_{pF}$. To have a general exposition, we here describe the case of running arbitrarily many inference iterations.

We can choose any effective strategy to generate the feature vector $\mathbf{d}_{pF}$ from the dependent messages. Here we discuss a simple example. According to (3.1), we define the feature vector $\mathbf{d}_{pF}$ as a $K$-dimensional vector which aggregates all dependent messages. In this case, $\mathbf{d}_{pF}$ is computed as:

With the definition of $\mathbf{d}_{pF}$ in (13) and ${\boldsymbol{\beta}}_{F\to p}$ in (12), it clearly shows that the message estimation requires evaluating a sequence of message estimators. Another example is to concatenate all dependent messages to construct the feature vector $\mathbf{d}_{pF}$.

There are different strategies to formulate the message estimators in different iterations. The simple strategy is using the same message estimator across all inference iteration. In this case the message estimator becomes a recursive function, and thus the CNN based estimator becomes a recurrent neural network (RNN). Another strategy is to formulate different estimator for each inference iteration.

2 Details for message estimator networks

We formulate the estimator $M_{F}$ as a CNN, thus the estimation is the network outputs:

We can consider any possible strategies for implementing ${\boldsymbol{z}}_{pF}$ with CNNs. For example, we here describe a strategy which is analogous to the network design in . We denote by $C^{(1)}$ as a fully convolutional network (FCNN) for convolutional feature generation, and $C^{(2)}$ as a traditional fully connected network for message estimation.

For a general CNN based potential function in conventional CRFs, the potential network is usually required to have a large number of output units (exponential in the order of the potentials). For example, it requires $K^{2}$ ($K$ is the number of classes) outputs for the pairwise potentials . A large number of output units would significantly increase the number of network parameters. It leads to expensive computations and tend to over-fit the training data. In contrast, for learning our CNN message estimator, we only need to formulate $K$ output units for the network. Clearly it is more scalable in the cases of a large number of classes.

3 Training CNN message estimators

Our goal is to estimate the variable marginals in (3), which can be re-written with the estimators:

Here $Z_{p}$ is the normalizer. The ideal variable marginal has the probability of $1$ for the ground truth class and for the remaining classes. Here we consider the cross entropy loss between the ideal marginal and the estimated marginal.

in which $\hat{y}_{p}$ is the ground truth label for the variable node $p$. Given a set of $N$ training images and label masks, the optimization problem for learning the message estimator network is:

The work in propose to learn the variable-to-factor message (${\boldsymbol{\beta}}_{p\to F}$). Unlike their approach, we aim to learn the factor-to-variable message (${\boldsymbol{\beta}}_{F\to p}$), for which we are able to naturally formulate the variable marginals, which is the ultimate goal for prediction, as the training objective. Moreover, for learning ${\boldsymbol{\beta}}_{p\to F}$ in their approach, the message estimator will depend on all neighboring nodes (connected by any factors). Given that variable nodes will have different number of neighboring nodes, they only consider a fixed number of neighboring nodes (e.g., $20$) and concatenate their features to generate a fixed-length feature vector for classification. In our case for learning ${\boldsymbol{\beta}}_{F\to p}$, the message estimator only depends a fixed number of neighboring nodes (connected by one factor), thus we do not have this problem. Most importantly, they learn message estimators by training traditional probabilistic classifiers (e.g., simple logistic regressors) with hand-craft features, and in contrast, we train deep CNNs in an end-to-end learning style without using hand-craft features.

4 Message learning with inference-time budgets

One advantage of message learning is that we are able to explicitly incorporate the expected number of inference iteration into the learning procedure. The number of inference iteration defines the learning sequence of message estimators. This is particular useful if we aim to learn the estimators which are able to have high-quality prediction for only running a few number of inference iterations. In contrast, the conventional potential function learning in CRFs are not able to directly incorporate the expected number of inference iterations.

We are particularly interested in learning message estimators for using only one message passing iteration, for which the inference can be very fast. In this case it might be preferable to have large-range neighborhood connections, for which the large range interaction can be captured by running one inference pass.

Experiments

We evaluate the proposed CNN message learning method for semantic image segmentation. We use the publicly available PASCAL VOC 2012 dataset . There are 20 object categories and one background category in the dataset. Its contains $1464$ images in the training set, $1449$ images in the “val” set and $1456$ images in the test set. Following the common setting in , the training set is augmented to $10582$ images by including the extra annotations provided in for the VOC images. We use intersection-over-union (IoU) score to evaluate the segmentation performance. For the learning and prediction of our method, we only use one message passing iteration.

The recent work in (referred to as ContextDCRF) learns multi-scale fully convolutional CNNs (FCNNs) for unary and pairwise potential functions to capture contextual information. We follow this CRF learning method and replace the potential functions by the proposed message estimators. We consider $2$ types of spatial relations for constructing the pairwise connections of variable nodes. One is the “surrounding” spatial relation, for which one node is connected to its surround nodes. The other one is the “above/below” spatial relation, for which one node is connected to the nodes that lie above. For the pairwise connections, the neighborhood size is defined by a range box. We learn one type of unary message estimator and $3$ types of pairwise message estimators in total. One type of pairwise message estimator is for the “surrounding” spatial relations, and the other two are for the “above/below” spatial relations. We formulate one network for one type of message estimator.

We formulate our message estimators as multi-scale FCNNs, for which we apply a similar network configuration as in . The network $C^{(1)}$ (see Sec.3.2 for details) has $6$ convolution blocks and $C^{(2)}$ has 2 fully connected layers (with $K$ output units). Our networks are initialized using the VGG-16 model . We train all layers using back-propagation.

We first evaluate our method on the VOC 2012 “val” set. We compare with several recent CNN based methods with available results on the “val” set. Results are shown in Table 1. Our method achieves the best performance. As mentioned, ContextDCRF learns CNN based potential functions in CRFs to capture contextual information. ContextDCRF follows a conventional CRF learning and prediction scheme: They first learn potentials and then perform inference based on the learned potentials to output final predictions. The result shows that learning the CNN message estimators is able to achieve similar performance compared to learning CNN potential functions in CRFs. Note that here we only use one message passing iteration for the training and prediction, the inference time cost is almost negligible. Hence our method enables much more efficient inference.

To further improve the performance, we perform simple data augmentation in training. We generate extra $4$ scales ($[0.8,0.9,1.1,1.2]$) of the training images and their flipped images for training. This result is denoted by “ours+” in the result table.

We further evaluate our method on the VOC 2012 test set. We compare with recent state-of-the-art CNN methods with competitive performance. The results are described in Table 3. Since the ground truth labels are not available for the test set, we evaluate our method through the VOC evaluation server. We achieve impressive performance on the test set: $73.4$ IoU score The result link provided by VOC evaluation server: http://host.robots.ox.ac.uk:8080/anonymous/DBD0SI.html, which is the so far best performance compared to the methods that use the same augmented VOC dataset (marked as “VOC extra” in the table). These results validate the effectiveness of direct message learning with CNNs.

We also include the comparison with the methods which are trained on the much larger COCO dataset (around $133$K training images). Our performance is comparable with these methods, while our method uses much less number of training images.

The results for each category is shown in Table 2. We compare with several recent methods which transfer layers from the same VGG-16 model and use the same training data. Our method performs the best for most categories.

Conclusion

We have proposed a new deep message learning framework for structured CRF prediction. Learning deep message estimators for the message passing inference reveals a new direction for learning deep structured model. Learning CNN message estimators is efficient, which does not involve expensive inference steps for gradient calculation. The network output dimension for message estimation is the same as the number of classes, which does not increase with the order of the potentials, and thus CNN message learning has less network parameters and is more scalable in the number of classes compared to conventional potential function learning. Our impressive performance for semantic segmentation demonstrates the effectiveness and usefulness of the proposed deep message learning. Our framework is general and can be readily applied to other structured prediction applications.

References