CapsNet再升級!堆棧式膠囊自編碼器面世

2017 年，Geoffrey Hinton 在論文《Dynamic Routing Between Capsules》中提出 CapsNet 引起了極大的關(guān)注，同時(shí)也提供了一個(gè)全新的研究的方向。今日，CapsNet 的作者 Sara Sabour、Hinton 老爺子聯(lián)合牛津大學(xué)的研究者提出了膠囊網(wǎng)絡(luò)的改進(jìn)版本——堆棧式膠囊自編碼器。這種膠囊自編碼器可以無監(jiān)督地學(xué)習(xí)圖像中的特征，并在無監(jiān)督分類任務(wù)取得最佳或接近最佳的表現(xiàn)。這也是膠囊網(wǎng)絡(luò)第一次在無監(jiān)督領(lǐng)域取得新的突破。

綜述

一個(gè)目標(biāo)可以被看做是一組相互關(guān)聯(lián)的部件按照幾何學(xué)形式組合的結(jié)果。利用這種幾何關(guān)系去重建目標(biāo)的系統(tǒng)應(yīng)當(dāng)對視點(diǎn)的變化具有魯棒性，因?yàn)槠浔举|(zhì)的幾何關(guān)系不應(yīng)隨著觀察視角的變化而發(fā)生改變。

本文中，研究人員描述了一種無監(jiān)督的膠囊網(wǎng)絡(luò)。其中，觀察組成目標(biāo)所有部件的神經(jīng)編碼器被用來推斷目標(biāo)膠囊的存在和姿態(tài)。編碼器通過解碼器的反向傳播方法訓(xùn)練。

訓(xùn)練中，解碼器使用姿態(tài)預(yù)測來預(yù)測每個(gè)已發(fā)現(xiàn)部件的姿態(tài)。這些部件是直接從圖像中被發(fā)現(xiàn)的，同樣也是使用神經(jīng)編碼器，該編碼器推斷這些部件及它們的仿射變換。

而對應(yīng)的解碼器將每個(gè)圖像像素建模為由仿射變換部件做出的預(yù)測混合。研究人員從目標(biāo)和目標(biāo)部件的膠囊中學(xué)習(xí)無標(biāo)簽數(shù)據(jù)，然后將這些目標(biāo)膠囊的存在向量進(jìn)行聚類。

得知這些聚類的名稱時(shí)，研究人員在 SVHN 和 MNIST 數(shù)據(jù)集上獲得了當(dāng)前最佳的無監(jiān)督分類結(jié)果，準(zhǔn)確率分別為 55% 和 98.5%。

本文提出了堆棧式膠囊自編碼器（SCAE），該編碼器包含兩個(gè)階段。在第一階段，部件膠囊自編碼器（PCAE）將圖像分割為組成部分，推斷其姿態(tài)，并將每個(gè)圖像像素重建為變換組件模板的像素混合。

在第二階段，目標(biāo)膠囊自編碼器（OCAE）嘗試將發(fā)現(xiàn)的部件及其姿態(tài)安排在一個(gè)更小的目標(biāo)集合中。這個(gè)目標(biāo)集合對每個(gè)部件進(jìn)行預(yù)測，從而解釋每個(gè)部件的姿態(tài)。通過將它們的姿態(tài)——目標(biāo)-觀察者關(guān)系（OV）和相關(guān)的目標(biāo)-部件關(guān)系（OP）相乘，每個(gè)目標(biāo)膠囊都會貢獻(xiàn)這些混合的一部分。

堆棧式膠囊自編碼器的工作原理

堆棧式膠囊自編碼器在使用未標(biāo)注數(shù)據(jù)訓(xùn)練時(shí)捕捉所有目標(biāo)和它們部件之間的空間關(guān)系。目標(biāo)膠囊存在概率的向量傾向于組成緊密的聚類。

當(dāng)給每個(gè)聚類一個(gè)分類時(shí)，其可以在無監(jiān)督分類任務(wù)上達(dá)到當(dāng)前最佳效果，如 SVHN 數(shù)據(jù)集上的 55% 和 MNIST 數(shù)據(jù)集上的 98.5%。以上結(jié)果還可以分別提升到 67% 和 99%，而且只需學(xué)習(xí)不到 300 個(gè)參數(shù)。

模型架構(gòu)

堆棧式膠囊自編碼器的結(jié)構(gòu)

