[DL輪読会]Set Transformer: A Framework for Attention-based Permutation-Invariant Nerural Network

DEEP LEARNING JP [DL Papers] Set Transformer: A Framework for Attention- 1 based Permutation-Invariant Neural Network Takumi Ohkuma, Nakayama Lab M1 http://deeplearning.jp/

• http://deeplearning.jp/

自己紹介  大熊拓海（オオクマ  東京大学 情報理工学系研究科 創造情報学専攻 中山研究室 M1  専門はfew-shot learning関連  タクミ） 現在の研究テーマはfew-shot learningにおける適切なトレーニングデータ選 択について

書誌情報 題名：Set Transformer: A Framework for Attentionbased Permutation-Invariant Neural Network  出典：ICML 2019   著者：Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungiin Choi, Yee Whye Teh

概要  NLP(自然言語処理)で用いられるTransformerのモデルを応用して、 Inputとして「集合」を受け取れるモデルを提案した研究。  単純なPooling等では困難な、集合の要素間の関係性を考慮した特徴 量の獲得が可能。  集合を扱う上で大切になる、「順序普遍性」と「任意サイズの集合 の取り扱い」について言及している。  様々なタスクに対して、通常のPoolingなどで集合を処理した場合よ り高いスコアを実現した。

目次 1. Inputとしての集合 2. Attention Set Transformer 4. 実験 3. 5. 結論

目次 1. Inputとしての集合 2. Attention Set Transformer 4. 実験 3. 5. 結論

Inputとしての集合  数値や画像、音の波形等の「インスタンス」単体ではなく、それらを集めた「集 合」を入力として扱いたいことがある。  例として  3D物体認識 ー「複数の視点から物体捉えた画像群」  系列データ整序 ー「系列データ」  Few-shot Learning ー「各クラスのデータ集合」  これらの例では各データごとの情報に加え、「集合」としての情報にも利用価値 がある。  上記の例からもわかるように、Inputとして「集合」を効率良く扱えるモデルがあ れば、様々なタスクの役に立つ。

集合の処理に必要なこと  とはいえ、Inputとして集合を扱うのはそんなに簡単ではない。  集合としての扱いを難しくしている理由が主に以下の二つで ある。 1. 任意のサイズの集合への対応が求められる 2. 順序普遍性が求められる

任意のサイズの集合への対応  入力として集合を受け取るには、任意のサイズの集合に対応しな ければならない。  例として集合の中の最大値を返すモデルを考えると、入力として {1, 4, 6}, {3}, {3, 5, 8, 1, 7, 4} 等の異なる大きさの集合が入力され、 それらに対して等しく答えを返さなければならない。  実用面でも、常に集合サイズが同じであるという状況は考えにく く、可変サイズ集合への対応は必要不可欠。  通常のMLPやCNNではこのような状況に対応することは難しい。

順序普遍性  順序普遍性とは、集合の要素の入力の順番によって出力が変化し ないことである。  集合は基本的に要素の「順序」の情報を持っていない。そのため 要素の並び順によって出力が異なるのは良くない。  例として集合の中の最大値を返すモデルを考えると、入力{1, 4, 6}, と{6, 4, 1}に対する出力は同じでなければならない。  集合を要素ごとにRNNで順番に処理する方法も考えられるが、こ の方法は順序普遍性を満たさない

関連研究   "Deep sets" Zaheer et al. (2017) [2]  順序普遍性を持つ集合の特徴量の抽出を可能にする研究。可変サイズにも対応。  弱点としては集合の特徴量の抽出が比較的単純な計算に基づいており表現力に乏しい 事と、特徴量を固定サイズに圧縮する機構が無いので別途Poolingなどをしなければな らない事。 “Robust attentional aggregation of deep feature sets for multi-view 3D reconstruction” Yang et al. (2018) [3]  Dot-product attentionを用いることで、可変サイズに対応しかつ順序普遍性を保った まま特徴量を固定サイズに圧縮することができる。  特徴量の重み付き和を取っているので、要素間の非線形な関係性に対する対応力は低 い

目次 1. Inputとしての集合 2. Attention Set Transformer 4. 実験 3. 5. 結論

Attentionとは  Attentionとは基本的に「どこに着目するか」を計算して重み付けをする機構であ る。  NLP分野で頻繁に用いられ、特に同分野のブレイクスルーの一つとなった “Attention Is All You Need [1]” は有名。  本研究で用いられるSet Transformer もモデルは基本的に[1]を参考にしており、モ デルに関する本資料の細かい部分がわかりにくければ[1]に関する解説資料も参考 にされたし。(非常に有名な論文なので沢山あると思われる)

