[Journal club] Hyena Hierarchy: Towards Larger Convolutional Language Models

慶應義塾⼤学 杉浦孔明研究室 M1 和田唯我 / Yuiga Wada

Transcript

1. Hyena Hierarchy: Towards Larger
   Convolutional Language Models
   Michael Poli1, Stefano Massaroli2, Eric Nguyen1, Daniel Y. Fu1, Tri Dao1,
   Stephen Baccus1, Yoshua Bengio2, Stefano Ermon1, Christopher Ré1
   (1Stanford University 2Mila and Université de Montréal)
   慶應義塾⼤学
   杉浦孔明研究室 M1 和⽥唯我
   Michael Poli et al., “Hyena Hierarchy: Towards Larger Convolutional Language Models”, arXiv preprint arXiv:2302.10866.
   

   View Slide

2. 概要
   2
   ü 背景
   • Transformerは強⼒だが系列⻑に対してquadraticな計算量が掛かる
   • Efficient Transformerは計算量を改善するものの性能が不⼗分
   ü 提案⼿法
   • Subquadraticな計算量でTransformerに匹敵するモデルHyenaを提案
   • Attention-freeであり，SSMベースのH3やGSSを⼀般化したモデル
   ü 結果
   • WikiText103やThe Pileにおいて80%の学習コストでTransformerに匹敵する結果
   • ImageNetにおいてもVision Transformerと同程度のaccuracy
   

   View Slide

3. 関連研究: H3 (Hungry Hungry Hippos)
   3
   ü SSMは①前⽅にあるトークンの記憶と②トークン間の⽐較が不得意
   ü H3 (Hungry Hungry Hippos) [Fu+, ICLR23]
   • ⾔語モデリングにおいて強⼒かつ⾼速なSSMレイヤ
   ① shift演算によりトークンを記憶
   ② Attention-Likeに乗算することでトークン間の関係を記憶 (like 𝑄𝐾!𝑉)
   → AttentionとH3を交互に配置したHybrid-H3はTransformerと同程度のperplexity
   課題: H3単体だと未だ性能に改善の余地あり / Hybrid-H3はAttention-Freeでない
   

   View Slide

4. 背景: Attentionの性能に匹敵するsubquadraticな演算が望まれている
   4
   • Transformerは強⼒だが O 𝑑𝐿" の計算量が掛かる
   • 系列⻑に対してquadraticな計算量 → ⼊⼒が⻑くなると計算コスト⼤
   • Efficient Transformer: 計算量を改善するが，Transformerとの性能のギャップ
   • 例: Routing Transformer[Roy, TACL21], Reformer[Kitaev+, ICLR20], Linformer[Wang+, 20]
   Attentionの性能に匹敵する
   subquadraticな演算が望まれる
   Reformer[Kitaev+, ICLR20]
   

   View Slide

5. 背景: Attentionの３つの特性
   5
   ü Data control
   • ⼊⼒に応じて異なる線形関数(重み)を表現できる
   → Attention Mapは⼊⼒により動的に変化
   ü Sublinear parameter scaling
   • パラメタ数が系列⻑に対しsublinear(2次以下)でスケーリング
   → パラメタをFFN等の他のブロックに割り振ることが可能
   ü Unrestricted context
   • コンテキストに対する制約が存在しない (e.g., 局所性)
   → ⼊⼒における任意の2点間の依存関係を近似することが可能
   本論⽂の⽬的: これら３つの特性を満たす subquadratic な演算の設計
   

   View Slide

6. 畳み込みはフィルタ ℎ の最適化⽅法によって２つに⼤別可能
   6
   • Explicit Convolutions
   • フィルタ ℎ を直接最適化する⽅法 (e.g., CNN)
   • CNNは信号処理において有限インパルス応答 (FIR) といえる
   • FIRのパラメタ数はカーネルサイズ 𝑀 (= step size)に関して線形にスケーリング
   • Implicit Convolutions
   • フィルタ ℎ を間接的に最適化する⽅法 (e.g., SSM)
   • パラメタ数がstep sizeに依存しないようフィルタ ℎ をstep 𝑡 で記述 (ℎ#
   ≔ 𝛾$
   (𝑡))
   SSM (reccurent) SSM (convolution)
   

   View Slide

