はじめに

Unsupervised Object Segmentation by Redrawingを読んだのでメモ．タイトルの通りunsupervisedなobject segmentationの論文．

Method

セグメンテーションのプロセス $\text{F}$ を定義する． $\text{F}$ は画像 $\mathbf{I}\in\mathbb{R}^{W\times H\times C}$ が与えられた時，各ピクセルに $n$ 個のラベルのどれかを割り当てる関数とする．ここでの課題は $\mathbf{I}$ が与えられた時に正しい分割を行う $\text{F}$ を見つけること．ここでは教師ありと違い，分割された各領域が満たすべき性質を定義し，その定義に従って分割する必要がある．今回はこれをエネルギー関数を定義することによって実現する．

Generative Process

今回は画像は次の3つの生成過程に従って作られると考える．

composition step : 画像中の異なる領域を定義
drawing step : 各領域毎に独立にピクセルの生成
assembling step : 領域を統合して画像を生成

自分の理解ではGANでよくやられるラベルからの画像生成とほぼ同じものを仮定している．

まず $n-1$ 個の物体と背景，すなわち $n$ 個の領域からなる画像を考える． $\mathbf{M}^k\in\{0,1\}^{W\times M}$ を $k$ 番目の物体に対応する2値のマスク(すなわちsegment)とし，仮に $(x,y)$ 座標のピクセルが $k$ 番目の物体に所属する場合 $\mathbf{M}_{x,y}^k=1$ となる．各ピクセルはただ一つの領域に所属することを仮定するため $\sum_{k=1}^n\mathbf{M}_{x,y}^k=1$ が成り立つ．

各物体のピクセルの値を $\mathbf{V}^k\in\mathbb{R}^{W\times H\times C}$ とすると，画像はマスクとピクセル値の要素積で与えられる．

$\displaystyle \mathbf{I}\leftarrow\sum_{k=1}^n\mathbf{M}^k\odot\mathbf{V}^k$

ここまで定義した文字を使って生成過程を表現すると，

composition step : マスクの事前分布 $p(\mathbf{M})$ に従ってマスク $\mathbf{M}^k$ を生成
drwaing step : $p(\mathbf{V}^k|\mathbf{M}^k, k)$ に従って物体領域毎に独立にピクセル値を選択
assembling step : マスクとピクセル値の要素積により画像を生成

From Generative Process to Object Segmentation

ここからが手法の本題の話．目的は物体のマスクを得ることで画像を生成することではないので，ここでは $\text{F}$ を通して入力画像 $\mathbf{I}$ から直接マスクを生成することを考える．すなわち生成過程としてはcomposition stepが $\text{F}(\mathbf{I})\rightarrow\mathbf{M}_1,\dots,\mathbf{M}_n$ と書き換わる．

ここでは生成過程に従って得られた出力画像が学習データの分布に一致するようにするためGANの学習法を用いて $\text{F},p(\mathbf{V}^k|\mathbf{M}^k,k)$ を学習する． $\text{D}:\mathbb{R}^{W\times H\times C}\rightarrow\mathbb{R}$ をdiscriminatorとし， $\text{G}_\text{F}(\mathbf{I},\mathbf{z}_1,\dots,\mathbf{z}_n)$ をgeneratorとする．generatorは入力として $\text{F}$ と $\mathbf{I}$ と $p(\mathbf{z})$ から独立にサンプリングされたベクトルの集合 $\mathbf{z}_1,\dots,\mathbf{z}_n$ が与えられる．今回は異なる領域のピクセル値は独立であることを仮定しているためgeneratorは $n$ 個に分解される， $\text{G}_k(\mathbf{M}^k,\mathbf{z}_k)$ ．最終的に生成過程は次のようになる．

composition step : $\mathbf{M}^1,\dots,\mathbf{M}^n\leftarrow\text{F}(\mathbf{I})$
drawing step : $\mathbf{V}^k\leftarrow\text{G}_k(\mathbf{M}^k,\mathbf{z}_k)\:\text{for}\: k\in\{1,\dots,n\}$
assembling step : $\text{G}_\text{F}(\mathbf{I},\mathbf{z}_1,\dots,\mathbf{z}_n)=\sum_{k=1}^n\mathbf{M}^k\odot\mathbf{V}^k$

