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 a meta algorithm for models involving a discrete bottleneck by making use of a specific target distribution with respect to which a max-likelihood loss is optimized in order to avoid having to compute the high variance score function gradients of the original objective.  This target distribution is characterized by being a local improvement to the overall objective.  To avoid marginalization on the discrete variables, the paper further consider approximations to the proposed I-MLE framework using the perturb-and-MAP approach, which repeatedly solves MAP subroutines for noise-perturbed parameters to approximate the marginals. In doing this, it also proposes a decomposition of Gumbel noise terms into a sum of Gamma distributions.  The ideas are evaluated empirically both on synthetic problems as well as benchmarks involving discrete VAEs and other combinatorial problems requiring flexible discrete latent variable modeling.