Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions
Authors: Mathias Niepert
Pasquale Minervini
Luca Franceschi
Venue: NeurIPS 2021 Poster (2021)
Keywords: discrete-continuous learning
gradient estimation
combinatorial optimization
exponential family distribution

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.

Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions
Mathias Niepert
Pasquale Minervini University College London p wmimmmviii@uul... ac.
Luca Franceschi
NEC Laboratories Europe mathias niepert
Istituto Italiano di Tecnologia University College London ucablfr@ucl. ac uk
Abstract
Combining discrete probability distributions and combinatorial optimization prob- lems 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 fam- ily distributions and differentiable neural components. I-MLE is widely applicable it only requires the ability 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 novel class of noise distributions for approximating marginals via Berturb-and-Ma Moreover, we show that I-MLE simplifies maximum likelihood estimation when used in some recently studied learning settings that involve combinatorial solvers. Experiments several datasets suggest that 1-MLE is competitive with and often outperforms existing approaches which rely on mooblem-specine relaxations.
Introduction
While deep neural networks excel at perceptual tasks, they tend 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 meciiinn-mmking [Misić and Perakis 2020). could benefit from general-purpose learning algorithms amenable integrating symbolic and neural computation. Motivated by these considerations, there is growing interest in end-to-end learnable models incorporating discrete components that allow, e.g., to sample from discrete latent distributions [Jang al. 2017 Paulus et al. 2020) or solve combinatorial optimization problems [Pogančie et al.) 2019 Mandi et al. 2020 Discrete energy-baat mased models (EBMs) [LeCun 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.
whose inductive biases
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 almost everywhere. The standard approach revolves around problem-specific smooth relaxations, which allow one fall back (stochastic) backpropagation. These strategies, however, require tailor-made relaxations, presuppose access to the constraints and are, therefore, not always feasible tractable for large state spaces. Moreover, reverting discrete outputs test time may cause unexpected behavior. In other situations, discrete outputs are required at training time because one has make one number of discrete choices, such as accessing discrete memory or deciding on an action in a game.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

With this paper, we take a step towards the vision of general-purpose ing systems. Specifically, we consider settings where the discrete component(s), embedded larger computational graph, are discrete random variables from the constrained exponential family" Grounded in concepts from Maximum Likelihood Estimation (MI.E) and perturbation-baxe 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 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 novel family of noise per- turbations well-suited for approximating marginals via perturb-ante and-MAP. I-MLE is ceneral-purposs as it only requires the ability to compute most probable states and not faithful samples probabilistic inference. In summary, we make the following contributions:
purpose algorithms for hybrid learn-
1. We propose implicit maximum likelihood estimation (I- (I-MLE) as framework for computing gradients with respect to the parameters of discrete exponential family distributions;
2. We show that this framework 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 families of target distributions and family of noise- distributiisssnns for Gumbel-max (perturb-andy aaob--ndd-AA and) 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 httpss: /tttps://githulte com/ nec http://eccreeeerrcttt----iiie
2 Problem Statement and Motivation
We consider models described by the equations
p(z; ()
0
7
0 = (x), Z 0(2;0), y y= fu(2, (1) where æ E X' and E y denote feature inputs and target outputs, h. : X O and fu 2 +
h
fu
are smooth parameterized maps, and p(2; 0) is a discrete probability distribution.
Figure 1: Illustration of the addressed leamiing prob-
lem. z the discrete (latent) structure.
Given a set of examples D = ((t,, 1 we are concerned with learning the parameters
w = (w. u) of [T] by finding approximate solutions of min E, (6;;),;:)//N The training error typically defined as:
= E E)p((=)) (2) 0) with θ= h
(2)
where /: 3 x R+ is a point-wise loss function. Fig. illustrates the setting. For example, an interesting instance of ( and (2) arises in learning to explain user reviews [Chen et al. 2018 where the task is infer target sentiment score (e.g. w.r.t. the quality of a product) from a review while also providing 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; 0) belongs to the (constrained) discrete exponential family, which we now formally introduce.
Let Z be vector of discrete random variables over a state space Z and C S 2 be the set of states Let 0 O m be a real-valued parameter vector.
that satisfy given set of linear constraints P
'This includes integer linear programs natural link that we outline in Example For the sake of simplicity we assume Z C (0, 1" The set C is the integral polytope spanned by the given, problem-specifie linear constraints.

---

The probability mass function (PMF) of discrete constrained exponential family r.V. is:
p(z; = { exp ((z, 0)/r A(O)) if 2 E C, otherwise.
(3)
0
Here, (+-)) is the inner product and the temperature, which, if not mentioned otherwise, is assumed to be 1. A(0) is the log- log-partition function defined as A(0) = log (Exec exp ((z, 0)/r)). We call (z. 0) the weight of the state z. The marginals (expected value, mean) of the r.v.s Z are defined as (0) Ez=p(=:0)(3) Finally, the most probable or Maximum A-Posteriori (MAP) states are defined MA == arg max (z, 0). The family of probability distributions we define here captures a broad range of settings and subsumes probability distributions such as positive Markoy 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 (TLP) identifying the ILP cost coefficients with the distribution's 's parameters O.
Example 1 (Categorical Variables). An m-way (one-hot) categorical variable corresponds to p(2; 0) exp (z, 0) A(0)). subject the constraint (z. 1) = 1, where 1 is vector of ones.
As C = [e where e; is the i-th vector of the canonical hase, 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 0: the softmax. This facilitates natural relaxation that involves using u0)) in place of = (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 7 by the constraint (z, 1) = k.
Here, closed-form expression for the marginals does not exist: sampling from this distribution requires computing the (x) = O(m²) weights (if l S 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 min (2, c), where C is an integral polytope and c E R is vector of cost coefficients. and let z" (c) be the set of its solutions. We can associate to the ILP the family (indexed by T > 0) of probability distributions p( 0) 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 maxzec (z 0) = z" (c) and for T + 0 one has that Pr(Z E 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. that respect the structure of C. Let (æt, ÿ) E D be training example and é h (+)). The gradient of L w.r.t u is given by V u L( ÿ: w) = Eg (2) fly ()) with y = fu f2), which may be estimated by drawing one or more samples from p. Regarding L. one has
V.I w) (
w)
where the major challenge is to compute VoL. 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(2; 0) or approximations of the derivative of the reparamcterization map. Our proposal departs from these two routes and instead involves the formulation of an implicit maximum likelihood estimation problem. In nutshell, I-MLE
The derivations are adaptable other types of losses defined over the outputs of Eq.
3

---

Algorithm Instance of I-MLE with perturbation--assed implicit differentiation.
function FORWARDPASS(O)
function BACKWARDPASSCVV ( (z), 9). A)
// Sample from the noise distribution p(e)
load 0, e, and & from the forward pass
p(E)
// Compute target distribution parameters
// Compute a MAP state of perturbed
0' XV = (()
- MAP(@ + c)
// Single sample I-MLE gradient estimate (oL 00) & MAP(O" + c)
save 0, e, and 2 for the backward pass return
return 0')
is a (biased) estimator that replaces Vol. L in Eq. (4) with VoL, where L is an implicitly defined MLE objective and V is an estimator of the gradient.
We now focus deriving the (implicit) MLE objective L. Let assume we can, for any given y. construct an exponential family distribution g(2; that, ideally, is such that
Bz~a(=:07) [*( 8)] s Bz~p((:0) [2) 9)]
(5)
We will call q the target distribution. The idea is that. by making p more similar to q can (iteratively) reduce the model loss L( ÿ w) To this purpose, we define as the MLE objective #
between the model distribution with parameters O and the target distribution q with parameters 0':
0(),0') -Eggg(::00 [log p(2; 0)] E =(= (A(0) (2]
(6)
Now, exploiting the fact that VoA(0) = =u(0). we can compute the gradient of C as
VoL(0,0)) (u(0) - Ex~g(=:00
(0 14(0),
(7)
that is, the difference between the marginals of the current distribution p and the marginals of the target distribution g, also equivalent to the gradient of the KI. 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 [ and are often used directly to approximate ( or to compute perturb-and- MAP approximations, where (0) Ee-p(e) MAP E) where p(E) is an appropriate noise distribution with domain Imm. 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 [ef. |Papandreou and Yuille [201]] J),
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 standard IIP solver in the case of Example 3) and rephrase Eq. ① as
0 h (2:), = MAP ( e) with p(E), fulz)
(8)
With Eq. ( in place, the general expression for the I-MLE estimator is V.l(a,y:)) Ooho (+) VeC(O, 0") with 0 (.t) where, for N
1 Σ MAP E.) MAP(O' + Ee)],, with p(e) for i E VoL(0,0)) =
S}. (I.
(9)
If the states of both the distributions and q are binary vectors, VoL(O, 0') E [-1, 1] and when
Vo 0') {-1,0, 1)m the following, we discuss the problem of constructing families 1 =
of target distributions We will also analyze under what assumptions the inequality of Eq. 5 holds.
3.1 Target Distributions via Perturbation-based Implicit Differentiation
The efficacy of the I-MLE estimator hinges on proper choice of g. hyperparameter of our framework. In this section we derive and motivate a class of seneral-purpose target distributions, rooted in perturbation-base implicit differentiation (PID):
<(; 0') = p(z; 0 AV & (z), Ø)) with z = MAP(O + e) and e ~ p(E).
(10)
4 We expand on this in Appendix A where we also review the classic MI.F sctup (Murphy 2012 Ch. 9].
'Thi is known the perceptron learning rule in standard MLE.

---

where 0 = (æ), (#: 9) E 'D is a data point, and À > 0 is hyperparameter that controls the perturbation intensity.
To motivate Eq. 10,, consider the setting where the inputs to / are the marginals of p(z; 0) (rather than discrete perturb-and-MA.. redefine the training error L of Eq. (2) accordingly. A seminal result by Domke [2010] shows that, in this case, we can obtain VoL by perturbation-base differentiation as:
samples in Eq. (8), that is, y fu. (u())) with 0 = ha (d),, and
{ [u()) - (0- 1::0-1000 AV L(2, 91 9:
Vol(a, (: lim X-2
(11)
where VnL = (1. yly.. The expression inside the limit may be interpreted as the gradient of an implicit MLE objective (see Eq. (7) between the distribution with (current) parame- ters 0 and p with parameters perturbed in the negative direction of the downstream gradient V L. Now, we can adapt (IT) to our setting of Eqq (8) by resorting to the straight-through estimator (STE) assumption [Bengio et al. 2013). Here, the STE assumption translates into reparameterizing z as a V L and we approximate
function of f and approximating duz= R I. Then, Vul. = PuzTVzl I. ~
Eq. as:
1 VoC L(,
VoL (: X (u(0) - (0 AV-L(2, [)))]
(12)
A
for some À > O. From Eq. 12 we derive (TO) by taking a single sample estimator of V.l. (with perturb-and-MAR sampling) and by incorporating the constant 1/A into a global learning rate. I- MLE with PID target distributions may be seen distributions. Instead of using the gradients V= L to backpropagate directly, I-MLE uses them to construct 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 p(c) óo(e{p [ the I-MLE estimator also recovers a recently proposed gradient estimation rule to differentiate through black- box combinatorial optimization problems (Pogančié et al. 2019]. I-MLE unifies existing gradient estimation rules in one framework. Algorithm| shows the pseudo-code of the algorithm implementing Eq. ( for S = 1, using the PID) target distribution of Eq. (TO). The simplicity of the code also demonstrates that instances of I-MLE can easily be implemented as layer.
way to generalize the STE more complex
We will resume the discussion about target distributions in Section 4 where we analyze more closely the setup of Example Next, we focus on the perturb-and-Mr strategies and derive a class of noise distributions that is particularly apt the settings we consider in this work.
3.2 A Novel Family of Perturb-and-MAP Noise Distributions
When p is a complex high- ligh-dimensional distribution, obtaining Monte Carlo estimates of the gradient in Eq. ( requires approximate sampling. In this paper, rely on cotturbattiinbbseed 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. Pauluus [e.g.Pauluus al. 2220). First, we start by revisiting classic result by Paaaanreouu and Yuille (201T] which theoretically motivates the perturb-and-M approach (also known as the Gumbel-max trick), which we generalize here consider also the temperature parameter T.
Proposition Let p(2; 0) be a discrete exponential family distribution with integer polytope C and temperature т, and let (z, 0) be the unnormalized weight of each z E C. Moreover, let be such that, for all 2 E C, (z. 0) = (z, 0) + <(z) with each e(z) sampled i.i.d. from Gumbel r). Then we have that Pr(MAP(() 2) = p(2;0) 0).
All proofs can be found in Appendix B The proposition states that if we can perturb the weights (z, 0) of each z E C' with independent Gumbel(0,, r) noise, then obtaining MAP states from the perturbed model is equivalent sampling from p(z; 0) temperature T. For complex exponential
$8 is Dirac delta centered around this is equivalent to approximating the marginals with MAP. "Note that the temperature here is different to the temperature of the Gumbel softmax trick (aang et al 2017 which scales both the sum of the logits and the samples from Gumbel(0, 1).
5

---

distributions, perturbing the weights (z, 0) for each state E C is at least as expensive as computing the marginals exactly. Henee, one usually resorts local perturbations of each [Ø], (the i-th entry of the vector 0) 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 & E N+ Gumbel distribution can be written as finite sum of k i.i.d. (implicitly defined) random variables.
Lemma 1. Let X
Gumbel 7) and let K N+. Define the Sum-of- Gammaa distribution as
SoG(K, T. = T
Σ (Gamma(1/xx
log(s)
(13)
where s E N+ and Gamma(a, 3) is the Gamma distribution with shape and scale 3. and let SoG(k, 7) lim (KG(K,, T, S). Then we have that X Ei=1 E,, with E, SoG(K, T).
Based on Lemma I we can show that for exponential family distributions where every z E C has design perturbations of (z, following Gumbel distribution. Theorem 1. Let p(z; 0) be a discrete exponential family distribution with integer polytope and temperature т. Assume that if z E C then (z, 1) = 1: for some constant E N Let 0 be the
exactly k non-zero entries
perturbation obtained by [0]; = [0]; + E, with Ej ~ SoG(k, from Eq. 13
Then, Vz E C we
have that (z, 0) (z, 0) + E(z), with e(z) Gumbel(O, 7).
Many problems such as k-subset selection, travel- ing salesman, spanning tree, and graph matching strictly satisfy the assumption of Theorem [I We can, however, also apply the strategy in cases where the variance of (Z, 1) is small (e.g. short- est weighted path). The Lum-of-Gamme pertur- bations provide a more fine fine-grained approach noise perturbations. For
0.
0.4
0.3
0.3
0.2
0.2
0.1
0.1
0.0
-5
15
0.0
10
Figure 2: Histograms for 10k samples where cach sample is (left) the sum of 5 e; ~ Gumbel(O, 1) or (right) the sum of 5 1,,
k = 1, we ob- tain the standard Gumbel perturbations. In con- trast to the standard Gun 1) noise, the pro- posed local Sum-of-Gamme, perturbations result in weights' perturbations that follow the Gum- bel distribution. Fig. [ shows histograms of 10k samples, where each sample is either the sum of 5 samples from Gumbel(O, 1) (the standard approach) or the sum of SoG(5, 1, 10) = Y/2 Ei= [Gamma(1/5, 5/1) - log((0)). While we still cannot sample faithfully from p(2; 0) the perturbations are not independent, we can counteract the problem of partially dependent perturbations by increasing the temperature 7 and, therefore, the variance of the noise distribution. We explore and verify the importance of tuning T empirically. In the appendix, also show that the infinite series from Lemma] can be well approximated by a finite sum using convergence results for the Mascherna series [Mortici 2010].
5 samples from
4 Target Distributions for Combinatorial Optimization Problems
In this section, we explore the setting where the discrete computational component arises from 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. (L) (where fu is the identity) and the supervision is available in terms of examples of either optimal solutions [e.g. Pogančić et al. 2019 or optimal cost coefficients (conditioned on the inputs) [e.g. Ellachtoub and Grigas 2020]. We have seen in Example that we can naturally associate to each IL. a probability distribution (see Eq. 3 with 0 given by the negative cost coefficients e of the ILP and C the integral polytope. Letting + 0 is equivalent to taking the MAP the forward pass. Furthermore, in Section 3.1 we showed that the I-MLE framework subsumes a recently propose method by [Pogančić 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 g(2;0) 0') be the distribution p(2; 0'), with parameters
0
if | =0 otherwise.
[0
=
(14)

---

In the first CO setting, we observe training data D { 8,)ر where y3 E C and the loss l measures a distance between discrete 2, p(2; 0,) with 0 (@)) and given optimal solution of the ILP ÿj. An example is the Hamming loss ε (Pogančic et al., 2019] defined as (z y) o (1 y) + (1 =)), where a denotes the Hadamard (or entry- product.
Fact 1. If one uses fH. then I-MLE with the target distribution of Eq. (14) and p(c) = 6n is equivalent to the perceptron-rule estimator of the MLE objective between pl (z (ă, and Dz-
It follows that the method by Pogančić et al. (011 returns, for a large enough X, the maximum- likelihood gradients (scaled by 1/2) approximated by the perceptron rule. The proof' are given in Appendix
In the second CO setting, observe training data D ((t, -- where è, is the optimal cost
conditioned on input zz.. Here, various authors [e.g. Elmachtoub and Grigas 2020 Mandi et al.
2020 Mandi and Guns 2020] use as point-wise loss the regret (n(0,c) c) = c (z ! (c(c) where
(0) is state sampled from p(z; 0) (possibly with temperature т + 0, that is, a MAP state) and 2 (c) E x" (c) is an optimal state for C.
Fact 2. If one uses l then I-MLE with the target distribution of Eq. (14) is equivalent to the perturb-and-MAR estimator of the MLE objective between p(z; h.p (+z)) and p(z; -c,).
This last result also implies that when using the target distribution q from (4) in conjunction with the regret, I-MLE performs -ikeehonl 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 defined by Eq. [4), we have that {(2, 9) = 0 for 2 E MAP( Therefore, {(2, D) = 0 S (x)p(=;0) [e((2, f)). meaning that the inequality of Eq. (5) is satisfied for O.
5 Related Work
Several papers address the gradient estimation problem for discrete r.v.s, many resorting to relaxations. Maddison et al. [2017), J. Jang et al. [2017] propose the numbel-softmaal distribution to relax categorical r.v.s: Paulus et al. [2200] study extensions to more complex probability distributions. The concrete distribution (the Gumbel-softma distribution) is only directly applicable to categorical variables. For more complex distributions, one has to come up with tailor-mad made relaxations or use the straight- through or score function estimators (see for instance Kim al. [016), Grover al.) (2019) In our experiments, we compare with the Gumbel-softtax 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 (1-subset) VAE trained with the Gumbel-softnaa gradient estimator, Tucker et al. [017]. Grathwohl al.) [ develop parameterized control variates (the former was named REBAR) based on continuous relaxations for the cooreffuncionn estimator. In contrast, we focus explicitly 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 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] ct al. [2118) is an approach to structured prediction and latent variables, replacing the exponential distribution (specifically, the softmax) with sparser distribution. Similar to work, it only presupposes the availability of a MAP oracle. LP-SparseMAPP [Niculae and Martins 2020 is an
extension that uses relaxation of the optimization problem rather than 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 Wiilder et al. 2019 Elmachtoub and
Grigas 2020, Ferber et al. 2020 Mandi and Guns 2020) relax ILPs by adding LT, L or log-barrier regularization terms and differentiate through the KKT conditions deriving from the application of the cutting plane or the interior-poin methods. These approaches are conceptually linked to
techniques for differentiating through smooth programs [Amos and Kolter 2017 Donti et al. 2017
7

---

Agrawal et al. 2019 Chen et al. 2020 Domke 2012
Franceschi et al. [018]] that arise not only
in modelling but also in hyperparameter optimization and meta-learning. Pogančié et al. (2019). Rolinek et al. 2220 Berthet et al. [020) propose methods that are not tied to a specific IIP solver. As we saw above, the former two, originally derived from 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 CO- related 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 [Bnnii et al. 2020 and references therein.
Direct Loss Minimization [DLM, McAllester et al. 2010 Sonn et al., 2016] is also related to our work, but the assumption there is that examples of optimal states 전 are given. Lorberbom al. [019] 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-ande and-MAP [Papandreou and Yuille, 2011]. The theory of ppeturb-and -nd--nd--a- MAP used to derive general ceneral-purpose upper bounds for log oog-partiiion Shpakova and Bach, 2016].
functions Hazan and Jaakkola 2012
6 Experiments
The set of experiments can be divided into three parts. First, we analyze and compare the behavior of I-MLEE with (i) the score function and (ii) the ttraight-through estimator using toy problem. Second, we explore the latent variable setting where both h, and fu in Eq. ① 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 E {0, 1310. We set the loss to L(O) = Eg-p(n,)) [I] 2 - -bl1", where b is a fixed vector sampled from (O, 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.
(PaM) with Gumbel and SoG(1, 5, 10) noise, respectively. The SFE uses faithful samples and exact marginals (which is fcasible 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. (0)) and the learning rate) and exhibit smaller variance (see Fig. 6). Additional details and results in Appendix C.I
[ For STE and I-MLE, we use Beturb-and-Mar
Learning to Explain. The BEER ADVOCADVO dataset (McAuley et al. 2012] consists of free-text reviews and ratings for 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. she 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 val 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 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
'Hyperparameters are optimized against L for all methods independently, Statisties are over 100 runs. We found STE slightly better with Gumbel rather than SoG noise. SFE failed with all tested PaM strategies.

---

Test MSE
Subset Precision
2.5
Method
2.0
I-MLL
Mean
Std. Dev.
Mean
Std. Dev.
Opiimalty
1.5
k: 10
1.0
L2X (1 = 0.1)
6.68
1.08
26.65
9.39
0.5
SoftSub (1 = 0.5)
2.67
0.14
44,44
2.27
0.9
STE (T 30)
4,44
0.09
38.93
0.14
20
40
Optimization
[-MLE MAP
4.08
0.91
14.55
0.04
I-MLE Gumbel
2.68
0.10
39.28
2.62
1-MLE 5oG Gum. (means)
I-MLE (r = 30)
2.71
0.10
47.98
2.26
0.0
Leaming rate
0.2
0.4
L2X 0.1)
5.75
0.30
33.63
6.91
0.6
SoftSub (f = 0.5)
2.57
0.12
54.06
6.29
I-MLE (T 5)
2.62
0.05
54.76
2.50
-1.0
A: 15
Lambda
1.XX (f = 0.1)
7.71
0.64
23.49
10.93
Figure 3: Top: Gradient-based optimiza-
SoftSub (f = 0.5)
2.52
0.07
37.78
1.71
tion of with various estimators, Bottom: Mean difference of the final value of T. be- tween I-MLE with SoG or Gumbel p(c). varying À and the learning rate (blue better SoG).
1-MLE (T = 30)
2.91
0.18
39.56
2.07
Table 1: Detailed results for the aspect AROMA. Test MSE and subset precision, both x100, for E (5, 10, 15].
cct 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 0 of the distribution p(z; 0) over k-hot baselines L2X (Chen] et al. 2018) and SoftSub [Xie and Ermon 2019]. We also compare the straight- through estimator (STE) with ofm-of-Gamma (SoG) perturbations. We used the standard hyperparameter settings of Chen al. 2018 and choose the temperature parameter E {0.1, 0.5, 1.0, 2.0). For I-MLE choose À E 1O 10², 10"). while for both I-MLE and STE we choose T E [k: 2k, 3k) based on the validation MSE We used the standard Adam settings. We trained separate models for each aspect
binary latent masks with & E 15, 10, 15). We compare relaxation-base
using MSE as point-wise loss
Table D lists detailed results for the aspect AROMA. I-MLE's MSE values are competitive with those of the best baseline, and its subset precision 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 coan-of-Gamma perturbation introduced in this paper outperforms the standard local Gumbel perturbations. More details and results can be found the appendix.
Discrete Variational Auto-Encoder. We evaluate various perturbation strategies for 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 probability distribution over k-subsets of (or top-k assignments too) binary vectors of length 20. The special case = 1 is cquivalent to categorical variable with 20 categories. For k 1, use I-MLE using the class of PID target distributions of Eq. (10) and compare various perturb-and anduur---n--APPP 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 12-784). We do
temperature annealing. Using Eq. ( with S =
@ either (u 1) perturbations (the standard approach or Sum-of-Gaman temperature of 10. We run 100 epochs and record the loss 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
(SoG) perturbations
"ncreasing the temperature of Gumbel(O, r) samples increased the test loss.

---

Discrete 10 subset VAE (1 MLE)
Discrete 10 sabbett VAE (STE)
Discrete 1-subeet VAE (GSMT)
BCE lase KL. diveggence
UCE Jues + KL. divergerre
250
BCE bos 4. KL. diverpaoo
E-LE Gumbaa (ouns)
STE (Gumbell) (SoG)
Gumbel-SMT (STE)
120
1MLE SoG (ouns
120
Chmbee-SMT (reesedd
1.00-
100+
210
80-
0
25
30 Rpocha
75
100
100 0
25
50
75
100
Epochs
Rpocis
Figure 4: Plots of the sum of the binary reconstruction loss and the KL. divergence function of the number of epochs (lower is better). (Left) Discrete 10-subset VAE trained with I-MLE with À = 10 (I-MLE). (Center) Discrete 10-subset VAE trained with the straight-through estimator (STE). (Right) Discrete 1-subset VAE using the Gumbel softmax trick (GSMT). The down-up up down artifact is due to temperature annealing. Sum-of- (SoG) perturbations have the lowest test loss for the 10-subset VAEs. For À = 10 and SoG perturbations, the test loss similar to that of the categorical (1-subset) VAE trained with the Gumbel softmax trick.
STE, which does 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 Combi- natorial Solvers. In these ex- periments, proposed by Pogančid al. [2019], the training datasets consists of 10,000 examples of randomly generated images of ter-
Results for the Warcraft shortest path task. Reported is the accuracy, i.c. percentage of paths with the optimal costs. Standard deviations are over five runs.
Table 2:
K
I-MLE (1-T1)
I-MLE (M-M)
BB
DPO
12
97.2 1 0.5
95.2 1 0.3
0.3
95.2 1 0.7
94.8 1 0.3
rain maps from the Warcraft tile
18
95.8 1 0.7
94.4 1 0.5
94.7 0.4
92.3 1 0.8
set [Guyomarch] 2017]. Each ex- ample has an underlying K x K grid whose cells represent terrains
24
94.3
1 1.0
93.2 1 0.2
0.3
91.5 0.4
30
93.6
0.4
93.7
0.6
93.6 91.5 0.5
with a fixed cost. The shortest (minimum cost) path between the top-left and bottom-right cell in indicator matrix and serves as the target output. An image of the terrain
the grid is encoded map is presented to CNN, which produces K x K matrix of vertex costs. These costs are then given to Dijkstra's 's algorithm (the MAP solver) compute the shortest path. We closely follow the evaluation protocol of Pogančie et al. [019]. We considered two instantiations of I-MLE: one derived from Fact (M-M in Table 2 using EH and one derived from Fact (14-p1 using εR, with p(c) = 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 (Pogančié al. (2019) (BE and Berthet
[2020) I-MLE obtains results comparable to (BB) with M-M and is more accurate with μ-p We believe that the ft-14 advantage may partially due to an implicit form of data augmentation since we know from Fact 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 (H1-4), and additional details are available in Table
(DPO). The results are listed in Table 2
7 Conclusions
I-MLE is an efficient, simple-to-implementt, and general-purpose framework for learning hybrid models. I-MLE is competitive with relaxation-basse 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 V is not available, studying the properties (including the bias) of the proposed estimator, developing adaptive strategies for and X, and integrating and testing I-MLE in several challenging application domains.
Noto that this is the same using I-MLE with PID target distribution form Eq. 10 and p(c) 80.
10

---

References
A. Agrawal, B. Amos, Barratt, Boyd, S. Diamond, and Z. Kolter. Differentiable convex optimization layers. arXiv preprint arXiv: /2.390000240, 2019.
B. Amos and J. Z. Kolter. Optnet: Differentiable optimization as layer in neural networks. In International Conference on Machine Learning, pages 136-145. PMLR, 2017.
Y. Bengio, N. Léonard, and A. Courville. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: MxXiv:/300.. 343,2 2013.
Y. Bengio, A. Lodi, and A. Prouvost. Machine learning for combinatorial optimization: method- ological tour d' d'horizon. European Journal of Operational Research, 2020.
Q Berthet, M. Blondel, O. Teboul, M. Cuturi, J. Vert, and R. Bach. Learning with differentiable pertubed optimizers. In NeurIPS, 2020.
J. Chen, I.. Song, M. Wainwright, and M. Jordan. Learning to explain: An perspective on model interpretation. In International Conference on Machine Learning, pages 883-892. PMLR, 2018.
X. Chen, Y. Zhang, C. Reisinger, and L. Song. Understanding deep architecture with reasoning layer. Advances in Neural Information Processing Systems, 33, 2020.
G. M. Correia, V. Niculae, W. Aziz, and A. F. Martins. Efficient marginalization of discrete and structured latent variables via sparsity. Advances in Neural Information Processing Systems, 2020.
Domke. Implicit differentiation by perturbation. In Advances Neural Information Processing Systems 23, pages 523-531. 2010.
J. Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics, pages 318-326. PMLR, 2012.
P. L. Donti, B. Amos, and J. Z. Kolter. Task-based end-to-end model learning in stochastic optimiza- tion. Advances in Neural Information Processing Systems, 2017.
A. N. Elmachtoub and Grigas. Smart "predict, then optimize", 2020.
R. Evans and E. Grefenstette. Learning explanatory rules from noisy data. Journal of Artificial Intelligence Research, 61:1-64, 2018.
A. Ferber, B. Wilder, B. Dilkina, and M. Tambe. Mipaal: Mixed integer program as a layer. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 1504-1511, 2020.
I. Franceschi, P. Frasconi, S. Salzo, R. Grazzi, and M. Pontil. Bilevel programming for hyperparam- eter optimization and meta-learning. In International Conference on Machine Learning, pages 1568-1577. PMLR, 2018.
I. Franceschi, M. Niepert, M. Pontil, and X. He. Learning discrete structures for graph neural networks. In International conference on machine learning, pages 1972-1982. PMLR, 2019.
W. Grathwohl, D. Choi, Y. Wu, G. Roeder, and D. Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. ICLR, 2018.
A. Grover, E. Wang, A. Zweig, and S. Ermon. Stochastic optimization of sorting networks via 2019.
continuous relaxations. arXiv preprint
J. Guyomarch. Wareraft II Open-Source Map http:///gtthbb.cm////r/aaaaaaaa 2017.
D. Hafner, T. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020.
T. Hazan and T. Jaakkola. On the partition function and random maximum a-posteriori perturbations. In Proceedings of the 29th International Coference on International Conference on Machine Learning, pages 1667-1674, 2012.
11

---

K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In CVPR, pages 770-778. IEEE Computer Society, 2016.
E. Jang, S. Gu, and B. Poole. Categorical reparameterization with gumbel-softmax. ICLR, 2017.
N. L. Johnson and N. Balakrishnan. Advances in the theory and practice of statistics, 1998.
C. Kim, A. Sabharwal, and S. Ermon. Exact sampling with integer linear programs and random perturbations. In Thirtieth AAAI Conference on Artificial Intelligence. 2016.
W. Kool, H. van Hoof, and M. Welling. Estimating gradients for discrete random variables by sampling without replacement. ICLR, 2020.
Y. LeCun, S. Chopra, R. Hadsell, M. Ranzato, and F. Huang. A tutorial energy-based learning. Predicting structured data, 1(0), 2006.
T. Lei, R. Barzilay, and Jaakkola. Rationalizing neural predictions. In EMNLP, 2016.
R. Liu, J. Regier, N. Tripuraneni, M. Jordan, and J. Mcauliffe. Rao-blackwellized stochastic gradients for discrete distributions. In International Conference on Machine Learning, pages 4023 4031-4031. PMLR, 2019.
G. Lorberbom, A. Gane, T. Jaakkola, and T. Hazan. Direct optimization through argmax for discrete variational auto-encoder. In Advances in Neural Information Processing Systems, pages 6203-6214 2019.
C. J. Maddison, A. Mnih, and Y. W. Teh. The concrete distribution: A continuous relaxation of discrete random variables. ICLR, 2017.
J. Mandi and T. Guns. Interior point solving for 1p-based prediction-optimiistion. In Advances in Neural Information Processing Systems, 2020.
Mandi, E. Demirović, P. J. Stuckey, and Guns. Smart predict-and-optimizz for hard combina- torial optimization problems. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pages 1603-1610, 2020.
D. A. McAllester, T. Hazan, and J. Keshet. Direct loss minimization for structured prediction. In Advances in Neural Information Processing Systems, volume 1, page 3, 2010.
McAuley, Leskovee, and Jurafsky. Learning attitudes and attributes from multi-aspect aspect reviews. 2012 IEEE 12th International Conference on Data Mining. pages 1020-1025, 2012.
V. V. Misić and G. Perakis. Data analytics in operations management: A review. Manufacturing & Service Operations Management,
2020.
C. Mortici. Fast convergences towards Euler Masch-Mas constant. Computational & Applied Mathematics, 29, 00 2010.
K. Murphy. Machine learming: probabilistic perspective. MIT press, 2012.
V. Niculae and A. F. Martins. Lp-sparsemap: Differentiable relaxed optimization for sparse structured prediction. In ICML, 2020.
V. Niculae, A. F. T. Martins, M. Blondel, and C. Cardie. Sparsemap: Differentiable sparse structured inference. In ICML, 2018
G. Papandreou and A. L. Yuille. Perturb-and-map random fields: Using discrete optimization to learn and sample from energy models. In 2 International Conference Computer Vision, pages 193-200, 2011.
M. B. Paulus, D. Choi, D. Tarlow, A. Krause, and C. J. Maddison. Gradient estimation with stochastic softmax tricks. arXiv preprint arXiv:2006. 08063, 2020.
M. Pogančić, Paulus, Musil, G. Martius, and M. Rolinek. Differentiation of blackbox combinatorial solvers. In International Conference Learning Representations, 2019.
12

---

L. D. Raedt, K. Kersting, S. Natarajan, and D. Poole. Statistical relational artificial intelligence: Logic, probability, and computation. Synthesis Lectures on Artificial Intelligence and Machine Learning, 10(2):1-189, 2016.
M. Rolínek, P. Swoboda, D. Zietlow, A. Paulus, V. Musil, and G. Martius. Deep graph matching via blackbox differentiation of combinatorial solvers. In ECCV, 2020.
J. Schulman, N. Heess, T. Weber, and P. Abbeel. Gradient estimation using stochastic computation graphs. In Proceedings of the 28th International Conference on Neural Information Processing Systems- Volume 2, pages 3528-3536, 2015.
T. Shpakova and F. Bach. Parameter learning for sugerrooro distributions. In Advances in Neural Information Processing Systems, volume 29, pages 3234-3242, 2016.
Y. Song, A. Schwing, Urtasun, et al. Training deep neural networks via direct loss minimization. In International Conference on Machine Learning, pages 2169-2177, 2016.
G. Tucker, A. Mnih, C. J. Maddison, D. Lawson, and J. Sohl-Dickstein Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. Advances in Neural Information Processing Systems, 2017.
M. Wainwright and M. Jordan. Graphical models, families,,, and variational inference. Now Publishers Inc, 2008.
P.-W. Wang, P. Donti, B. Wilder, and Z. Kolter. Satnet: Bridging deep learning and logical reasoning using a differentiable satisfiability solver. In International Conference on Machine Learning, pages 6545 6554. PMLR, 2019.
B. Wilder, B. Dilkina, and M. Tambe. Melding the data-decisions pipeline: Decision-focde focused learning for combinatorial optimization. Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 1658-1665, 2019.
Xi'an. x?
Which pdf of leads to gumbel distribution of the finite-size average of Cross Validated, 2016. URL https://stats wttessssttcecceeaaseeeeeeeeee com/q/21487 httmps///tatsstaacecccaagn..moo
(version: 2016-05-27).
S. M. Xie and S. Ermon. Reparameterizable subset sampling via continuous relaxations. In IJCAI, 2019.
13


Reviews:
[{"review_id": "jmlh-ZeLSVu", "reviewer": "Reviewer_JsU3", "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.\n\nThe 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. \n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper considers an approximation to maximum likelihood learning in models that involve a discrete latent variable ($z$) bottleneck between the input features ($x$) and the output ($y$), where $z$ is modeled as being from the discrete exponential family where a parametrized function approximator predicts the natural family parameters for the exponential family.  The last stage of the model maps the realizations of this discrete random variable to the model output. It is assumed that we are given a loss function that depends only on $z, y$, with an overall loss for a given $x, y$ pair being the expectation over the $z$ with a distribution defined by the natural parameters specified by $x$. \n\nTraining the model involves learning $\\omega = (u, v)$, where the first stage parameters are $v$ and the final stage params are $u$. While a standard score function estimator approach may be enough for learning $u$ (since the partials wrt $u$ are an expectation of gradients, whose distribution does not depend on $u$), the variance for this approach is too high to obtain gradients wrt $v$.  To address this, a target distribution decoupled from $\\theta$ is considered, which is only required to strictly improve on the original objective and then the goal is to optimize a maximum likelihood gradient wrt this target distribution. The resulting objective is called the implicit-MLE loss, but clearly a key aspect of this involves defining the target distribution to learn from. \n\nQuestions/comments for the authors to clarify:\n\n- A key property that characterizes these target distributions is given by the Equation between lines 116 and 117 (I would strongly encourage the authors to give this Equation a reference as this seems very critical to me).  \n$$E_{\\hat{z} \\sim q(z; \\theta')} l(f_u(\\hat{z}), \\hat{y})  \\le E_{\\hat{z} \\sim p(z; \\theta)} l(f_u(\\hat{z}), \\hat{y})$$\nBut the above equation does not seem to be actually verified in any of the proposals constructed by the authors both in Sections 3.1 or in Section 4. Could the authors please elaborate? \n\n- In Eq (8), the target $q(z, \\theta')$ definition seems ambiguous as the parameters on the RHS depend on $z$, which is the variable whose distribution is being specified. Could you please clarify the precise definition for $q$ in this Equation?\n\n- In Eq (12), it appears that target distribution $q$ has the same form as $p$, except that the parameters are $\\theta'$ instead of $\\theta$. If that is the case, isn't the solution to the maximum likelihood problem simply $\\theta'$? \n\n- The authors propose the core idea as a fairly general recipe/meta-algorithm (implicit-MLE), however its specific instantiation via the perturbation based implicit differentiation seems to be an equally important piece of the proposal, and I'm not able to assess which of these is more important even though the presentation and the title suggests that its the former.  \n\n- In figure 1, the output needs to be $y$ rather than $L$?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper is not directly relevant for societal impact issues due to its theoretical nature. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "T2YiGuXzPBH", "reviewer": "Reviewer_MvPH", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\n- In an attempt to optimize computation graphs with discrete latent variables, the paper revisits the classic MAP-and-perturb scheme for sampling via optimization. While these strategies are known, their use and extension in an end-to-end pipeline is novel to this work. Moreover, this scheme presents an interesting departure from the predominant strategies: score function estimation with control variates and continuous relaxations+reparameterization.\n\n  While I am generally in favor of fresh takes, my general worry is that in doing so, the paper adds more complexity to both the forward and backward pass. In the forward pass, there is a need for a (black-box) MAP solver and the backward pass relies on a new biased estimate of the gradient. There are also additional hyperparameters, such as lambda or the choice of q more generally, that are not needed for alternative baseline gradient estimators. Having said that, the authors are able to demonstrate good performance, thereby lending credibility that the extra machinery is perhaps worthy of further exploration.\n\n- The writing is mostly clear, except for a key piece of intuition: \nIn L119-120, the authors say that \"The intuition is that making p more similar to q will also reduce the loss\" --- why is that the case? does it hold universally for all q's that satisfy the equation following L116?\n\n- The results in Tables 1 and 2 use different q's (eq. 8 and 12). What would the performance be if the q's for Tables 1 and 2 were switched?\n\n- I would also move Algorithm 1 to the main paper. It summarizes the algorithm well enough to aid clarity.\n\n- The estimator name, I-MLE, is somewhat too general. The connections with MLE hold only for very specific choice of q's and loss functions (sec 4). Or is there a broader connection that holds more generally?\n\n- Could the SoG noise distribution also improve the default Gumbel-softmax relaxation? In general, it could be insightful to report results from a SoG-softmax relaxation to assess the general-purpose utility of this noise distribution.\n\n- Empirically, how does this method scale with the number of categories for the discrete VAE?\n\n- There are some missing related works:\n\nKim, C., Sabharwal, A. and Ermon, S.Exact sampling with integer linear programs and random perturbations. AAAI, 2016 -- Guarantees for MAP and perturb based methods\n\nGrover, A., Wang E., Zweig A., and Ermon, S. \"Stochastic optimization of sorting networks via continuous relaxations.\", ICLR 2019  -- Continuous relaxation for combinatorial discrete distributions\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "HdPyqwNBaCT", "reviewer": "Reviewer_pyap", "summary": "The paper proposes a method to compute gradients while training stochastic models with discrete latent variables. The proposed method assumes that latents are drawn from a discrete exponential family distribution: it defines a target distribution over the discrete latents, and uses the target distribution to compute a likelihood which can then be minimised to learn the parameters of the model.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The novel approach to estimating gradients for models with discrete latents, as well as the evidence that it outperforms score-function estimators appears to be compelling. Overall the paper was clear and well-written.\n\nHowever, sections 3 and 4 were quite hard to follow since the discussion jumps back and forth between the gradient estimator / loss function, and the perturb-and-map distributions. For example, $z=MAP(\\theta + \\epsilon)$ is already defined and used extensively in the introduction to section 3 and in section 3.1, before perturb-and-MAP distributions are discussed in detail and section 4. These sub-sections could be re-arranged and re-written for clarity.\n\nMoreover, in section 3 the loss $l(f_u(z), y)$ is replaced with $L(\\theta, \\theta')$. The text states that it is intuitive that minimising the latter results in minimising the former.  The text also suggests that the construction of $q(z, \\theta')$, the target distribution is a hyperparameter. However, from section 3.1, it seems that the replacement of the loss is possible only when the target distribution $q$ is specifically constructed such that the $l(f_u(z), y)$ is minimised, and that this in turn, hinges on sampling the latent variables using the perturb-and-MAP technique. From this it seems that the \"hyperparameter\" really is the choice of the noise distribution for the perturbation step, and that the construction of the loss / target distribution is otherwise non-trivial. If this is indeed the case, this should be made clear in the text.\n\nThe conclusion claims that I-MLE outperforms relaxation-based techiniques -- this is not obvious from the results.  Although the Gumbel distribution is used in conjunction with the MLE estimator, this does not appear to be the same as the approach in Tucker et al (REBAR) or using concrete relaxed distributions as in Maddison et al. It would be useful to have comparisons between I-MLE and these methods as well.\n\nMinor comments:\n1) Figure 1 shows $f_u$ mapping $z$ to $L$ rather than $y$\n1) Figure 3 top plot shows optimisation gap versus optimisation steps. Is the optimisation gap the same as the loss function? If not, this performance measure should be explained.\n2) Example 3 in section 2 reads more like a theorem -- it would be good to have a clear explanation.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": " ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "BOd5pWbV938", "reviewer": "Reviewer_e8wB", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The problem of differentiating through discrete representations (e.g., representations of discrete probability distributions) is certainly very relevant.\n\nThe paper borrows some ideas from methods such as perturb-and-MAP, and applies them to the setting described in eq. (2) (input feature map -- discrete distribution -- output feature map). The authors contrast their method to be potentially more widely applicable than, say, methods based on smooth relaxations. While this may be true for some classes of problems, I find the requirement to be able to perform MAP inference to be quite a big constraint. Efficient MAP inference is only possible for quite a limited number of discrete distribution families.\n\nI think the biggest problem with this paper is its exposition. This may be due to the fact that I am not an expert is some of the topics covered, but I find that section 3 could have been presented in a much clearer way. As it is, it is quite hard to judge the specific contributions of the paper. The experimental section is better, but even there some parts could have been motivated in a better way (e.g., the VAE application).\n\nFurther comments:\n- How is this setup related to energy-based models (e.g., see [1])? While I see some basic differences, the form of distributions, and some of the problems considered (e.g., computing gradients of discrete representations) may be very similar.\n- Table 1 only reports only results on the \"aroma\" aspect? How do the other aspects compare?\n\n&nbsp;\n### References\n[1] Song and Kingma. How to Train Your Energy-Based Models. ArXiv, 2021.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents new methods for gradient estimation in computation graphs with certain types of discrete stochastic variables. Reviewers all voted accept (scores 6, 6, 6, 8). The consensus was that the paper has fresh ideas and there is evidence in favor of the method working well, but it also has some tradeoffs, e.g., in terms of added complexity in the computations. Several reviewers raised (non-fatal) concerns about clarity, which the authors are encouraged to address in the final revision. The meta-reviewer recommends accept.
Decision: Accept (Poster)