対照学習（NNCLR） の encoder に Attention をつけてみた。

概要

　近年、自己教師あり学習の１つである、対照学習が注目を集めています。土木の分野でも後々注目しそうです。

　対照学習は、ある1枚の画像に対して、ランダムなデータ拡張を行い2枚の画像を出力させます。データ拡張が行われた2枚の画像をエンコーダー（CNN）に入力し、2つの特徴ベクトルを出力させます。この2枚の画像は、元々は同じ画像であるので、この2つのベクトルは類似しているはずです。それらが類似するように学習を行っていきます。

　もちろん、1枚の画像だけでは類似度を判断できないので、ミニバッチに含まれる他の画像を類似していないとして学習させます。

　データ拡張されても元々は同じ画像なので、エンコーダー（CNN）を通して得られた特徴マップは似ていると考えられます。よって、Attention（注意機構）をエンコーダーの最終層に張り付けて、データ拡張されても、同じ場所に Attention が反応すると考えられます。

　今回の記事は、以上について説明した内容になります。対照学習に関する良い記事は、たくさんあるので他に譲ります。

論文は以下を引用しています。

プログラムは、以下の keras のブログを参考に作成できます。

NNCLR (Nearest Neighbor Contrastive Learning of Visual Representations)

以下の図は、NNCLR の概要図であり、論文から引用した。

　最初に、ミニバッチ内の画像に対しランダムにデータ拡張を行い、2つの画像セットを作る。それらを（同じ）エンコーダーに入力して、特徴ベクトルの組を出力させる。

　NNCLRの先行研究である SimCLR では、2つの画像セットから得られた特徴ベクトルの類似度を計算し、誤差関数を計算し最適化を行う。同じ画像ならばポジティブペア（類似している画像）として、そのほかのミニバッチ内の画像ならばネガティブペア（類似していない画像）として類似度・誤差関数を評価する。
　
　この記事において、類似度とは、2つの特徴ベクトルの距離が近いか、2つの特徴ベクトルの向きが揃っているか（cos 類似度）である。

　NNCLR では、 support set （ベクトルの集合）を準備する。片方の特徴ベクトルに対して、support set の中から最も類似しているベクトルを選ぶ。それをもう片方の特徴ベクトルと比較することで類似度・誤差関数の計算を行う。 SimCLR などの自己教師あり学習では、データ拡張に依存するので複雑なデータ拡張が必要である。NNCLR では、support set から選ぶことで複雑なデータ拡張を避けることができる。

　誤差関数は、$i$ 番目の画像に対して、片方の特徴ベクトルを $z_i$ とし、もう片方のポジティブペアにおける特徴ベクトルを $z_i^{+}$ とすると、

\begin{align}
\mathcal{L}_i = - \log \frac{\exp\left(\mbox{NN}(z_i,Q) \cdot z_i^{+} /\tau \right) }{\sum_{k=1}^n \exp\left(\mbox{NN}(z_i,Q) \cdot z_k^{+} /\tau  \right) }
\end{align}

と書ける。$\tau$ は Softmax Temperature であり、$z_k^{+}$ は $k$ 番目の画像のポジティブペアである。分母の和の $n$ は、バッチサイズである。（特徴ベクトルは規格化されているので、分子の内積が大きくなることは、ポジティブペアの類似度が高いことを意味している。）

　$\mbox{NN}(z,Q) $ は、Nearest Neighbor 演算子であり、特徴ベクトル $z$ に最も近い（最も類似している）ベクトルを、support set $Q$ から選ぶ。

\begin{align}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator*{\argmax}{arg\,max}
\mbox{NN}(z,Q) = \argmin_{q \in Q } ||z-q||_2
\end{align}

　実際の計算では、特徴ベクトル $z$ は正規化される。つまり、二乗ノルムは

\begin{align}
& (z-q)^2 = z^2 - 2 z\cdot q + q^2 = 2 -2 z\cdot q \\
\end{align}

と書けるので、二乗ノルムが最小となる $q$ を選ぶことは、内積が最大となる $q$ を選ぶことである。（keras の実装はこれを使っている。）

\begin{align}
\therefore & \ \ \argmin_{q \in Q } ||z-q||_2 =\argmax_{q \in Q } z\cdot q
\end{align}

最終的に、誤差関数は

\begin{align}
\mathcal{L} = \frac{1}{n}\sum_{i=1}^n \mathcal{L}_i
\end{align}

と計算される。

　

Softmax 関数は使わない？

