Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions

Abstract:
Combining discrete probability distributions and combinatorial optimization problems with neural network components has numerous applications but poses several challenges. We propose Implicit Maximum Likelihood Estimation (I-MLE), a framework for end-to-end learning of models combining discrete exponential family distributions and differentiable neural components. I-MLE is widely applicable as it only requires the ability to compute the most probable states and does not rely on smooth relaxations. The framework encompasses several approaches such as perturbation-based implicit differentiation and recent methods to differentiate through black-box combinatorial solvers. We introduce a novel class of noise distributions for approximating marginals via perturb-and-MAP. Moreover, we show that I-MLE simplifies to maximum likelihood estimation when used in some recently studied learning settings that involve combinatorial solvers. Experiments on several datasets suggest that I-MLE is competitive with and often outperforms existing approaches which rely on problem-specific relaxations.

heading: 1 Introduction
text: 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

heading: 2 Problem Statement and Motivation
text: 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.

heading: 3 The Implicit Maximum Likelihood Estimator
text: 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.

heading: 3.1 Target Distributions via Perturbation-based Implicit Differentiation
text: The efficacy of the I-MLE estimator hinges on a proper choice of q, a hyperparameter of our framework. In this section we derive and motivate a class of general-purpose target distributions, rooted in perturbation-based implicit differentiation (PID):
q(z;θ′) = p(z;θ − λ∇z`(fu(z), ŷ)) with z = MAP(θ + ) and ∼ ρ( ), (10) 4We expand on this in Appendix A where we also review the classic MLE setup [Murphy, 2012, Ch. 9]. 5This is known as the perceptron learning rule in standard MLE.
where θ = hv(x̂), (x̂, ŷ) ∈ D is a data point, and λ > 0 is a hyperparameter that controls the perturbation intensity.
To motivate Eq. (10), consider the setting where the inputs to f are the marginals of p(z;θ) (rather than discrete perturb-and-MAP samples as in Eq. (8)), that is, y = fu(µ(θ)) with θ = hv(x̂), and redefine the training error L of Eq. (2) accordingly. A seminal result by Domke [2010] shows that, in this case, we can obtain∇θL by perturbation-based differentiation as:
∇θL(x̂, ŷ;ω) = lim λ→0
{ 1
λ [µ(θ)− µ (θ − λ∇µL(x̂, ŷ;ω))]
} , (11)
where ∇µL = ∂µfu(µ)ᵀ∇y`(y, ŷ). The expression inside the limit may be interpreted as the gradient of an implicit MLE objective (see Eq. (7)) between the distribution p with (current) parameters θ and p with parameters perturbed in the negative direction of the downstream gradient ∇µL. Now, we can adapt (11) to our setting of Eq. (8) by resorting to the straight-through estimator (STE) assumption [Bengio et al., 2013]. Here, the STE assumption translates into reparameterizing z as a function of µ and approximating ∂µz ≈ I . Then, ∇µL = ∂µzᵀ∇zL ≈ ∇zL and we approximate Eq. (11) as:
∇θL(x̂, ŷ;ω) ≈ 1
λ [µ(θ)− µ (θ − λ∇zL(x̂, ŷ;ω))] =
1 λ ∇θL(θ,θ − λ∇zL(x̂, ŷ;ω)), (12)
for some λ > 0. From Eq. (12) we derive (10) by taking a single sample estimator of ∇zL (with perturb-and-MAP sampling) and by incorporating the constant 1/λ into a global learning rate. IMLE with PID target distributions may be seen as a way to generalize the STE to more complex distributions. Instead of using the gradients ∇zL to backpropagate directly, I-MLE uses them to construct a target distribution q. With that, it defines an implicit maximum likelihood objective, whose gradient (estimator) propagates the supervisory signal upstream, critically, taking the constraints into account. When using Eq. (10) with ρ( ) = δ0( )6, the I-MLE estimator also recovers a recently proposed gradient estimation rule to differentiate through black-box combinatorial optimization problems [Pogančić et al., 2019]. I-MLE unifies existing gradient estimation rules in one framework. Algorithm 1 shows the pseudo-code of the algorithm implementing Eq. (9) for S = 1, using the PID target distribution of Eq. (10). The simplicity of the code also demonstrates that instances of I-MLE can easily be implemented as a layer.
We will resume the discussion about target distributions in Section 4, where we analyze more closely the setup of Example 3. Next, we focus on the perturb-and-MAP strategies and derive a class of noise distributions that is particularly apt to the settings we consider in this work.

heading: 3.2 A Novel Family of Perturb-and-MAP Noise Distributions
text: When p is a complex high-dimensional distribution, obtaining Monte Carlo estimates of the gradient in Eq. (7) requires approximate sampling. In this paper, we rely on perturbation-based sampling, also known as perturb and MAP [Papandreou and Yuille, 2011]. In this Section we propose a novel way to design tailored noise perturbations. While the proposed family of noise distributions works with I-MLE, the results of this section are of independent interest and can also be used in other (relaxed) perturb-and-MAP based gradient estimators [e.g. Paulus et al., 2020]. First, we start by revisiting a classic result by Papandreou and Yuille [2011] which theoretically motivates the perturb-and-MAP approach (also known as the Gumbel-max trick), which we generalize here to consider also the temperature parameter τ . Proposition 1. Let p(z;θ) be a discrete exponential family distribution with integer polytope C and temperature τ , and let 〈z,θ〉 be the unnormalized weight of each z ∈ C. Moreover, let θ̃ be such that, for all z ∈ C, 〈z, θ̃〉 = 〈z,θ〉+ (z) with each (z) sampled i.i.d. from Gumbel(0, τ). Then we have that Pr(MAP(θ̃) = z) = p(z;θ).
All proofs can be found in Appendix B. The proposition states that if we can perturb the weights 〈z,θ〉 of each z ∈ C with independent Gumbel(0, τ) noise, then obtaining MAP states from the perturbed model is equivalent to sampling from p(z;θ) at temperature7 τ . For complex exponential
6δ0 is the Dirac delta centered around 0 – this is equivalent to approximating the marginals with MAP. 7Note that the temperature here is different to the temperature of the Gumbel softmax trick [Jang et al., 2017]
which scales both the sum of the logits and the samples from Gumbel(0, 1).
distributions, perturbing the weights 〈z,θ〉 for each state z ∈ C is at least as expensive as computing the marginals exactly. Hence, one usually resorts to local perturbations of each [θ]i (the i-th entry of the vector θ) with Gumbel noise. Fortunately, we can prove that, for a large class of distributions, it is possible to design more suitable local perturbations. First, we show that, for any κ ∈ N+, a Gumbel distribution can be written as a finite sum of κ i.i.d. (implicitly defined) random variables. Lemma 1. Let X ∼ Gumbel(0, τ) and let κ ∈ N+. Define the Sum-of-Gamma distribution as
SoG(κ, τ, s) := τ
κ { s∑ i=1 {Gamma(1/κ, κ/i)} − log(s) } , (13)
where s ∈ N+ and Gamma(α, β) is the Gamma distribution with shape α and scale β, and let SoG(κ, τ) := lims→∞ SoG(κ, τ, s). Then we have that X ∼ ∑κ j=1 j , with j ∼ SoG(κ, τ).
Based on Lemma 1, we can show that for exponential family distributions where every z ∈ C has exactly k non-zero entries we can design perturbations of 〈z,θ〉 following a Gumbel distribution. Theorem 1. Let p(z;θ) be a discrete exponential family distribution with integer polytope C and temperature τ . Assume that if z ∈ C then 〈z,1〉 = k for some constant k ∈ N+. Let θ̃ be the perturbation obtained by [θ̃]j = [θ]j + j with j ∼ SoG(k, τ) from Eq. (13). Then, ∀z ∈ C we have that 〈z, θ̃〉 = 〈z,θ〉+ (z), with (z) ∼ Gumbel(0, τ).
Many problems such as k-subset selection, traveling salesman, spanning tree, and graph matching strictly satisfy the assumption of Theorem 1. We can, however, also apply the strategy in cases where the variance of 〈Z,1〉 is small (e.g. shortest weighted path). The Sum-of-Gamma perturbations provide a more fine-grained approach to noise perturbations. For τ = κ = 1, we obtain the standard Gumbel perturbations. In contrast to the standard Gumbel(0, 1) noise, the pro-
posed local Sum-of-Gamma perturbations result in weights’ perturbations that follow the Gumbel distribution. Fig. 2 shows histograms of 10k samples, where each sample is either the sum of 5 samples from Gumbel(0, 1) (the standard approach) or the sum of k = 5 samples from SoG(5, 1, 10) = 15 ∑10 i=1{Gamma(1/5, 5/i) − log(10)}. While we still cannot sample faithfully from p(z;θ) as the perturbations are not independent, we can counteract the problem of partially dependent perturbations by increasing the temperature τ and, therefore, the variance of the noise distribution. We explore and verify the importance of tuning τ empirically. In the appendix, we also show that the infinite series from Lemma 1 can be well approximated by a finite sum using convergence results for the Euler-Mascheroni series [Mortici, 2010].

heading: 4 Target Distributions for Combinatorial Optimization Problems
text: In this section, we explore the setting where the discrete computational component arises from a combinatorial optimization (CO) problem, specifically an integer linear program (ILP). Many authors have recently considered the setup where the CO component occupies the last layer of the model defined by Eq. (1) (where fu is the identity) and the supervision is available in terms of examples of either optimal solutions [e.g. Pogančić et al., 2019] or optimal cost coefficients (conditioned on the inputs) [e.g. Elmachtoub and Grigas, 2020]. We have seen in Example 3 that we can naturally associate to each ILP a probability distribution (see Eq. (3)) with θ given by the negative cost coefficients c of the ILP and C the integral polytope. Letting τ → 0 is equivalent to taking the MAP in the forward pass. Furthermore, in Section 3.1 we showed that the I-MLE framework subsumes a recently propose method by Pogančić et al. [2019]. Here, instead, we show that, for a certain choice of the target distribution, I-MLE estimates the gradient of an explicit maximum likelihood learning loss L where the data distribution is ascribed to either (examples of) optimal solutions or optimal cost coefficients.
Let q(z;θ′) be the distribution p(z;θ′), with parameters
[θ′]i := { [θ]i if [∇zL]i = 0 −[∇zL]i otherwise. (14)
In the first CO setting, we observe training data D = {(x̂j , ŷj)}Nj=1 where ŷj ∈ C and the loss ` measures a distance between a discrete ẑj ∼ p(z;θj) with θj = hv(x̂) and a given optimal solution of the ILP ŷj . An example is the Hamming loss `H [Pogančić et al., 2019] defined as `H(z,y) = z ◦ (1− y) + y ◦ (1− z), where ◦ denotes the Hadamard (or entry-wise) product. Fact 1. If one uses `H , then I-MLE with the target distribution of Eq. (14) and ρ( ) = δ0 is equivalent to the perceptron-rule estimator of the MLE objective between p(z;hv(x̂j)) and ŷj .
It follows that the method by Pogančić et al. [2019] returns, for a large enough λ, the maximumlikelihood gradients (scaled by 1/λ) approximated by the perceptron rule. The proofs are given in Appendix B.
In the second CO setting, we observe training data D = {(x̂j , ĉj)}Nj=1, where ĉj is the optimal cost conditioned on input x̂j . Here, various authors [e.g. Elmachtoub and Grigas, 2020, Mandi et al., 2020, Mandi and Guns, 2020] use as point-wise loss the regret `R(θ, c) = c> (z(θ)− ẑ∗(c)) where z(θ) is a state sampled from p(z;θ) (possibly with temperature τ → 0, that is, a MAP state) and ẑ∗(c) ∈ z∗(c) is an optimal state for c. Fact 2. If one uses `R then I-MLE with the target distribution of Eq. (14) is equivalent to the perturb-and-MAP estimator of the MLE objective between p(z;hv(x̂j)) and p(z;−ĉj).
This last result also implies that when using the target distribution q from (14) in conjunction with the regret, I-MLE performs maximum-likelihood learning minimizing the KL divergence between the current distribution and the distribution whose parameters are the optimal cost.
Moreover, both facts imply that, when sampling from the MAP states of the distribution q defined by Eq. (14), we have that `(ẑ, ŷ) = 0 for ẑ ∈ MAP(θ′). Therefore, `(ẑ, ŷ) = 0 ≤ Eẑ∼p(z;θ) [`((ẑ, ŷ)], meaning that the inequality of Eq. (5) is satisfied for τ → 0.

heading: 5 Related Work
text: Several papers address the gradient estimation problem for discrete r.v.s, many resorting to relaxations. Maddison et al. [2017], Jang et al. [2017] propose the Gumbel-softmax distribution to relax categorical r.v.s; Paulus et al. [2020] study extensions to more complex probability distributions. The concrete distribution (the Gumbel-softmax distribution) is only directly applicable to categorical variables. For more complex distributions, one has to come up with tailor-made relaxations or use the straightthrough or score function estimators (see for instance Kim et al. [2016], Grover et al. [2019]). In our experiments, we compare with the Gumbel-softmax estimator in Figure 4 (left and right). We show that the k-subset VAE trained with I-MLE achieves loss values that are similar to those of the categorical (1-subset) VAE trained with the Gumbel-softmax gradient estimator. Tucker et al. [2017], Grathwohl et al. [2018] develop parameterized control variates (the former was named REBAR) based on continuous relaxations for the score-function estimator. In contrast, we focus explicitly on problems where only discrete samples are used during training. Moreover, REBAR is tailored to categorical distributions. I-MLE is intended for models with complex distributions (e.g. those with with many constraints).
Approaches that do not rely on relaxations are specific to certain distributions [Bengio et al., 2013, Franceschi et al., 2019, Liu et al., 2019] or assume knowledge of C [Kool et al., 2020]. We provide a general-purpose framework that does not require access to the linear constraints and the corresponding integer polytope C. Experiments in the next section show that while I-MLE only requires a MAP solver, it is competitive and sometimes outperforms tailor-made relaxations. SparseMAP [Niculae et al., 2018] is an approach to structured prediction and latent variables, replacing the exponential distribution (specifically, the softmax) with a sparser distribution. Similar to our work, it only presupposes the availability of a MAP oracle. LP-SparseMAP [Niculae and Martins, 2020] is an extension that uses a relaxation of the optimization problem rather than a MAP solver. Sparsity can also be exploited for efficient marginal inference in latent variable models [Correia et al., 2020].
A series of works about differentiating through CO problems [Wilder et al., 2019, Elmachtoub and Grigas, 2020, Ferber et al., 2020, Mandi and Guns, 2020] relax ILPs by adding L1, L2 or log-barrier regularization terms and differentiate through the KKT conditions deriving from the application of the cutting plane or the interior-point methods. These approaches are conceptually linked to techniques for differentiating through smooth programs [Amos and Kolter, 2017, Donti et al., 2017,
Agrawal et al., 2019, Chen et al., 2020, Domke, 2012, Franceschi et al., 2018] that arise not only in modelling but also in hyperparameter optimization and meta-learning. Pogančić et al. [2019], Rolínek et al. [2020], Berthet et al. [2020] propose methods that are not tied to a specific ILP solver. As we saw above, the former two, originally derived from a continuous interpolation argument, may be interpreted as special instantiations of I-MLE. The latter addresses the theory of perturbed optimizers and discusses perturb and MAP in the context of the Fenchel-Young loss. All the COrelated works assume that either optimal costs or solutions are given as training data, while I-MLE may be also applied in the absence of such supervision by making use of implicitly generated target distributions. Other authors focus on devising differentiable relaxations for specific CO problems such as SAT [Evans and Grefenstette, 2018] or MaxSAT [Wang et al., 2019]. Machine learning intersects with CO also in other contexts, e.g. in learning heuristics to improve the performances of CO solvers or differentiable models such as GNNs to “replace” them; see Bengio et al. [2020] and references therein.
Direct Loss Minimization [DLM, McAllester et al., 2010, Song et al., 2016] is also related to our work, but the assumption there is that examples of optimal states ẑ are given. Lorberbom et al. [2019] extend the DLM framework to discrete VAEs using coupled perturbations. Their approach is tailored to VAEs and not general-purpose. Under a methodological viewpoint, I-MLE inherits from classical MLE [Wainwright and Jordan, 2008] and perturb-and-MAP [Papandreou and Yuille, 2011]. The theory of perturb-and-MAP was used to derive general-purpose upper bounds for log-partition functions [Hazan and Jaakkola, 2012, Shpakova and Bach, 2016].

heading: 6 Experiments
text: The set of experiments can be divided into three parts. First, we analyze and compare the behavior of I-MLE with (i) the score function and (ii) the straight-through estimator using a toy problem. Second, we explore the latent variable setting where both hv and fu in Eq. (1) are neural networks and the optimal structure is not available during training. Finally, we address the problem of differentiating through black-box combinatorial optimization problems, where we use the target distribution derived in Section 4. More experimental details for available in the appendix.
Synthetic Experiments. We conducted a series of experiments with a tractable 5-subset distribution (see Example 2) where z ∈ {0, 1}10. We set the loss to L(θ) = Eẑ∼p(z;θ)[‖ẑ − b‖2], where b is a fixed vector sampled from N (0, I). In Fig. 3 (Top), we plot optimization curves with means and standard deviations, comparing the proposed estimator with the straight-through (STE) and the score function (SFE) estimators. 8 For STE and I-MLE, we use Perturb-and-MAP (PaM) with Gumbel and SoG(1, 5, 10) noise, respectively. The SFE uses faithful samples and exact marginals (which is feasible only when m is very small) and converges much more slowly than the other methods, while the STE converges to worse solutions than those found using I-MLE. Fig. 3 (Bottom) shows the benefits of using SoG rather than Gumbel perturbations with I-MLE. While the best configurations for both are comparable, SoG noise achieves in average (over 100 runs) strictly better final values of L for more than 50% of the tested configurations (varying λ from Eq. (10) and the learning rate) and exhibit smaller variance (see Fig. 6). Additional details and results in Appendix C.1.
Learning to Explain. The BEERADVOCATE dataset [McAuley et al., 2012] consists of free-text reviews and ratings for 4 different aspects of beer: appearance, aroma, palate, and taste. Each sentence in the test set has annotations providing the words that best describe the various aspects. Following the experimental setup of recent work [Paulus et al., 2020], we address the problem introduced by the L2X paper [Chen et al., 2018] of learning a distribution over k-subsets of words that best explain a given aspect rating. The complexity of the MAP problem for the k-subset distribution is linear in k. sThe training set has 80k reviews for the aspect APPEARANCE and 70k reviews for all other aspects. Since the original dataset [McAuley et al., 2012] did not provide separate validation and test sets, we compute 10 different evenly sized validation/test splits of the 10k held out set and compute mean and standard deviation over 10 models, each trained on one split. Subset precision was computed using a subset of 993 annotated reviews. We use pre-trained word embeddings from Lei et al. [2016]. Prior work used non-standard neural networks for which an implementation is not available [Paulus
8Hyperparameters are optimized against L for all methods independently. Statistics are over 100 runs. We found STE slightly better with Gumbel rather than SoG noise. SFE failed with all tested PaM strategies.
et al., 2020]. Instead, we used the neural network from the L2X paper with 4 convolutional and one dense layer. This neural network outputs the parameters θ of the distribution p(z;θ) over k-hot binary latent masks with k ∈ {5, 10, 15}. We compare to relaxation-based baselines L2X [Chen et al., 2018] and SoftSub [Xie and Ermon, 2019]. We also compare the straight-through estimator (STE) with Sum-of-Gamma (SoG) perturbations. We used the standard hyperparameter settings of Chen et al. [2018] and choose the temperature parameter t ∈ {0.1, 0.5, 1.0, 2.0}. For I-MLE we choose λ ∈ {101, 102, 103}, while for both I-MLE and STE we choose τ ∈ {k, 2k, 3k} based on the validation MSE. We used the standard Adam settings. We trained separate models for each aspect using MSE as point-wise loss `.
Table 1 lists detailed results for the aspect AROMA. I-MLE’s MSE values are competitive with those of the best baseline, and its subset precision is significantly higher than all other methods (for τ = 30). Using only MAP as the approximation of the marginals leads to poor results. This shows that using the tailored perturbations with tuned temperature is crucial to achieve state of the art results. The Sum-of-Gamma perturbation introduced in this paper outperforms the standard local Gumbel perturbations. More details and results can be found in the appendix.
Discrete Variational Auto-Encoder. We evaluate various perturbation strategies for a discrete k-subset Variational Auto-Encoder (VAE) and compare them to the straight-through estimator (STE) and the Gumbel-softmax trick. The latent variables model a probability distribution over k-subsets of (or top-k assignments too) binary vectors of length 20. The special case of k = 1 is equivalent to a categorical variable with 20 categories. For k > 1, we use I-MLE using the class of PID target distributions of Eq. (10) and compare various perturb-and-MAP noise sampling strategies. The experimental setup is similar to those used in prior work on the Gumbel softmax tricks [Jang et al., 2017]. The loss is the sum of the reconstruction losses (binary cross-entropy loss on output pixels) and the KL divergence between the marginals of the variables and the uniform distribution. The encoding and decoding functions of the VAE consist of three dense layers (encoding: 512-256-20x20; decoding: 256-512-784). We do not use temperature annealing. Using Eq. (9) with S = 1, we use either Gumbel(0, 1) perturbations (the standard approach)9 or Sum-of-Gamma (SoG) perturbations at a temperature of τ = 10. We run 100 epochs and record the loss on the test data. The difference in training time is negligible. Fig. 4 shows that using the SoG noise distribution is beneficial. The test loss using the SoG perturbations is lower despite the perturbations having higher variance and, therefore, samples of the model being more diverse. This shows that using perturbations of the weights that follow a proper Gumbel distribution is indeed beneficial. I-MLE significantly outperforms the
9Increasing the temperature τ of Gumbel(0, τ) samples increased the test loss.
STE, which does not work in this setting and is competitive with the Gumbel-Softmax trick for the 1-subset (categorical) distribution where marginals can be computed in closed form.
Differentiating through Combinatorial Solvers. In these experiments, proposed by Pogančić et al. [2019], the training datasets consists of 10,000 examples of randomly generated images of terrain maps from the Warcraft II tile set [Guyomarch, 2017]. Each example has an underlying K ×K grid whose cells represent terrains
with a fixed cost. The shortest (minimum cost) path between the top-left and bottom-right cell in the grid is encoded as an indicator matrix and serves as the target output. An image of the terrain map is presented to a CNN, which produces a K ×K matrix of vertex costs. These costs are then given to Dijkstra’s algorithm (the MAP solver) to compute the shortest path. We closely follow the evaluation protocol of Pogančić et al. [2019]. We considered two instantiations of I-MLE: one derived from Fact 1 (M-M in Table 2) using `H and one derived from Fact 2 (µ-µ) using `R, with ρ( ) = SoG(k, 1, 10) where k is the empirical mean of the path lengths (different for each grid size K). We compare with the method proposed by Pogančić et al. [2019] (BB10) and Berthet et al. [2020] (DPO). The results are listed in Table 2. I-MLE obtains results comparable to (BB) with M-M and is more accurate with µ-µ. We believe that the µ-µ advantage may be partially due to an implicit form of data augmentation since we know from Fact 2 that, by using I-MLE, we obtain samples from the distribution whose parameters are the optimal cost. Training dynamics, showing faster convergence of I-MLE (µ-µ), and additional details are available in Table 4.

heading: 7 Conclusions
text: I-MLE is an efficient, simple-to-implement, and general-purpose framework for learning hybrid models. I-MLE is competitive with relaxation-based approaches for discrete latent-variable models and with approaches to backpropagate through CO solvers. Moreover, we showed empirically that I-MLE outperforms the straight-through estimator. A limitation of the work is its dependency on computing MAP states which is, in general, an NP-hard problem (although for many interesting cases there are efficient algorithms). Future work includes devising target distributions when∇zL is not available, studying the properties (including the bias) of the proposed estimator, developing adaptive strategies for τ and λ, and integrating and testing I-MLE in several challenging application domains.
10Note that this is the same as using I-MLE with PID target distribution form Eq. (10) and ρ( ) = δ0.

Review summaries:
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.  

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.

The authors propose a method for differentiating models that combine neural network components with a discrete distribution. The method estimates gradients by computing most probable states akin to perturb-and-MAP. Simulations on three problems show comparable performance or slight improvement over existing methods.