经典神经网络非常灵活，但是某些任务并不适合它们。特别是一些算术运算证明对神经网络具有挑战性。这就是为什么Trask等人在他们的论文“ 神经算术逻辑单元（Neural Arithmetic Logic Units）”中引入了两个新模块，这些模块意味着在某些算术任务上表现良好。在这篇文章中，我将描述论文要解决的问题， 提出的解决方案，并使用PyTorch讨论实现细节和结果再现。

恒等函数

对于神经网络来说，这听起来像一个简单的任务，对吧？如果我们的测试集与我们的训练集在同一范围内，那的确是。然而，一旦我们的测试数据超出训练范围，大多数神经网络都会失败。这种失败揭示了一个事实：多层感知器理论上能够代表任何连续函数，但是架构，训练数据和学习计划的选择将严重影响MLP最终学习的函数。Trask等人用下图有效地说明了这种失败：

MLP学习恒等函数失败

为了克服这些限制，论文建议使用两个新模块：神经累加器（NAC）和神经算术逻辑单元（NALU）。让我们来看看它们是什么以及它们背后的数学。

NAC和NALU

让我们从NAC开始：

NAC

我们可以看到，NAC类似于线性变换，但是我们使用的权值矩阵是两个已知矩阵的结果，在元素相乘之前，我们先应用tanh和sigmoid

为什么要这么麻烦呢？这如何帮助我们的网络正确学习算术函数呢？我们来看看下图：

tanh（x）* sigmoid（y）

我们可以看到权重向量的值偏向0,1和-1。这意味着NAC的输出是输入向量的加法和减法，而不是scalings。这允许NAC无限期地堆叠在一起并且仍然保持输入表示的一致scaling。这对于学习处理加法和减法的算术运算显然是有帮助的，但对于乘法和除法来说，这是不够的。

这就是NALU的用武之地：

NALU

NALU使用NAC作为子单元，并添加两个额外的信息流。第一个使用NAC在log-space中操作，第二个用作学习门（learned gate），计算标准NAC和log-space NAC之间的加权和。理论上，这应该允许NALU学习乘法和除法以及加法和减法等函数。

Python代码实现

NAC和NALU的实施相当简单。首先，我们将从NAC开始，Python代码如下：

import torch from torch.nn.parameter import Parameter from torch.nn import functional as F from torch.nn.modules import Module class NAC(Module): def __init__(self, n_in, n_out): super().__init__() self.W_hat = Parameter(torch.Tensor(n_out, n_in)) self.M_hat = Parameter(torch.Tensor(n_out, n_in)) self.reset_parameters() def reset_parameters(self): init.kaiming_uniform_(self.W_hat) init.kaiming_uniform_(self.M_hat) def forward(self, input): weights = torch.tanh(self.W_hat) * torch.sigmoid(self.M_hat) return F.linear(input, weights) 

在初始化步骤中，我们定义两个学习矩阵W_hat和M_hat。在forward步骤中，我们将W_hat（通过tanh）和M_hat（通过sigmoid）进行相乘。然后，我们取结果矩阵并使用它和单元格的输入进行线性变换。你有了一个功能完整的NAC模块!

在我们转向NALU实现之前，重要的是要强调上述代码中只有两个学习矩阵的事实。传递到线性变换中的矩阵没有经过训练，它只是学习矩阵W_hat和M_hat的数学运算的结果。当NALU论文首次出现时，许多PyTorch实现都将最终权重作为参数，将NAC转换为简单的线性层。

我们来看看NALU的Python实现：

import torch from torch.nn.parameter import Parameter from torch.nn import functional as F from torch.nn.modules import Module from NAC import NAC class NALU(Module): def __init__(self, n_in, n_out): super().__init__() self.NAC = NAC(n_in, n_out) self.G = Parameter(torch.Tensor(1, n_in)) self.eps = 1e-6 self.reset_parameters() def reset_parameters(self): init.kaiming_uniform_(self.G) def forward(self, input): g = torch.sigmoid(F.linear(input, self.G)) y1 = g * self.NAC(input) y2 = (1 - g) * torch.exp(self.NAC(torch.log(torch.abs(input) + self.eps))) return y1 + y2 

我们从定义NAC和学习门矩阵开始，并初始化它们。在前进的步骤中，我们首先通过对输入和gate矩阵G进行线性变换来计算gate值，然后在返回加权和之前计算标准的NAC输出和log-space NAC输出。重要是要注意，在计算log – space NAC时，我们需要将小的ε值传递给log函数，以便避免任何未定义的行为。

现在我们知道了NALU背后的理论以及如何在PyTorch中实现它们。我们将尝试使用基于kevinzakka代码的脚本复制论文的函数学习部分的结果，在这里包括FauxNALU。函数学习测试的想法是，我们将提供随机生成的输入数据，这些数据符合一个特定的函数(加法、减法、除法、平方根等)，然后测试我们的MLP是否能够正确学习这个算术函数。Python代码如下：

