OpenAI发布高度优化的GPU计算内核—块稀疏GPU内核
2017年12月13日 由 yuxiangyu 发表
313230
0
深度学习领域的模型架构和算法的发展在很大程度上受到GPU能否高效实现初等变换的限制。其中一个问题是缺乏GPU不能高效执行稀疏线性操作,我们现在正在发布高度优化的GPU计算内核实现一些稀疏模式(附带初步研究结果)。从初步的结果来看希望很大,但有不确定性,我们邀请社区加入我们,研究这些内核解锁的体系结构的极限。
稠密层(左)可以替换为稀疏并且宽的层(中)或稀疏并且深的层(右),而它们的计算时间几乎相同。
稀疏权矩阵与密集权矩阵相反,它具有大量值为零的项。我们希望稀疏权重矩阵作为模型的构建模块,因为矩阵乘法和稀疏块卷积的计算成本仅与非零块的数量成正比。例如,在给定的参数预算和计算预算中,神经网络的训练比其他方法更宽也更深,,例如具有数万个隐藏单元的LSTM 。(目前训练的最大LSTM只有数千个隐藏单元。)
内核
稠密权重矩阵(左)和块稀疏权重矩阵(中)的可视化,其中空白部分表示权重为零。
这个内核允许在完全连接层和卷积层中有效地使用块稀疏权重(如上所示)。对于卷积层,内核的输入和输出特征维度都允许有稀疏性; 空间维度上的连通性不受影响。稀疏性是在块级别(右上图)中定义的,并且已针对8x8(例如本例中)16x16或32x32的块大小进行了优化。在这个块级别中,稀疏模式是完全可配置的。由于内核计算时跳过值为零的块,所以计算成本只与非零权重的数量成正比,而不是与输入或输出特征的数量成正比。存储参数的成本也只与非零权重的数量成比例。
不同水平的稀疏的加速系数,与cuBLAS相比,在宽的状态下(12288个隐藏单元),块大小为32×32,小批量大小为32。在使用CUDA 8的NVIDIA Titan X Pascal GPU上进行比较。相对于cuSPARSE的加速在测试的稀疏水平上事实上更大。
使用内核
下面我们展示用于在Tensorflow中执行稀疏矩阵乘法的示例代码。
from blocksparse.matmul import BlocksparseMatMul
import tensorflow as tf
import numpy as np
hidden_size = 4096
block_size = 32
minibatch_size = 64
# Create a (random) sparsity pattern
sparsity = np.random.randint(2, size=(hidden_size//block_size,hidden_size//block_size))
# Initialize the sparse matrix multiplication object
bsmm = BlocksparseMatMul(sparsity, block_size=block_size)
# Input to graph
x = tf.placeholder(tf.float32, shape=[None, hidden_size])
# Initialize block-sparse weights
w = tf.get_variable("w", bsmm.w_shape, dtype=tf.float32)
# Block-sparse matrix multiplication
y = bsmm(x, w)
# Run
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
result = sess.run([y], feed_dict = {x: np.ones((minibatch_size,hidden_size), dtype='float32')})
print(result)
小世界的LSTM
块稀疏内核的一个特别有趣的用途是使用它们来创建小世界神经网络。
小世界网络图的连接方式是,图中的任何两个节点都可以通过少量步骤连接(即使图中有数十亿个节点)。我们实现小世界连通性是为了在稀疏性很高的情况下,仍然可以通过网络迅速传播信息。其实人类的大脑也是
小世界的连接模式,这就引发了一个问题:让LSTM具有同样的特性能否提高性能。利用小世界稀疏连通性,我们很快训练了近2万个隐藏单元的LSTM,比具有相似参数计数的密集网络宽5倍,提高了文本生成建模的结果,以及半监督的情感分类。
参阅:https://s3-us-west-2.amazonaws.com/openai-assets/blocksparse/blocksparsepaper.pdf
在小世界图中,即使具有高度稀疏,节点也能够以少量步骤连接。上面的动图展示了从二维Watts-Strogatz小世界图中中心节点(像素)传播的激活(为了美观做了平滑的随机)。在图中,节点之间的平均路径长度小于5,类似于我们的LSTM实验中使用的Barabasi-Albert图。
情感表征学习
在我们的
情绪神经元实验中,我们使用了近似等效参数计数的LSTM,并比较了比较了具有密集权重矩阵与块稀疏变量的模型。稀疏模型在所有情感数据集上都优于稠密模型。我们的稀疏模型将文档级别IMDB数据集的最新状态从5.91%的错误率(
Miyato et al,2016)降低到5.01%。与
以前的结果相比,这个改进看起来很不错,之前只有在较短的句子级别的数据集上表现最好
对密集和稀疏生成模型的特征进行训练的线性模型的情感分类错误(%)大致相当于总参数计数。
压缩结果
通过使用稀疏和宽的LSTM,我们的实验位每字符的结果从1.059下降到1.048,且参数计数相同(约1亿)。具有块稀疏线性层的架构也可以根据用稠密连接的线性层获得的结果进行改善。我们对CIFAR-10自然图像的
PixelCNN ++模型进行了简单的修改。用稀疏内核代替常规的二维卷积核,在加深网络的同时保持超参数的其余部分不变,导致每维的位数从2.92降低到2.90,现在这个数据集的最新技术水平。
研究方向
这里我们列出一些对未来研究的建议。
- 神经网络中的大部分权重可以在训练结束后修剪。如果推理时这种修剪与这些内核一起使用能节省多少执行的时间?
- 在生物大脑中,除了连接强度之外,网络的部分稀疏结构在成长过程中就已确定。我们是否可以在人工神经网络中进行模仿,是否可以不仅使用梯度学习连接权重还学习最佳的稀疏结构?近期有两篇论文可以在这个方向起到作用:arxiv.org/abs/1711.02782 arxiv.org/abs/1712.01312
- 我们训练了有上万个的隐藏单位的LSTM,获得了更好的文字模型表现。那么,稀疏层使得训练具有大量权重矩阵的模型成为可能,而且与其较小的稠密网络相同数量的参数和相同的计算成本。那么应用在哪些领域将对性能产生最大的影响?