堆棧式膠囊自編碼器的兩個(gè)組成部分為：部件膠囊自編碼器（PCAE）和目標(biāo)膠囊自編碼器（OCAE）。在下文中，論文首先介紹了集群自編碼器（CCAE），通過一系列數(shù)學(xué)公式說明自編碼器如何分解圖像中的部件的過程，然后由此引出堆棧式膠囊自編碼器的兩個(gè)組成部分。

集群自編碼器

圖 2：使用集群自編碼器對不同形狀的點(diǎn)進(jìn)行聚類的示意圖。

論文首先介紹了集群自編碼器，通過這種結(jié)構(gòu)的數(shù)學(xué)原理，引出堆棧式膠囊自編碼器的結(jié)構(gòu)。令 {x_m | m = 1, . . . , M } 為一組二維的輸入點(diǎn)，每個(gè)點(diǎn)屬于一個(gè)集群（見圖2）。首先使用Set Transformer將所有的輸入點(diǎn)（相當(dāng)于部件膠囊）編碼進(jìn)k個(gè)目標(biāo)膠囊中，Set Transformer是一種基于注意力機(jī)制的、有置換不變性的編碼器h^cap (Lee et al., 2019) 。

一個(gè)目標(biāo)膠囊 k 包括一個(gè)膠囊特征向量 c_k（其存在概率 a_k ∈ [0, 1]）和一個(gè) 3 × 3 的目標(biāo)-觀察者（OV）關(guān)系矩陣。關(guān)系矩陣代表著目標(biāo)（集群）和觀察者之間關(guān)系的仿射變換。

需要注意的是，每個(gè)目標(biāo)膠囊每次只能代表一個(gè)目標(biāo)。每個(gè)目標(biāo)膠囊都使用一個(gè)獨(dú)立的多層感知機(jī) h_k^part 從膠囊特征向量 c_k 中預(yù)測 N ≤ M 個(gè)候選部件。

每個(gè)候選由條件概率 a_k,n ∈ [0, 1] (當(dāng)其存在)，一個(gè)關(guān)聯(lián)標(biāo)量的標(biāo)準(zhǔn)差λ_k,n，以及一個(gè) 3 × 3 的目標(biāo)-部件（OP）關(guān)系矩陣組成。這些代表著目標(biāo)膠囊和候選部件的仿射變換。

候選預(yù)測 μ_k,n 根據(jù)目標(biāo)膠囊 OV 和候選 OP 矩陣相乘得來。然后，研究人員將每個(gè)輸入部件建模為高斯混合模型，其中μ_k,n 和 λ_k,n 是各向同性組件的中心和標(biāo)準(zhǔn)差。其標(biāo)準(zhǔn)公式如下：

集群膠囊編碼器的公式。論文通過舉出集群膠囊編碼器的例子，用于說明目標(biāo)膠囊編碼器和它的區(qū)別。

部件膠囊自編碼器

如果要將圖像分解為組成部件的集合關(guān)系，就需要首先推斷圖像是由哪些部件組成的，同時(shí)也需要了解觀察者和這些部件之間的關(guān)系（稱之為他們的姿態(tài)）。

在本研究中，每個(gè)部件膠囊都有六個(gè)維度的自由姿態(tài)，一個(gè)存在變量，和一個(gè)獨(dú)特的特征。研究人員把部件發(fā)現(xiàn)問題視為自編碼：編碼器學(xué)習(xí)去推斷不同部件膠囊的姿態(tài)和存在，而解碼器學(xué)習(xí)每個(gè)部件的圖像模板。

模板對應(yīng)的部件是使用其姿態(tài)的仿射變換，而這些變換過的模板的像素點(diǎn)被用來為每個(gè)圖像像素創(chuàng)建單獨(dú)的混合模型。在部件膠囊自編碼器后是目標(biāo)膠囊自編碼器。

令 y ∈ [0, 1]^h×w×c 為圖像。研究人員將部件膠囊的數(shù)量限定在 M 之內(nèi)。對于每個(gè)部件膠囊，他們使用一個(gè)編碼器去推斷姿態(tài) x_m ∈ R^6，存在概率 d_m ∈ [0, 1]，以及特殊特征 z_m ∈ R^c_z。

雖然后者不會直接參與圖像重建，但是會將對應(yīng)部件的特殊信息提供給目標(biāo)膠囊自編碼器。他們會通過目標(biāo)膠囊自編碼器使用反向傳播微分的方式訓(xùn)練。

