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.

NeurIPS 2021 camera-ready; repo: https://github.com/nec-research/tf-imle

References in corpus (13)

• Categorical Reparameterization with Gumbel-Softmax
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
• Gradient Estimation Using Stochastic Computation Graphs
• REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models
• Task-based End-to-end Model Learning in Stochastic Optimization
• Deep Graph Matching via Blackbox Differentiation of Combinatorial Solvers
• Training Deep Neural Networks via Direct Loss Minimization
• On the Partition Function and Random Maximum A-Posteriori Perturbations
• Gradient Estimation with Stochastic Softmax Tricks
• Rao-Blackwellized Stochastic Gradients for Discrete Distributions
• Estimating Gradients for Discrete Random Variables by Sampling without Replacement
• Parameter Learning for Log-supermodular Distributions