作者：Marina Munkhoeva，Skolkovo 科技学院博士生，Alphabet 旗下 X 的 AI 研究员；Chase Roberts，Alphabet 旗下 X 的研究工程师；Stefan Leichenauer，Alphabet 旗下 X 的研究科学家

简介

在本篇文章中，我们将介绍 TensorNetwork，以及如何使用它来加速 TensorFlow 中的前馈神经网络。TensorNetwork 是一个开源库，它于 2019 年 6 月 发布，旨在促进使用张量网络进行计算。通常，人们问我们的第一个问题是：“什么是张量网络？”，紧随其后的是“为什么我应该关心张量网络？”。许多资源都回答了第一个问题（例如我们之前的 Google AI 博客文章），而这里我们将重点关注回答第二个问题。这个基本思想被称为“张量化”神经网络，它的根源可以追溯到 Novikov 等人于 2015 年发表的 论文。使用 TensorNetwork 库，可以轻松地实现此过程。下面我们将使用 Keras 和 TensorFlow 2.0 提供一个明确且具有教学意义的示例。

入门使用 TensorNetwork 很容易。可以使用 pip 安装该库

pip install tensornetwork

本文中将讨论的示例代码也可以在 Colab 中找到。

TN 层

增强 TN 的神经网络的基本思想是用 TN 层替换网络中的一个或多个层。可以将 TN 层视为原始层的压缩版本。当你从一个密集的、高度连接的层开始，并且该层具有很大的压缩潜力时，可以预期获得最佳结果。例如，考虑以下内容

Dense = tf.keras.layers.Dense
fc_model = tf.keras.Sequential(
    [
     tf.keras.Input(shape=(2,)),
     Dense(1024, activation=tf.nn.swish),
     Dense(1024, activation=tf.nn.swish),
     Dense(1, activation=None)])

该网络具有二维输入、一维输出，并包含两个隐藏层，每个隐藏层具有 1024 个神经元。参数总数（包括权重和偏差）为 (2+1)*1024 + (1024+1)*1024 + (1024 +1)*1 = 1,053,697。其中大部分参数，即 (1024+1)*1024 = 1,049,600 个参数，与第二个隐藏层相关联。我们将用一个 TN 层替换这些参数，该层的大小不到原始层大小的 1%。

隐藏层的 1,049,600 个参数包括以 1024×1024 矩阵排列的 1024×1024 = 1,048,576 个权重，以及 1024 个偏差。我们将用张量网络替换 1024×1024 的权重矩阵，从而节省大量参数。为了与网络的其余部分保持兼容，我们不改变 1024 个输入和 1024 个输出的事实。这样，我们可以将 TN 层视为原始层的直接替换。

通常，我们将层中的 1024 个输入视为形状为 (1024,) 的简单数组 X。这些输入乘以权重矩阵 W 以生成新的向量 Y。下一步是对 Y 应用非线性函数，但我们不会关注这一点，因为 TN 修改只影响层的线性部分。在张量网络表示法中，我们将 Y 的计算写成如下形式

通常，神经网络层中的权重 W 作为矩阵乘法作用于输入 X，产生输出 Y = WX。在这里，我们以图表形式说明了矩阵乘法。

在这一点上，我们可以讨论层的张量化。首先，我们将输入数组重新整形为类似 (32,32) 而不是 (1024,)。从图示来看，整形看起来像

将输入 X 从形状为 (1024,) 的向量重新整形为形状为 (32,32) 的数组是张量化过程的第一步。原则上，任何整形都是允许的，我们只是在我们的示例中选择了这种特别简单的整形。

现在，我们不应用 (1024,1024) 的权重矩阵 W，而是应用由两个核心组成的张量网络运算

张量化权重乘法的图表表示。输出 Y 可以根据需要重新整形为向量。

我们需要对连接两个核心的边的维数（通常称为键维数）进行选择。为简单起见，我们将其设为二维。然后每个核心有两个 32 维边以及单个 2 维边，这意味着形状为 (32,32,2)。键维数控制模型中的参数数量，通过适当选择键维数，在许多情况下，我们可以实现良好的参数减少率，而性能几乎没有下降。

使用 TN 层的结果是，我们用张量网络的 2*(32*32*2) = 4,096 个参数替换了全连接权重矩阵的 1,048,576 个权重。这是一个巨大的减少！即使考虑到其他层，模型总大小也减少到了 9,217 个参数，而原来的模型大小是 1,053,697 个参数。

