Efficient Inference in Fully Connected CRFs with Gaussian Edge Potentials

Intro

回顧一下經典，文章是早期CNN直接做分割效果不好的情況下，利用條件隨機場建模，解決細節分割問題的一篇文章，文章主要貢獻是提出了Mean Field Approximation來近似原來的全連線條件隨機場，同時提出一個Efficient演算法將整個更新過程變為線性時間。

Why use CRF

在早期分割領域，FCN直接分割的效果並不好，當時並沒有太多的結構解決這一問題，而效果不好的原因基本上來自於下采樣和網路深度過深導致的資訊損失，輸出層難以利用這樣的資訊擬合這樣的任務，文章給了一個圖來描述這樣的情況，其中圖(b)就是直接分割的效果：

The Fully Connected CRF Model

為了解決直接分割效果不好的問題，一般利用CRF對影象進行建模，圖(c)是基於patch的CRF，概率圖中依賴關係只和一定patch內的pixel有關，圖(d)是全連線CRF，即一個節點和圖中其他所有節點有關。

定義兩隨機場\(\mathbf{X} = \{X_1, X_2, \cdots X_N \}、\mathbf{I} = \{ I_1, I_2, \cdots I_N \}\)，其中X表示標籤變數，I表示影象資訊，\(X_i\)定義在\(\mathcal{L} = \{l_1, l_2, \cdots, l_k\}\)，\(I_i\)

則表示索引為i的畫素所攜帶的顏色資訊。

對於條件隨機場\((\mathbf{X},\mathbf{I})\)，其概率分佈定義為：

\[P(\mathbf{X}|\mathbf{I}) = \frac{1}{Z(\mathbf{I})} exp(-E(\mathbf{X}, \mathbf{I})) \]

其中E為能量函式，Z為歸一化項，分別定義如下：

\[E(\mathbf{x}|\mathbf{I})=\sum_{i} \psi_{u}\left(x_{i}|\mathbf{I}\right)+\sum_{i<j} \psi_{p}\left(x_{i}, x_{j} | \mathbf{I}\right) \]\[Z(\mathbf{I}) = \sum_{x} exp(-E(\mathbf{X}, \mathbf{I})) \]

其能量函式第一項為一元勢函式，描述的是某一節點的勢能，可以用神經網路的輸出層進行擬合；第二項為二元勢函式，描述的是節點i和無向圖中的其他所有和i有邊的節點的關係勢能，而這一項的定義如下：

\[\psi_{p}\left(x_{i}, x_{j}\right| \mathbf{I})=\mu\left(x_{i}, x_{j}\right) \underbrace{\sum_{m=1}^{K} w^{(m)} k^{(m)}\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right)}_{k\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right)} \]

其中\(k\left(\mathbf{f}*{i}, \mathbf{f}*{j}\right)=w^{(1)} \exp \left(-\frac{\left|p_{i}-p_{j}\right|^{2}}{2 \theta_{\alpha}^{2}}-\frac{\left|I_{i}-I_{j}\right|^{2}}{2 \theta_{\beta}^{2}}\right)+w^{(2)} \exp \left(-\frac{\left|p_{i}-p_{j}\right|^{2}}{2 \theta_{\gamma}^{2}}\right)\)為高斯函式，該公式中p表示位置，第一項描述的是顏色相似性和位置相似性，是基於“距離較近且顏色相似的節點可能屬於同一標籤類”這一先驗，第二項描述的是平滑性；f為人工設計的特徵，\(w\)為不同類別的權重，\(\mu\)是一指示函式，在兩者標籤不同時為1，否則為0，即對於標籤相同的情況，並不進行能量傳導，而標籤不通的情況才進行能量傳導. \(\theta\)為可學習的引數(\(w\)其實也是可學習的，但是在實驗中發現對\(w^{(2)}\))設定為1效果就挺好了，所以這一引數是超參。

關於為什麼一元勢函式要用神經網路擬合，而二元勢函式要在原圖上直接利用“人工設計”的特徵來描述，是因為CNN本身感受野有限，網路輸出feature的某一pixel位置只和其領域一定範圍內的畫素有關，無法建模全域性關係如全域性的二元關係，但描述該位置的資訊是足夠的；而用“人工設計”的特徵來描述二元勢函式則是根據顏色先驗和平滑先驗來的，神經網路難以對這兩個先驗進行建模描述，而我們卻可以很方便的通過影象本身運算得到量化描述（儘管需要優化建模的引數）。

