训练速度最高100倍提升！基于PyTorch实现的可微逻辑门网络开源

选自arXiv

作者：Felix Petersen等

机器之心编译

编辑：袁铭怿

本文通过逻辑门组合学习来探索面向机器学习任务的逻辑门网络。这些网络由 AND 和 XOR 等逻辑门组成，为了实现有效训练，本文提出可微逻辑门网络，一种结合了实值逻辑和网络连续参数化松弛的架构。

随着神经网络的成功应用，各项研究和机构也一直在致力于实现快速且高效的计算，特别是在推理时。对此，各种技术应运而生，包括降低计算精度，二进制和稀疏神经网络。本文中，来自斯坦福大学、萨尔茨堡大学等机构的研究者希望训练出一种不同的、在计算机领域被广泛应用的体系架构：逻辑门网络（logic (gate) networks）。

训练像逻辑门这样的离散组件网络所面临的问题是，它们是不可微的，因此，一般而言，不能通过梯度下降等标准方法进行优化。研究者提出了一种方法是无梯度优化方法，如演化训练（evolutionary training），它适用于小型模型，但不适用于大型模型。

在这项工作中，该研究探索了用于机器学习任务的逻辑门网络。这些网络由「AND」和「XOR」等逻辑门电路组成，可以快速执行任务。逻辑门网络的难点在于它通常是不可微的，不允许用梯度下降进行训练。因此，可微逻辑门网络的出现是为了进行有效的训练。由此产生的离散逻辑门网络实现了快速的推理速度，例如，在单个 CPU 核上每秒处理超过一百万张 MNIST 图像。这篇论文入选 NeurIPS 2022 。

论文地址：https://arxiv.org/pdf/2210.08277.pdf

项目地址：https://github.com/Felix-Petersen/difflogic

纽约大学计算机科学教授 Alfredo Canziani 表示：由逻辑门（如 AND 和 XOR）组成的可学习组合网络，允许非常快速的执行任务及硬件实现。

距离论文公开才过去三个月，论文作者 Felix Petersen 表示该研究的官方实现已经公布，他们发布了 difflogic 项目，这是一个基于 pytorch 实现的可微逻辑门网络库。不仅如此，经过优化现在的训练速度比最初的速度快 50-100 倍，因为该研究提供了高度优化的 CUDA 内核。

项目介绍

difflogic 是一个基于 Python 3.6 + 和 PyTorch 1.9.0 + 的库，基于逻辑门网络进行训练和推理。该库安装代码如下：

需要注意的是，使用 difflogic，还需要 CUDA、CUDA 工具包（用于编译）以及 torch>=1.9.0（匹配 CUDA 版本）。

下面给出了 MNIST 数据集的可微逻辑网络模型的定义示例：

该模型接收了 784 维的输入，并返回 k=10，对应于 MNIST 的 10 个类。该模型可以用 torch.nn.CrossEntropyLoss 进行训练，类似于在 PyTorch 中训练其他神经网络模型的方式。值得注意的是，Adam 优化器（torch. optimt .Adam）可用于训练，推荐的默认学习率是 0.01 而不是 0.001。最后，同样需要注意的是，与传统的 MLP 神经网络相比，逻辑门网络每层神经元的数量要高得多，因为后者非常稀疏。

为了深入了解这些模块的细节，下面是一些更详细的示例：

模型推理

在训练期间，模型应该保持在 PyTorch 训练模式，即.train ()，这种模式使得模型保持可微。现在有两种模式可以进行快速推理：

第一种选择是使用 PackBitsTensor。PackBitsTensors 允许在 GPU 上高效动态的执行训练好的逻辑门网络。

第二种选择是使用 CompiledLogicNet。CompiledLogicNet 允许在 CPU 上高效地执行固定训练的逻辑门网络。

下面给出一些实验示例，它们包含在 experiments 目录中。main.py 用于执行，main_baseline.py 是包含规则的神经网络基线。

MNIST

CIFAR-10