　近年の github や tensorflow・keras のプログラムコード付き解説ブログなどでは、Softmax 関数があまり見られない。
　
　例えば、$k$ 分類問題の考えたとき、$x$ が $i$ 番に分類される確率は、 Softmax 関数を使い

\begin{align}
P(\mathcal{C}_i|x) &= \frac{\exp\left(u_i \right) }{\sum_{j=1}^k \exp\left(u_j \right) } \\
 u_j &= \omega_j x +b_j
\end{align}

と書ける。$u$ はロジットであり、$\omega,b$ はパラメータである。 $x$ は、Softmax 関数が最大となるクラス

\begin{align}
\argmax_{i=1,2,\ldots,k} P(\mathcal{C}_i|x) 
\end{align}

に分類される。Softmax 関数は使わない場合、ロジットが最大となるクラス

\begin{align}
\argmax_{i=1,2,\ldots,k} u_i
\end{align}

に分類される。

　自然言語処理において、例えば、翻訳AIやチャットボットのように入力を文章として、出力も文章の場合を考える（このようなAIは、単語ごとに番号を割り当てる）。単語の数を $K$ 個とすると、文章を予測する際は、$K$ 個の単語（番号）の中から最適な単語（番号）を選ぶ。つまり、$K$ 分類問題である。単語の数は膨大なので、Softmax 関数 の分母を計算するのは、計算コストがかかる。

　引用している keras のブログでは、誤差関数は

\begin{align}
\mathcal{L}_i = - \log \left(\mbox{NN}(z_i,Q) \cdot z_i^{+} /\tau \right) 
\end{align}

\begin{align}
\mathcal{L} = \frac{1}{n}\sum_{i=1}^n \mathcal{L}_i
\end{align}

と計算している。

Attention

　今回の記事で用いる Attention は、以下の tensorflow のチュートリアルで使用される BahdanauAttention の RNN の入力部分を除いたレイヤーを用いる。

つまり、

class BahdanauAttention(f.keras.layers.Layer):
  def __init__(self, units):
    super(BahdanauAttention, self).__init__()
    self.W1 = tf.keras.layers.Dense(units)
    self.V = tf.keras.layers.Dense(1)

  def call(self, features):
    attention_hidden_layer = tf.nn.tanh(self.W1(features))
    score = self.V(attention_hidden_layer)
    attention_weights = tf.nn.softmax(score, axis=1)
    context_vector = attention_weights * features
    context_vector = tf.reduce_sum(context_vector, axis=1)
    return context_vector, attention_weights

である。

　そして、以下の tensorflow のチュートリアルで使用される Self Attention（ MultiHeadAttention） である。

　エンコーダーの畳み込み部分は keras のブログの通りで、最終層に attention を加える。BahdanauAttention の場合、コードで書けば以下のようになる。

class encoder(tf.keras.layers.Layer):
	def __init__(self,width):
		super(encoder, self).__init__()
		self.conv1 = tf.keras.layers.Conv2D(width, kernel_size=3, strides=2, activation="relu")
		self.conv2 = tf.keras.layers.Conv2D(width, kernel_size=3, strides=2, activation="relu")
		self.conv3 = tf.keras.layers.Conv2D(width, kernel_size=3, strides=2, activation="relu")
		self.conv4 = tf.keras.layers.Conv2D(width, kernel_size=3, strides=2, activation="relu")
		self.attention = BahdanauAttention(width)

	def call(self,x):
		attention_weights = {}
		x = self.conv1(x) 
		x = self.conv2(x) 
		x = self.conv3(x)
		x = self.conv4(x)
		# (b,h,w,c) の配列を (b,h*w,c) にする
		x = tf.reshape(x,(x.shape[0], -1, x.shape[3]))
		x ,block= self.attention(x)
		attention_weights['attention_weights'] = block

		return x,attention_weights

　畳み込みによって得られた2次元特徴マップを、1次元の配列にして、Attention weigh を計算し、1次元の配列の特徴マップに張り付ける。Attention weigh を可視化するときは、1次元の配列を2次元配列に戻し、入力画像のサイズにリスケールした。

　Self Attention の場合は、Attention を計算した後のテンソルは、チャネル方向の成分が残るので、Flatten をする必要がある。また、Self Attention は、2次元特徴マップが $N\times N$ の大きさの場合、1次元の配列にすると $N^2$ の配列となる。Attention weight は$N^2\times N^2$ のテンソルとなるので、Attention weigh を可視化するときは、$N\times N$ の配列を $N$個 作成し、それら $N$個 を平均して $N\times N$ の配列にしてから、入力画像のサイズにリスケールした。MultiHead Attention の場合は、Head ごとに平均を行った。

