Energy-Based Generative Cooperative Saliency Prediction

As a class-agnostic segmentation task, salient object detection has attracted a lot of attentions in the computer vision community for its close relationship to human visual perception. A salient region is a visually distinctive scene region that can be located rapidly and with little human effort. Salient object detection is commonly treated as a pixel-wise binary output of a deterministic prediction model in most recent works  (Wu, Su, and Huang 2019a; Qin et al. 2019; Wu, Su, and Huang 2019b; Wei, Wang, and Huang 2020; Wang et al. 2019; Xu et al. 2021). Despite the success of those recent models, the one-to-one deterministic mapping has prevented them from modeling the uncertainty of human salient object prediction, which is considered to be subjective (Itti, Koch, and Niebur 1998) and affected by biological factors (e.g., contrast sensitivity), contextual factors (e.g., task, experience, and interest), etc. In this way, it is more reasonable to represent visual saliency as a conditional probability distribution over a saliency map given an input image, and formulate the saliency prediction as a stochastic sampling process from the conditional distribution.

Generative models (Goodfellow et al. 2014; Kingma and Welling 2013; Xie et al. 2016, 2018a) have demonstrated their abilities to represent conditional distributions of high-dimensional data and produce multiple plausible outputs given the same input (Zhu et al. 2017; Xie et al. 2021b). In this work, we fit the saliency detection task into a generative framework, where the input image is the condition, and the goal is to generate multiple saliency maps, representing the “subjective nature” of human visual saliency. Zhang et al. (2020a, 2021) have used conditional variational auto-encoders (VAEs) (Kingma and Welling 2013; Sohn, Lee, and Yan 2015a), which are latent variable models (LVMs), to implicitly represent distributions of visual saliency. However, VAEs only learn a stochastic mapping from image domain to saliency domain, and lack an intrinsic cost function to evaluate the visual saliency output and guide the saliency prediction process. As to a prediction task, a cost function of solution is more reliable than a mapping function because the former is more generalizable than the latter.

In contrast, we propose to model the conditional distribution of visual saliency explicitly via an energy-based model (EBM) (Xie et al. 2016; Nijkamp et al. 2019), where the energy function defined on both image and saliency domains serves as a cost function of the saliency prediction. Given an input image, the saliency prediction can be achieved by performing sampling from the EBM via Markov chain Monte Carlo (MCMC) (Neal 2012) method, which is a gradient-based algorithm that searches local minima of the cost function of the EBM conditioned on the input image.

A typical high-dimensional EBM learns an energy function by MCMC-based maximum likelihood estimation (MLE), which commonly suffers from convergence difficulty and computational expensiveness of the MCMC process. Inspired by prior success of energy-based generative cooperative learning  (Xie et al. 2018a, 2021b), we propose the energy-based generative cooperative saliency prediction framework to tackle the saliency prediction task. Specifically, the framework consists of a conditional EBM whose energy function is parameterized by a bottom-top neural network, and a conditional LVM whose transformation function is parameterized by an encoder-decoder framework. The framework brings in an LVM as an ancestral sampler to initialize the MCMC computational process of the EBM for efficient sampling, so that the EBM can be learned efficiently. The EBM, in turn, refines the LVM’s generated samples via MCMC and feeds them back to the LVM, so that the LVM can learn its mapping function from the MCMC transition. Thus, the resulting cooperative saliency prediction process first generates an initial saliency map via a direct mapping and then refines the saliency map via an iterative process. This is a coarse-to-fine generative saliency detection, and corresponds to a fast-thinking and slow-thinking system (Xie et al. 2021b).

Moreover, based on the generative cooperative saliency prediction framework, we further propose a cooperative learning while cooperative recovering strategy for weakly supervised saliency learning, where each training image is associated with a partially observed annotation (e.g., scribble annotation (Zhang et al. 2020b)). At each learning iteration, the strategy has two sub-tasks: cooperative recovery and cooperative learning. As to the cooperative recovery sub-task, each incomplete saliency ground truth is firstly recovered in the low-dimensional latent space of the LVM via inference, and then refined by being pushed to the local mode of the cost landscape of the EBM via MCMC. For the cooperative learning sub-task, the recovered saliency maps are treated as pseudo labels to update the parameters of the framework as in the scenario of learning from complete data.