对于这个简单的示例，将双核心张量网络转换为 einsum 表达式并不困难。但是，对于具有更多核心和更复杂连接的张量网络，调试或扩展 einsum 会非常困难。因此，我们改为使用 TN 库以更面向对象的方式构建网络。在下面的代码示例中，我们将坚持使用简单的双核心示例，以方便理解，但在最后我们将回到讨论其他可能性。

TN 层的代码

以下是一些示例代码（也可以在 这个 Colab 中找到），用于在 Keras 中创建 TN 层，专门针对上面讨论的 1024×1024 案例

import tensorflow as tf
import tensornetwork as tn

class TNLayer(tf.keras.layers.Layer):
 
  def __init__(self):
    super(TNLayer, self).__init__()
    # Create the variables for the layer.
    self.a_var = tf.Variable(tf.random.normal(
            shape=(32, 32, 2), stddev=1.0/32.0),
             name="a", trainable=True)
    self.b_var = tf.Variable(tf.random.normal(shape=(32, 32, 2), stddev=1.0/32.0),
                             name="b", trainable=True)
    self.bias = tf.Variable(tf.zeros(shape=(32, 32)), name="bias", trainable=True)
 
  def call(self, inputs):
    # Define the contraction.
    # We break it out so we can parallelize a batch using
    # tf.vectorized_map (see below).
    def f(input_vec, a_var, b_var, bias_var):
      # Reshape to a matrix instead of a vector.
      input_vec = tf.reshape(input_vec, (32,32))
 
      # Now we create the network.
      a = tn.Node(a_var, backend="tensorflow")
      b = tn.Node(b_var, backend="tensorflow")
      x_node = tn.Node(input_vec, backend="tensorflow")
      a[1] ^ x_node[0]
      b[1] ^ x_node[1]
      a[2] ^ b[2]
 
      # The TN should now look like this
      #   |     |
      #   a --- b
      #    \   /
      #      x
 
      # Now we begin the contraction.
      c = a @ x_node
      result = (c @ b).tensor
 
      # To make the code shorter, we also could've used Ncon.
      # The above few lines of code is the same as this:
      # result = tn.ncon([x, a_var, b_var], [[1, 2], [-1, 1, 3], [-2, 2, 3]])
 
      # Finally, add bias.
      return result + bias_var
  
    # To deal with a batch of items, we can use the tf.vectorized_map
    # function.
    # https://tensorflowcn.cn/api_docs/python/tf/vectorized_map
    result = tf.vectorized_map(
        lambda vec: f(vec, self.a_var, self.b_var, self.bias), inputs)
    return tf.nn.relu(tf.reshape(result, (-1, 1024)))

在这个示例中，我们对层的大小进行了硬编码，但调整起来相当容易。创建了这个层之后，我们可以非常简单地将其用作 Keras 模型的一部分

tn_model = tf.keras.Sequential(
  [
    tf.keras.Input(shape=(2,)),
    Dense(1024, activation=tf.nn.relu),
    # Here use a TN layer instead of the dense layer.
    TNLayer(),
    Dense(1, activation=None)
  ]
)

可以使用 Keras 的fit方法像往常一样训练模型。

压缩 Transformer 以及更多

虽然我们上面讨论的简单示例很好地说明了张量化的概念，但这种特定模型在实践中可能并不实用。为了更现实地说明，我们最近尝试以非常类似的方式对 Transformer 模型进行张量化。我们查看的模型具有更大的密集层：原始模型具有 236M 个参数！我们使用四个核心的张量网络对 8 个 Transformer 块中的全连接层进行了张量化

我们用来对 Transformer 模型的全连接层进行张量化的四核心张量网络。

张量化后，模型缩小到了约 101M 个参数！除了更小之外，张量化后的模型也能够更快地生成英文句子。以下是展示性能差异的并排视频



还有许多其他类型的张量网络和张量分解可应用于神经网络中的各种类型层，我们这里不会尝试给出详尽的参考列表。但首先，可以查看 Khrulkov 等人 关于嵌入层中的张量网络，Ma 等人 关于注意力层，以及 Lebedev 等人 关于卷积层。

神经网络的张量化仍处于起步阶段。还有许多悬而未决的问题，还有许多实验需要尝试。本文中考虑的所有张量网络都是 张量列车 类型的，在物理学中被称为 MPO（也可以参见 这篇论文），但其他经过充分研究的张量网络，如 PEPS 和 MERA 也可用于此。 TensorNetwork 库 旨在处理任何可能的张量网络，特别是像我们这里展示的那样在机器学习的背景下进行处理。我们期待看到您使用它烹饪出来的所有新颖而令人兴奋的结果。