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

编辑：袁铭怿

本文通过逻辑门组合学习来探索面向机器学习任务的逻辑门网络。这些网络由 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 + 的库，基于逻辑门网络进行训练和推理。该库安装代码如下：

pip install difflogic

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

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

from difflogic import LogicLayer, GroupSumimport torch
model = torch.nn.Sequential(    torch.nn.Flatten(),    LogicLayer(784, 16_000),    LogicLayer(16_000, 16_000),    LogicLayer(16_000, 16_000),    LogicLayer(16_000, 16_000),    LogicLayer(16_000, 16_000),    GroupSum(k=10, tau=30))

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

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

layer = LogicLayer(    in_dim=784,             # number of inputs    out_dim=16_000,         # number of outputs    device='cuda',          # the device (cuda / cpu)    implementation='cuda',  # the implementation to be used (native cuda / vanilla pytorch)    connections='random',   # the method for the random initialization of the connections    grad_factor=1.1,        # for deep models (>6 layers), the grad_factor should be increased (e.g., 2) to avoid vanishing gradients)

模型推理

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

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

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

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

MNIST

python experiments/main.py  -bs 100 -t  10 --dataset mnist20x20 -ni 200_000 -ef 1_000 -k  8_000 -l 6 --compile_modelpython experiments/main.py  -bs 100 -t  30 --dataset mnist      -ni 200_000 -ef 1_000 -k 64_000 -l 6 --compile_model# Baselines:python experiments/main_baseline.py  -bs 100 --dataset mnist    -ni 200_000 -ef 1_000 -k  128 -l 3python experiments/main_baseline.py  -bs 100 --dataset mnist    -ni 200_000 -ef 1_000 -k 2048 -l 7

CIFAR-10

python experiments/main.py  -bs 100 -t 100 --dataset cifar-10-3-thresholds  -ni 200_000 -ef 1_000 -k    12_000 -l 4 --compile_modelpython experiments/main.py  -bs 100 -t 100 --dataset cifar-10-3-thresholds  -ni 200_000 -ef 1_000 -k   128_000 -l 4 --compile_modelpython experiments/main.py  -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k   256_000 -l 5python experiments/main.py  -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k   512_000 -l 5python experiments/main.py  -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k 1_024_000 -

想要了解更多内容，请参考原项目。