當(dāng)前條件下，不允許圖像中同一種類型的部件多次出現(xiàn)，從而導(dǎo)致部件膠囊不會在空間中被復(fù)制（盡管它們可能會）。然而，確實(shí)需要分辨出所有出現(xiàn)在圖像中的部件，因此編碼器會采用帶有從下到上（bottom-up）注意力機(jī)制的卷積神經(jīng)網(wǎng)絡(luò)。

對于每個(gè)膠囊 k，其預(yù)測一個(gè)特征矩陣 e^k，特征矩陣是 6（姿態(tài)）+1（存在）+c_z（特殊特征）的膠囊參數(shù)，其空間維度是 h_e × w_e，以及一個(gè)單通道注意力層 a_k。

最終，該膠囊的參數(shù)計(jì)算公式是。softmax 是對空間維度上的計(jì)算。這種計(jì)算有點(diǎn)類似于全局平均池化，但是允許一些空間點(diǎn)比其他點(diǎn)對最終結(jié)果的權(quán)重影響更大。研究人員將其稱為注意力池化（attention-based pooling）。

圖像的像素點(diǎn)被建模為獨(dú)立的高斯混合模型。對于每個(gè)像素點(diǎn)，研究人員采用其對應(yīng)的變換模板，并將其視為有著恒定方差的各向同性高斯組件的中心點(diǎn)。其混合概率對部件膠囊的存在概率和在該位置的色值函數(shù)（c 指的是圖像的通道數(shù)）都是成比例的。

部件膠囊自編碼器的公式推導(dǎo)過程

目標(biāo)膠囊自編碼器（OCAE）

下一步是從已經(jīng)發(fā)現(xiàn)的部件中尋找目標(biāo)。因此，需要使用相連的姿態(tài) x_m，特殊特征 z_m，以及平滑化的模板 T_m（通過將部件膠囊的特征進(jìn)行轉(zhuǎn)化）。這些將會成為目標(biāo)膠囊自編碼器的輸入，這里和集群自編碼器有一些不同。

首先，研究人員將部件膠囊的存在概率 d_m 輸入目標(biāo)膠囊自編碼器——由于平衡注意力機(jī)制，避免將缺失點(diǎn)考慮在內(nèi)。

其次，d_m 同時(shí)用于衡量部件膠囊的對數(shù)似然 cf。另外，除了特殊特征外，不對其他目標(biāo)膠囊自編碼器的輸入計(jì)算梯度，以便提升訓(xùn)練的穩(wěn)定性，并避免隱變量崩潰。

最后，通過部件膠囊自編碼器發(fā)現(xiàn)的部件有著獨(dú)立的特征（模板和特殊特征）。因此，每個(gè)部件姿態(tài)都可以被解釋為是目標(biāo)膠囊預(yù)測的獨(dú)立混合——即每個(gè)目標(biāo)膠囊都做出 M 個(gè)候選預(yù)測 V_k,1:M，或者對每個(gè)部件做出一個(gè)候選預(yù)測。

最終，部件膠囊的似然公式是：

圖 3：從MNIST（左）和SVHN（中）和CIFAR 10（右）學(xué)習(xí)到的模板。

圖 4：展示了膠囊自編碼器對MNIST數(shù)據(jù)集的重建過程。a）MNIST圖像；b）紅色的部件膠囊和綠色的目標(biāo)膠囊在重建中的組合；c）實(shí)際參與重建的被激活膠囊；d）根據(jù)圖像捕捉到的信息；e）部件的仿射變換，用于展示其重建圖像的過程。

模型性能評估

堆棧式膠囊自編碼器使用仿射變換，這樣可以使編碼器的輸入由一組較小的變換目標(biāo)或部件解釋。

無監(jiān)督分類評價(jià)

研究人員在 MNIST、SVHN 和 CIFAR 10 數(shù)據(jù)集上進(jìn)行了測試，并將目標(biāo)膠囊的存在打上類別標(biāo)簽。他們使用了多種評價(jià)方法。

在部件膠囊編碼器上，研究人員在 MNIST 數(shù)據(jù)集上使用了 24 個(gè)單通道，11 × 11 的模板，在 SVHN 和 CIFAR 10 上則分別使用了 32 個(gè) 3 通道，14 × 14 的模板。

對于后兩個(gè)數(shù)據(jù)集的圖像，研究人員進(jìn)行了 Sobel 過濾，作為重建的目標(biāo)。對于目標(biāo)膠囊編碼器，研究則分別使用了 24、32 和 64 個(gè)目標(biāo)膠囊。

表 1：運(yùn)行五次后取平均的無監(jiān)督分類結(jié)果和標(biāo)準(zhǔn)差。