In experiments, we demonstrate that our framework can not only achieve state-of-the-art performances in both fully supervised and weakly supervised saliency predictions, but also generate diverse saliency maps from one input image, indicating the success of modeling the uncertainty of saliency prediction. Furthermore, we show that the learned energy function in the EBM can serve as a cost function, which is useful to refine the results from other pre-trained saliency prediction models.

Our contributions can be summarized as below:

• •

  We study generative modeling of saliency prediction, and formulate it as a sampling process from a probabilistic model using EBM and LVM respectively, which are new angles to model and solve saliency prediction.

• •

  We propose a generative cooperative saliency prediction framework, which jointly trains the LVM predictor and the EBM predictor in a cooperative learning scheme to offer reliable and efficient saliency prediction.

• •

  We generalize our generative framework to the weakly supervised saliency prediction scenario, in which only incomplete annotations are provided, by proposing the cooperative learning while recovering algorithm, where we train the model and simultaneously recover the unlabeled areas of the incomplete saliency maps.

• •

  We provide strong empirical results in both fully supevised and weakly supervised settings to verify the effectiveness of our framework for saliency prediction.

We first briefly introduce existing fully supervised and weakly supervised saliency prediction models. We then review the family of generative cooperative models and other conditional deep generative frameworks.

Fully Supervised Saliency Prediction. Existing fully supervised saliency prediction models (Wang et al. 2018; Liu, Han, and Yang 2018; Wei, Wang, and Huang 2020; Liu et al. 2019; Qin et al. 2019; Wu, Su, and Huang 2019b, a; Wang et al. 2019, 2019; Wei et al. 2020; Xu et al. 2021) mainly focus on exploring image context information and generating structure-preserving predictions.Wu, Su, and Huang (2019b); Wang et al. (2019, 2019, 2018); Liu, Han, and Yang (2018); Liu et al. (2019); Wu, Su, and Huang (2019a); Xu et al. (2021) propose saliency prediction models by effectively integrating higher-level and lower-level features. Wei, Wang, and Huang (2020); Wei et al. (2020) propose an edge-aware loss term to penalize errors along object boundaries.Zhang et al. (2020a) present a stochastic RGB-D saliency detection network based on the conditional variational auto-encoder (Kingma and Welling 2013; Jimenez Rezende, Mohamed, and Wierstra 2014). In this paper, we introduce the conditional cooperative learning framework (Xie et al. 2018a, 2021b) to achieve probabilistic coarse-to-fine RGB saliency detection, where a coarse prediction is produced by a conditional latent variable model and then is refined by a conditional energy-based model. Our paper is the first work to use a deep energy-based generative framework for probabilistic saliency detection.

Weakly Supervised Saliency Prediction. Weakly supervised saliency prediction frameworks (Wang et al. 2017; Li, Xie, and Lin 2018a; Nguyen et al. 2019; Zhang et al. 2020b) attempt to learn predictive models from easy-to-obtain weak labels, including image-level labels (Wang et al. 2017; Li, Xie, and Lin 2018b), noisy labels (Nguyen et al. 2019; Zhang et al. 2018; Zhang, Han, and Zhang 2017) or scribble labels (Zhang et al. 2020b). In this paper, we also propose a cooperative learning while recovering strategy for weakly supervised saliency prediction, in which only scribble labels are provided and our model treats them as incomplete data and recovers them during learning.

Energy-Based Generative Cooperative Networks. Deep energy-based generative models (Xie et al. 2016), with energy functions parameterized by modern convolutional neural networks, are capable of modeling the probability density of high-dimensional data. They have been applied to image generation (Xie et al. 2016; Gao et al. 2018; Nijkamp et al. 2019; Du and Mordatch 2019; Grathwohl et al. 2020; Zhao, Xie, and Li 2021; Zheng, Xie, and Li 2021), video generation (Xie, Zhu, and Wu 2019), 3D volumetric shape generation (Xie et al. 2018b, 2020), and unordered point cloud generation (Xie et al. 2021a). The maximum likelihood learning of the energy-based model typically requires iterative MCMC sampling, which is computationally challenging. To relieve the computational burden of MCMC, the Generative Cooperative Networks (CoopNets) in Xie et al. (2018a) propose to learn a separate latent variable model (i.e. a generator) to serve as an efficient approximate sampler for training the energy-based model. Xie, Zheng, and Li (2021) propose a variant of CoopNets by replacing the generator with a variational auto-encoder (VAE) (Kingma and Welling 2014). Xie et al. (2021b) propose a conditional CoopNets for supervised image-to-image translation. Our paper proposes a conditional CoopNets for visual saliency prediction. Further, we generalize our model to the weakly supervised learning scenario by proposing a cooperative learning while recovering algorithm. In this way, we can learn from incomplete data for weakly supervised saliency prediction.

