はじめに

Momentum Contrast for Unsupervised Visual Representation Learningを読んだのでメモ． FAIRのKaiming Heの教師なし表現学習の論文．どうでもいいが，FAIRからの論文は研究のモチベーションから手法の提案まで非常に明快に書かれてて読みやすくて好き．一方，近年のcontrastive learningによる表現学習はほとんど流れを知らないのでこの記事の内容はとても読みにくいものになってしまった気がするので，こんな素人のブログより原著を読むことをおすすめする．

Contrastive Learning as Dictionaly Look-up

まず，contrastive learningとdictionary look-upの関係について考える．

エンコードされたクエリ $q$ とエンコードされた(辞書のkeyとして考えられる)サンプルの集合 $\{k_0,k_1,k_2,\dots\}$ を考える．仮定として $q$ には一つのkeyが対応するものとし，このkeyを $k_+$ とする． Contrastive lossは $q$ が $k_+$ に似ている(positive key)ときに小さく，異なる(negative keys)ときに大きな値を取る損失であり，内積を利用したcontrastive lossの一つであるInfoNCEは次のように定義される．

$\displaystyle \mathcal{L}_q=-\log\frac{\exp(q\cdot k_+/\tau)}{\sum_{l=0}^K\exp(q\cdot k_i/\tau)}$

ただし， $\tau$ はハイパーパラメータで，分母の和は $K$ 個のnegative keyと一つのpositive keyに関する足し合わせ． Contrastive lossは簡単に言えば， $K+1$ 個の中からpositive keyを分類する問題として，クエリとkeyを表現するためのencoder network $q=f_q(x^q),k=f_k(x^k)$ の教師なし学習のための目的関数と言える．この時， $x^q$ や $x^k$ は画像そのものや，画像のパッチ，パッチの集合など様々．また， $f_q$ と $f_k$ も様々で，一部を共有したり，全く別なものにしたりする．

Momentum Contrast

上記の解釈からcontrastive learningは画像のような高次の連続な空間での離散の辞書を作る方法と言える．辞書は，学習中にエンコーダのパラメータが更新されていくという点で動的なものである．ここでは，例えパラメータが更新されて辞書が変化していったとしても，辞書のkeyは可能な限りconsistentな値を保ち，豊富なnegative sampleの集合をカバーする巨大な辞書によって良い表現が獲得されるという仮説をおいている．この仮説からMomentum Contrast(MoCo)を提案する．

Dictionary as a queue

MoCo中心となるアイディアはデータサンプルのキューとして辞書を維持することで，エンコードされたkeyの再利用や辞書の大きさとミニバッチの大きさを分離を可能にする． MoCoの辞書サイズは(当然)ミニバッチよりも大きく，それゆえ十分な表現力を持つことができる．

辞書内のサンプルは動的に置き換わっていき，現在のミニバッチは辞書にenqueされ，最も古いミニバッチは削除される．よって，辞書は常に全データの部分集合で表現され，辞書を維持するための計算コストは保持するミニバッチ数として扱いやすい．

Momentum update

キューの利用はいくつかの利点もあるが，辞書の巨大化によりkeyエンコーダのパラメータの更新を扱いにくくしている側面もある．ナイーブには， $f_k$ を $f_q$ からコピーしてくるという手が考えられるが，これは実験的にはあまり良くなかった．これはkeyのconsistencyが保たれにくいためであるとい仮説を考え，ここでは移動平均でパラメータを更新する．

$f_k$ のパラメータを $\theta_k$ とし， $f_q$ を $\theta_q$ とすれば， $\theta_k$ は次のよう更新する．

$\displaystyle \theta_k\leftarrow m\theta_k+(1-m)\theta_q$

$m\in[0,1)$ はmomentum係数で $\theta_q$ はcontrastive lossを目的関数として，backpropにより更新される．この更新方法によりkeyのconsistencyが保たれやすくなり性能の向上へ繋がるとのこと．実験では $m$ は $0.999$ などの値にしたとのこと．

Relations to previous mechanisms

MoCoは辞書サイズが大きいこととmomentum updateによるconsistencyの保証以外は従来のcontrastive learningの枠組みと大差ない．ここでは論文の図2を使って比較をしており，end-to-endではbackpropが必要となる点からMoCo以上に計算リソースに辞書サイズが縛られる，memory bankは過去の計算結果からランダムにkeyをサンプルしてくるため辞書サイズは巨大にできるが，メモリバンクは更新がなかなか起こらないという問題がある．こう考えるとmomentum updateは半教師あり学習のmean teacherのモチベーションとほぼ等しい．つまり全サンプルの値を追跡していくのではなくモデルのパラメータに任せようというもの．

まとめ

細かい実装等に関する部分は割愛したが，PyTorchライクの擬似コード含め細かく書いてあるため再現は容易にできそう．実際すでに再現リポジトリがgithub上にいくつか見られる． billionスケールのデータで学習をして教師ありを超える成果を出しているのは機械学習の活用としてはとても夢の溢れる結果．

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

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

Momentum Contrast for Unsupervised Visual Representation Learningを読んだのでメモ