Attentionの基本構造  Attentionは基本的に”Query”, “Key”, “Value”の3要素を用いて計算される。  Queryを𝑄 ∊ ℝ𝑛×𝑑𝑞 、KeyををK ∊ ℝ𝑛𝑣 ×𝑑𝑞 、ValueをV ∊ ℝ𝑛𝑣 ×𝑑𝑣 として 𝐴𝑡𝑡 𝑄, 𝐾, 𝑉; ω = ω 𝑄𝐾 𝑇 𝑉 と表すことができる。  この式の意味は「queryベクトルとkeyベクトルの類似度を求め、 その正規化し た重みをvalueベクトルに適用して値を取り出す」と解釈でき、出力ω 𝑄𝐾 𝑇 𝑉は 𝑉の行ごとの重み付き和である。  活性化関数ωは行ごとのsoftmax関数が用いられることが多い。今回はsoftmax 関数を更に 𝑑𝑞 で割ってスケーリングする関数を活性化関数ωとして用いる。

Multi-Head Attention  CNNの特徴量が複数のチャンネルを持つように、Attentionも複数のチャン ネル (ヘッドと呼ばれる) を並列させて表現力を高めることができる。  具体的にはQ, K, Vをヘッドの数(右図のh)だけあらかじめ線形活性しておき、 個別にAttentionした後に合体させ(Concat)、最後にまた線形活性する。  数式で表すと、 𝑀𝑢𝑙𝑡𝑖ℎ𝑒𝑎𝑑 𝑄, 𝐾, 𝑉; ω = 𝑐𝑜𝑛𝑐𝑎𝑡 𝑂1 , 𝑂2, … , 𝑂ℎ 𝑊 𝑂 𝑄 𝑂𝑗 = 𝐴𝑡𝑡 𝑄𝑊𝑗 , 𝐾𝑊𝑗𝐾 , 𝑉𝑊𝑗𝑉 ; ω 𝑄 と表すことができ、各𝑊 𝑂 , 𝑊𝑗 , 𝑊𝑗𝐾 , 𝑊𝑗𝑉は学習可能なパラメータである。

Multi-Head Attention Block (MAB)  前スライドで定義したMulti-Head Attentionを用いてMulti-Head Attention (MAB)を定義する。ここの構造も基本的には[1]と同じである。 具体的には右図の構造である。  Source Target Attention と呼ばれるKeyとValueの入力が同じである Attentionを用いる。  数式で表すと 𝑀𝐴𝐵(𝑋, 𝑌) = 𝐿𝑎𝑦𝑒𝑟𝑁𝑜𝑟𝑚(𝐻 + 𝑟𝐹𝐹(𝐻)) 𝐻 = 𝐿𝑎𝑦𝑒𝑟𝑁𝑜𝑟𝑚(𝑋 + 𝑀𝑢𝑙𝑡𝑖ℎ𝑒𝑎𝑑 𝑋, 𝑌, 𝑌; ω ) とあらわされる。 𝐿𝑎𝑦𝑒𝑟𝑁𝑜𝑟𝑚はLayer Normalization、 𝑟𝐹𝐹は行ごとの線形活性を表す。  本研究のSet TransformerはこのMABを基準として構築される

目次 1. Inputとしての集合 2. Attention Set Transformer 4. 実験 3. 5. 結論

Set Transformer 概要  Set Transformerは前スライドで導入したMABを用いて可変サイズの集合を入力と して受け取り特徴量を抽出し (Encoder)、同様にMABを用いて特徴量を固定サイ ズに圧縮する (Pooling)ことで固定長の特徴ベクトルを得ることを目的とする。  詳細は後述する。  EncoderとPoolingを経て得られた固定長特徴量は、通常の特徴量と同様に様々な タスクに用いることができる。  可変サイズの集合のサイズを𝑛、集合の各要素の次元を𝑑である時、集合𝑋を𝑛 × 𝑑 行列として、 X ∊ ℝ𝑛×𝑑 と表すこととする。

