はじめに

SegSort: Segmentation by Discriminative Sorting of Segmentsを読んだのでメモ．

気持ち

従来のsemantic segmentationはピクセル毎のクラス分類問題として解決が試みられているが，人間はそうやっていないからもっと別なアプローチが必要だと言うもの．ここではピクセル毎のmetric learningを利用しようというもの．すなわち個々の画像内で異なる領域の表現の類似度は低く同じ領域の表現の類似度は高いかつ，異なる画像内の同一クラスの領域の表現の類似度は高いという前提に基づいたアプローチを取る．この方法が従来と異なる点は領域分割自体はクラスタリングによって行われ，セマンティックなラベル付けは学習データの集合からnearest neighborで行うという点．このnearest neighborによるラベル付けの方がより人間の認識方法に近いだろうという主張．

Method

基本的な枠組みは以下の3行程からなる．

任意のCNNモデル(DeepLab，FCN，PSPNetなど)でピクセル毎の埋め込みベクトルを計算
pixel sortingと呼ばれる方法でピクセルを細かいセグメントにクラスタリング
segment sortingと呼ばれる各セグメントの分離と結合のためのmetric learning

まず仮定として，CNNによって得られた埋め込みベクトルはvon Mises-Fisher分布(vMF)に従うとする．この仮定からpixel sortingはspherical K-Meansとして，segment sortingは最尤推定として定式化する．ただし推論時にはsegment sortingはk-nearest neighborで実行される．

$\mathcal{V}=\{\boldsymbol{v}_i\}=\{\phi(x_i)\}$ をピクセルの埋め込みベクトルの集合とする．ただし $\boldsymbol{v}_i$ はピクセル $x_i$ を中心としてCNN $\phi$ から作られる埋め込みベクトルを表す． $\mathcal{Z}=\{z_i\}$ をあるセグメントに属するピクセルの集合とし， $z_i=s$ はピクセル $i$ がセグメント $s$ に属することを意味する． $\Theta=\{\theta_{z_i}\}$ をセグメントの埋め込みベクトルの事前分布(ここではmixture vMF)のパラメータ集合とする．この論文での主要な目的関数は次のようになる．

$\displaystyle \min_{\phi,\mathcal{Z},\Theta}-\log P(\mathcal{V},\mathcal{Z}|\Theta)=\min_{\phi,\mathcal{Z},\Theta}-\sum_i\log\frac{1}{k}f_{z_i}(\boldsymbol{v}_i|\theta_{z_i})$

pixel sortingにおいて $\phi$ 固定のもとEMアルゴリズムによって $\mathcal{Z},\Theta$ を求める．segment sortingでは $\mathcal{Z},\Theta$ を固定して $\phi$ を更新する．そのため上記の目的関数はtwo-stageのEMアルゴリズムで最適化を行うこととなる．

Pixel Sorting

まずvMF分布とspherical K-Meansについて．セグメント内のピクセルの持つ $d$ 次元埋め込みベクトル $\boldsymbol{v}\in\mathbb{S}^{d-1}$ はvMF分布に従うとする．vMFの確率密度関数は次式．

$\displaystyle f(\boldsymbol{v}|\boldsymbol{\mu},\kappa)=C_d(\kappa)\exp(\kappa\boldsymbol{\mu}^\top\boldsymbol{v})$

$C_d(\kappa)=\frac{\kappa^{d/2-1}}{(2\pi)^{d/2}I_{d/2-1}(\kappa)}$ は正規化定数． $I_r(\cdot)$ は第１種変形ベッセル関数(modified Bessel function)． $\boldsymbol{\mu}=\sum_i\boldsymbol{v}_i/\|\sum_i\boldsymbol{v}_i\|$ は平均方向で $\kappa\geq0$ は集中度．ここでは $\kappa$ は定数と仮定し，コストの高い $C_d(\kappa)$ の計算を避ける．

画像中の $k$ 個のセグメントの埋め込みベクトルを $k$ 個の一様分布を事前分布とするvMF分布の混合モデルとして考える．

$\displaystyle f(\boldsymbol(v)|\Theta)=\sum_{s=1}^k\frac{1}{k}f_s(\boldsymbol{v}|\boldsymbol{\mu}_s,\kappa)$