$\text{F},\text{G}_k$ は微分可能であるため勾配ベースで次の目的関数の学習が可能．今回はhinge versionのadversarial lossを使うとのこと．

$\displaystyle \max_{\text{G}_\text{F}}\mathcal{L}_\text{G}=\mathbb{E}_{\mathbf{I}\sim p_{data},\mathbf{z}_1,\dots,\mathbf{z}_n\sim p(\mathbf{z})}[\text{D}(\text{G}_\text{F}(\mathbf{I},\mathbf{z}_1,\dots,\mathbf{z}_n))]\\ \max_\text{D}\mathcal{L}_\text{D}=\mathbb{E}_{\mathbf{I}\sim p_{data}}[\min(0,-1+\text{D}(\mathbf{I}))]+\mathbb{E}_{\mathbf{I}\sim p_{data},\mathbf{z}_1,\dots,\mathbf{z}_n\sim p(\mathbf{z})}[\min(0,-1-\text{D}(\text{G}_\text{F}(\mathbf{I},\mathbf{z}_1,\dots,\mathbf{z}_n)))]$

単純にこの学習をした場合，二つの理由から失敗する可能性がある．一つ目は $\mathbf{I}$ に対して意味のあるセグメントを抽出しない可能性がある．というのも $\mathbf{I}$ を完全に無視して画像を生成可能であるため．二つ目は $\text{F}$ が自明な解にたどり着く，すなわち単一の領域のみに分割(=分割しない)という結果を返してしまう可能性があるということ．なのでここにさらなる制約を追加する．

Constraining mask extraction by redrawing a single region

一つ目の制約はモデルが意味のある領域を抽出することを強いるもので，異なる物体のピクセルは独立に生成されるという仮定を利用したもの．やり方は単純である一つの領域を除いた全ての領域はピクセル値を生成するのではなく元画像のピクセル値を使うというもの．すなわち生成過程のdrawing stepを $\mathbf{V}^k\leftarrow\mathbf{I}\:\text{for}\:k\in\{1,\dots,n\}\backslash,\:\mathbf{V}^i\leftarrow\text{G}_i(\mathbf{M}^i,\mathbf{z}_i$ と修正する． $i$ は一様分布からランダムにサンプリングされる．

Conservation of Region Information

二つ目の制約は潜在ベクトル $\mathbf{z}_i$ から生成された領域 $i$ が与えられた時，最終的な画像 $\text{G}_\text{F}(\mathbf{I},\mathbf{z}_i,i)$ が $\mathbf{z}_i$ に関する情報を含んでいることを強いるもので， $\text{F}$ が空領域を抽出することを防ぐ．実際 $i$ が空，すなわち全ての $x,y$ に関して $\mathbf{M}_{x,y}^i=0$ の場合 $\mathbf{z}_i$ の情報は生成された画像に含まれない．この制約はInfoGANの相互情報量の最大化と同じ方法で実現される．

$\delta_k$ を $\mathbf{I}$ が与えられた時 $\mathbf{z}_k$ の値を推論する関数とする．上記二つの制約を踏まえた損失関数は次のようになる．

$\displaystyle \max_{\text{G}_\text{F},\delta}\mathcal{L}_\text{G}=\mathbb{E}_{\mathbf{I}\sim p_{data},i\sim\mathcal{U}(n),\mathbf{z}_i\sim p(z)}[\text{D}(\text{G}_\text{F}(\mathbf{I},\mathbf{z}_i,i))-\lambda_z\|\delta_i(\text{G}_\text{F}(\mathbf{I},\mathbf{z}_i,i))-\mathbf{z}_i\|_2^2]\\ \max_\text{D}\mathcal{L}_\text{D}=\mathbb{E}_{\mathbf{I}\sim p_{data}}[\min(0,-1+\text{D}(\mathbf{I})]+\mathbb{E}_{\mathbf{I}\sim p_{data},i\sim\mathcal{U}(n),\mathbf{z}_i\sim p(\mathbf{z})}[\min(0,-1-\text{D}(\text{G}_\text{F}(\mathbf{I},\mathbf{z}_i,i))]$