文 / ,AI 定向实习生与 Chase ,研究工程师以及 ,研究员,来自 X
Colab
简介
在本文中,我们要讨论 (TN) 以及如何用其来增强 中的前馈神经网络。
是 19 年 6 月发布的一个开源库,用于简化张量网络中的计算。通常人们首先会问,“什么是张量网络”,紧接着就会问“为什么我要关心张量网络?”。第一个问题可以靠查阅很多现有资源来解决(例如我们之前的 AI 文章),现在我们来重点回答第二个问题。其基本思路称作将神经网络“张量化”,来源于 et. al 2015 年的一篇论文。可利用 库直接实现此过程。下面我们给出一个通俗易懂的教学示例,帮助您理解如何使用 Keras 和 2.0。
易于上手。该库可通过 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)])
此网络为二维输入keras.layers,一维输出,并包含两个隐藏层,每层 1024 个神经元。参数总量(包含权重和偏置)为 (2+1)*1024 + (1024+1)*1024 + (1024 +1)*1 = 1,053,697。这些参数中的大部分((1024+1)*1024 = 1,049,600 个)与第二个隐藏层关联。我们要将其替换为大小不足原始大小 1% 的 TN 层。
隐藏层中的 1,049,600 个参数包含排列在 1024*1024 矩阵中的 1024*1024 = 1,048,576 个权重,与 1024 个偏置。我们要通过将 1024*1024 的权重矩阵替换为一个张量网络来节省空间。为保持与其余网络的兼容性,我们不会改变原本的 1024 个输入和 1024 个输出。这样我们就能把 TN 层视为原始层的直接替代。
通常,我们将 1024 个层输入视为一个简单的 X 形数组 (1024,)。这些输入乘以权重矩阵 W 产生一个新的向量 Y。
下一步是将 Y 非线性化,但我们的重点不在这里,因为 TN 的修改仅影响到层的线性部分。在张量网络表示法中,我们将 Y 的计算记录如下:
通常,神经网络层的权重 W 充当输入 X 的矩阵乘数,生成输出 Y = WX。这里我们用图解法展示该矩阵乘数
这时,我们可以讨论层的张量化。首先,我们将输入数组重组为 (32,32) 这样的形式,以代替原来的 (1024,)。重组的图片展示如下
将输入 X 从 (1024,) 的向量形式重组为 (32,32) 的向量形式,这是向量化过程的第一步。从原则上来说,任何重组都可以,在示例中我们只是选择了这种特别简单的形式
现在,我们不再使用 (1024,1024) 加权矩阵 W,而是使用由两个核心部分组成的张量网络运算:
张量化权重乘数的图形化表示:可根据需要将输出 Y 重组为一个向量
我们需要选择连接到两个核心部分的连线的维度(通常称为 连接维度)。方便起见,将其设为二维。然后,每个核心部分均有两个三十二维连线和一个二维连线,代表形式为 (32,32,2)。连接维度控制着模型中的参数数量keras.layers,通过选择合适的连接维度,我们能大幅缩减参数数量,且许多情况下对性能的影响很小或几乎没有。
使用 TN 层的结果是我们将稠密权重矩阵中的 1,048,576 个权重替换为张量网络中的 2*(32*32*2) = 4096 个参数。实现了大幅缩减!即便算上其他层,总模型大小也降至 9,217 个参数,而原来有 1,053,697 个。
在这个简单示例中,将两个核心的张量网络转换为 表达式并不太难。但是,对于核心更多、连接也更复杂的张量网络,调试或扩展 可能十分困难。因此,我们改为使用 TN 库,以面向对象的方式构建网络。
在以下代码示例中,我们仍然沿用上面的两个核心简单示例,以便于理解,但最后我们将回来讨论其他各种可能的情况。
TN 层代码
以下是在 Keras 中创建 TN 层的一些示例代码(仍在此Colab 中),专门针对以上讨论的 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_mapfunction.
# https://tensorflow.google.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进行训练。
的压缩及更多内容
上面讨论的简单示例是张量化思想的完美演示,但实际中这种特殊模型可能并不适用。针对更加实际的情况,我们最近试验了以非常类似的方法对一个 模型进行张量化。我们看到的模型中有一些稠密层,这些层的规模要大得多:原始模型的参数高达2.36亿个!我们通过包含四个核心的张量网络将 8 个 块中的全连接层进行张量化:
我们用来张量化 模型中完全互连层的四核心张量网络
张量化之后,模型缩减为约1.01亿个参数!除了规模大幅变小之外,张量化模型生成英语句子的速度也快了很多。下面的视频显示了性能差异:
还有其他类型的一些张量网络和张量分解适用于神经网络中各种类型的层,在这里我们不再一一详述。但首先,请查看 et al. 的论文了解嵌入层中的张量网络,查看 Ma et al. 的论文了解注意力层,并查看 et al. 的论文了解卷积层。
神经网络的张量化研究才刚刚起步。仍有许多问题有待解决,许多试验有待尝试。本文中讨论的所有张量网络均为张量训练类型,在物理学中称为 MPO(另请参见此论文),但也可以使用其他研究较为成熟的张量网络,例如 PEPS 和 MERA。 库旨在处理任何可能的张量网络,在本文所述的机器学习领域中尤为适用。我们期待看到您采用此技术取得令人欣喜的全新成果。
如果您想详细了解 本文提及 的相关内容,请参阅以下文档。这些文档深入探讨了这篇文章中提及的许多主题: