torch-imle

这是一个简洁且独立的PyTorch库，实现了我们在NeurIPS 2021论文《Implicit MLE: Backpropagating Through Discrete Exponential Family Distributions》中提出的I-MLE梯度估计器。

该仓库包含一个库，用于将任何组合黑盒求解器转换为可微分层。NeurIPS论文中所有实验的复现代码都可在NEC欧洲实验室的官方仓库中找到。

概述

隐式最大似然估计（I-MLE）使得在标准深度学习架构中包含离散组合优化算法（如Dijkstra算法或整数线性规划求解器）成为可能。I-MLE的核心思想是定义一个隐式的最大似然目标，其梯度用于更新模型的上游参数。每个I-MLE实例需要两个要素：

一种从复杂且难以处理的分布中近似采样的方法，该分布由组合求解器在解空间上诱导，其中最优解具有最高的概率质量。为此，我们使用扰动映射（又称Gumbel-max技巧），并提出了一种针对特定问题的新型噪声扰动家族。
一种计算替代经验分布的方法：普通MLE减少当前分布与经验分布之间的KL散度。由于在我们的设置中无法获得经验分布，我们必须设计替代经验分布。这里我们提出了两种广泛适用且实践效果良好的替代分布家族。

示例

例如，让我们考虑一个简单游戏的地图，任务是找到从左上角到右下角的最短路径。较暗的区域成本较高，较亮的区域成本较低。在中间，你可以看到当我们使用提出的伽玛噪声分布之和来采样路径时会发生什么。在右侧，你可以看到每个格子的边际概率结果（每个格子成为采样路径一部分的概率）。

[图片1] [图片2] [图片3]

梯度和学习

假设最优最短路径是左侧的路径。从随机权重开始，模型可以通过梯度下降学习产生将导致最优最短路径的权重，方法是最小化生成路径与黄金路径之间的汉明损失。这里我们展示了训练过程中产生的路径（中间）和相应的地图权重（右侧）。

输入噪声温度设置为0.0，目标噪声温度设置为0.0：

[图片4] [图片5] [图片6]

输入噪声温度设置为1.0，目标噪声温度设置为1.0：

[图片7] [图片8] [图片9]

输入噪声温度设置为2.0，目标噪声温度设置为2.0：

[图片10] [图片11] [图片12]

输入噪声温度设置为5.0，目标噪声温度设置为5.0：

[图片13] [图片14] [图片15]

输入噪声温度设置为5.0，目标噪声温度设置为0.0：

[图片16] [图片17] [图片18]

所有动画都由这个脚本生成。

代码

使用这个库非常简单 -- 请参考这个示例。假设我们有一个实现黑盒组合求解器（如Dijkstra算法）的方法：

import numpy as np

import torch
from torch import Tensor

def torch_solver(weights_batch: Tensor) -> Tensor:
    weights_batch = weights_batch.detach().cpu().numpy()
    y_batch = np.asarray([solver(w) for w in list(weights_batch)])
    return torch.tensor(y_batch, requires_grad=False)

我们可以通过以下方式获得相应的分布和梯度：

from imle.wrapper import imle
from imle.target import TargetDistribution
from imle.noise import SumOfGammaNoiseDistribution

target_distribution = TargetDistribution(alpha=0.0, beta=10.0)
noise_distribution = SumOfGammaNoiseDistribution(k=k, nb_iterations=100)

def torch_solver(weights_batch: Tensor) -> Tensor:
    weights_batch = weights_batch.detach().cpu().numpy()
    y_batch = np.asarray([solver(w) for w in list(weights_batch)])
    return torch.tensor(y_batch, requires_grad=False)

imle_solver = imle(torch_solver,
                   target_distribution=target_distribution,
                    noise_distribution=noise_distribution,
                    nb_samples=10,
                    input_noise_temperature=input_noise_temperature,
                    target_noise_temperature=target_noise_temperature)

或者，使用简单的函数注解：

@imle(target_distribution=target_distribution,
      noise_distribution=noise_distribution,
      nb_samples=10,
      input_noise_temperature=input_noise_temperature,
      target_noise_temperature=target_noise_temperature)
def imle_solver(weights_batch: Tensor) -> Tensor:
    return torch_solver(weights_batch)

使用I-MLE的论文

Patrick Betz, Mathias Niepert, Pasquale Minervini, 和 Heiner Stuckenschmidt：《Backpropagating through Markov Logic Networks》，NeSy'20/21 @ IJCLR：第15届神经符号学习与推理国际研讨会

参考文献

@inproceedings{niepert21imle,
  author    = {Mathias Niepert and
               Pasquale Minervini and
               Luca Franceschi},
  title     = {Implicit {MLE:} Backpropagating Through Discrete Exponential Family
               Distributions},
  booktitle = {NeurIPS},
  series    = {Proceedings of Machine Learning Research},
  publisher = {{PMLR}},
  year      = {2021}
}