MoCo：モメンタム更新のPyTorch実装+注意点

MoCoとは?

MoCoは自己教師あり学習の手法の一つで，特に表現学習に使われる．教師ラベルなしで大量のデータから特徴量を学習することが目的．

Momentum updateの仕方

PyTorchでモメンタム更新をする方法

model : クエリエンコーダ

model_key: キーエンコーダ

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

with torch.no_grad():
    for p_q, p_k in zip(model.parameters(), model_key.parameters(), strict=False):
        p_k.data.mul_(momentum).add_(p_q.data, alpha=1.0 - momentum)

実装上の注意点

事前学習の損失計算は専用のMLPヘッドを介して行う

事前学習: 類似度損失の計算は, $g(e(x))$に対して行う．

下流タスク: MLPを捨てて，クラス予測ヘッドを使う $h(e(x))$ ．

• $x$: 入力
• $e$: エンコーダ
• $g$: MLP
• $h$: クラス予測ヘッド

なぜこうするのか？

• 類似度の損失計算にエンコーダの特徴$e(x)$を直接使うと，下流タスクに有用な特徴が失われやすいから
  • SimCLRの4.2節でこれに関する実験がある
    • 事前学習における，MLPの有無でImageNetの転移性能が10%ほど変化

SimCLR↓

queueを介したvalidationデータのリーク

validationで学習時のqueueを絶対に更新しない．
valでもtrainと同様のフローにすると，queueにvalサンプルの特徴が入ってしまうので🙅

grad accumulationを併用したときのkey encoderの更新

モメンタム更新を行うタイミングは，パラメータ更新が発生したあとだけにする．

これを考慮しないで毎バッチで更新してしまうと，
実質のmomentumが小さくなってしまうので学習安定性が低下する．（key encoderが想定よりも早くquery encoderに追従する．）