7. 畳み込みはフィルタ ℎ の最適化⽅法によって２つに⼤別可能
   7
   • Explicit Convolutions
   • フィルタ ℎ を直接最適化する⽅法 (e.g., CNN)
   • CNNは信号処理において有限インパルス応答 (FIR) といえる
   • FIRのパラメタ数はカーネルサイズ 𝑀 (= step size)に関して線形にスケーリング
   • Implicit Convolutions
   • フィルタ ℎ を間接的に最適化する⽅法 (e.g., SSM)
   • パラメタ数がstep sizeに依存しないようフィルタ ℎ を時間 𝑡 で記述 (ℎ#
   ≔ 𝛾$
   (𝑡))
   SSM (reccurent) SSM (convolution)
   SSMのフィルタはstep sizeである 𝑀に依存せず，学習可能パラメタ 𝐴, 𝐵, 𝐶, 𝐷 で記述可
   

   View Slide

8. 畳み込み / H3はToeplitz⾏列の⾏列積で記述できる
   8
   • ⼀般的な畳み込みにおける 𝑡 成分
   • ℎ!"#
   が繰り返し使⽤されるので，対⾓線に沿って
   値が⼀定である Toeplitz⾏列 S$
   により記述可能
   • H3 [Fu+, ICLR23]
   𝜓, 𝜙 : SSMでparametrizeされたフィルタ
   𝐷!
   , 𝐷"
   : 𝑞, 𝑘 の要素で構成された対⾓⾏列
   → SSMはImplicit Convolutionsの類なので，
   H3はToeplitz⾏列と対⾓⾏列の積で記述可能
   

   View Slide

9. 提案⼿法: Hyena
   9
   ü 提案⼿法: Hyena
   • H3を⼀般化 + ⻑い畳み込みフィルタによってAttentionを代替
   • Hyena MatricesとHyena Filtersによって構成
   • Hyena Matrices
   • Toeplitz⾏列と対⾓⾏列の積を𝑁 回に⼀般化 (H3では2回)
   

   View Slide

10. 提案⼿法: Hyena
    10
    ü 提案⼿法: Hyena
    • H3を⼀般化 + ⻑い畳み込みフィルタによってAttentionを代替
    • Hyena MatricesとHyena Filtersによって構成
    • Hyena Matrices
    • Toeplitz⾏列と対⾓⾏列の積を𝑁 回に⼀般化 (H3では2回)
    ⼊⼒に応じて重みが動的に変化
    → Data controlled matrix
    

    View Slide

11. 提案⼿法: Hyena
    11
    • Hyena Filters
    • ⼊⼒と同じ⻑さの⻑いフィルタによる畳み込み (FFN部分のみ学習可能)
    • Window: 指数関数的に減衰する関数は⾼周波フィルタと相性が良い
    • Positional Encoding: ⾼周波数成分の学習が安定化 [Basri+, PMLR20]
    • FFN: S4[Gu+, ICLR22]のフィルタを近似的に学習することが可能
    

    View Slide

12. 提案⼿法: Hyena
    12
    • Hyena Filters
    • ⼊⼒と同じ⻑さの⻑いフィルタによる畳み込み (FFN部分のみ学習可能)
    • Window: 指数関数的に減衰する関数は⾼周波フィルタと相性が良い
    • Positional Encoding: ⾼周波数成分の学習が安定化 [Basri+, PMLR20]
    • FFN: S4[Gu+, ICLR22]のフィルタを近似的に学習することが可能
    • パラメタ数が系列⻑に対して強依存していない
    → Sublinear parameter scaling
    • ⼊⼒と同系列⻑のフィルタを扱うため，任意の⼆点間の依存関係を近似可能
    → Unrestricted context
    H3ではSSMによってフィルタが規定されたが，HyenaではHyena Filterを使⽤
    

    View Slide

13. Hyena Algorithm
    13
    ⼊⼒ 𝑢 に対して，Attentionにおける
    Q = 𝑢𝑊!
    , 𝐾 = 𝑢𝑊"
    , 𝑉 = 𝑢𝑊#
    と同じ要領でProjection (Linear + DwConv)
    

    View Slide

14. Hyena Algorithm
    14
    系列⻑ 𝐿 とPositional Encodingの次元 𝐷 から
    Hyena Filterを⽣成
    

    View Slide

15. Hyena Algorithm
    15
    𝐷$
    𝑆%
    をFFTにより計算
    

    View Slide

