【DL輪読会】Simplifying Transformer Blocks

DEEP LEARNING JP [DL Papers] Simplifying Transformer Blocks Jeong Seong Cheol, M1, Matsuo Lab, The University of Tokyo http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 • Publish date(ArXiv): 3 Nov 2023 • Paper: https://arxiv.org/abs/2311.01906 • Github(official): https://github.com/bobby-he/simplified_transformers

• https://arxiv.org/abs/2311.01906
• https://github.com/bobby-he/simplified_transformers

背景&目的 • 複雑化するTransformer • 同一構造(Transformer Block)の積み上げ • Transformerの各要素を削除することで単純化を目指す • 各要素を削除する複数の先行研究の組み合わせと信号伝播理論(signal propagation theory)の観点からアプローチ

Transformer Block Transformer Blockの入出力 • Query, Key, Value • Self-Attention • Multi Head Attention • Feed Forward Network • Skip connections • Layer normalisation

アプローチ Dataset • CodeParrot datasets Model(Baseline) • 18-block 768 width causal decoder onlyの GPT-likeなモデル Baselineから4つの要素を性能維持したままそれぞれ消せるか検証 • • • • Attention sub-blockのSkip connection ValueとProjectionのパラメータ Layer normalization MLPのSkip connection

(先出し)本研究の結論 Baselineから4つの要素をそれぞれ消せるか? • • • • Attention sub-blockのSkip connection→yes ValueとProjectionのパラメータ→yes MLPのSkip connection →yes Layer normalization → No(消せるが微小な性能劣化&下位タスクで劣化が激しい,原因不明のまま) 単純化されたTransformer baseline

検証1：Attention sub-blockのSkip connectionは消せるか？→YES Transformer Blockの入出力 = 0にする • 先行研究ではこれを消すとrank collapseが起き[1]，うまく学習できない[2] • そこでAttentionの計算にHeら[3] の変更を加える(Value-SkipInit) • 他のトークンに比べてより自分自身に注意を向けるようにする • さらにNociら[4]の変更を加える(Shaped Attention)

検証2：ValueとProjectionのパラメータは消せるか？ →YES 𝑾𝑽, 𝑾𝑷が単位行列になる (導出過程が長すぎるので詳細省略) Valueが入力Norm(X)そのものになる (導出過程が長すぎるので詳細省略)

検証3：MLPのSkip connectionは消せるか？ →YES そこでMHAとMLPを並列化する技術(Wang & Komatsuzaki, 2021) ．PaLM(2022)にも 使われている．そもそもMLPにSkip connectionのない設計になっている． 複数の先行研究で困難である ことがわかっている

検証3：MLPのSkip connectionは消せるか？ →YES = 0にするとskip connectionが消える

検証4： Layer normalizationは消せるか？→No • 著者ら「これまでのモデルの修正により，信号伝播理論(signal propagation theory)の観点からLayer normalizationはなくて良い」 • しかし，実際に消してみると学習スピードの劣化がおきたり，後の fine-tuningによる下流タスクで計算中にNanが発生するバグがおきる など． • 著者ら「お手上げ(原因不明のまま)」 ”normalisation layers have some beneficial properties for training speed beyond what is captured by signal propagation theory”

実験結果 単純化されたTransformer baseline • • • • 既存のTransformerと単純化されたTransformerを比較 15%早い学習スループット -15%の省メモリを実現 実験が多すぎるので省略(すみません)

まとめ：本研究の結論 Baselineから4つの要素をそれぞれ消せるか? • • • • Attention sub-blockのSkip connection→Yes ValueとProjectionのパラメータ→Yes MLPのSkip connection →Yes Layer normalization → No(消せるが微小な性能劣化&下位タスクで劣化が激しい,原因不明のまま) 単純化されたTransformer baseline

感想 • Transformer周辺の論文をがっつり読んだことがなかったため，引用文 献が多く，たくさんの先行研究を知ることができた • しかし，abstractにnormalizationの消去に成功しているみたいな表現は ミスリードかなとも思った • タイトルに惹かれて選んでしまった