Conditional Deep Generative Models. Our framework belongs to the family of conditional generative models, which include conditional generative adversarial networks (CGANs) (Mirza and Osindero 2014) and conditional variational auto-encoders (CVAEs) (Sohn, Lee, and Yan 2015a). Different from existing CGANs (Luc et al. 2016; Zhang et al. 2018; Xue et al. 2017; Pan et al. 2017; Yu and Cai 2018; Hung et al. 2018; Souly, Spampinato, and Shah 2017), which train a conditional discriminator and a conditional generator in an adversarial manner, or CVAEs (Kohl et al. 2018; Zhang et al. 2020a, 2021), in which a conditional generator is trained with an approximate inference network, our model learns a conditional generator with a conditional energy-based model via MCMC teaching. Specifically, our model allows an additional refinement for the generator during prediction, which is lacking in both CGANs and CVAEs.

We will first present two types of generative modeling of saliecny prediction, i.e., the energy-based model (EBM) and the latent variable model (LVM). Then, we propose a novel generative saliency prediction framework, in which the EBM and the LVM are jointly trained in a generative cooperative manner, such that they can help each other for better saliency prediction in terms of computational efficiency and prediction accuracy. The latter aims to generate a coarse but fast prediction, and the former serves as a fine saliency predictor. The resulting model is a coarse-to-fine saliency prediction framework.

In Section 3, the framework is trained from fully-observed training data. In this section, we want to show that our generative framework can be modified to handle the scenario in which each image $X_{i}$ only has a partial pixel-wise annotation $Y^{\prime}_{i}$, e.g., scribble annotation (Zhang et al. 2020b). Since the saliency map for each training image is incomplete, directly applying the algorithm to the incomplete training data can lead to a failure of learning the distribution of saliency given an image. However, generative models are good at data recovery, therefore they can learn to recover the incomplete data. In our framework, we will leverage the recovery powers of both EBM and LVM to deal with the incomplete data in our cooperative learning algorithm, and this will lead to a novel weakly supervised saliency prediction framework.

To learn from incomplete data, our algorithm alternates the cooperative learning step and the cooperative recovery step. Both steps need a cooperation between EBM and LVM. The cooperative learning step is the same as the one used for fully observed data, except that it treats the recovered saliency maps, which are generated from the cooperative recovery step, as training data in each iteration. The following is the cooperative recovery step, which consists of two sub-steps driven by the LVM and the EBM respectively:

(i) Recovery by LVM in Latent Space. Given an image $X_{i}$ and its incomplete saliency map $Y^{\prime}_{i}$, the recovery of the missing part of $Y^{\prime}_{i}$ can be achieved by first inferring the latent vector $h^{\prime}_{i}$ based on the partially observed saliency information via $h^{\prime}_{i}\sim p_{\alpha}(h|Y^{\prime}_{i},X_{i})$, and then generating $\hat{Y}^{\prime}_{i}=G_{\alpha}(X_{i},h^{\prime}_{i})$ with the inferred latent vector $h^{\prime}_{i}$. Let $O_{i}$ be a binary mask, with the same size as $Y^{\prime}$, indicating the locations of visible annotations in $Y^{\prime}_{i}$. $O_{i}$ varies for different $Y^{\prime}_{i}$ and can be extracted from $Y^{\prime}_{i}$. The Langevin dynamics for recovery iterates the same step in Eq. (7) except that $\frac{\partial}{\partial h}\log p_{\alpha}(Y^{\prime},h_{\tau}|X)=\frac{1}{\sigma^{2}}(O\circ(Y-G_{\alpha}(X,h_{\tau})))\frac{\partial}{\partial h}G_{\alpha}(X,h_{\tau})-h_{\tau}$, where $\circ$ denotes element-wise matrix multiplication operation.

