Set Prediction without Imposing Structure as Conditional Density Estimation

1 INTRODUCTION . This paper strives for set prediction . Making multiple predictions with intricate interactions is essential in a variety of applications . Examples include predicting the set of attributes given an image ( Rezatofighi et al. , 2020 ) , detecting all pedestrians in video footage ( Wang et al. , 2018 ) or predicting the future state for a group of molecules ( Noé et al. , 2020 ) . Because of their unordered nature , sets constitute a challenge for both the choice of machine learning model and training objective . Models that violate permutation invariance suffer from lower performance , due to the additional difficulty of needing to learn it . Similarly , loss functions should be indifferent to permutations in both the ground-truth and predictions . Additional ambiguity in the target set exacerbates the problem of defining a suitable set loss . We propose Deep Energy-based Set Prediction ( DESP ) to address the permutation symmetries in both the model and loss function , with a focus on situations where multiple plausible predictions exist . DESP respects the permutation symmetry , by training a permutation invariant energy-based model with a likelihood-based objective . In the literature , assignment-based set distances are applied as loss functions ( Zhang et al. , 2019 ; Kosiorek et al. , 2020 ) . Examples include the Chamfer loss ( Fan et al. , 2017 ) and the Hungarian loss ( Kuhn , 1955 ) . Both compare individual elements in the predicted set to their assigned groundtruth counterpart and vice-versa . While they guarantee permutation invariance , they also introduce a structure over sets , in the form of a metric space . Choosing the wrong set distance can result in implausible predictions , due to interpolations in the set space for underdefined problems . For example , Fan et al . ( 2017 ) observe different set distances to lead to trade-offs between fine-grained shape reconstruction and compactness , for 3d reconstruction from RGB images . As an additional shortcoming , optimizing for a set loss during training poses a limitation on the family of learnable data distributions . More specifically , conditional multi-modal distributions over sets can not be learned by minimizing an assignment-based set loss during training . To overcome the challenges of imposed structure and multi-modal distributions , we propose to view set prediction as a conditional density estimation problem , where P ( Y |x ) denotes the distribution for the target set Y given observed features x . In this work we focus on distributions taking the form of deep energy-based models ( Ngiam et al. , 2011 ; Zhai et al. , 2016 ; Belanger & McCallum , 2016 ) : Pθ ( Y |x ) = 1 Z ( x ; θ ) exp ( −Eθ ( x , Y ) ) , ( 1 ) with Z as the partition function and Eθ the energy function with parameters θ . The expressiveness of neural networks ( Cybenko , 1989 ) allows for learning multi-modal densities Pθ ( Y |x ) . This sets the approach apart from forward-processing models , that either require conditional independence assumptions ( Rezatofighi et al. , 2017 ) , or an order on the predictions , when applying the chain rule ( Vinyals et al. , 2016 ) . Energy-based prediction is regarded as a non-linear combinatorial optimization problem ( LeCun et al. , 2006 ) : Ŷ = argmin Y Eθ ( x , Y ) , ( 2 ) which is typically approximated by gradient descent for deep energy-based models ( Belanger & McCallum , 2016 ; Belanger et al. , 2017 ) . We replace the deterministic gradient descent with a stochastically augmented prediction algorithm , to account for multiple plausible predictions . We show that our stochastic version outperforms standard gradient descent for set prediction tasks . Our main contribution is DESP , a training and prediction framework for set prediction , that removes the limitations imposed by assignment-based set losses . Sampling plays a key role in DESP . For training , sampling approximates the intractable model gradients , while during prediction , sampling introduces stochasticity . We show the generality of our framework by adapting recently proposed permutation invariant neural networks as set prediction deep energy-based models . We demonstrate that our approach ( i ) learns multi-modal distributions over sets ( ii ) makes multiple plausible predictions ( iii ) generalizes over different deep energy-based model architectures and ( iv ) is competitive even in non-stochastic settings , without requiring problem specific loss-engineering . 2 DEEP ENERGY BASED SET PREDICTION . 2.1 TRAINING . Our goal is to train a deep energy based model for set prediction , such that all plausible sets are captured by the model . Regression models with a target in the Rd space , that are trained with a root mean-square error ( RMSE ) loss , implicitly assume a Gaussian distribution over the target . Analog to the RMSE , assignment-based set losses assume a uni-modal distribution over the set space . Training with the negative log-likelihood ( NLL ) circumvents the issues of assignment-based set losses . Notably , NLL does not necessitate explicit element-wise comparisons , but treats the set holistically . We reformulate the NLL for the training data distribution PD as : E ( x , Y ) ∼PD [ − log ( Pθ ( Y |x ) ) ] = E ( x , Y ) ∼PD [ Eθ ( x , Y ) ] + Ex∼PD [ log ( Z ( x ; θ ) ) ] . ( 3 ) The gradient of the left summand is approximated by sampling a mini-batch of n tuples { ( xi , Y +i ) } i=0 .. n from the training set . The gradient of the right summand is approximated by solely sampling input features { xi } i=0 .. m. Directly evaluating ∂∂θ log ( Z ( x ; θ ) ) is intractable ; instead we approximate the gradient by sampling { Y −j } j=0 .. k from the model distribution : ∂ ∂θ log ( Z ( x ; θ ) ) = −EY ∼Pθ [ ∂ ∂θ Eθ ( x , Y ) ] ≈ − k∑ j=0 ∂ ∂θ Eθ ( x , Y − j ) . ( 4 ) The resulting approximate NLL objective is equivalent to contrasting the energy value for real and synthesized targets , with the former being minimized and the latter maximized . The objective is reminiscent of the discriminator ’ s loss in generative adversarial networks ( Goodfellow et al. , 2014 ) , where a real sample is contrasted to a sample synthesized by the generator network . In practice , setting k=1 suffices . The Langevin MCMC algorithm allows for efficient sampling from high dimensional spaces ( Geman & Geman , 1984 ; Neal et al. , 2011 ) . Access to the derivative of the unnormalized density function provides sufficient information for sampling . We apply the following modified transition function and keep only the last sample : Y ( t+1 ) = Y ( t ) − ∂Eθ ( x , Y ( t ) ) ∂Y +U ( t ) , ( 5 ) with U ( t ) ∼ N ( 0 , I ) , > 0 , Y ( 0 ) ∼ N ( 0 , I ) a sample from a fixed initial distribution and Y ( T ) the final sample . The proper formulation of the Langevin MCMC algorithm multiplies the gradient in Equation 5 by a factor and further requires a Metropolis-Hastings acceptance step ( Neal , 1993 ) . We forgo both of these components in favor of increased efficiency , but at the cost of forfeiting theoretical guarantees for desirable properties such as not being trapped in a subset of the sampling space , i.e. , ergodicity . Discarding all but the last sample Y ( T ) of each chain constitutes a non typical usage that undermines the usual importance of ergodicity . Notably , this weakens the hard to meet requirement for the sampler to mix between multiple modes in a single MCMC chain , making it sufficient for independently sampled chains to find different local modes . Although the fixed cutoff at T and missing Metropolis-Hastings update result in a biased sampler , previous works have demonstrated the feasibility of training generative models on images with similar Langevin MCMC methods ( Xie et al. , 2016 ; 2018 ; Nijkamp et al. , 2020 ; Du & Mordatch , 2019 ; Grathwohl et al. , 2019 ) . The model density from Equation 1 approaches the data distribution PD while training , leading to an increased ability in distinguishing between synthesized sets Y − from real sets Y + . This in turn enhances the samples Y − to be closer to the ground-truth , making it harder for the model to discriminate between real and fake . In practice , it is necessary to smooth out the data distribution . Otherwise , the deep energy-based model would be required to fit a distribution with zero density everywhere except the training examples . Any gradient based sampling and prediction algorithm would be rendered useless . Additional Gaussian distributed noise on the data samples Y + alleviates this issue and facilitates stable training . 2.2 PREDICTION . Prediction from an energy-based viewpoint corresponds to finding the set with the lowest energy value . One approach addresses this intractable optimization problem by approximating a local minimum via gradient descent ( Belanger & McCallum , 2016 ; Belanger et al. , 2017 ) . Learning a multimodal distribution is clearly not sufficient , as the deterministic gradient descent algorithm would not be able to cover all possible sets . This would make the learning process pointless , except for a single local minimum in the energy function . We propose to augment the gradient descent optimizer with additional Gaussian noise during the first n steps : Y ( t+1 ) = Y ( t ) − ∂ ∂Y Eθ ( x , Y ( t ) ) +U ( t ) , for t ≤ S , ( 6a ) Y ( t+1 ) = Y ( t ) − ∂ ∂Y Eθ ( x , Y ( t ) ) , for S < t ≤ T. ( 6b ) For simplicity we choose the same maximum number of steps T , both for training and prediction . One interpretation of the prediction procedure is : 1 . Langevin MCMC sample Y ( S ) based on the energy Eθ and 2 . Refine the sample via gradient descent , such that Y ( T ) is a local minimum of Eθ that is close to Y ( S ) . Note that the partial derivative ∂∂Y Eθ ( x , Y ( t ) ) is not stochastic and can be computed independent of a mini-batch . Thus the sole source of randomness lies with the addition of U , resulting in a prediction procedure that allows for different predictions given the same observation . From the set prediction point of view , the noise term addresses an optimization problem that is specific to set functions . Commonly used set neural networks ( Zaheer et al. , 2017 ) , require permutation invariant pooling operators . Examples include sum or mean pooling . Both of these result in identical partial gradients for identical elements : ∂ ∂yi Eθ ( x , Y ) = ∂ ∂yj Eθ ( x , Y ) , ( 7 ) where yi and yj are two different elements in Y with identical value , i.e. , yi=yj . Although we consider set , not multi-set prediction ; in practice the set Y needs to be stored as a tensor of numbers with limited precision . For the purpose of successfully sampling Y from Eθ , we restrict the parameters θ to energy functions with numerically stable derivatives . Specifically , the difference in the gradients of two elements in Y is limited by the difference between the same two elements . This poses the additional difficulty for the optimizer of separating different elements that are too close , next to the original task of moving the element to the correct position . It is reasonable to assume several elements in close vicinity for problems where the set size is much larger than the number of features . The independently sampled noise term helps disambiguate such proximal elements and speeds up the optimization procedure . A naive alternative would be to solely initialize the set Y ( 0 ) with the constraint of a minimal distance between each element . While this approach addresses the problem at step t=0 , it is ignored in the subsequent steps t > 0 , where two elements may have collapsed . Our proposed prediction procedure adds independently sampled noise at several steps ; thus removing some of the responsibility , for separating elements , from the gradient-based optimizer .

Summary: Authors propose a new method for formulating set prediction tasks. They propose to use a noisy energy-based model with langevin mcmc + noisy startup as their model. The can approximate the gradient of the likelihood function by computing the enery of ground truth pairs and energy of synthesized pairs where the target is sampled from the model distribution.