$\Theta=\{\boldsymbol{\mu}_1,\dots,\boldsymbol{\mu}_k,\kappa\}$ は各vMFのパラメータの集合． $z_i$ をピクセルの埋め込みベクトル $\boldsymbol{v}_i$ の所属セグメント $s$ を表す隠れ変数． $\mathcal{V}=\{\boldsymbol{v}_1,\dots,\boldsymbol{v}_k\}$ を埋め込みベクトルの集合， $\mathcal{Z}=\{z_1,\dots,z_k\}$ を対応する隠れ変数の集合とする．混合vMFの対数尤度は次のように与えられる．

$\displaystyle \log P(\mathcal{V},\mathcal{Z}|\Theta)=\sum_i\log\frac{1}{k}f_{z_i}(\boldsymbol{v}_i|\boldsymbol{\mu}_{z_i},\kappa)$

EMアルゴリズムを用いてこの対数尤度の最尤推定を行う．E-step(尤度最大化)で事後分布に $z_i=s$ を代入する．

$\displaystyle p(z_i=s|\boldsymbol{v}_i,\Theta)=\frac{f_s(\boldsymbol{v}_i|\Theta)}{\sum_{l=1}^kf_l(\boldsymbol{v}_l|\Theta)}$

ここではK-Meansによるクラスタリングを行うため， $z_i$ は $z_i=\mathrm{argmax}_sp(z_i=s|\boldsymbol{v}_i,\Theta)=\mathrm{argmax}_s\boldsymbol{\mu}_s^\top\boldsymbol{v}_i$ として割り当てが決まる．ここでセグメント $c$ 内のピクセルの集合を $\mathcal{R}_c$ とする．従ってK-Meansのハードアサインメントにより $i\in\mathcal{R}_c$ である場合に $p(z_i=c|\boldsymbol{v}_i,\Theta)=1$ となりそれ以外は0となる．

M-step(期待値最大化)は次の式で実行される．

$\displaystyle \hat{\boldsymbol{\mu}}_c=\frac{\sum_i\boldsymbol{v}_ip(z_i=c|\boldsymbol{v}_i,\Theta)}{\|\sum_i\boldsymbol{v}_ip(z_i=c|\boldsymbol{v}_i,\Theta)\|}=\frac{\sum_{i\in\mathcal{R}_c}\boldsymbol{v}_i}{\|\sum_{i\in\mathcal{R}_c}\boldsymbol{v}_i\|}$

これはセグメント $c$ 内のピクセルの埋め込みベクトルの平均方向を表す．よってspherical K-Meansは $\mathcal{Z}$ (E-step)と $\Theta$ (M-step)の交互最適化によって行われる．この最適化の繰り返し回数はプラクティカルに固定回数(論文では10回)で良いとのこと．

このK-Meansクラスタリングに埋め込みベクトルのみを使うと，得られたクラスタ内のピクセルが非連結かつ散らばることがある．セグメントは連結成分であってほしいため埋め込みベクトルとピクセルの座標値を結合させたベクトルをK-Meansに利用するとのこと．

通常上記の処理で得られたセグメントは，そのセグメントに所属するピクセルのセマンティックなラベルが同一であるとは限らない(すなわち一つのセグメントが複数の物体の領域をまたぐ場合がある)．なので教師ありの設定では真のラベルを使ってセグメントを分割するとのこと(論文のFigure 3)．

自分の理解ではCNNの特徴マップに対してSLIC(厳密には異なるが)によるスーパーピクセルを作るという感じ．

Segment Sorting

pixel sortngでは埋め込みベクトル $\boldsymbol{v}_i$ を作るCNNのパラメータ $\phi$ は固定だった．ここでは $\phi$ のパラメータの最適化を考える．

定数 $\kappa$ に対して $\phi$ をフリーパラメータとし， $z_i=s$ の事後確率は次のように計算する．

$\displaystyle p_\phi(z_i=s|\boldsymbol{v}_i,\Theta)=\frac{f_s(\boldsymbol{v}_i|\Theta)}{\sum_{l=1}^kf_l(\boldsymbol{v}_l|\Theta)}=\frac{\exp(\kappa\boldsymbol{\mu}_s^\top\boldsymbol{v}_i)}{\sum_{l=1}^k\exp(\kappa\boldsymbol{\mu}_l^\top\boldsymbol{v}_i)}$

