Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions

1 Introduction . While deep neural networks excel at perceptual tasks , they tend to generalize poorly whenever the problem at hand requires some level of symbolic manipulation or reasoning , or exhibit some ( known ) algorithmic structure . Logic , relations , and explanations , as well as decision processes , frequently find natural abstractions in discrete structures , ill-captured by the continuous mappings of standard neural nets . Several application domains , ranging from relational and explainable ML to discrete decision-making [ Mišić and Perakis , 2020 ] , could benefit from general-purpose learning algorithms whose inductive biases are more amenable to integrating symbolic and neural computation . Motivated by these considerations , there is a growing interest in end-to-end learnable models incorporating discrete components that allow , e.g. , to sample from discrete latent distributions [ Jang et al. , 2017 , Paulus et al. , 2020 ] or solve combinatorial optimization problems [ Pogančić et al. , 2019 , Mandi et al. , 2020 ] . Discrete energy-based models ( EBMs ) [ LeCun et al. , 2006 ] and discrete world models [ Hafner et al. , 2020 ] are additional examples of neural network based models that require the ability to backpropagate through discrete probability distributions . For complex discrete distributions , it is intractable to compute the exact gradients of the expected loss . For combinatorial optimization problems , the loss is discontinuous , and the gradients are zero almost everywhere . The standard approach revolves around problem-specific smooth relaxations , which allow one to fall back to ( stochastic ) backpropagation . These strategies , however , require tailor-made relaxations , presuppose access to the constraints and are , therefore , not always feasible nor tractable for large state spaces . Moreover , reverting to discrete outputs at test time may cause unexpected behavior . In other situations , discrete outputs are required at training time because one has to make one of a number of discrete choices , such as accessing discrete memory or deciding on an action in a game . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . With this paper , we take a step towards the vision of general-purpose algorithms for hybrid learning systems . Specifically , we consider settings where the discrete component ( s ) , embedded in a larger computational graph , are discrete random variables from the constrained exponential family1 . Grounded in concepts from Maximum Likelihood Estimation ( MLE ) and perturbation-based implicit differentiation , we propose Implicit Maximum Likelihood Estimation ( I-MLE ) . To approximate the gradients of the discrete distributions ’ parameters , I-MLE computes , at each update step , a target distribution q that depends on the loss incurred from the discrete output in the forward pass . In the backward pass , we approximate maximum likelihood gradients by treating q as the empirical distribution . We propose ways to derive target distributions and introduce a novel family of noise perturbations well-suited for approximating marginals via perturb-and-MAP . I-MLE is general-purpose as it only requires the ability to compute most probable states and not faithful samples or probabilistic inference . In summary , we make the following contributions : 1 . We propose implicit maximum likelihood estimation ( I-MLE ) as a framework for computing gradients with respect to the parameters of discrete exponential family distributions ; 2 . We show that this framework is useful for backpropagating gradients through both discrete probability distributions and discrete combinatorial optimization problems ; 3 . I-MLE requires two ingredients : a family of target distribution q and a method to sample from complex discrete distributions . We propose two families of target distributions and a family of noise-distributions for Gumbel-max ( perturb-and-MAP ) based sampling . 4 . We show that I-MLE simplifies to explicit maximum-likelihood learning when used in some recently studied learning settings involving combinatorial optimization solvers . 5 . Extensive experimental results suggest that I-MLE is flexible and competitive compared to the straight-through and relaxation-based estimators . Instances of the I-MLE framework can be easily integrated into modern deep learning pipelines , allowing one to readily utilize several types of discrete layers with minimal effort . We provide implementations and Python notebooks at https : //github.com/nec-research/tf-imle 2 Problem Statement and Motivation . We consider models described by the equations θ = hv ( x ) , z ∼ p ( z ; θ ) , y = fu ( z ) , ( 1 ) where x ∈ X and y ∈ Y denote feature inputs and target outputs , hv : X → Θ and fu : Z → Y are smooth parameterized maps , and p ( z ; θ ) is a discrete probability distribution . Given a set of examples D = { ( x̂j , ŷj ) } Nj=1 , we are concerned with learning the parameters ω = ( v , u ) of ( 1 ) by finding approximate solutions of minω ∑ j L ( x̂j , ŷj ; ω ) /N . The training error L is typically defined as : L ( x̂ , ŷ ; ω ) = Eẑ∼p ( z ; θ̂ ) [ ` ( fu ( ẑ ) , ŷ ) ] with θ̂ = hv ( x̂ ) , ( 2 ) where ` : Y × Y → R+ is a point-wise loss function . Fig . 1 illustrates the setting . For example , an interesting instance of ( 1 ) and ( 2 ) arises in learning to explain user reviews [ Chen et al. , 2018 ] where the task is to infer a target sentiment score ( e.g . w.r.t . the quality of a product ) from a review while also providing a concise explanation of the predicted score by selecting a subset of exactly k words ( cf . Example 2 ) . In Section 6 , we present experiments precisely in this setting . As anticipated in the introduction , we restrict the discussion to instances in which p ( z ; θ ) belongs to the ( constrained ) discrete exponential family , which we now formally introduce . Let Z be a vector of discrete random variables over a state space Z and let C ⊆ Z be the set of states that satisfy a given set of linear constraints.2 Let θ ∈ Θ ⊆ Rm be a real-valued parameter vector . 1This includes integer linear programs via a natural link that we outline in Example 3 . 2For the sake of simplicity we assume Z ⊆ { 0 , 1 } m. The set C is the integral polytope spanned by the given , problem-specific , linear constraints . The probability mass function ( PMF ) of a discrete constrained exponential family r.v . is : p ( z ; θ ) = { exp ( 〈z , θ〉/τ −A ( θ ) ) if z ∈ C , 0 otherwise . ( 3 ) Here , 〈· , ·〉 is the inner product and τ the temperature , which , if not mentioned otherwise , is assumed to be 1 . A ( θ ) is the log-partition function defined as A ( θ ) = log ( ∑ z∈C exp ( 〈z , θ〉/τ ) ) . We call 〈z , θ〉 the weight of the state z . The marginals ( expected value , mean ) of the r.v.s Z are defined as µ ( θ ) : = Eẑ∼p ( z ; θ ) [ ẑ ] . Finally , the most probable or Maximum A-Posteriori ( MAP ) states are defined as MAP ( θ ) : = arg maxz∈C 〈z , θ〉 . The family of probability distributions we define here captures a broad range of settings and subsumes probability distributions such as positive Markov random fields and statistical relational formalisms [ Wainwright and Jordan , 2008 , Raedt et al. , 2016 ] . We now discuss some examples which we will use in the experiments . Crucially , in Example 3 we establish the link between the constrained exponential family and integer linear programming ( ILP ) identifying the ILP cost coefficients with the distribution ’ s parameters θ . Example 1 ( Categorical Variables ) . An m-way ( one-hot ) categorical variable corresponds to p ( z ; θ ) = exp ( 〈z , θ〉 −A ( θ ) ) , subject to the constraint 〈z,1〉 = 1 , where 1 is a vector of ones . As C = { ei } mi=1 , where ei is the i-th vector of the canonical base , the parameters of the above distribution coincide with the weights , which are often called logits in this context . The marginals µ coincide with the PMF and can be expressed through a closed-form smooth function of θ : the softmax . This facilitates a natural relaxation that involves using µ ( θ ) in place of z [ Jang et al. , 2017 ] . The convenient properties of the categorical distribution , however , quickly disappear even for slightly more complex distributions , as the following example shows . Example 2 ( k-subset Selection ) . Assume we want to sample binary m-dimensional vectors with k ones . This amounts to replacing the constraint in Example 1 by the constraint 〈z,1〉 = k. Here , a closed-form expression for the marginals does not exist : sampling from this distribution requires computing the ( m k ) = O ( mk ) weights ( if k ≤ m/2 ) . Computing MAP states instead takes time linear in m. Example 3 ( Integer Linear Programs ) . Consider the combinatorial optimization problem given by the integer linear program arg minz∈C〈z , c〉 , where C is an integral polytope and c ∈ Rm is a vector of cost coefficients , and let z∗ ( c ) be the set of its solutions . We can associate to the ILP the family ( indexed by τ > 0 ) of probability distributions p ( z ; θ ) from ( 3 ) , with C the ILP polytope and θ = −c . Then , for every τ > 0 , the solutions of the ILP correspond to the MAP states : MAP ( θ ) = arg maxz∈C〈z , θ〉 = z∗ ( c ) and for τ → 0 one has that Pr ( Z ∈ z∗ ( c ) ) → 1 . Many problems of practical interest can be expressed as ILPs , such as finding shortest paths , planning and scheduling problems , and inference in propositional logic . 3 The Implicit Maximum Likelihood Estimator . In this section , we develop and motivate a family of general-purpose gradient estimators for Eq . ( 2 ) that respect the structure of C . 3 Let ( x̂ , ŷ ) ∈ D be a training example and ẑ ∼ p ( z ; hv ( x̂ ) ) . The gradient of L w.r.t . u is given by∇uL ( x̂ , ŷ ; ω ) = Eẑ [ ∂ufu ( ẑ ) ᵀ∇y ` ( y , ŷ ) ] with y = fu ( ẑ ) , which may be estimated by drawing one or more samples from p. Regarding ∇vL , one has ∇vL ( x̂ , ŷ ; ω ) = ∂vhv ( x̂ ) ᵀ∇θL ( x̂ , ŷ ; ω ) , ( 4 ) where the major challenge is to compute∇θL . A standard approach is to employ the score function estimator ( SFE ) which typically suffers from high variance . Whenever a pathwise derivative estimator ( PDE ) is available it is usually the preferred choice [ Schulman et al. , 2015 ] . In our setting , however , the PDE is not readily applicable since z is discrete and , therefore , every ( exact ) reparameterization path would be discontinuous . Various authors developed ( biased ) adaptations of the PDE for discrete r.v.s ( see Section 5 ) . These involve either smooth approximations of p ( z ; θ ) or approximations of the derivative of the reparameterization map . Our proposal departs from these two routes and instead involves the formulation of an implicit maximum likelihood estimation problem . In a nutshell , I-MLE 3The derivations are adaptable to other types of losses defined over the outputs of Eq . ( 1 ) . Algorithm 1 Instance of I-MLE with perturbation-based implicit differentiation . function FORWARDPASS ( θ ) // Sample from the noise distribution ρ ( ) ∼ ρ ( ) // Compute a MAP state of perturbed θ ẑ = MAP ( θ + ) save θ , , and ẑ for the backward pass return ẑ function BACKWARDPASS ( ∇z ` ( fu ( z ) , ŷ ) , λ ) load θ , , and ẑ from the forward pass // Compute target distribution parameters θ′ = θ − λ∇z ` ( fu ( z ) , ŷ ) // Single sample I-MLE gradient estimate ∇̂θL ( θ̂ , θ̂′ ) = ẑ − MAP ( θ′ + ) return ∇̂θL ( θ̂ , θ̂′ ) is a ( biased ) estimator that replaces∇θL in Eq . ( 4 ) with ∇̂θL , where L is an implicitly defined MLE objective and ∇̂ is an estimator of the gradient . We now focus on deriving the ( implicit ) MLE objective L. Let us assume we can , for any given ŷ , construct an exponential family distribution q ( z ; θ′ ) that , ideally , is such that Eẑ∼q ( z ; θ′ ) [ ` ( fu ( ẑ ) , ŷ ) ] ≤ Eẑ∼p ( z ; θ ) [ ` ( fu ( ẑ ) , ŷ ) ] . ( 5 ) We will call q the target distribution . The idea is that , by making p more similar to q we can ( iteratively ) reduce the model loss L ( x̂ , ŷ ; ω ) . To this purpose , we define L as the MLE objective4 between the model distribution p with parameters θ and the target distribution q with parameters θ′ : L ( θ , θ′ ) : = −Eẑ∼q ( z ; θ′ ) [ log p ( ẑ ; θ ) ] = Eẑ∼q ( z ; θ′ ) [ A ( θ ) − 〈ẑ , θ〉 ] ( 6 ) Now , exploiting the fact that∇θA ( θ ) = µ ( θ ) , we can compute the gradient of L as ∇θL ( θ , θ′ ) = µ ( θ ) − Eẑ∼q ( z ; θ′ ) [ ẑ ] = µ ( θ ) − µ ( θ′ ) , ( 7 ) that is , the difference between the marginals of the current distribution p and the marginals of the target distribution q , also equivalent to the gradient of the KL divergence between p and q . We will not use Eq . ( 7 ) directly , as computing the marginals is , in general , a # P-hard problem and scales poorly with the dimensionality m. MAP states are typically less expensive to compute ( e.g . see Example 2 ) and are often used directly to approximate µ ( θ ) 5 or to compute perturb-andMAP approximations , where µ ( θ ) ≈ E ∼ρ ( ) MAP ( θ + ) where ∼ ρ ( ) is an appropriate noise distribution with domain Rm . In this work we follow – and explore in more detail in Section 3.2 – the latter approach ( also referred to as the Gumbel-max trick [ cf . Papandreou and Yuille , 2011 ] ) , a strategy that retains most of the computational advantages of the pure MAP approximation but may be less crude . Henceforth , we only assume access to an algorithm to compute MAP states ( such as a standard ILP solver in the case of Example 3 ) and rephrase Eq . ( 1 ) as θ = hv ( x ) , z = MAP ( θ + ) with ∼ p ( ) , y = fu ( z ) . ( 8 ) With Eq . ( 8 ) in place , the general expression for the I-MLE estimator is ∇̂vL ( x , y ; ω ) = ∂vhv ( x̂ ) ᵀ∇̂θL ( θ , θ′ ) with θ = hv ( x̂ ) where , for S ∈ N+ : ∇̂θL ( θ , θ′ ) = 1 S S∑ i=1 [ MAP ( θ + i ) − MAP ( θ′ + i ) ] , with i ∼ ρ ( ) for i ∈ { 1 , . . . , S } . ( 9 ) If the states of both the distributions p and q are binary vectors , ∇̂θL ( θ , θ′ ) ∈ [ −1 , 1 ] m and when S = 1 ∇̂θL ( θ , θ′ ) ∈ { −1 , 0 , 1 } m. In the following , we discuss the problem of constructing families of target distributions q . We will also analyze under what assumptions the inequality of Eq . ( 5 ) holds .

Summary: The paper proposes I-MLE, a class of gradient estimation techniques for computational graphs involving sampling from discrete exponential family distributions. The techniques are built on top of the perturb-and-MAP strategies for sampling from such distributions, extended to account for more general noise perturbation strategies via Sum of Gamma distributions in the forward pass and (biased) gradient estimators that combine straight-through-estimation with perturbation-based differentiation in the backward pass. Finally, the paper provides empirical evaluation on real and synthetic benchmarks that indicates strong and competitive performance of the proposed approach.