對於總畫素點數為N的一張圖，其邊的數量可以達到\(N^2\)，上式的二元勢函式的每個節點則需要進行\(N^2\)的運算，本文提出了mean field來近似fully connected crf。

Mean Field Approximation

對於\(P(\mathbf{X}|\mathbf{I})\)，文章希望求的一個分佈\(Q(\mathbf{X}|\mathbf{I})\)來近似\(P(\mathbf{X}|\mathbf{I})\)，文章中是用\(P(\mathbf{X})\)和\(Q(\mathbf{X})\)來表示，省略了後面的條件量，表達的意思是一致的。為了和文章保持一致，後面沿用文章的寫法。

其中，對於\(Q(\mathbf{X})\)需要滿足\(Q(\mathbf{X})=\prod_{i} Q_{i}\left(X_{i}\right)\)，因此只需要最大化分佈P和分佈Q之間的KL散度，就可以實現對P的近似了。

對於

\[P(\mathbf{X})=\frac{1}{Z} \tilde{P}(\mathbf{X})=\frac{1}{Z} \exp \left(-\sum_{i} \psi_{u}\left(x_{i}\right)-\sum_{i<j} \psi_{p}\left(x_{i}, x_{j}\right)\right) = \frac{1}{Z} \exp \left(-E(\mathbf{X})\right) \]

這裡原文應該把exp裡的正負搞錯了，糾正一下。

P和Q的KL散度為：

\[\begin{aligned}\mathbf{D}(Q \| P) &=\sum_{\mathbf{x}} Q(\mathbf{x}) \log \left(\frac{Q(\mathbf{x})}{P(\mathbf{x})}\right) \\&=-\sum_{\mathbf{x}} Q(\mathbf{x}) \log P(\mathbf{x})+\sum_{\mathbf{x}} Q(\mathbf{x}) \log Q(\mathbf{x}) \\&=-\mathbf{E}_{\mathbf{U} \sim Q}[\log P(\mathbf{U})]+\mathbf{E}_{\mathbf{U} \sim Q}[\log Q(\mathbf{U})] \\&=-\mathbf{E}_{\mathbf{U} \sim Q}[\log \tilde{P}(\mathbf{U})]+\mathbf{E}_{\mathbf{U} \sim Q}[\log Z]+\sum_{i} \mathbf{E}_{U_{i} \sim Q}\left[\log Q\left(U_{i}\right)\right] \\&=\mathbf{E}_{\mathbf{U} \sim Q}[E(\mathbf{U})]+\sum_{i} \mathbf{E}_{U_{i} \sim Q_{i}}\left[\log Q_{i}\left(U_{i}\right)\right]+\log Z\end{aligned} \]

第三行到第四行是進行了代換；第四行到第五行是能量函式代換且Z只和P有關，因此可以從期望中提取出來。

此外，Q的約束條件為：

\[\sum_{x_i} Q_i(x_i) = 1 \]

這是因為邊際概率和為1。

根據拉格朗日乘數法，構造拉氏函式

\[L = \mathbf{E}_{\mathbf{U} \sim Q}[E(\mathbf{U})]+\sum_{i} \mathbf{E}_{U_{i} \sim Q_{i}}\left[\log Q_{i}\left(U_{i}\right)\right]+\log Z + \lambda \sum_{x_i} (Q_i(x_i) - 1) \]

因此：

\[\begin{equation} \begin{aligned} \frac{\partial L}{\partial Q_i(x_i)} &= \frac{\partial L}{\partial Q_i(x_i)} \left ( \sum_i \mathbf{E}_{U_i \sim Q_i}\left[ \psi_{u}\left(U_{i}\right) \right ] + \sum_{i<j} \mathbf{E}_{U_i \sim Q_i, U_j \sim Q_j} \left[ \psi_{p}\left(U_{i}, U_{j}\right )\right ] \right. \\ &+ \left.\sum_{i} \mathbf{E}_{U_{i} \sim Q_{i}}\left[\log Q_{i}\left(U_{i}\right)\right]+\log Z + \lambda \sum_{x_i} (\log Q_i(x_i) - 1) \right) \\ &= \underline{\psi_u(x_i)} + \underline{\frac{\partial L}{\partial Q_i(x_i)}\left(\frac{1}{2} \sum_i Q_i(U_i)\sum_{i\neq j}\mathbf{E}_{U_j \sim Q_j} \left [ \psi_p(U_i, U_j) \right] \right)} \\ &+ \underline{(\log Q_i(x_i) +1)} + \underline{0} + \underline{\lambda} \\ & =\underline{\psi_u(x_i)} + \underline{\left(\frac{1}{2} \sum_{i\neq j}\mathbf{E}_{U_j \sim Q_j} \left [ \psi_p(x_i, U_j) \right] \right)} + \underline{(\log Q_i(x_i) +1)} + \underline{0} + \underline{\lambda} \\ \end{aligned} \end{equation} \]