$\boldsymbol{\mu}$ はM-stepで計算された平均方向ベクトルで $\boldsymbol{v}$ がCNNによって作られた埋め込みベクトル．それぞれ長さが1に正規化されているためその内積はコサイン類似度を表す．よって上記の式はコサイン類似度に対するソフトマックスとなっている．ただし $\kappa$ が定数として入っているのである意味ではGumbel Softmaxといえる．意味としてはあるピクセルの埋め込みベクトルと各セグメントの平均方向ベクトルとの類似度を基準とした所属確率を表現している．

パラメータ $\phi$ の推定は負の対数尤度の最小化によって行われる．

$\displaystyle L_\text{vMF}^i=-\log p_\phi(c|\boldsymbol{v}_i,\Theta)=-\log\frac{\exp(\kappa\boldsymbol{\mu}_c^\top\boldsymbol{v}_i)}{\sum_{l=1}^k\exp(\kappa\boldsymbol{\mu}_l^\top\boldsymbol{v}_i)}$

この負の対数尤度の最小化は，K-Meansで得られたプロトタイプのセグメントに対してピクセルの埋め込みベクトルが近づくことと，各埋め込みベクトルが他の全てのプロトタイプセグメントから離れることを要請する．この負の対数尤度の最小化の嬉しいことは真のラベルを必要としないことで，CNNを教師なしで学習することが可能となる．

真のラベルの利用のためにはNeighborhood Components Analysisにおけるsoft neighborhood assignmentsについて考える．soft neighborhood assignmentsはこの課題においては埋め込みベクトル $\boldsymbol{v}_i$ に関して同一カテゴリの任意の他のセグメント(ここでの他のセグメントは同一画像内に限らない．これをこの文脈での近傍 $c^+$ として定義する)の平均方向ベクトルに関する確率を高めることで，その確率は次のように定義される．

$\displaystyle p_\phi'(z_i=c^+|\boldsymbol{v}_i,\Theta)=\frac{f_{c^+}(\boldsymbol{v}_i|\Theta)}{\sum_{l\neq c}f_l(\boldsymbol{v}_l|\Theta)}=\frac{\exp(\kappa\boldsymbol{\mu}_{c^+}^\top\boldsymbol{v}_i)}{\sum_{l\neq c}\exp(\kappa\boldsymbol{\mu}_l^\top\boldsymbol{v}_i)},\\ p_\phi'(z_i=c|\boldsymbol{v}_i,\Theta)=0$

ピクセル $i$ に関する近傍のセグメントの集合を $C^+_i\{c^+\}$ として定義する．上記を踏まえて最終的な学習のための損失関数は次のようになる．

$\displaystyle L_\text{vMF-N}^i=-\log\sum_{s\in C^+_i}p_\phi'(z_i=s|\boldsymbol{v}_i,\Theta)\\ =-\log\frac{\sum_{s\in C^+_i}\exp(\kappa\boldsymbol{\mu}_s^\top\boldsymbol{v}_i)}{\sum_{l\neq c}\exp(\kappa\boldsymbol{\mu}_l^\top\boldsymbol{v}_i)}$

この損失関数の最小化は正しい近傍のプロトタイプに割り当てられるピクセルの数を最大化することを意味する．この損失関数の最小化において真のラベルはミニバッチ内のピクセル $i$ に関する同一クラスのセグメントの集合 $C^+_i$ の探索に用いられる．仮にミニバッチ内に同一クラスのセグメントが存在しない場合には $L_\text{vMF}^i$ を損失として扱う．

ここでの損失関数の利点は近傍（同一クラス）セグメントと非近傍（異なるクラス）セグメントの存在で，セグメントの事例が増えれば増えるほど最適化で得られる解の性質が向上する．そのため最適化のプロセスの中で損失はバッチ内(論文中でbatchと書いてあるがおそらくミニバッチのこと)の全セグメントを利用して行われる．これは従来の画像ごとの損失の計算と異なるスキームとなっており独自性の一つ．また，2時刻前までに使ったバッチをメモリに保存しておいて損失計算時に利用するということも行う．これらの方法により多数のセグメントの事例を用意できて最適化が捗るとのこと．

Inference via K-Nearest Neighbor Retrieval

学習後，学習データの全セグメントをプロトタイプとして保存しておく．推論時にはK-Meansによるクラスタリングを行なったのち，各クラスタリング結果のセグメントと保存されたセグメントとのk-nearest neighbor searchによってラベル付けを行う．個人的にこの方法の面白いところはnon-parametericな手法であるためクラス数が固定されないという点．学習データから作られたプロトタイプのラベル付け規則を変えることで原理的には任意の粒度でクラス分類を行うことができる．