(ii) Recovery by EBM in Data Space. With the initial recovered result $\hat{Y}^{\prime}$ generated by the coarse saliency predictor $p_{\alpha}$, the fine saliency predictor $p_{\theta}$ can further refine the result by running a finite-step Langevin dynamics, which is initialized with $\hat{Y}^{\prime}$, to obtain $\tilde{Y}^{\prime}$. The underlying principle is that the initial recovery $\hat{Y}^{\prime}$ might be just around one local mode of the energy function. A few steps of Langevin dynamics (i.e., stochastic gradient descent) toward $p_{\theta}$, starting from $\hat{Y}_{i}^{\prime}$, will push $\hat{Y}_{i}^{\prime}$ to its nearby low energy mode, which might correspond to its complete version $Y_{i}$.

Cooperative Learning and Recovering. At each iteration $t$, we perform the above cooperative recovery of the incomplete saliency maps $\{Y^{\prime}\}_{i=1}^{n}$ via $p_{\theta^{(t)}}$ and $p_{\alpha^{(t)}}$, while learning $p_{\theta^{(t+1)}}$ and $p_{\alpha^{(t+1)}}$ from $\{X_{i},\tilde{Y}_{i}^{\prime(t)}\}_{i=1}^{n}$, where $\{\tilde{Y}_{i}^{\prime(t)}\}_{i=1}^{n}$ are the recovered saliency maps at iteration $t$. The parameters $\theta$ are still updated via Eq. (5) except that we replace $Y_{i}$ by $\tilde{Y}^{\prime}_{i}$. That is, at each iteration, we use the recovered $\tilde{Y}^{\prime}_{i}$, as well as the synthesized $\tilde{Y}_{i}$, to compute the gradient of the log-likelihood, which is denoted by $\Delta\theta(\{\tilde{Y}^{\prime}_{i}\},\{\tilde{Y}_{i}\})$. The algorithm simultaneously performs (i) cooperative recovering of missing annotations of each training example; (ii) cooperative sampling to generate annotations; (iii) cooperative learning of the two models by updating parameters with both recovered annotations and generated annotations. See Algorithm 2 for a detailed description of the learning while recovering algorithm.

We present the details of architecture designs of the LVM and the EBM, as well as the hyper-parameters below.

Latent Variable Model: The LVM $G_{\alpha}(X,h)$, using the ResNet50 (He et al. 2016) as an encoder backbone, maps an image $X$ and a latent vector $h$ to a saliency map $\hat{Y}$. Specifically, we adopt the decoder from the MiDaS (Ranftl et al. 2020) for its simplicity, which gradually aggregates the higher level features with lower level features via residual connections. We introduce the latent vector $h$ to the bottleneck of the LVM by concatenating the tiled $h$ with the highest level features of the encoder backbone, and then feed them to a $3\times 3$ convolutional layer to obtain a feature map with the same size as the original highest level feature map of the encoder. The latent-vector-aware feature map is then fed to the decoder from Ranftl et al. (2020) to generate a final saliency map. As shown in Eq. (8), the parameters of the LVM are updated with the revised predictions $\{\tilde{Y}\}$ provided by the EBM. Thus, immature $\{\tilde{Y}\}$ in the early stage of the cooperative learning might bring in fluctuation in training the LVM, which in turn affects the convergence of the MCMC samples $\{\tilde{Y}\}$. To stabilize the cooperative training, especially in the early stage, we let the LVM learn from not only $\{\tilde{Y}\}$ but also $\{Y\}$. Specifically, we add an extra loss for the LVM as $\lambda\mathcal{L}_{ce}(G_{\alpha}(X,\tilde{h}),Y)$, where $\lambda$ linearly decreases to 0 during training, and $\mathcal{L}_{ce}$ is the cross-entropy loss.

Energy-Based Model: The energy function $U_{\theta}(Y,X)$ is parameterized by a neural network that maps the channel-wise concatenation of $X$ and $Y$ to a scalar. Let cksl-n denote a k$\times$k Convolution-BatchNorm-ReLU layer with n filters and a stride of l. Let fc-n be a fully connected layer with n filters. The $U_{\theta}(Y,X)$ is our framework consists of the following layers: c3s1-32, c4s2-64, c4s2-128, c4s2-256, c4s1-1, fc-100.

