[Journal club] LambdaNetworks: Modeling Long-Range Interactions Without Attention

1. -BNCEB/FUXPSLT
   .PEFMJOH
   -POH3BOHF
   *OUFSBDUJPOT
   8JUIPVU
   "UUFOUJPO 慶應義塾大学 杉浦孔明研究室 飯田 紡 Irwan Bello (Google) ICLR 2021 Bello, I. (2021). Lambdanetworks: Modeling long-range interactions without attention. arXiv preprint arXiv:2102.08602.
2. എܠɿ"UUFOUJPOͷܭࢉྔ͸ଟ͍ Self-Attention シーケンス長 ! -> "(!!) 画像をFlattenして入力する場合 画像全体を入力は困難（256×256 → 256") わずかな部分しか入力不可 画像全体との関係を捉えられない 1 I have a pen I 0.5 0.2 0.1 0.2 have 0.1 0.4 0.1 0.4 a 0.2 0.1 0.5 0.2 pen 0.2 0.2 0.1 0.5
3. ؔ࿈ݚڀ Attentionの計算量削減手法は汎用的 画像特有の構造を考慮しつつ画像全体を入力したい 2 Linear Transformer [Katharopoulos+, ICML20] Transformerを低ランク近似することで計算量削減 [Shazeer, arXiv preprint19] Queryを分割することでattentionの計算量削減 ViT [Dosovitskiy+, arXiv preprint21] パッチ入力により画像にTransformerを適用 パッチは高解像度の画像やdetectionに応用できるか不明
4. ఏҊख๏ɿ-BNCEB
   /FUXPSLT 3 š 入力 Inputs: ! ∈ ℝ!×# Context: $ ∈ ℝ$×# ! = $のときSelf-attentionと同等 š 出力 & ∈ ℝ!×# ': Query, (: Key, ): Value ' = !*% , ( = $*& , ) = $*' + Content Lambda Position Lambdas * ! + , - . / 0# = 2# $3# = 2% + 2# & $ 3# 5 = 2' , 2! , ⋯ , 2( 2# = 2% + 2# &
5. ఏҊख๏ɿ$POUFOU
   -BNCEB Attentionに対応 Key, Valueを集約後Queryに作用 → 計算量削減 4 + Position Lambdas Content Lambda * ! + - . /
6. "UUFOUJPOͱ$POUFOU
   -BNCEBͷൺֱ 5 +(×* *+×* /(×, -!×# -+×, .+×* -!×# -!×! 8(×+ 9 8(×+ soft max ,(×* !, #の積でAttention map を作成 画像データでは一般に データ長 $, % : 大 &, ' ∶ 小さくできる（ハイパラ） 大きな ! × : を使用せずに ! × ; の出力できないか？ Attention Layer
7. "UUFOUJPOͱ$POUFOU
   -BNCEBͷൺֱ 6 画像データでは一般に データ長 :, ! : 大 ;, < ∶ 小さくできる（ハイパラ） -, .の積で5%を作成 ! × : → k × ; Content Lambda +(×* *+×* /(×, -!×# -+×, .+×* -!×# -!×! ,(×* soft max ? -+×, 5% ,×* ※正確には.$と同じ/×0の Position Lambdasを足したものを1にかける
8. "UUFOUJPOͱ$POUFOU
   -BNCEBͷൺֱ 7 /(×, -+×, .+×* 8(×+ 9 8(×+ soft max ,(×* /(×, -+×, .+×* ,(×* soft max ? -+×, 5% ,×* " !:(< + ; ) " ! + : <; Attention Layer Content Lambda 2×3, 3×5行列の積は6(235) 時間計算量 空間計算量 " !: " <; ((*+)などもありえますが 後にO +. しか使わないので省略
9. ఏҊख๏ɿ1PTJUJPO
   -BNCEBT Positional Encodingに対応 ポジション1つ(1〜,)につき -×/の行列0を使用 0は全てのデータで同じ 時間計算量 1(,-/3) 空間計算量 1(,-/) 8 + Content Lambda Position Lambdas * ! + , - . / . +×* @' +×, 5' & ,×* 5! & ,×* 5( & ,×* … @( +×, …
10. ܭࢉྔͷൺֱ 9 時間計算量 空間計算量 Attention 6(9:; / + 0 ) 6(9:;) Content only 6( : + ; /0) 6(9/0) Position only 6(9:;/0) 6(:;/) Lambda 6(9:;/0) 6(:;/ + 9:/0) 入力サイズ %×& , $×& バッチサイズ * !, #の大きさ ' 時間計算量 ほぼ等しい（工夫で高速化可能） 空間計算量 Positionは全部同じ → *に非依存 = 大きな%$がバッチサイズ非依存 → バッチサイズを大きくできる 画像サイズ224 × 224, バッチサイズ 128 Attention: 120 GB Lambda: 0.63 GB
11. ߴ଎Խͷ޻෉ š Multi-Query Queryをℎ 個に分割 ! = [!! , !" , ⋯ , !# ] 並列にlambdaを計算して結合 0 = 1 !! , 1 !" , ⋯ , 1 !# % → $ # より時間計算量 3 %$&'( # に（Key, Valueはそのまま） š Lambda Convolution 計算量の多いPosition LambdasをCNNに置き換え 4) &×' = Conv2D(<) 11
12. ఆྔత݁Ռɿ-BNCEB͕BUUFOUJPO΍$//Λ্ճΔ 12 ResNet50の3 × 3Convを Layerに置き換え ImageNetの 224 ×224画像を使用 最も少ないパラメータ数で Top-1 accuracy最大を達成
13. ఆྔత݁Ռɿܭࢉྔগ͔ͭਫ਼౓ྑ 13 Lambda Layerは • メモリ使用量 少 • Throughput 大 全レイヤーでpos lambdasを 共有しても精度は落ちない Attentionでは画像全体を使 用できない（Global ~） 9 ∶ バッチサイズ, ℎ ∶ ヘッド/クエリ数, : ∶ inputの長さ ; ∶ contextの長さ, / ∶ Q, Kのdepth, ? ∶ レイヤー数 画像全体 軸ごと クロップ
14. ఆྔత݁Ռɿগͳ͍ύϥϝʔλ͔ͭߴ଎ͰେܕϞσϧʹ ඖఢ͢Δਫ਼౓Λୡ੒ 14 横軸：Trainにかかる時間 ≒ モデルサイズ LambdaResNet420, 画像サイズ320 LambdaResNet350, 画像サイズ320 350 epochの訓練をしたとき 3〜4倍高速でEfficientNetに匹敵 ViTには劣るものの、非常に高速
15. "CMBUJPO
    4UVEJFT 16 Content, Position片方のみを使用 Positionが精度に大きく貢献 ConvとLambdaの併用 前方にL → Throughput低下 後方にL → Throughputそのまま 後ろ2つL → 精度向上かつ高速
16. ·ͱΊ 17 背景 Attentionの計算量削減 提案 先にKey, ValueをまとめるLambda Networks 結果 高速かつ大規模モデルに匹敵する精度