[论文解读] Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions

Mathias Niepert, Pasquale Minervini+1 more|arXiv (Cornell University)|Jun 3, 2021

Machine Learning and Algorithms被引用 16

一句话总结

本文提出隐式最大似然估计（I-MLE），一种通用框架，通过利用损失函数和最可能状态导出的目标分布来近似梯度，实现对离散指数族分布的反向传播。I-MLE 在性能上优于基线方法（如直通估计器），且在使用新型伽马分布和噪声分布进行扰动-最大后验（perturb-and-MAP）采样时，与基于松弛的方法具有竞争力。

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.

研究动机与目标

• 解决具有离散潜在变量或组合优化组件的神经网络训练问题，标准反向传播因输出不可微或不连续而失效。
• 开发一种无需特定问题松弛且不依赖完整概率推理或忠实采样的通用学习框架。
• 实现结合离散指数族分布与可微神经组件的模型的端到端训练，尤其适用于涉及组合求解器或离散潜在变量的场景。
• 通过仅依赖计算最可能状态（MAP 推断）的能力来简化梯度计算，使方法具备可扩展性和广泛适用性。

提出的方法

• I-MLE 通过构建一个目标分布 $ q $ 来计算梯度，该分布近似前向传播过程中观察到的最可能状态的经验分布。
• 该框架采用基于扰动的隐式微分，通过向参数添加噪声以生成样本，并利用扰动-最大后验策略估计梯度。
• 提出了一类新型的伽马分布和（Sum-of-Gamma, SoG）噪声分布，以更好地近似 Gumbel-max 采样中的边缘分布，从而提升梯度质量。
• 提出了两种类型的目标分布：一种基于 MAP 解（M-M），另一种基于损失加权分布的均值（μ-μ），两者均用于定义隐式梯度估计器。
• 该方法避免了平滑松弛，且无需访问完整概率分布，仅依赖 MAP 推断和损失反馈。
• I-MLE 在涉及组合求解器的特定场景下可退化为标准最大似然估计，从而验证了其理论一致性。

实验结果

研究问题

• RQ1 我们能否在不依赖问题特定平滑松弛或梯度近似的情况下，训练包含离散组件的混合神经模型？
• RQ2 当梯度在几乎所有位置均为零时，如何有效实现对离散指数族分布的反向传播？
• RQ3 在离散潜在变量模型中，何种噪声分布能为扰动-最大后验采样提供更优的梯度估计？
• RQ4 I-MLE 是否能在离散潜在变量学习中实现与基于松弛的方法（如 Gumbel-Softmax）相当或更优的性能？
• RQ5 I-MLE 是否能有效实现涉及黑箱组合求解器（如最短路径计算）的模型的端到端训练？

主要发现

• 在离散 VAE 中，I-MLE 显著优于直通估计器（STE），在 10-子集和 1-子集任务中均实现了更低的测试损失。
• 在 10-子集 VAE 中，I-MLE 使用伽马分布和（SoG）噪声分布时，尽管样本方差更高，但测试损失仍低于 Gumbel-Softmax。
• 在《魔兽争霸》最短路径任务中，I-MLE 使用 μ-μ 目标分布时，在 K=12 时达到 97.2% 的准确率，优于 BB 和 DPO 方法。
• I-MLE 使用 μ-μ 配置时，在 K=30 时达到 93.7% 的准确率，匹配或超过 BB 和 DPO 在所有网格尺寸下的表现。
• 该方法在训练动态中表现出更快的收敛速度，表明其优化稳定性优于基线方法。
• 在涉及组合求解器的场景中，I-MLE 简化为显式最大似然估计，确认了其理论一致性。