機械学習とかコンピュータビジョンとか

CVやMLに関する勉強のメモ書き。

Neural Nearest Neighbors Networksを読んだのでメモ

はじめに

Neural Nearest Neighbors Networksを読んだのでメモ.differentiableなKNN selection ruleを提案するというもの.NIPS 2018でgithubコードも公開されている.

Differentiable k-Nearest Neighbors

クエリとなるアイテムをq,データベースの候補を(x_i)_{i\in I},距離関数をd(\cdot,\cdot)とする.qはデータベースの中にないものと仮定し,dはクエリからの距離に従ってデータベースのアイテムのランキングを生成する.ここで[tes:\pi_q:I\rightarrow I]をqの距離にしたがってソートを行うpermutationと定義する.

\displaystyle
\pi_q(i)\lt\pi_q(I')\Rightarrow d(q,x_i)\leq d(q,x_{i'}),\:\forall i,i'\in I

するとqのKNNはpermutation \pi_qにおいて最初のk個のアイテムの集合として与えられる.

\displaystyle
\mathrm{KNN}(q)=\{x_i|\pi_q(i)\leq k\}

KNNのselection ruleはdeterministicなため微分不可能で,距離に関する微分を伝播することができない.そこでこの問題をone-hotベクトルの連続値への緩和を使って解消する.

KNN rule as limit distribution

ここではKNN selection ruleをカテゴリカル分布の極限分布としてみなすことで緩和する.\mathrm{Cat}(w^1|\alpha^1,t)をデータベースのアイテムのインデックスIに関するカテゴリカル分布とする.クエリアイテムとの距離d(q,x_i)\alpha_i^1とし,温度パラメータtを導入すれば,i\in I番目のアイテムを選択する確率w_1は次のように与えられる.

\displaystyle
\mathbb{P}[ w^1=i|\alpha^1,t]\equiv\mathrm{Cat}(\alpha^1,t)=\frac{\exp(\alpha_i^1/t)}{\sum_{i'\in I}\exp(\alpha_{i'}^1/t)}\\
\mathrm{where}\:\alpha_i^1\equiv -d(q,x_i)

Gumbel softmaxやconcrete distributionで見たような式だが,ここではアイテム間の距離によって決定的にサンプリングされるアイテムが決まるためgumbel分布からのサンプリングがない.t\rightarrow 0の時w^1はone-hotベクトルとなる.この式を1-NNのstochastic relaxationとしてみなすことができるため,これをkの場合に拡張することでKNNを微分可能な形に緩和することができる.そのために条件付き分布\mathrm{Cat}(w^{j+1}|\alpha^{j+1},t)へと拡張する.ここで\alpha^{j+1}は次のように計算される.

\displaystyle
\alpha_i^{j+1}\equiv\alpha_i^j+\log(1-w_i^j)

つまり一度選択されたアイテムは選択される確率が0に近くなるようクエリアイテムとの距離を離すというもの.実際w_i^jがone-hotベクトルの場合には\alpha_i^{j+1}w^j=iにおいて-\inftyになる.これによってj+1番目のサンプルのインデックスは次のように得られる.

\displaystyle
\mathbb{P}[w^{j+1}=i|\alpha^{j+1},t]\equiv\mathrm{Cat}(\alpha^{j+1},t)=\frac{\exp(\alpha_i^{j+1}/t)}{\sum_{i'\in I}\exp(\alpha_{i'}^{j+1}/t)}

Index vector w^jからqのstochastic nearest neighbors \{X^1,\dots,X^k\}を次のように定義する.

\displaystyle
X^j\equiv\sum_{i\in I}w_i^jx_i

温度t\rightarrow 0の時w^j_iがone-hotベクトルとなりk nearest neighborsと一致する.ただ基本的にはone-hotベクトルではないので多分mixup的な形になるはず(というかone-hotベクトルに近づけば近づくほど勾配が消失していくので意味がなくなる).

Neural Nearest Neighbors Block

今までの議論を踏まえてニューラルネットの一つの層としてneural nearest neighbors block (\mathrm{N}^3 blocks)を提案する.\mathrm{N}^3 blocksは二つの部分からなり,一つはembedding networkで,もう一つは連続緩和されたnearest neighborsを計算する部分がある.

Embedding network

Embedding networkはf_EでparameterizeされたCNNを使って入力から特徴ベクトルを得る部分.得られた特徴をEとし,ユークリッド距離を用いてD_{ij}=d(E_i,E_j)というpairwise distance matrix Dを得る.さらに,ここでは温度tを計算するもう一つのネットワークT=f_T(Y)を用意する.ただし,f_Ef_Tは出力層以外は重み共有している.

Continuous nearest neighbors selection

距離行列D,温度テンソルT,入力の特徴Yからk continuous nearest neighbors feature volumes N_1,\dots,N_kを計算する.nearest neighborsとYは入力の次元が同じなのでelement-wiseな操作によってまとめることができるが今回はチャネル方向に結合することで情報をまとめたとのこと.

\mathrm{N}^3 block for image data

基本的に\mathrm{N}^3 blockは入力のドメインに関係なく適用可能だが,画像データにカニsてはちょっとした修正が必要.従来画像におけるnon-local methodはパッチレベルで行われ,パッチレベルの処理には様々な利点が存在する.そのため画像に\mathrm{N}^3 blockを適用する際にはim2colを使ってパッチレベルで計算するとのこと.

まとめ

基本的にはdeterministicなGumbel softmaxという感じ.実験ではdenoisingや超解像などの逆問題を解くのに使っていたが,いろいろなことに使えそう.ただ,複雑なことをしようとすると距離(というかembedding network)に何らかのpriorを持たせないと学習は難しそう.