Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: 
- 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.

  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.

- The writing is mostly clear, except for a key piece of intuition: 
In 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?

- 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?

- I would also move Algorithm 1 to the main paper. It summarizes the algorithm well enough to aid clarity.

- 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?

- 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.

- Empirically, how does this method scale with the number of categories for the discrete VAE?

- There are some missing related works:

Kim, C., Sabharwal, A. and Ermon, S.Exact sampling with integer linear programs and random perturbations. AAAI, 2016 -- Guarantees for MAP and perturb based methods

Grover, A., Wang E., Zweig A., and Ermon, S. "Stochastic optimization of sorting networks via continuous relaxations.", ICLR 2019  -- Continuous relaxation for combinatorial discrete distributions

limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper proposes a meta algorithm for models involving a discrete bottleneck by making use of a specific target distribution with respect to which a max-likelihood loss is optimized in order to avoid having to compute the high variance score function gradients of the original objective.  This target distribution is characterized by being a local improvement to the overall objective.  To avoid marginalization on the discrete variables, the paper further consider approximations to the proposed I-MLE framework using the perturb-and-MAP approach, which repeatedly solves MAP subroutines for noise-perturbed parameters to approximate the marginals. In doing this, it also proposes a decomposition of Gumbel noise terms into a sum of Gamma distributions.

The ideas are evaluated empirically both on synthetic problems as well as benchmarks involving discrete VAEs and other combinatorial problems requiring flexible discrete latent variable modeling. 

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$. 

Training 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. 

Questions/comments for the authors to clarify:

- 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).  
$$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})$$
But 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? 

- 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?

- 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'$? 

- 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.  

- In figure 1, the output needs to be $y$ rather than $L$?

limitations_and_societal_impact: The paper is not directly relevant for societal impact issues due to its theoretical nature.
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: 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.
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.

However, 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.

Moreover, 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.

The 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.

Minor comments:
1) Figure 1 shows $f_u$ mapping $z$ to $L$ rather than $y$
1) 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.
2) Example 3 in section 2 reads more like a theorem -- it would be good to have a clear explanation.

limitations_and_societal_impact:
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: The problem of differentiating through discrete representations (e.g., representations of discrete probability distributions) is certainly very relevant.

The 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.

I 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).

Further comments:
- 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.
- Table 1 only reports only results on the "aroma" aspect? How do the other aspects compare?

&nbsp;
### References
[1] Song and Kingma. How to Train Your Energy-Based Models. ArXiv, 2021.
limitations_and_societal_impact: Yes.
review_title: null

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)