每一項對應的求導我用下劃線標出來了，方便區分；第三行是因為將\(i<j\)換成了\(i\neq j\)，所以加上了一個\(\frac{1}{2}\)，之所以換成不等號，是和前面說到的指示函式對應。

接下來讓其導數等於0，求解Q：

\[\psi_u(x_i) + \left(\frac{1}{2} \sum_{i\neq j}\mathbf{E}_{U_j \sim Q_j} \left [ \psi_p(x_i, U_j) \right] \right) + (\log Q_i(x_i) +1) + 0 + \lambda = 0 \]\[\log Q_{i}\left(x_{i}\right)=-\psi_{u}\left(x_{i}\right)-\frac{1}{2} \sum_{j \neq i} \mathbf{E}_{U_{j} \sim Q_{j}}\left[\psi_{p}\left(x_{i}, U_{j}\right)\right]-1-\lambda \]

所以：

\[\begin{equation} \begin{aligned} Q_{i}\left(x_{i}\right)&=\exp \left\{-\psi_{u}\left(x_{i}\right)-\frac{1}{2} \sum_{j \neq i} \mathbf{E}_{U_{j} \sim Q_{j}}\left[\psi_{p}\left(x_{i}, U_{j}\right)\right]-1-\lambda\right\} \\ & =\frac{1}{Z_{i}} \exp \left\{-\psi_{u}\left(x_{i}\right)-\sum_{j \neq i} \mathbf{E}_{U_{j} \sim \mathrm{Q}_{j}}\left[\psi_{p}\left(x_{i}, U_{j}\right)\right]\right\} \end{aligned} \end{equation} \]

之所以把\(\frac{1}{2}\)去掉這裡我理解是後面的期望項展開是含有可學習引數\(w\)的，因此前面的係數並不重要；後面的常數項被歸到normalize項Z裡，是因為Q的求和等於1，最終是需要normalize的，可以直接將其歸進去而不影響結果。

將相關的變數帶入：

\[\begin{aligned}Q_{i}\left(x_{i}=l\right) &=\frac{1}{Z_{i}} \exp \left\{-\psi_{u}\left(x_{i}\right)-\sum_{j \neq i} \mathbf{E}_{U_{j} \sim Q_{j}}\left[\mu\left(l, U_{j}\right) \sum_{m=1}^{K} w^{(m)} k^{(m)}\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right)\right]\right\} \\&=\frac{1}{Z_{i}} \exp \left\{-\psi_{u}\left(x_{i}\right)-\sum_{m=1}^{K} w^{(m)} \sum_{j \neq i} \mathbf{E}_{U_{j} \sim Q_{j}}\left[\mu\left(l, U_{j}\right) k^{(m)}\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right)\right]\right\} \\&=\frac{1}{Z_{i}} \exp \left\{-\psi_{u}\left(x_{i}\right)-\sum_{m=1}^{K} w^{(m)} \sum_{j \neq i} \sum_{l^{\prime} \in \mathcal{L}} Q_{j}\left(l^{\prime}\right) \mu\left(l, l^{\prime}\right) k^{(m)}\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right)\right\} \\&=\frac{1}{Z_{i}} \exp \left\{-\psi_{u}\left(x_{i}\right)-\sum_{l^{\prime} \in \mathcal{L}} \mu\left(l, l^{\prime}\right) \sum_{m=1}^{K} w^{(m)} \sum_{j \neq i} k^{(m)}\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right) Q_{j}\left(l^{\prime}\right)\right\}\end{aligned} \]

很不幸的是，對於\(w^{(1)}、\theta\)的學習，並沒有利用梯度來學習，主要是速度上的考慮，可能還不如直接暴力搜尋快和效果好；文章是利用grid search來搜尋的，即預設一些值，利用暴力法挨個嘗試在一個固定的驗證集上選擇最佳引數。

再貼一下文章中收斂的一個大致效果：

