機器之心原創
作者：蔣思源
上週 Geoffrey Hinton 等人公開了那篇備受關注的 NIPS 論文，而後很多研究者與開發者都閱讀了該論文並作出了一定的程式碼實現。機器之心在本文中將詳細解釋該論文提出的結構與過程，並藉助 GitHub 上熱烈討論的專案完成了 CapsNet 的 TensorFlow 實現，並提供了主體架構的程式碼註釋。
本文是機器之心的第三個 GitHub 專案，旨在解釋 CapsNet 的網路架構與實現。為了解釋 CapsNet，我們將從卷積層與卷積機制開始，從工程實踐的角度解釋卷積操作的過程與輸出，這對進一步理解 Capsule 層的處理十分有利，後面我們將基於對 Capsule 層的理解解釋 Geoffrey Hinton 等人最近提出來的 CapsNet 架構。最後我們會根據 naturomics 的實現進行測試與解釋。
機器之心 GitHub 專案地址：https://github.com/jiqizhixin/ML-Tutorial-Experiment
卷積層與卷積機制
這一部分主要是為不太瞭解卷積機制具體過程的讀者準備，因為 CapsNet 的前面兩層本質上還是傳統的卷積操作。若讀者已經瞭解基本的卷積操作，那麼可以跳過這一章節直接閱讀 Capsule 層的結構與過程。
若要解釋卷積神經網路，我們先要知道為什麼卷積在影像上能比全連線網路有更好的效能，以下分別展示了全連線網路和卷積網路一般的架構：
我們知道全連線網路前一層的每個神經元（或單元）都會與後一層中每個神經元相連，連線的強弱可以通過相對應的權重控制。而所有連線權重就是該全連線神經網路希望學到的。上圖可知卷積神經網路也是由一層一層的神經元組織起來的，只不過全連線網路相鄰兩層的神經元都有連線，所以可以將相同層的神經元排列為一列，這樣可以方便顯示連線結構。而卷積網路相連兩層之間只有部分神經元相連，為了展示每一層神經元的維度，我們一般會將每一個卷積層的結點組織為一個三維張量。
全連線網路處理影像最大的問題是每層之間的引數或權重太多了，主要是因為兩層間的神經元都有連線。若使用一個隱藏層為 500 個單元的全連線網路（784×500×10）識別 MNIST 手寫數字，那麼引數的數量為 28×28×500+5000+510=397510 個引數，這大大限制了網路層級的加深。
而對於卷積網路來說，每一個單元都只會和上一層部分單元相連線。一般每個卷積層的單元都可以組織成一個三維張量，即矩陣沿第三個方向增加一維資料。例如 Cifar-10 資料集的輸入層就可以組織成 32×32×3 的三維張量，其中 32×32 代表圖片的尺寸或畫素數量，而 3 代表 RGB 三色通道。
卷積層
卷積層試圖將神經網路中的每一小塊進行更加深入的分析，從而得出抽象程度更高的特徵。一般來說通過卷積層處理的神經元結點矩陣會變得更深，即神經元的組織在第三個維度上會增加。
下圖展示了卷積核或濾波器（filter）將當前層級上的一個子結點張量轉化為下一層神經網路上的一個長和寬都為 1，深度不限的結點矩陣。下圖輸入是一個 32×32×3 的張量，中間的小長方體為卷積核，一般可以為 3×3 或 5×5 等，且因為要計算乘積，那麼卷積核的第三個維度必須和其處理的影像深度（即輸入張量第三個維度 3）相等。最右邊的矩形體的深度為 5，即前面使用了五個卷積核執行卷積操作。這五個卷積核有不同的權重，但每一個卷積層使用一個卷積核的權重是一樣的，所以下圖五層特徵中每一層特徵都是通過一個卷積核得出來的，也就是該層共享了權重。
卷積操作
可能入門讀者對卷積的具體過程還是不夠了解，下面我們可以討論卷積操作的具體過程。如下所示，該圖展示了卷積的具體操作過程。首先我們的輸入為 5×5×3 的張量，即 x[:, :, 0 : 3]。其次我們有兩個 3×3 的卷積核，即 W0 和 W1，第三個維度必須和輸入張量的第三個維度相等，所以一般只用兩個維度描述一個卷積核。最後卷積操作輸出 3×3×2 的張量，其中 o[:, :, 0] 為第一個卷積核 W0 的卷積輸出，o[:, :, 1] 為第二個卷積核的輸出。因為輸入張量使用了 Padding，即每一個通道的輸入影像周圍加 0，且卷積核移動的步幅為 2，則每個卷積核輸出的維度為 3×3（即 (7-3)/2）。
在上圖中，卷積核會與輸入張量對應相乘相加，然後再加上偏置項就等於輸出張量中對應位置的值。例如使用卷積和 W0 對輸入張量（深度為 3 可看作影像擁有的 RGB 三個通道）做卷積，卷積和三個層級將對應輸入張量的三個層級做乘積累計。w0[:, :, 0] 乘以 x[:, :, 0] 左上角的九個元素為 1*0+1*0-1*0-1*0+0*0+1*1-1*0-1*0+0*1=1，同理 w0[:, :, 1] 乘以 x[:, :, 1] 左上角九個元素為-1、w0[:, :, 2] 乘以 x[:, :, 2] 左上角九個元素為 0，這三個值相加再加上偏置項 b0 就等於最右邊輸出張量 o[:, :, 0] 的左上角第一個元素，即 1-1+0+1=1。
隨著卷積核移動一個步長，我們可以計算出輸出矩陣移動一個元素的值。注意但卷積核在輸入張量上移動的時候，卷積核權重是相同的，也就是說這一層共享了相同的權重，即 o[:, :, 0] 和 o[:, :, 1] 分別共享了一組權重。這裡之所以強調權重的共享，不僅因為它是卷積層核心的屬性，同時還有利於我們在後面理解 CapsNet 的 PrimaryCaps 層。
卷積還有很多性質沒有解釋，例如最大池化選取一個濾波器內數值最大的值代表該區域的特徵以減少輸出張量的尺寸，Inception 模組將多組卷積核並聯地對輸入張量進行處理，然後再將並聯處理得到的多個輸出張量按序串聯地組成一個很深的輸出張量作為 Inception 模組的輸出等。讀者也可以繼續閱讀機器之心關於卷積的文章進一步瞭解。最後，我們提供了一個簡單的實現展示卷積操作的計算過程：
這一段程式碼執行了卷積操作和平均池化，它的輸出如下：
Capsule 層與動態路由
這一部分主要是解釋 Capsule 層與動態路由（DynamicRouting）機制的大概原理，這一部分基於我們對 Hinton 原論文的理解完成，並採用了知乎 SIY.Z、Debarko De 等人的觀點。文末將給出更多的參考資料，讀者可進一步閱讀以瞭解更多。
前面我們已經知道卷積通過權重共享和區域性連線可以減少很多引數，此外共享卷積核權重可以使影像上的內容不受位置的影響。例如 Cifar-10 中的影像為 32×32×3，而由 16 個尺寸為 5×5 的卷積核（或表述深度為 16）所構成的卷積層，其引數共有 5*5*3*16+16=1216 個。但這這樣的卷積層單元還是太簡單了，它們也不能表徵複雜的概念。
例如當影像進行一些旋轉、變形或朝向不同的方向，那麼 CNN 本身是無法處理這些圖片的。當然這個問題可以在訓練中新增相同影像的不同變形而得到解決。在 CNN 中每一層都以非常細微的方式理解影像，因為我們卷積核的感受野一般使用 3×3 或 5×5 等畫素級的操作來理解影像，所以卷積層總是嘗試理解區域性的特徵與資訊。而當我們由前面低階特徵組合成後面複雜與抽象的特徵時，我們很可能需要使用池化操作來減少輸出張量或特徵圖的尺寸，而這種操作實際上會丟失一些資訊，比如說位置資訊。
而等變對映（Equivariance）可以幫助 CNN 理解旋轉或比例等屬性變換，並相應地調整自己，這樣影像空間中的位置等屬性資訊就不會丟失。而 Geoffrey Hinton 等人提出的 CapsNet 使用向量代替標量點，因此能獲取更多的資訊。此外，我們感覺 Capsule 使用向量作為輸入與輸出是這篇論文的亮點。
Capsule 層
在論文中，Geoffrey Hinton 介紹 Capsule 為：「Capsule 是一組神經元，其輸入輸出向量表示特定實體型別的例項化引數（即特定物體、概念實體等出現的概率與某些屬性）。我們使用輸入輸出向量的長度表徵實體存在的概率，向量的方向表示例項化引數（即實體的某些圖形屬性）。同一層級的 capsule 通過變換矩陣對更高階別的 capsule 的例項化引數進行預測。當多個預測一致時（本論文使用動態路由使預測一致），更高階別的 capsule 將變得活躍。」
Capsule 中的神經元的啟用情況表示了影像中存在的特定實體的各種性質。這些性質可以包含很多種不同的引數，例如姿勢（位置，大小，方向）、變形、速度、反射率，色彩、紋理等等。而輸入輸出向量的長度表示了某個實體出現的概率，所以它的值必須在 0 到 1 之間。
為了實現這種壓縮，並完成 Capsule 層級的啟用功能，Hinton 等人使用了一個被稱為「squashing」的非線性函式。該非線性函式確保短向量的長度能夠縮短到幾乎等於零，而長向量的長度壓縮到接近但不超過 1 的情況。以下是該非線性函式的表示式：
其中 v_j 為 Capsule j 的輸出向量，s_j 為上一層所有 Capsule 輸出到當前層 Capsule j 的向量加權和，簡單說 s_j 就為 Capsule j 的輸入向量。該非線性函式可以分為兩部分，即
和
，前一部分是輸入向量 s_j 的縮放尺度，第二部分是輸入向量 s_j 的單位向量，該非線性函式既保留了輸入向量的方向，又將輸入向量的長度壓縮到區間 [0,1) 內。s_j 向量為零向量時 v_j 能取到 0，而 s_j 無窮大時 v_j 無限逼近 1。該非線性函式可以看作是對向量長度的一種壓縮和重分配，因此也可以看作是一種輸入向量後「啟用」輸出向量的方式。
那麼如上所述，Capsule 的輸入向量就相當於經典神經網路神經元的標量輸入，而該向量的計算就相當於兩層 Capsule 間的傳播與連線方式。輸入向量的計算分為兩個階段，即線性組合和 Routing，這一過程可以用以下公式表示：
其中 u_j|i hat 為 u_i 的線性組合，這一點可以看作是一般全連線網路前一層神經元以不同強弱的連線輸出到後一層某個神經元。只不過 Capsule 相對於一般神經網路每個結點都有一組神經元（以生成向量），即 u_j|i hat 表示上一層第 i 個 Capsule 的輸出向量和對應的權重向量相乘（W_ij 表示向量而不是元素）而得出的預測向量。u_j|i hat 也可以理解為在前一層為第 i 個 Capsule 的情況下連線到後一層第 j 個 Capsule 的強度。
在確定 u_j|i hat 後，我們需要使用 Routing 進行第二個階段的分配以計算輸出結點 s_j，這一過程就涉及到使用動態路由（dynamic routing）迭代地更新 c_ij。通過 Routing 就能獲取下一層 Capsule 的輸入 s_j，然後將 s_j 投入「Squashing」非線性函式後就能得出下一層 Capsule 的輸出。後面我們會重點解釋 Routing 演算法，但整個 Capsule 層及它們間傳播的過程已經完成了。
所以整個層級間的傳播與分配可以分為兩個部分，第一部分是下圖 u_i 與 u_j|i hat 間的線性組合，第二部分是 u_j|i hat 與 s_j 之間的 Routing 過程。若讀者對傳播過程仍然不是太明晰，那麼可以看以下兩層 Capsule 單元間的傳播過程，該圖是根據我們對傳播過程的理解而繪製的：
Capsule 層級結構圖
如上所示，該圖展示了 Capsule 的層級結構與動態 Routing 的過程。最下面的層級 u_i 共有兩個 Capsule 單元，該層級傳遞到下一層級 v_j 共有四個 Capsule。u_1 和 u_2 是一個向量，即含有一組神經元的 Capsule 單元，它們分別與不同的權重 W_ij（同樣是向量）相乘得出 u_j|i hat。例如 u_1 與 W_12 相乘得出預測向量 u_2|1 hat。隨後該預測向量和對應的「耦合係數」c_ij 相乘並傳入特定的後一層 Capsule 單元。不同 Capsule 單元的輸入 s_j 是所有可能傳入該單元的加權和，即所有可能傳入的預測向量與耦合係數的乘積和。隨後我們就得到了不同的輸入向量 s_j，將該輸入向量投入到「squashing」非線性函式就能得出後一層 Capsule 單元的輸出向量 v_j。然後我們可以利用該輸出向量 v_j 和對應預測向量 u_j|i hat 的乘積更新耦合係數 c_ij，這樣的迭代更新不需要應用反向傳播。
Dynamic Routing 演算法
因為按照 Hinton 的思想，找到最好的處理路徑就等價於正確處理了影像，所以在 Capsule 中加入 Routing 機制可以找到一組係數 c_ij，它們能令預測向量 u_j|i hat 最符合輸出向量 v_j，即最符合輸出的輸入向量，這樣我們就找到了最好的路徑。
按照原論文所述，c_ij 為耦合係數（coupling coefficients），該係數由動態 Routing 過程迭代地更新與確定。Capsule i 和後一層級所有 Capsule 間的耦合係數和為 1，即圖四 c_11+c_12+c_13+c_14=1。此外，該耦合係數由「routing softmax」決定，且 softmax 函式中的 logits b_ij 初始化為 0，耦合係數 c_ij 的 softmax 計算方式為：
b_ij 依賴於兩個 Capsule 的位置與型別，但不依賴於當前的輸入影像。我們可以通過測量後面層級中每一個 Capsule j 的當前輸出 v_j 和 前面層級 Capsule i 的預測向量間的一致性，然後藉助該測量的一致性迭代地更新耦合係數。本論文簡單地通過內積度量這種一致性，即
，這一部分也就涉及到使用 Routing 更新耦合係數。
Routing 過程就是上圖 4 右邊表述的更新過程，我們會計算 v_j 與 u_j|i hat 的乘積並將它與原來的 b_ij 相加而更新 b_ij，然後利用 softmax(b_ij) 更新 c_ij 而進一步修正了後一層的 Capsule 輸入 s_j。當輸出新的 v_j 後又可以迭代地更新 c_ij，這樣我們不需要反向傳播而直接通過計算輸入與輸出的一致性更新引數。
該 Routing 演算法更具體的更新過程可以檢視以下虛擬碼：
對於所有在 l 層的 Capsule i 和在 l+1 層的 Capsule j，先初始化 b_ij 等於零。然後迭代 r 次，每次先根據 b_i 計算 c_i，然後在利用 c_ij 與 u_j|i hat 計算 s_j 與 v_j。利用計算出來的 v_j 更新 b_ij 以進入下一個迭代迴圈更新 c_ij。該 Routing 演算法十分容易收斂，基本上通過 3 次迭代就能有不錯的效果。
CapsNet 架構
Hinton 等人實現了一個簡單的 CapsNet 架構，該架構由兩個卷積層和一個全連線層組成，其中第一個為一般的卷積層，第二個卷積相當於為 Capsule 層做準備，並且該層的輸出為向量，所以它的維度要比一般的卷積層再高一個維度。最後就是通過向量的輸入與 Routing 過程等構建出 10 個 v_j 向量，每一個向量的長度都直接表示某個類別的概率。
以下是 CapsNet 的整體架構：
第一個卷積層使用了 256 個 9×9 卷積核，步幅為 1，且使用了 ReLU 啟用函式。該卷積操作應該沒有使用 Padding，輸出的張量才能是 20×20×256。此外，CapsNet 的卷積核感受野使用的是 9×9，相比於其它 3×3 或 5×5 的要大一些，這個能是因為較大的感受野在 CNN 層次較少的情況下能感受的資訊越多。這兩層間的權值數量應該為 9×9×256+256=20992。
隨後，第二個卷積層開始作為 Capsule 層的輸入而構建相應的張量結構。我們可以從上圖看出第二層卷積操作後生成的張量維度為 6×6×8×32，那麼我們該如何理解這個張量呢？雲夢居客在知乎上給出了一個十分形象且有意思的解釋，如前面章節所述，如果我們先考慮 32 個（32 channel）9×9 的卷積核在步幅為 2 的情況下做卷積，那麼實際上得到的是傳統的 6×6×32 的張量，即等價於 6×6×1×32。
因為傳統卷積操作每次計算的輸出都是一個標量，而 PrimaryCaps 的輸出需要是一個長度為 8 的向量，因此傳統卷積下的三維輸出張量 6×6×1×32 就需要變化為四維輸出張量 6×6×8×32。如下所示，其實我們可以將第二個卷積層看作對維度為 20×20×256 的輸入張量執行 8 次不同權重的 Conv2d 操作，每次 Conv2d 都執行帶 32 個 9×9 卷積核、步幅為 2 的卷積操作。
由於每次卷積操作都會產生一個 6×6×1×32 的張量，一共會產生 8 個類似的張量，那麼將這 8 個張量（即 Capsule 輸入向量的 8 個分量）在第三個維度上合併在一起就成了 6×6×8×32。從上可知 PrimaryCaps 就相當於一個深度為 32 的普通卷積層，只不過每一層由以前的標量值變成了長度為 8 的向量。
此外，結合 Hinton 等人給出的 Capsule 定義，它就相當於一組常見的神經元，這些神經元封裝在一起形成了新的單元。在本論文討論的 CapsNet 架構中，我們將 8 個卷積單元封裝在一起成為了一個新的 Caosule 單元。PrimaryCaps 層的卷積計算都沒有使用 ReLU 等啟用函式，它們以向量的方式預備輸入到下一層 Capsule 單元中。
PrimaryCaps 每一個向量的分量層級是共享卷積權重的，即獲取 6×6 張量的卷積核權重為相同的 9×9 個。這樣該卷積層的引數數量為 9×9×256×8×32+8×32=5308672，其中第二部分 8×32 為偏置項引數數量。
第三層 DigitCaps 在第二層輸出的向量基礎上進行傳播與 Routing 更新。第二層共輸出 6×6×32=1152 個向量，每一個向量的維度為 8，即第 i 層共有 1152 個 Capsule 單元。而第三層 j 有 10 個標準的 Capsule 單元，每個 Capsule 的輸出向量有 16 個元素。前一層的 Capsule 單元數是 1152 個，那麼 w_ij 將有 1152×10 個，且每一個 w_ij 的維度為 8×16。當 u_i 與對應的 w_ij 相乘得到預測向量後，我們會有 1152×10 個耦合係數 c_ij，對應加權求和後會得到 10 個 16×1 的輸入向量。將該輸入向量輸入到「squashing」非線性函式中求得最終的輸出向量 v_j，其中 v_j 的長度就表示識別為某個類別的概率。
DigitCaps 層與 PrimaryCaps 層之間的引數包含兩類，即 W_ij 和 c_ij。所有 W_ij 的引數數量應該是 6×6×32×10×8×16=1474560，c_ij 的引數數量為 6×6×32×10×16=184320，此外還應該有 2×1152×10=23040 個偏置項引數，不過原論文並沒有明確指出這些偏置項。最後小編計算出該三層 CapsNet 一共有 5537024 個引數，這並不包括後面的全連線重構網路引數。（算錯了不要怪小編呦～）
損失函式與最優化
前面我們已經瞭解 DigitCaps 層輸出向量的長度即某個類別的概率，那麼我們該如何構建損失函式，並根據該損失函式迭代地更新整個網路？前面我們耦合係數 c_ij 是通過一致性 Routing 進行更新的，他並不需要根據損失函式更新，但整個網路其它的卷積引數和 Capsule 內的 W_ij 都需要根據損失函式進行更新。一般我們就可以對損失函式直接使用標準的反向傳播更新這些引數，而在原論文中，作者採用了 SVM 中常用的 Margin loss，該損失函式的表示式為：
其中 c 是分類類別，T_c 為分類的指示函式（c 存在為 1，c 不存在為 0），m+ 為上邊界，m- 為下邊界。此外，v_c 的模即向量的 L2 距離。
因為例項化向量的長度來表示 Capsule 要表徵的實體是否存在，所以當且僅當圖片裡出現屬於類別 k 的手寫數字時，我們希望類別 k 的最頂層 Capsule 的輸出向量長度很大（在本論文 CapsNet 中為 DigitCaps 層的輸出）。為了允許一張圖裡有多個數字，我們對每一個表徵數字 k 的 Capsule 分別給出單獨的 Margin loss。
構建完損失函式，我們就能愉快地使用反向傳播了。
重構與表徵
重構即我們希望利用預測的類別重新構建出該類別代表的實際影像，例如我們前面部分的模型預測出該圖片屬於一個類別，然後後面重構網路會將該預測的類別資訊重新構建成一張圖片。
前面我們假設過 Capsule 的向量可以表徵一個例項，那麼如果我們將一個向量投入到後面的重構網路中，它應該能重構出一個完整的影像。因此，Hinton 等人使用額外的重構損失（reconstruction loss）來促進 DigitCaps 層對輸入數字圖片進行編碼。下圖展示了整個重構網路的的架構：
我們在訓練期間，除了特定的 Capsule 輸出向量，我們需要矇住其它所有的輸出向量。然後，使用該輸出向量重構手寫數字影像。DigitCaps 層的輸出向量被饋送至包含 3 個全連線層的解碼器中，並以上圖所示的方式構建。這一過程的損失函式通過計算 FC Sigmoid 層的輸出畫素點與原始影像畫素點間的歐幾里德距離而構建。Hinton 等人還按 0.0005 的比例縮小重構損失，以使它不會主導訓練過程中的 Margin loss。
Capsule 輸出向量的重構與表徵除了能提升模型的準確度以外，還能提升模型的可解釋性，因為我們能修正需要重構向量中的某個或某些分量而觀察重構後的影像變化情況，這有助於我們理解 Capsule 層的輸出結果。
以上就是本論文構建的 CapsNet 架構，當然 Hinton 還描述了很多試驗結果與發現，感興趣的讀者可以查閱論文的後一部分。
CapsNet 的 TensorFlow 實現
以下定義構建 CapsNet 後面兩層的方法。在 CapsNet 架構中，我們能訪問該類中的物件和方法構建 PrimaryCaps 層和 DigitCaps 層。
下面是整個 CapsNet 的架構與推斷過程程式碼，我們需要從 MNIST 抽出影像並投入到以下定義的方法中，該批量的影像將先通過三層 CapsNet 網路輸出 10 個類別向量，每個向量有 16 個元素，且每個類別向量的長度為輸出影像是該類別的概率。隨後，我們會將一個向量投入到重構網路中構建出該向量所代表的影像。
以上是該網路的主體程式碼，更多程式碼請檢視 naturomics 的 GitHub 地址，或機器之心的 GitHub 地址，我們上傳的是帶註釋的程式碼，希望能幫助初學者更加理解 CapsNet 的過程與架構。以下是上面我們定義 CapsNet 的主體計算圖，即 TensorFlow 中的靜態計算圖：
我們也迭代訓練了大概 3 萬多步，不過因為使用的是 CPU，所以我們將批量大小調整為 8 以減少單次迭代的計算壓力，以下是我們訓練時的損失情況，最上面是 Margin loss，下面還有重構損失和總損失：
最後放上兩張由 DigitCaps 層輸出向量重構出的對應影像：
我們只是初步地探索了 CapsNet，它還存在很多的可能性，例如它以向量的形式應該能獲取非常多的影像資訊，這種優勢是否能在其它大型資料集或平面 3D 影像資料集中進一步展現出非凡的表徵力？而且第二層 PrimaryCaps 的引數非常多，就像一組橫向並聯的卷積結構以產生向量（類似 Inception 模組，但要寬地多），我們是否能通過某種方式的共享進一步減少該層級的引數？還有當前 Routing 過程的效果至少在 MNIST 資料集中並不好，它僅僅只能展示存在這個概念，那麼我們能否找到更加高效的 Routing 演算法？此外，Capsule 是否能擴充套件到其他神經網路結構如迴圈或門控單元？這些可能都是我們存在的疑惑，但向前走，時間總會給我們答案的。
歡迎大家留言討論，本文在機器之心網站上將持續更新與修正。
參考資料
- 原論文：Dynamic Routing Between Capsules（https://arxiv.org/abs/1710.09829）
- 知乎討論地址：https://www.zhihu.com/question/67287444/answer/251241736
- naturomics 實現地址（TensorFlow）：https://github.com/naturomics/CapsNet-Tensorflow
- XifengGuo 實現地址（Keras）：https://github.com/XifengGuo/CapsNet-Keras
- leftthomas 實現地址（Pytorch）：https://github.com/leftthomas/CapsNet
本文為機器之心原創，轉載請聯絡本公眾號獲得授權。