选自arXiv

作者:Felix Petersen等

机器之心编译

编辑:袁铭怿

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

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

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

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

打开网易新闻 查看精彩图片

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

纽约大学计算机科学教授 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 torchmodel = 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 -

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