Set Attention Block  Set TransformerのEncoderはMABの2つの入力の両方に、入力の集合 X ∊ ℝ𝑛×𝑑 を代入したSet Attention Block (SAB) が用いられる。  数式で表すと、 𝑆𝐴𝐵(𝑋) = 𝑀𝐴𝐵(𝑋, 𝑋) と表される。  Self Attentionと同じと言った方が伝わるかもしれない。  入出力の次元は ℝ𝑛×𝑑 → ℝ𝑛×𝑑 と変化しない。  計算量はMulti-head Attentionの計算でX𝑊𝑗𝑄 (𝑋𝑊𝑗𝐾 )𝑇 等を計算するので、 O(𝑛2 )の計算量となる。

Induce Set Attention Block (ISAB)  SABは計算量が𝑛の2乗オーダーとなるので、入力となる集合サイズが大きい時 計算量が爆発してしまう。これを防止するために代わりに用いることができる のがInduce Set Attention Block (ISAB) である。  ISABはMABを右図の様に二段重ねにしたものであり注目すべきは一段目の入力 の 𝐼 ∊ ℝ𝑚×𝑑であり、これは学習可能パラメータである。  ISABは処理の中で𝑛に依存する行列同士の積が出て来ず、𝐼は入力集合の大きさ にかかわらず一定なので計算量はO(𝑚𝑛)である。𝑛に関して計算量が線形オー ダーなので、大きな集合を扱うのに適している。  ISABを数式で表すと 𝐼𝑆𝐴𝐵𝑚 𝑋 = 𝑀𝐴𝐵 𝑋, 𝐻 ∊ ℝ𝑛×𝑑 𝐻 = 𝑀𝐴𝐵(𝐼, 𝑋) ∊ ℝ𝑚×𝑑 と表される。

Pooling by Multihead Attention (PMA)  前スライドで導入したSAB及びISABによって抽出された特徴量を固定サイズに 圧縮するための、Pooling by Multihead Attention (PMA) を導入する。  PMAは、入力 Z ∊ ℝ𝑛×𝑑 、学習可能パラメータ S ∊ ℝ𝑘×𝑑 を用いて 𝑃𝑀𝐴𝑘 (𝑍) = 𝑀𝐴𝐵(𝑆, 𝑟𝐹𝐹(𝑍)) ∊ ℝ𝑘×𝑑 と表される。  行列Sの行数を表す𝑘は定数なので、入力サイズ𝑛に寄らない固定長の出力を得る ことができる。𝑘 = 1としてベクトルの出力を得る場合が多いが、 𝑘 > 1として さらにSAB等を重ねることも可能。  (meanやmax 等の)通常のPoolingが入力要素に関係なく単純に平均や最大値を 取るのに対し、PMAはAttention機構を持つのでより柔軟なPoolingが可能と なっている。

Set Transformerの全体構造  SAB, ISAB, PAMを単層もしくは多層に重ねてSet Transformer全体を構築する。 例1 例2 FC FC 𝑃𝑀𝐴1 SAB or ISAB SAB or ISAB SAB or ISAB SAB or ISAB SAB or ISAB SAB or ISAB SAB or ISAB 𝑃𝑀𝐴4

目次 1. Inputとしての集合 2. Attention Set Transformer 4. 実験 3. 5. 結論

条件設定  本論文では、複数のタスクに対して Set Transformerの有効性検証している。 そのために提案手法と既存手法を組み合わせて以下の条件で比較実験を行う。  rFF + Pooling：通常の row-wise feedforward + Pooling。  rFFp-mean/max + Pooling：rFFの代わりに[2]の順序普遍特徴量を用いる。  rFF + Dotprod：Poolingの代わりに[3]のAttentionによる重み付き和を用いる。  SAB (ISAB) + Pooling (ours)：rFFの代わりにSAB (ISAB)を用いる。  rFF + PMA (ours)：Poolingの代わりにPMAを用いる。  SAB (ISAB) + PMA (ours)：提案手法をFullで用いる。

実験１ (最大値検出)  単純なトイタスクで、与えられた数値集合の最大値を返すというタスク。  データセットは、[0, 100]の実数を1～10個独立に選び集合とする。  このタスクは「rFF + Pooling (mean/sum/max)」と「SAB + PMA (ours)」で実験。  結果（平均絶対値誤差）は以下のとおりである。  最大値を取るというタスクなので、maxpoolingを用いたときの性能が最も良いが、特に タスク特化していない提案手法もそれに肉薄する精度を実現している。

