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

随着神经网络的成功应用，各项研究和机构也一直在致力于实现快速且高效的计算，特别是在推理时。对此，各种技术应运而生，包括降低计算精度，二进制和稀疏神经网络。本文中，来自斯坦福大学、萨尔茨堡大学等机构的研究者希望训练出一种不同的、在计算机领域被广泛应用的体系架构：逻辑门网络（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, GroupSum
import 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_model
python 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 3
python 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_model
python experiments/main.py  -bs 100 -t 100 --dataset cifar-10-3-thresholds  -ni 200_000 -ef 1_000 -k   128_000 -l 4 --compile_model
python experiments/main.py  -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k   256_000 -l 5
python experiments/main.py  -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k   512_000 -l 5
python experiments/main.py  -bs 100 -t 100 --dataset cifar-10-31-thresholds -ni 200_000 -ef 1_000 -k 1_024_000 -

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

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

项目介绍

模型推理

MNIST

CIFAR-10

作者信息

文章信息

上一篇

下一篇