こんにちは、CCCMKホールディングス TECH LAB三浦です。
ふと気が付くと近所に小型のスポーツジムがたくさんできていました。スポーツジムというと、プールがあって、サウナがあって・・・と、とても大規模な施設の印象があったのですが、小型のスポーツジムにはそういった設備はなくて、トレーニング用の機械がいくつか設置されていて、ホテルの中に時々あるトレーニングルームのような感じです。先日試しにそのうちの1つに行ってみたのですが、散歩や買い物の途中にちょっと寄って体を動かして帰って来られる気軽さがとてもいいな、と思いました。今までこういうスポーツジムの使い方って考えたことがなかったので、すごく新鮮な体験でした。
さて、現在のLarge Language Model(LLM)において、Transformerというモデル構造は欠かせないものになりました。Transformerは自然言語処理の領域で誕生し、今なおその領域で大きな貢献をしているのですが、自然言語にとどまらず画像解析の領域においてもTransformerを利用する研究が行われており、その中でもVision Transformer(ViT)は特に有名で、ViTを元にした様々な派生モデルも誕生しています。
画像処理においてはConvolutional Neural Network(CNN)という構造が長く使われており、私もCNNベースのモデルを使うことがほとんどです。CNNは使われてきた期間が長いためドキュメントや実装例も豊富で、扱いやすいことが理由です。一方、ViT関連の情報も増えてきており、ちゃんと扱えるようになって最新の情報をキャッチアップ出来るようになっておきたいな、と考えるようになりました。
そこで最近拝読しているのがこちらの書籍です。
タイトル: Computer Vision Libraryシリーズ Vision Transformer入門 著者: 片岡裕雄 監修，山本晋太郎，徳永匡臣，箕浦大晃，邱玥（QIU YUE），品川政太朗 著 出版: 2022年9月17日紙版発売, 2022年9月15日電子版発売 出版社: 技術評論社
ViTの実装例やCNNとの比較、様々な画像解析タスクに対応するViT派生モデルの紹介など、読んでいるとViTって想像以上に面白いんだな、と感じました。様々なViT派生モデルの中で、特に自分が興味を持ったのがSwin Transformerというモデルです。ResNetやVGGといったCNNベースのモデルが画像分類をはじめ、物体検出などの様々なタスクのバックボーンとして利用出来るように、このSwin Transformerも様々な画像解析タスクのバックボーンとして利用することが出来る点に、魅力を感じました。
今回はSwin Transformerについて、論文と公式の実装を見ながら調べてみたのでまとめてみたいと思います。
Vision Transformer(ViT)について
Swin Transformerを理解する前に、そもそもVision Transformer(ViT)ってどんなモデルなのかについて、簡単にまとめてみたいと思います。ViTはこちらの論文で提案されたモデルです。
Title: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale Authors:Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby Submit:Submitted on 22 Oct 2020 (v1), last revised 3 Jun 2021 arXivURL:https://arxiv.org/abs/2010.11929
この論文の"1 INTRODUCTION"から"2 RELATED WORK"でいくつか触れられているのですが、論文が書かれた当時、自然言語で成功したTransformerをどうやって画像に適用するのか、様々な方法が模索されていたことが分かります。たとえば画像を構成するピクセルをテキストのトークンのようにTransformerのself-attentionにかける方法です。これは自然な発想ですが、テキストに含まれるトークンに比べ、画像に含まれるピクセルの数は膨大なため、計算コストを考えると非現実です。また、畳み込み処理で抽出した局所的な特徴データをTransformerに入力する、というCNNと併用したアプローチもあったようです。
ViTはCNNを使わずにTransformerの構造だけで画像を取り扱うことを実現しました。その方法は画像を決まったサイズの格子(パッチ)に分割し、そのパッチをテキストと同じくトークンとして扱います。パッチは埋め込み処理によってベクトル化され、画像における位置情報を表す位置埋め込みを加算された後、Transformerに入力されます。
ViTからSwin Transformerへ
Swin Transformerは2021年の論文で提案したモデルです。
Title: Swin Transformer: Hierarchical Vision Transformer using Shifted Windows Authors:Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo Submit:Submitted on 25 Mar 2021 (v1), last revised 17 Aug 2021 arXivURL:https://arxiv.org/abs/2103.14030
Swin Transformerの公式の実装がGithubのレポジトリに公開されています。
ViTは高解像度の画像を扱おうとするとself-attentionにかかる計算コストが急激に高まってしまうため、扱う画像が高解像度な傾向にある物体検出やセグメンテーションなどの用途には扱えないという課題がありました。
Swin Transformerではこの課題に対し、複数のパッチを囲んだ矩形のwindowを設けてその中のパッチに限定してself-attentionを計算することで計算コストを削減することを可能にしています。しかしwindowの内部だけのself-attentionでは、画像の局所的な情報しか捉えることが出来ず、大域的な特徴を捉えることが出来なくなります。そこでSwin Transformerではwindowの切り方を2パターン用意します。1つは等間隔で区切ったパターンで、もう1つはそれをwindowのサイズの約半分、左上にずらして作成したパターン(Shifted windows)です。Swin Transformerの"Swin"はShifted windowsに由来しているようです。公式の実装ではPyTorchのtorch.roll()を使ってずらす処理を実現しています。この2つのwindowパターンを繰り返しながらself-attentionを計算すると、より広い範囲の情報を取り込んだself-attentionの計算が可能になります。
また、ネットワークが深くなるのに従って、PatchMergingという処理を施し、周囲4つのパッチを1つのパッチにまとめていきます。1つのwindowに含まれるパッチの数は不変に保ちながら画像のより大域的な特徴を取得出来るようになっています。この処理はCNNの畳み込み処理になんとなく似ているな、と感じました。
このように同じ数のパッチが含まれるwindow内部でself-attentionの処理を制限することで、入力画像のサイズが大きくなったとしてもwindowの数を増やせば対応することが出来、そのため計算コストの増加が線形増加に抑えられるという点がSwin Transformerの特徴です。そのためViTでは扱えなかったより高解像度の画像をSwin Transformerでは処理することが出来るようになります。
Shifted Window based Self-Attention
Swin Transformerのself-attention計算は等間隔に区切られたwindowパターン(W-MSA)と、それをずらしたShifted windowsパターン(SW-MSA)で行います。SW-MSAの方は先の図(Figure 2)を見ると分かるように、W-MSAよりも扱うwindowの数が2.25倍に増えてしまっています(先の図だとW-MSA:4からSW-MSA:9に増加)。
これに対しcyclic shiftという方法を取ると、windowの数をW-MSAと同じ数に保ちながらSW-MSAの計算が可能になっています・・・とのことなのですが、ここを理解するのに苦戦してしまいました。
cyclic shiftを理解してみる
画像処理関連の内容は、文章だけ読んでもなかなかイメージしづらいことがあります。そういう時はソースコードを見るとちょっと理解できることがあります。GithubのSwin Transformerのリポジトリには、cyclic shiftに関する質問がいくつかissueとして挙がっており、そちらがcyclic shiftを理解するのにとても参考になりました。
cyclic shiftについて図で表すと次のようになります。
なんだかパズルみたいですね。これでW-MSAと同じ形状のwindow内に細かいwindowを全て含めることが出来ました。あとはこのwindow(上の図だと4つ)ごとにself-attentionを計算するのですが、このままでは異なる領域間でself-attentionの計算をしてしまうことになります。これを防ぐため、self-attention計算時にattention-MASKを使用します。こちらもイメージを図にしてみました。
このようにして2つのwindowのパターンでも同じような計算コストで対応することが出来る点がSwin Transformerの面白いアプローチです。理解するのに大分苦労したのですが、手続き自体はかなりシンプルだと思います。
まとめ
ということで、今回はViTの話からそれを改良したSwin Transformerについて調べたことをまとめてみました。本当はもう少しSwin Transformerについて調べ、モデルも動かしてみたかったところですがcyclic shiftの理解で大分詰まってしまい、そこまでたどり着くことが出来ませんでした・・・。今後は動かしてみて、CNNベースのものと比較をしてみたいと思います！