文章中還提到指示函式的梯度學習，前面我們將其認為是不可學習的，將結果根據標籤是否一致進行一刀切，顯然比較hard，文章後面認為指示函式本身可以作為可學習引數進行學習，而這一塊可以利用梯度方法進行學習。

對於訓練影象\(\mathcal{I}\)和其對應的分割標註\(\mathcal{T}\)，對於學習引數\(\mu\)，其損失表示為：

\[\begin{aligned}\ell\left(\mu: \mathcal{T}^{(n)}, \mathcal{I}^{(n)}\right) &=\log P\left(\mathbf{X}=\mathcal{T}^{(n)} \mid \mathcal{I}^{(n)}, \mu\right) \\&=-E\left(\mathcal{T}^{(n)} \mid \mathcal{I}^{(n)}, \mu\right)-\log Z\left(\mathcal{I}^{(n)}, \mu\right)\end{aligned} \]

其導數：

\[\begin{aligned} \frac{\partial}{\partial \mu_{a, b}} \ell\left(\mu: \mathcal{T}^{(n)}, \mathcal{I}^{(n)}\right) &=\frac{\partial}{\partial \mu_{b, a}} \ell\left(\mu: \mathcal{T}^{(n)}, \mathcal{I}^{(n)}\right) \\ &=-\frac{\partial}{\partial \mu_{a, b}} E\left(\mathcal{T}^{(n)} \mid \mathcal{I}^{(n)}, \mu\right)-\frac{\partial}{\partial \mu_{a, b}} \log Z\left(\mathcal{I}^{(n)}, \mu\right) \\ &=-\sum_{m} \frac{1}{2} \sum_{i \neq j} k^{(m)}\left(\mathcal{T}_{i}^{(n)}, \mathcal{T}_{j}^{(n)}\right) g(\mu_{a,b})-\frac{1}{Z} \frac{\partial}{\partial \mu_{a, b}} Z\left(\mathcal{I}^{(n)}, \mu\right) \end{aligned} \]

其中，

\[g(\mu_{a,b}) = 1_{a=\mathcal{T}_{i}^{(n)}} 1_{b=\mathcal{T}_{j}^{(n)}}+1_{b=\mathcal{T}_{i}^{(n)}} 1_{a=\mathcal{T}_{j}^{(n)}} \]\[\begin{aligned} \frac{1}{Z} \frac{\partial}{\partial \mu_{a, b}} Z\left(\mathcal{I}^{(n)}, \mu\right) &=\frac{1}{Z} \sum_{\mathbf{X}} \frac{\partial}{\partial \mu_{a, b}} \tilde{P}\left(\mathbf{X} \mid \mathcal{I}^{(n)}, \mu\right) \\ &=\frac{1}{Z} \sum_{\mathbf{X}} \frac{\partial}{\partial \mu_{a, b}} \exp \left(-E\left(\mathbf{X} \mid \mathcal{I}^{(n)}, \mu\right)\right) \\ &=-\sum_{\mathbf{X}} \frac{1}{Z} \exp \left(-E\left(\mathbf{X} \mid \mathcal{I}^{(n)}, \mu\right)\right) \frac{\partial}{\partial \mu_{a, b}} E\left(\mathbf{X} \mid \mathcal{I}^{(n)}, \mu\right) \\ &=-\sum_{\mathbf{X}} P(\mathbf{X}) \sum_{m} w^{(m)} \frac{1}{2} \sum_{i \neq j} k^{(m)}\left(\mathcal{T}_{i}^{(n)}, \mathcal{T}_{j}^{(n)}\right)\left(1_{a=\mathcal{T}_{i}^{(n)}} 1_{b=\mathcal{T}_{j}^{(n)}}+\right. \left.1_{b=\mathcal{T}_{i}^{(n)}} 1_{a=\mathcal{T}_{j}^{(n)}}\right) \\ &=-\sum_{\mathbf{X}} P(\mathbf{X}) \sum_{m} w^{(m)} \frac{1}{2}\left(\sum_{i \neq j} k^{(m)}\left(\mathcal{T}_{i}^{(n)}, \mathcal{T}_{j}^{(n)}\right) 1_{a=\mathcal{T}_{i}^{(n)}} 1_{b=\mathcal{T}_{j}^{(n)}}+\right. \left.\sum_{j \neq i} k^{(m)}\left(\mathcal{T}_{j}^{(n)}, \mathcal{T}_{i}^{(n)}\right) 1_{a=\mathcal{T}_{j}^{(n)}} 1_{b=\mathcal{T}_{i}^{(n)}}\right)\\ &=-\sum_{\mathbf{X}} P(\mathbf{X}) \sum_{m} w^{(m)} \sum_{i \neq j} k^{(m)}\left(\mathcal{T}_{i}^{(n)}, \mathcal{T}_{j}^{(n)}\right) 1_{a=\mathcal{T}_{i}^{(n)}} 1_{b=\mathcal{T}_{j}^{(n)}} \end{aligned} \]

