Learning Generative Vision Transformer with Energy-Based Latent Space for Saliency Prediction

1 Introduction . In the field of computer vision , salient object detection [ 64 , 65 , 16 , 17 , 5 , 89 ] ( SOD ) or visual saliency prediction , which aims at highlighting objects more attentive than the surrounding areas in images , has achieved significant performance improvement with the deep convolutional neural network revolution . Given a set of training images along with their saliency annotations , the conventional SOD models seek to learn a deterministic one-to-one mapping from image domain to saliency domain . Two main issues exist in the above conventional deep saliency prediction framework : ( i ) the convolution operation based on sliding window makes the deep saliency prediction model less effective in modeling the global contrast of the image , which is essential for salient object detection [ 7 ] ; ( ii ) the one-to-one deterministic mapping mechanism makes the current framework not only impossible to represent the pixel-wise uncertainty in predicting salient objects , but also hard to handle incomplete data in a weakly supervised scenario [ 89 ] . Besides , given an image , the saliency output of a human is subjective , therefore , a stochastic generative model is more natural than a deterministic model for representing saliency prediction . Although [ 85 ] introduces a conditional variational autoencoder ( CVAE ) [ 56 ] for RGB-D salient object detection , the potential posterior collapse problem [ 23 ] makes the stochastic predictions less effective in generating meaningful uncertainty estimation . Motivated by the above two issues , we propose a novel framework , the generative vision transformer , for salient object detection , where a vision transformer structure [ 40 ] is used as a backbone and latent variables are introduced in designing our generative framework . On the one hand , transformers [ 60 ] have proven to be very effective in long-range dependency modeling , and are capable of modeling 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . various scopes of object context information with the multi-head self-attention module . With such an architecture , we can achieve global context modeling for effective salient object detection . On the other hand , the latent variables account for randomness and uncertainty in modeling the mapping from image domain to saliency domain , and also enable the model to produce stochastic saliency predictions for uncertainty estimation . Therefore , the proposed model is a latent variable transformer . Nowadays , there are two types of generative models that have been widely used , namely the variational autoencoder ( VAE ) [ 31 ] and the generative adversarial network ( GAN ) [ 20 ] , which correspond to two different generative learning strategies to train latent variable models . To train a top-down latent variable generator , the VAE introduces an extra encoder to approximate the intractable posterior distribution of the latent variables , and trains the generator via a perturbation of maximum likelihood ; while the GAN introduces a discriminator to distinguish between generated samples and real data , and trains the generator to fool the discriminator . [ 22 , 69 ] present the third learning strategy , namely alternating back-propagation ( ABP ) , to train the generator with latent variables being directly sampled from the true posterior distribution by using a gradient-based Markov chain Monte Carlo ( MCMC ) [ 38 ] , e.g. , Langevin dynamics [ 43 , 66 , 12 ] . All the three generative models define the prior distribution of the latent variables as a simple non-informative isotropic Gaussian distribution , which is less expressive in capturing meaningful latent representation of the data . In this paper , we investigate generative modeling and learning of the vision transformer . We construct a generative model for salient object detection in the form of a top-down conditional latent variable model . Specifically , we propose a generative vision transformer by adding latent variables into the traditional deterministic vision transformer , and assume the latent variables follow an informative trainable energy-based prior distribution [ 47 , 48 ] . Following [ 72 ] , we parameterize the energy function of the energy-based model ( EBM ) by a deep net . Instead of using variational learning or adversarial learning , we jointly train the parameters of the EBM prior and the transformer network by maximum likelihood estimation ( MLE ) . The MLE algorithm relies on MCMC sampling to evaluate the intractable prior and posterior distributions of the latent variables . Experimental results on RGB and RGB-D salient object detections [ 64 , 68 , 85 , 16 ] show that the generative framework equipped with the EBM prior and the transformer-based non-linear mapping is powerful in representing the conditional distribution of object saliency given an image , leading to more reasonable uncertainty estimation as shown in Figure 1 , where stochastic saliency prediction is provided by a learned model and the visualization of the pixel-wise uncertainty is presented . We summarize our main contributions and novelties as follows : ( i ) we propose a novel top-down generative vision transformer network with an energy-based prior distribution defined on latent space for salient object detection ; ( ii ) we jointly train the vision transformer network and the energy-based prior model by an MCMC-based maximum likelihood estimation , without relying on any extra assisting network for adversarial learning or variational learning ; ( iii ) we achieve new benchmark results for both RGB and RGB-D salient object detections , and obtain meaningful uncertainty maps that are highly consistent with human perception for saliency prediction . 2 Related Work . Salient object detection : The main goal of the existing deep fully-supervised RGB image-based salient object detection models [ 67 , 41 , 68 , 64 , 61 , 87 , 65 ] is to achieve structure preserving saliency prediction , by either sophisticated feature aggregation [ 67 , 61 ] , auxiliary edge detection [ 68 , 52 , 65 ] , or resorting to structure-aware loss functions [ 41 , 64 ] . With extra depth information , RGB-D salient object detection models [ 53 , 85 , 5 , 27 , 93 , 90 , 16 , 28 , 51 , 92 , 59 , 86 ] mainly focus on effective multi-modal modeling . Our paper solves the same problems , i.e. , RGB and RGB-D salient object detection , by developing a new generative transformer-based framework . Vision transformers : The breakthroughs of the Transformer networks [ 60 ] in natural language processing ( NLP ) domain have sparked the interest of the computer vision community in developing vision transformers for different computer vision tasks , such as image classification [ 10 , 40 ] , object detection [ 4 , 63 , 6 , 40 ] , image segmentation [ 96 , 54 , 63 , 40 ] , object tracking [ 80 , 81 ] , pose estimation [ 42 , 58 ] , etc . Among them , DPT [ 54 ] adopts a U-shape structure and uses ViT [ 10 ] as an encoder to perform semantic segmentation and monocular depth estimation . Swin [ 40 ] presents a hierarchical transformer with a shifted windowing scheme to achieve an efficient transformer network with high resolution images as input . Different from the above vision transformers that mainly focus on discriminative modeling and learning , our paper emphasizes generative modeling and learning of the vision transformer by involving latent variables and MCMC inference . Dense prediction with generative models : VAEs have been successfully applied to image segmentation [ 3 , 32 ] . For saliency prediction , [ 34 ] adopts a VAE to model the image background , and separates salient objects from the background through reconstruction residuals . [ 85 , 84 ] design CVAEs to model the subjective nature of saliency . GAN-based methods can be divided into two categories , namely fully-supervised and semi-supervised settings . The former [ 21 , 33 ] uses the discriminator to distinguish model predictions from ground truths , while the latter [ 57 , 26 ] uses the GAN to explore the contribution of unlabeled data . [ 88 ] uses a cooperative learning framework [ 71 , 74 ] for generative saliency prediction . [ 84 ] trains a single top-down generator in the ABP framework for RGB-D saliency prediction . Our model generalizes [ 84 ] by replacing the simple Gaussian prior by a learnable EBM prior and adopting a vision transformer-based generator for salient object prediction . Energy-based models : Recent works have shown strong performance of data space EBMs [ 72 , 44 ] in modeling high-dimensional complex dependencies , such as images [ 97 , 95 , 18 , 11 , 19 ] , videos [ 78 , 79 ] , 3D shapes [ 75 , 76 ] , and point clouds [ 73 ] , and also demonstrated the effectiveness of latent space EBMs [ 47 ] in improving the model expressivity for text [ 48 ] , image [ 47 ] , and trajectory [ 49 ] generation . Our paper also learns a latent space EBM as the prior model but builds the EBM on top of a vision transformer generator for image-conditioned saliency map prediction . 3 Generative Vision Transformer with Energy-Based Latent Space . 3.1 Model . We formulate the supervised saliency prediction problem as a conditional generative learning problem . Let I ∈ Rh×w×3 be an observed RGB image , s ∈ Rh×w×1 be the saliency map , and z ∈ R1×1×d be the d-dimensional vector of latent variables , where h× w d. Consider the following generative model to predict a saliency map s from an image I , s = Tθ ( I , z ) + , z ∼ pα ( z ) , ∼ N ( 0 , σ2 I ) , ( 1 ) where Tθ is the non-linear mapping process from [ z , I ] to s with parameters θ , pα ( z ) is the prior distribution with parameters α , and ∼ N ( 0 , σ2 I ) is the observation residual of saliency with σ being given . Due to the stochasticity of the latent variables z , given an image I , its saliency map is also stochastic . Such a probabilistic model is in accord with the uncertainty of the image saliency . Following [ 47 ] , the prior pα ( z ) is not assumed to be a simple isotropic Gaussian distribution as GAN [ 20 ] , VAE [ 31 , 56 ] or ABP [ 22 ] . Specifically , it is in the form of the energy-based correction or exponential tilting [ 72 ] of an isotropic Gaussian reference distribution p0 ( z ) = N ( 0 , σ2zI ) , i.e. , pα ( z ) = 1 Z ( α ) exp [ −Uα ( z ) ] p0 ( z ) ∝ exp [ −Uα ( z ) − 1 2σ2z ||z||2 ] , ( 2 ) where Eα ( z ) = Uα ( z ) + 12σ2z ||z|| 2 is the energy function that maps the latent variables z to a scalar , and Uα ( z ) is parameterized by a multi-layer perceptron ( MLP ) with trainable parameters α . The standard deviation σz is a hyperparameter . Z ( α ) = ∫ exp [ −Uα ( z ) ] p0 ( z ) dz is the intractable normalizing constant that resolves the requirement for a probability distribution to have a total probability equal to one . pα ( z ) is an informative prior distribution in our model and its parameters α need to be estimated along with the non-linear mapping function Tθ from the training data . The mapping function Tθ is parameterized by a vision transformer [ 40 ] with self-attention mechanism , which encodes the input image I and then decodes it along with the vector of latent variables z to the saliency map s , thus , pθ ( s|I , z ) = N ( Tθ ( I , z ) , σ2 I ) . The resulting generative model is a conditional directed graphical model that combines the EBM prior [ 47 ] and the vision transformer [ 40 ] .

Summary: This paper proposes a method based on generative vision transformer with latent variables  following an informative energy-based prior for salient object detection. Both the vision transformer network and the energy based prior model are jointly trained  via Markov chain Monte Carlo (MCMC)-based maximum likelihood estimation.  Extensive experiments reveal both high performance model and meaningful confidence map scores, as well as estimation of uncertainty maps.