近日,OpenAI 在其一篇題為《Block-Sparse GPU Kernels》的博文中釋出了一個低階別神經網路架構的高度優化 GPU 核心,並且這個神經網路帶有「塊稀疏」(block-sparse)權重。根據已選的稀疏性,這些核心可以比 cuBLAS 或 cuSPARSE 執行快幾個數量級,並在文字情感分析與文字、影象生成建模方面取得了當前最優結果。機器之心對博文進行了編譯,GitHub 連結及論文摘要請見文中。
- GitHub 地址:https://github.com/openai/blocksparse
- 論文地址:https://s3-us-west-2.amazonaws.com/openai-assets/blocksparse/blocksparsepaper.pdf
深度學習模型架構和演算法的發展在很大程度上受限於基本操作的高效 GPU 實現的可用性。這其中的一個問題是缺乏稀疏線性操作的高效 GPU 實現,它也正是今天我們釋出的內容,以及用其實現一些稀疏模式的初步結果。這些初步結果很有潛在價值但還存在不確定性,因此我們想邀請社群加入進來,共同解決這一難題。
密集層(左)可由寬而稀疏的層(中)或者深而稀疏的層(右)替代,並同時近似地保持計算時間。
與密集權重矩陣相反,稀疏權重矩陣具有大量為零的條目。稀疏權重矩陣作為模型的構建模組很有吸引力,因為矩陣乘法和稀疏塊卷積的計算成本僅與非零塊的數量成比例。例如,稀疏使得某些神經網路的訓練成為可能,這些神經網路比給定相同引數和計算預算的其它網路更寬、更深,例如具有數萬個隱藏單元的 LSTM(目前訓練的最大的 LSTM 只有數千個隱藏單元)。
核心
密集權重矩陣(左)和塊稀疏(中)權重矩陣的視覺化,其中白色代表零值權重。
核心允許在全連線層和卷積層高效地使用塊稀疏權重(如上所示)。對於卷積層,核心允許輸入和輸出特徵維度的稀疏性;連線性在空間維度上不受影響。稀疏性是在塊的層面被定義的(如上右圖所示),並已經優化為 8x8(本例項中)、 16x16 或 32x32 的塊大小。在塊的層面上,稀疏性模式是完全可配置的。由於核心跳過零值塊的計算,所以計算成本只與非零權重的數量成正比,而不是輸入/輸出特徵的數量。儲存引數的成本也只與非零權重的數量成正比。
不同級別的稀疏性相對於 cuBLAS 的加速效果,其使用了寬狀態(12288 個隱藏單元)、32x32 的塊大小和 32 的批尺寸。該對比實驗使用英偉達 Titan X Pascal GPU 和 CUDA 8 完成。和 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)
小世界(Small-World)LSTM
塊稀疏核心的一個很有趣的應用是用於建立小世界(small-world)神經網路。小世界圖的連線方式是圖中的任意兩個節點之間僅通過少量幾步就能連線,即使圖中有數百萬個節點。我們實現小世界連線的動機是,即使具有高度的稀疏度,我們仍然希望資訊在網路中能快速地傳播。大腦中的連線模式就是小世界的,這使我們好奇同樣的性質能不能用於提高 LSTM 的效能。我們使用小世界稀疏連線高效地訓練了擁有大約兩萬個隱藏單元(與擁有相似引數數量的密集網路相比具有 5 倍的寬度)的 LSTM,提升了文字生成建模和半監督情緒分類的結果,詳情參見我們的論文。
在小世界圖中,小世界圖的連線方式是圖中的任意兩個節點之間僅通過少量幾步就能連線,即使圖中有數百萬個節點。以上動圖展示了一個二維 Watts-Strogatz 小世界圖上的中心節點(畫素)的啟用擴散現象(隨機平滑灰度僅僅是為了審美效果)。在這個圖中,節點之間的平均路徑長度小於 5,和我們在 LSTM 實驗中使用的 Barabasi-Albert 圖相似。
情緒表徵學習
通過使用在我們的情緒神經實驗(Unsupervised Sentiment Neuron,https://blog.openai.com/unsupervised-sentiment-neuron/)中的設定,我們使用近似相等的引數數量訓練 LSTM,並比較了使用密集權重矩陣的模型和一個塊稀疏的變體。稀疏模型在所有的情緒資料集上的效能都超越了相對的密集模型。我們的稀疏模型將在文件級 IMDB 資料集上的當前最佳效能從 5.91% 誤差率(Miyato 等,2016)降低到了 5.01%。相對於我們之前的結果(Unsupervised Sentiment Neuron,僅在短句級資料集上獲得最佳結果),這個提升是很有潛在價值的。
在密集特徵上訓練的線性模型和稀疏生成模型的情緒分類錯誤率(%),兩者擁有近似相等的引數總數。
壓縮結果
通過使用稀疏和寬的 LSTM,在我們實驗中的每個特徵佔用的位元數從 1.059 下降到了 1.048(在相等引數數量的前提下,即大約 1 億個引數)。擁有塊稀疏線性層的架構也可以提高相對於使用密集連線線性層的結果。我們對在 CIFAR-10 自然影象上訓練的 PixelCNN++模型做了一個簡單的修改。用稀疏核心替代了常規的 2D 卷積核,同時增加網路深度,但保持其它超引數不變,使每個維度佔用的位元數從 2.92 下降到 2.90,達到了在該資料集上的當前最佳。
研究方向
我們在這裡提幾個對未來研究的建議。
大多數神經網路的權重(連線)在訓練完成之後都可以修剪掉。如果能結合剪枝技術和我們的稀疏核心,能獲得怎樣的加速效果呢?
在生物大腦中,除了連線強度以外,網路的稀疏結構部分地在成長過程中確定。是否可以在人工神經網路中實現類似的行為,即不僅將梯度用於學習連線強度,還用於尋找最優的稀疏結構。最新的一篇論文提出了一種方法用於學習塊稀疏 RNN(https://arxiv.org/abs/1711.02782)。我們近日也提出了一種演算法用於在神經網路中進行 L0 正則化,可用於實現這一目標。
我們訓練了擁有成千上萬個隱藏單元的 LSTM,可以更好地處理文字。更一般地說,稀疏層使得訓練擁有大型權重矩陣(但引數數量、計算開銷和相對的密集模型相等)的模型成為可能。這種模型能為哪些領域帶來顯著的進步呢?
論文:GPU Kernels for Block-SparseWeights
論文連結:https://s3-us-west-2.amazonaws.com/openai-assets/blocksparse/blocksparsepaper.pdf
摘要:我們正在釋出一個低階別神經網路架構(帶有塊稀疏(block-sparse)權重)的高度優化 GPU 核心,它允許在權重矩陣中帶有靈活可配置的塊稀疏性模式的線性層(包括卷積層)的高效評估與微分。我們發現根據已選稀疏性,這些核心可以比 cuBLAS 或 cuSPARSE 執行快幾個數量級。我們藉助這些核心取得了文字情感分析與文字、影象生成建模方面的當前最優結果。通過公開發布這些核心,我們希望可以促進模型和演算法設計的進一步發展。