医療系AIエンジニアの技術メモ

ディープラーニング（主に画像系）の技術メモブログです

Attention Augmented Convolutional Networks

論文メモ

スポンサーリンク

f:id:y_kurashina:20190506002354j:plain
Fig1. Attention-augmented convolution

論文URL

https://arxiv.org/pdf/1904.09925.pdf

2019年4月公開

ポイント

  • 画像の特徴量マップを算出するAttention Augmentation(AA)ブロックをGoogle Brainが提案
    • 機械翻訳等のシーケンスデータを扱うTransfomerに、空間並進普遍性を考慮する項を追加
      • Transformerでは注意を向けるべき場所を抽出する辞書(クエリ、キー、バリュー)を学習を通して構築する
    • SE blockのような従来型の注意機構では、畳み込み層の出力特徴量マップの補正重みを計算していたのに対し、AAブロックでは特徴量マップそのものを出力する
      • 論文中の実験では、畳み込み層とAAブロックの両方の特徴量マップをconcatして使用する方が高精度となっていたが、AAブロックのみで画像分類することも可能
  • 性能
    • ResNet50によるImageNet画像分類では精度が1.3%改善
    • RetinaNetによるCOCO物体検出ではmAPが1.4%改善
    • モデルのパラメータ数や計算コストはオリジナルのResNetと同程度
      • 特徴量マップのチャンネル数を変えない場合、畳み込み層のチャンネルを減らしAAブロックの特徴量を増やすほどパラメータ数と計算コストを減らすことができる
    • AAブロックのメモリ使用量は大きい
      • (height \times width) ^ 2に比例している

Attention Augmentation(AA)ブロック

  • 数式に用いる変数
    • B, H, W, F _ {in} : 入力特徴量マップのバッチサイズ、行数、列数、チャンネル数
    • N _ h : ヘッドの個数。ヘッドごとに異なるクエリ、キー、バリューを計算する。
    • d _ k ^ h : h番目ヘッドのクエリおよびキー特徴量の深さ(＝チャンネル数)
    • d _ v ^ h : h番目ヘッドのバリュー特徴量の深さ
    • X : 入力テンソル。サイズはB \times H \times W \times F _ {in}
    • Q ^ h, K ^ h, V ^ h : h番目ヘッドのクエリ、キー、バリュー
    • S _ H ^ {h,rel}, S _ W ^{h,rel} : h番目ヘッドの行方向、列方向の空間並進普遍性考慮項
    • O _ h : h番目ヘッドの出力特徴量マップ
    • MHA : multihead-attention 。AAブロックの出力特徴量マップ
    • 補足
      • 論文ではヘッドごとに異なる項の一部にしかhが付記されていないが、明確化のため関連する全ての項にhを追加
  • クエリの計算
    • Q ^ h = XW _ q ^ h
      • W _ q ^ h : Xをクエリ特徴量に変換する学習パラメータ。サイズは F _ {in} \times d _ k ^ h
      • 特徴量マップのチャンネル数を変換する演算となっており、1x1の畳み込みで計算できる
      • Q ^ hのサイズはB \times H \times W \times d _ k ^ h
  • キーの計算
    • K ^ h = XW _ k ^ h
      • W _ k ^ h : Xをキー特徴量に変換する学習パラメータ。サイズは F _ {in} \times d _ k ^ h
      • K ^ hのサイズはB \times H \times W \times d _ k ^ h
  • バリューの計算
    • V ^ h = XW _ v ^ h
      • W _ v ^ h : Xをバリュー特徴量に変換する学習パラメータ。サイズは F _ {in} \times d _ v ^ h
      • V ^ hのサイズはB \times H \times W \times d _ v ^ h
  • 行方向の空間並進普遍性考慮項
    • S _ H ^ {h,rel} \lbrack i,j \rbrack = q _ i ^ h (r _ {ij} ^ H) ^ T
      • i : 空間位置。行列方向成分をflat化したベクトル上のインデックスと考えれば良い。
      • j : 位置iに対する相対行。相対行の範囲は(-H+1) ～ (+H-1)
      • q _ i ^ h : Q ^ hの位置iにおける成分
      • r _ {ij} ^ H : i,jにおけるクエリ補正項。学習パラメータ。
        • サイズは(2H-1) \times d _ k ^ h
        • 全ヘッドの全列で共通の補正項となっている
        • q _ i ^ h (r _ {ij} ^ H) ^ Tの計算時に画像範囲外となるr _ {ij} ^ Hは使用しない
          • 例：一行目画素であれば、相対行0 ～ (+H-1)のr _ {ij} ^ H成分が使用される
      • S _ H ^ {h,rel}のサイズはB \times H ^2 \times W ^2
        • 全画素について他の画素との相対的位置に基づく補正量が保持されている
  • 列方向の空間並進普遍性考慮項
    • 上記、行方向の普遍性考慮項において行と列を入れ替えて考えれば良い
  • ヘッドの出力特徴量マップ
    • O _ h = Softmax ( \frac{Q ^ h (K ^ h) ^ T}{ \sqrt{d _ k ^ h }}+S _ H ^ {h,rel} + S _ W ^ {h,rel} ) V
      • 上記の式からS _ H ^ {h,rel} , S _ W ^ {h,rel}を取り除くとTransformerと同じ形式のブロックになる。
    • Softmax ()項のサイズはB \times H ^2 \times W ^2
    • O _ hのサイズはB \times H \times W \times d _ v ^ h
  • AAブロックの出力特徴量マップ
    • MHA(X) = Concat \lbrack O _ 1 , ... , O _ {N _ h} \rbrack W ^ O
      • W ^ O : サイズが d _ v \times d _ vの学習パラメータ
      • d _ v = \sum _ {h} d _ v ^ h
    • MHA(X)のサイズはB \times H \times W \times d _ v

AA-ResNetによるImageNet画像分類性能

  • オリジナルのResNetに対してはもちろんAttentionを行っているSE-ResNetよりも高性能

f:id:y_kurashina:20190506172224j:plain