Implementation Details: We train our model with a maximum of 30 epochs. Each image is rescaled to $352\times 352$. We set the number of dimensions of the latent space as $d=8$. The number of Langevin steps is $K=5$ and the Langevin step sizes for EBM and LVM are 0.4 and 0.1. The learning rates of the LVM and EBM are initialized to $5\times 10^{-5}$ and $10^{-3}$ respectively. We use Adam optimizer with momentum 0.9 and decrease the learning rates by 10% after every 20 epochs. It takes 20 hours to train the model with a batch size of 7 using a single NVIDIA GeForce RTX 2080Ti GPU.

We conduct a series of experiments to test the performances of the proposed generative cooperative frameworks for saliency prediction. We start from experiment setup.

Datasets: We use the DUTS dataset (Wang et al. 2017) to train the fully supervised model, and S-DUTS (Zhang et al. 2020b) dataset with scribble annotations to train the weakly supervised model. Testing images include (1) DUTS testing dataset, (2) ECSSD (Yan et al. 2013), (3) DUT (Yang et al. 2013), (4) HKU-IS (Li and Yu 2015), (5) PASCAL-S (Li et al. 2014) and (6)SOD dataset (Movahedi and Elder 2010).

Compared methods: We compare our method against state-of-the-art fully supervised saliency detection methods, e.g., PoolNet (Liu et al. 2019), BASNet (Qin et al. 2019), SCRN (Wu, Su, and Huang 2019b), F3Net (Wei, Wang, and Huang 2020), ITSD (Zhou et al. 2020), LDF (Wei et al. 2020), UCNet+ (Zhang et al. 2021) and PAKRN (Xu et al. 2021). UCNet+ (Zhang et al. 2021) is the only generative framework. We also compare our weakly supervised solution with the scribble saliency detection models, e.g., SSAL (Zhang et al. 2020b) and SCWS (Yu et al. 2021).

Evaluation Metrics: We evaluate performance of our models and compared methods with four saliency evaluation metrics, including Mean Absolute Error ($\mathcal{M}$), mean F-measure ($F_{\beta}$), mean E-measure ($E_{\xi}$) (Fan et al. 2018) and S-measure ($S_{\alpha}$) (Fan et al. 2017).

In this paper, we propose a novel energy-based generative saliency prediction framework based on the conditional generative cooperative network, where a conditional latent variable model and an conditional energy-based model are jointly trained in a cooperative learning scheme to achieve a coarse-to-fine saliency prediction. The latent variable model serves as a coarse saliency predictor that provides a fast initial saliency prediction, while the energy-based model serves as a fine saliency predictor that further refines the initial output by the Langevin revision. Even though each of the models can represent the conditional probability distribution of saliency, the cooperative representation and training can offer the best of both worlds. Moreover, we propose a cooperative learning while recovering strategy and apply the model to the weakly supervised saliency detection scenario, in which partial annotations (e.g., scribble annotations) are provided for training. As to the cooperative recovery part of the proposed strategy, the latent variable model serves as a fast but coarse saliency recoverer that provides an initial recovery of the missing annotations from the latent space via inference process, while the energy-based model serves as a slow but fine saliency recoverer that refines the initial recovery results by Langevin dynamics. Combining these two types of recovery schemes leads to a coarse-to-fine recoverer. Further, we find that the learned energy function in the energy-based model can serve as a refinement module, which can be easily plugged into the existing pre-trained saliency prediction models. The energy function is the potential cost function trained from the saliency prediction task. In comparison to a mapping function from image to saliency, the cost function captures the criterion to measure the quality of the saliency given an image, and is more generalizable so that it can be used to refine other saliency predictions. Extensive results exhibit that, compared with both conventional deterministic mapping methods and alternative uncertainty estimation methods, our framework can lead to both accurate saliency predictions for computer vision tasks and reliable uncertainty maps indicating the model confidence in performing saliency prediction from an image. As to a broader impact, the proposed computational framework might also benefit the researchers in the field of computational neuroscience who investigate human attentional mechanisms. The proposed coarse-to-fine saliency prediction model and recovery model may shed light on a clear path toward the understanding of relationship between visual signals and human saliency.