然後再將P換成我們估計的Q：

\[\begin{aligned}&\frac{1}{Z} \frac{\partial}{\partial \mu_{a, b}} Z\left(\mathcal{I}^{(n)}, \mu\right) \\&\approx \sum_{\mathbf{X}} Q(\mathbf{X}) \sum_{m} w^{(m)} \sum_{i \neq j} k^{(m)}\left(\mathcal{T}_{i}^{(n)}, \mathcal{T}_{j}^{(n)}\right) 1_{a=X_{i}} 1_{b=X_{j}} \\&=\sum_{m} w^{(m)} \sum_{i \neq j} k^{(m)}\left(\mathcal{T}_{i}^{(n)}, \mathcal{T}_{j}^{(n)}\right) \sum_{\mathbf{X}} Q\left(\mathbf{X} /\left\{X_{i}, X_{j}\right\}\right) 1_{a=X_{i}} Q_{i}\left(X_{i}\right) 1_{b=X_{j}} Q_{j}\left(X_{j}\right) \\&=\sum_{m} w^{(m)} \sum_{i \neq j} k^{(m)}\left(\mathcal{T}_{i}^{(n)}, \mathcal{T}_{j}^{(n)}\right) Q_{i}(a) Q_{j}(b)\end{aligned} \]

最後得到

\[\begin{aligned}\frac{\partial}{\partial \mu(a, b)} \ell_{n}\left(\mu: \mathcal{I}^{(n)}, \mathcal{T}^{(n)}\right) \approx \sum_{m} w^{(m)}(&-\sum_{i} \mathcal{T}_{i}^{(n)}(a) \sum_{j \neq i} k^{(m)}\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right) \mathcal{T}_{j}^{(n)}(b) \\&\left.+\sum_{i} Q_{i}(a) \sum_{j \neq i} k^{(m)}\left(\mathbf{f}_{i}, \mathbf{f}_{j}\right) Q_{i}(b)\right)\end{aligned} \]

Algorithm

之前我們得到：

\[Q_{i}\left(x_{i}\right)=\frac{1}{Z_{i}} \exp \left\{-\psi_{u}\left(x_{i}\right)-\sum_{j \neq i} \mathbf{E}_{U_{j} \sim \mathrm{Q}_{j}}\left[\psi_{p}\left(x_{i}, U_{j}\right)\right]\right\} \]

不考慮\(\mu\)的學習的情況下，公式左右兩邊都是有Q的，因此可以先初始化一個Q，然後不斷迭代代入，直到P和Q的KL散度小於一定閾值。

其演算法流程如下：

該演算法在Message passing階段的複雜度是\(O(N^2)\)的，利用mean field approximation之後絳為\(O(N)\).

而複雜度的降低來自於對\(k^{(m)}(f_i, f_j)\)用卷積替代，卷積核的大小並非全域性(不然不和原來沒區別了)，而是利用截斷，設定了一個閾值選擇卷積核的大小，即只連線高斯分佈附近的節點。

替換之後的演算法：

可以看到除了用卷積處理，還加上了下采樣和上取樣，這裡先下采樣是為了減少運算，卷積過程等價於低通濾波，低通濾波本身就會過濾掉高頻部分，試想如果不進行任何取樣，那message passing的過程其實就是過濾高頻部分，但是加了下采樣會使得計算量降低，這裡可能會有個疑問，直接下采樣不是有可能造成混疊嗎，一般我們應該先低通濾波再去下采樣的呀？其實這裡用到了Nobel恆等式，也就是先低通濾波再下采樣和先下采樣再低通濾波是等價的，簡單的解釋是先低通濾波再去下采樣，最終留下來的點其實有一些是沒用到的，所以可以事先進行取樣選擇需要的點，再進行濾波，更為詳細的理解可以在The Scientist and Engineer's Guide to Digital Signal Processing找到，大概在p40；而上取樣部分僅僅是為了恢復形狀。