import math import random import numpy as np import matplotlib.pyplot as plt import torch import torch.nn as nn import torch.nn.functional as F from models import MLP, NAC, NALU, FauxNALU NORMALIZE = True NUM_LAYERS = 2 HIDDEN_DIM = 2 LEARNING_RATE = 1e-3 NUM_ITERS = int(1e5) RANGE = [5, 10] ARITHMETIC_FUNCTIONS = { 'add': lambda x, y: x + y, 'sub': lambda x, y: x - y, 'mul': lambda x, y: x * y, 'div': lambda x, y: x / y, 'squared': lambda x, y: torch.pow(x, 2), 'root': lambda x, y: torch.sqrt(x), } def generate_data(num_train, num_test, dim, num_sum, fn, support): data = torch.FloatTensor(dim).uniform_(*support).unsqueeze_(1) X, y = [], [] for i in range(num_train + num_test): idx_a = random.sample(range(dim), num_sum) idx_b = random.sample([x for x in range(dim) if x not in idx_a], num_sum) a, b = data[idx_a].sum(), data[idx_b].sum() X.append([a, b]) y.append(fn(a, b)) X = torch.FloatTensor(X) y = torch.FloatTensor(y).unsqueeze_(1) indices = list(range(num_train + num_test)) np.random.shuffle(indices) X_train, y_train = X[indices[num_test:]], y[indices[num_test:]] X_test, y_test = X[indices[:num_test]], y[indices[:num_test]] return X_train, y_train, X_test, y_test def train(model, optimizer, data, target, num_iters): for i in range(num_iters): out = model(data) loss = F.mse_loss(out, target) mea = torch.mean(torch.abs(target - out)) optimizer.zero_grad() loss.backward() optimizer.step() if i % 1000 == 0: print("\t{}/{}: loss: {:.7f} - mea: {:.7f}".format( i+1, num_iters, loss.item(), mea.item()) ) def test(model, data, target): with torch.no_grad(): out = model(data) return torch.abs(target - out) def main(): save_dir = './results/' models = [ MLP( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, activation='relu6', ), MLP( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, activation='none', ), NAC( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, ), NALU( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1 ), FauxNALU( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1 ), ] results = {} for fn_str, fn in ARITHMETIC_FUNCTIONS.items(): results[fn_str] = [] # dataset X_train, y_train, X_test, y_test = generate_data( num_train=500, num_test=50, dim=100, num_sum=5, fn=fn, support=RANGE, ) # random model random_mse = [] for i in range(100): net = MLP( num_layers=NUM_LAYERS, in_dim=2, hidden_dim=HIDDEN_DIM, out_dim=1, activation='relu6', ) mse = test(net, X_test, y_test) random_mse.append(mse.mean().item()) results[fn_str].append(np.mean(random_mse)) # others for net in models: optim = torch.optim.RMSprop(net.parameters(), lr=LEARNING_RATE) train(net, optim, X_train, y_train, NUM_ITERS) mse = test(net, X_test, y_test).mean().item() results[fn_str].append(mse) with open(save_dir + "interpolation.txt", "w") as f: f.write("Relu6\tNone\tNAC\tNALU\tfauxNALU\n") for k, v in results.items(): rand = results[k][0] mses = [100.0*x/rand for x in results[k][1:]] if NORMALIZE: f.write("{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\n".format(*mses)) else: f.write("{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\t{:.3f}\n".format(*results[k][1:])) if __name__ == '__main__': main() 

结果：

这些结果有几点：

• NAC在加法和减法方面的表现令人惊叹，但在乘法、除法和幂函数方面却很糟糕。这是有意义的，因为同样的事情使得它对加法（没有任意的输入的扩展）非常强大，当它必须学习其他函数时，会极大地限制它。
• 除了减法之外，NALU在所有事情上表现都很好。鉴于NAC和FauxNALU在减法上都表现良好，这表明在减法用例中，对数门和NAC在某种程度上相互干扰。
• FauxNALU的结果是最有趣的。它的表现与None MLP（仅使用线性度）一样，表明学习的门本身无助于学习这些函数。在学习分区和平方根时，NALU果断地击败了FauxNALU，这表明对数空间NAC和标准NAC之间的结合确实发挥了作用。

最后

NAC和NALU肯定是有趣的模块，尽管他们对我们看到的问题的表现是喜忧参半。只有时间和进一步的实验才能说明这些专业模块是否会成为机器学习从业者常规工具箱的一部分。有一件事是肯定的：我们可以通过巧妙的权重矩阵定义和门的使用来设计具有特定偏差的模块的想法。