実験２：要素の種類カウント(1)  集合として画像のデータセットOmniglot[4]を6枚から10枚集めた集合を入力として、 何種類のクラスが集合内に含まれているかカウントする。  結果としては、提案手法「SAB + PMA」が最も高い。  PMAのみを用いた場合は殆ど効果が無く、逆にSABのみでも高いスコアが出ている。 Omniglotの画像例

実験2：要素数の種類のカウント(2)  SABとISABの精度を比較する。  ISABはパラメータ𝐼 ∊ ℝ𝑚×𝑑の次元数𝑚を変化させ て精度を比較する(右図では𝑛)。  次元数が大きくなると精度が向上するがSABを用 いた場合「SAB + PMA」よりも低い。  集合の大きさは6から10なので、ISABで𝑚 が6から 10くらいの場合と比較すると、明らかにSABより も低いので、入力される集合のサイズが小さいと わかっているときはSABを用いた方がよさそう。

実験3：混合ガウス分布  混合ガウス分布のパラメータ推定には一般的にEMアルゴリズムが用いられるが、Set Transformerを用いてダイレクトにパラメータを推定することができる。  具体的には、 𝑛 𝑘 𝑙𝑜𝑔𝑝 𝑋; θ = ෍ 𝑙𝑜𝑔 ෍ π𝑗 𝑁(𝑥𝑖 ; μ𝑗 , 𝑑𝑖𝑎𝑔(σ𝑗2)) 𝑖=1 𝑗=1 のパラメータθ = {π𝑗 , μ𝑗 , σ𝑗 }𝑘𝑗=1にデータをフィットさせる。  用いたデータセットは、人為的に生成した4峰の2D混合ガウス分布と、CIFAR-100で ある。CIFAR-100に関してはpretrained-VGG[5]で抽出した特徴量を用いた。

実験3：混合ガウス分布(2)  Set Transformer等の手法に加えて、OracleとしてEMアルゴリズムを収束するまで用いた結果 (LL0/data, ARI0)、さらに各手法の後に1ループだけEMアルゴリズムを用いた結果を (LL1/data, ARI1)記載する。  結果はSyntheticが人為的に生成したガウス分布で平均尤度で評価。CIFAR-100はadjusted rand index (ARI) で評価している。  Set Transformer+EMアルゴリズム1サイクルで、収束するまで行ったEMアルゴリズムを上 回った。

実験４：例外検出  属性ラベルを持つ顔のデータセットCelebAを用いて例外検出を行う。  各集合ごとにランダムで属性を二つ選び、その両方の属性を持つデータを7つ、 持たないものを1つ合わせて集合とし、その一つを予測するモデルを学習。  Set Transformerが従来の手法を上回った。

実験５：Point Cloud 分類  3DモデルデータModelNet40[6]のクラス分類  3Dモデル形式としては、多数の頂点を3次元空間に配置して3Dモデルを構成する Point Cloudという形式である。  入力は、各データの頂点のうちの{100, 1000, 5000}点選択して集合を作る。  集合サイズが5000のとき、SABを用いると計算に非常に時間がかかってしまう （数十分）ので、より高速（数秒）なISABのみを用いる。  集合が小さい (100点) のとき、提案手法はより効果を発揮する。

目次 1. Inputとしての集合 2. Attention Set Transformer 4. 実験 3. 5. 結論

結論  自然言語処理のTransformerモデルを応用することで、順序普遍かつ任意のサイ ズの集合の特徴量を抽出する手法を提案した。  従来手法と比較して、Attentionを用いることでより要素間の関係性を柔軟に捉 えることができており、結果として精度の向上につながっている。  集合サイズが大きいと計算量が爆発するという弱点を、ISABによって克服。  Future work として、Set Transformerをメタ学習等に応用するという方向性が考 えられる。

参考文献 [1] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems. 2017. [2] Zaheer, Manzil, et al. "Deep sets." Advances in neural information processing systems. 2017. [3] Yang, Bo, et al. "Robust attentional aggregation of deep feature sets for multi-view 3D reconstruction." International Journal of Computer Vision 128.1 (2020): 53-73. [4] Lake, B. M., et al. “Human-level concept learning through probabilistic program induction.” Science, 350(6266):1332–1338, 2015 [5] Karen Simonyan et al. “Very deep convolutional networks for large-scale image recognition” In proc of ICLR 2014 [6] Chang, A. X., et al “An information-rich 3D model repository.”, arXiv:1512.03012, 2015.