16. 定量的結果: Attention-freeであるにも拘らず，Transformerと同程度のPerplexity
    16
    ⽐較モデル
    • Hyena-3: 𝑁 = 3 としたHyena (12レイヤ)
    • Hyena-3-slim: depth⼤・width⼩ (18レイヤ)
    • Efficient Transformer
    • Performer[Choromanski+, ICLR21]
    • Reformer[Kitaev+, ICLR20]
    • Linear Attention[Katharopoulos+, ICML20]
    WikiText103 におけるperplexity
    各モデルのパラメタについて
    HyenaはTransformerおよび
    Hybrid H3と同程度のperplexity
    

    View Slide

17. 定量的結果: ImageNetやCIFAR10においてもViTと同程度のaccuracy
    17
    • データセット: ImageNet / CIFAR10
    • Hyena-ViT: S4ND[Nguyen+, NeurIPS22]に則り設計
    • AttentionをHyenaに代替 / CLSトークンとPEを削除
    • 8台のA100で学習
    → ViT[Dosovitskiy+, ICLR21]と⽐較して同程度のaccuracy
    "-ISO”
    = isotropic
    

    View Slide

18. 定量的結果: Hyenaは性能を維持しつつ⾼速に動作
    18
    • HyenaはFLOPs20%減でGPT[Brown+, NeurIPS20]と同程度のperplexity
    • 系列⻑64KでAttention / FlashAttention[Dao+, NeurIPS22]よりも100倍⾼速
    The PileにおけるperplexityとFLOPs HyenaとFlash AttentionのRuntime
    

    View Slide

19. まとめ
    20
    ü 背景
    • Transformerは強⼒だが系列⻑に対してquadraticな計算量が掛かる
    • Efficient Transformerは計算量を改善するものの性能が不⼗分
    ü 提案⼿法
    • Subquadraticな計算量でTransformerに匹敵するモデルHyenaを提案
    • Attention-freeであり，SSMベースのH3やGSSを⼀般化したモデル
    ü 結果
    • WikiText103やThe Pileにおいて80%の学習コストでTransformerに匹敵する結果
    • ImageNetにおいてもVision Transformerと同程度のaccuracy
    

    View Slide

20. Appendix: Hyena Filterの実装
    21
    https://github.com/HazyResearch/safari/blob/main/standalone_hyena.py
    

    View Slide

21. Appendix: Hyena Operatorの実装
    22
    https://github.com/HazyResearch/safari/blob/main/standalone_hyena.py
    

    View Slide

22. Appendix: H3がHyena-2と等価であることの証明 (1/3)
    23
    

    View Slide

23. Appendix: H3がHyena-2と等価であることの証明 (2/3)
    24
    

    View Slide

24. Appendix: H3がHyena-2と等価であることの証明 (3/3)
    25
    

    View Slide

25. Appendix: 状態空間モデルについて
    26
    • LSSL[Gu+, NeurIPS21]における定式化
    • ⼊⼒ 𝑢 𝑡 ，状態 𝑥 𝑡 ，出⼒ 𝑦 𝑡 に対して以下のように定義
    • GBTにより離散化 (GBT; generalized bilinear transform)
    連続空間
    離散空間
    𝜶はハイパラ, 𝑨, 𝑩, 𝑪, 𝑫 は学習可能パラメタ (BPなどで学習)
    

    View Slide

26. Appendix: SSM + 機械学習の研究は盛んに⾏われている
    27
    • H3までの系譜は以下の図の通り
    • HiPPO→LSSL→S4→H3→Hyena
    • 特にS4は⾮常に重要な研究であり，後続の研究が盛んに⾏われている
    HiPPO
    [Gu+, NeurIPS20]
    LSSL
    [Gu+, NeurIPS21]
    S4
    [Gu+, ICLR22]
    H3
    [Fu+, ICLR23]
    S4D
    [Gu+, NeurIPS22]
    MEGA [Ma+, ICLR23]
    S5 [Smith+, ICLR23]
    

    View Slide

27. Appendix: H3 (Hungry Hungry Hippos)
    28
    ü SSMは①前⽅にあるトークンの記憶と②トークン間の⽐較が不得意
    • ⼀⽅，Attentionは𝑄𝐾! によりトークン間の関係を記憶可能(②) &
    softmax 𝑄𝐾! 𝑉 によりトークン⾃体を直接記憶可能(①)
    ü 提案⼿法: H3 (Hungry Hungry Hippos)
    • この⼆つの難点を乗り越える新たなSSMとしてH3を提案
    • 上記考察に基づき，Q, K, V によって
    Attention-Likeに設計
    • またGPU上のFFT, iFFT等の⾼速化⼿法
    としてFlashConvを提案 (詳細は省略)
    ⇒ ⾼速かつ⾔語に強いSSMを実現
    

    View Slide