　あとは、keras のブログの通り NNCLR の計算を行った。

stl10 データセットの結果

　
　Keras のブログで使われた stl10 データセットを使い、学習を行った。

　stl10 データセットは、ラベル無し画像が 100,000 枚、ラベル付き訓練データが 5,000 枚、テストデータが 8,000 枚からなるデータセットである。画像サイズは、$96\times 96 \times 3$ であり、ラベル付きデータにおけるクラス数は $10$ である。

　自己教師あり学習は、 ラベル無し画像が 100,000 枚で学習を行った。そのうち、8割のデータを訓練に、残りの2割のデータを検証データとした。
　

Bahdanau Attention

　Bahdanau Attention の結果を確認する。
　以下の左図は、誤差関数の値であり、概ね収束していることが確認できる。右図は、正答率である。

　contrastive accuracy は、support set から得られた特徴ベクトルの組ともう片方の特徴ベクトル組から内積を計算する。つまり、ポジティブペアとネガティブペアの類似度を計算する。ポジティブペアの類似度（内積）が、最も高ければ正解、そうでなければ不正解とする。

　correlation accuracy は、support set から得られた特徴ベクトルの組ともう片方の特徴ベクトル組から相関関数を計算する。ポジティブペアの相関が、最も高ければ正解、そうでなければ不正解とする。

　次に、ラベル付き訓練データが 5,000 枚から 100枚 選び、類似度（内積）の計算を行った。以下の図は、support set から得られた特徴ベクトルの組ともう片方の特徴ベクトル組から内積を計算した図である。

　
対角成分（ポジティブペア）が高くなっていることが確認できる。

　Attention weight の値を確認する。以下の図は、左から本来の画像・データ拡張された画像1・データ拡張された画像2・データ拡張された画像1から得られたAttention weight・データ拡張された画像2から得られたAttention weightである。

概ね、画像内の似たような場所の Attention weight が高くなっていることが確認できる。

Self Attention

　Self Attention の結果を確認する。
　以下の左図は、誤差関数の値であり、概ね収束していることが確認できる。右図は、正答率である。 Bahdanau Attention と比べ correlation accuracy は、低い値となった。

　次に、ラベル付き訓練データが 5,000 枚から 100枚 選び、類似度（内積）の計算を行った。以下の図は、support set から得られた特徴ベクトルの組ともう片方の特徴ベクトル組から内積を計算した図である。

　
対角成分（ポジティブペア）が高くなっていることが確認できる。

　Attention weight の値を確認する。以下の図は、左から本来の画像・データ拡張された画像1・データ拡張された画像2・データ拡張された画像1から得られたAttention weight・データ拡張された画像2から得られたAttention weightである。

概ね、画像内の似たような場所の Attention weight が高くなっていることが確認できる。

深層距離学習または教師無し学習のように使える？

　学習方法や誤差関数が、深層距離学習と似ている。K-means 法を使い、特徴ベクトルをクラスタリングすれば、概ね、同じ物体が写っている画像を（教師無し学習のように）分類してくれると考えた。自明なことかもしれないが、上手く分類してくれない。

　以下の図は、t-SNE を使い特徴ベクトルを2次元に次元削除した結果である。図中の番号は、画像のラベルである。

　自己教師あり学習は、自身の画像を教師としているので、他の画像との関係は指定していないので教師無し学習のように分類してくれないだろう。

　例えば、赤いリンゴ・青いリンゴ・赤い車・青い車 の画像を分類することを考える。色で分類を行うならば、（赤いリンゴ、赤い車）と（青いリンゴ、青い車）と分けられる。物体で分類するならば、（赤いリンゴ、青いリンゴ）と（赤い車、青い車）と分けられる。したがって、何で分類すればよいか指定することは、画像の特徴を指定していることになる。通常の分類問題や深層距離学習では、その特徴を指定しているのだろう。

　

最後に

　NNCLR の エンコーダー部分に Attention を追加して可視化を行った。異なるデータ拡張をしても、概ね、同じ部分に反応することが確認できた。

　通常の自己教師あり学習の使い方である、大量のラベル無しデータセットを使い自己教師あり学習を行い、少数のラベルありデータセットを使い finetuning を行うことである。

　今後は、実データをもとに NNCLR を使ってみたい。