高階関数認知のための合成可能な操作としての自己注意
著者: Asher Bond (asher.bond@distillative.ai)
要旨
自己注意 [1] は Transformer アーキテクチャとそのスケール化された後継 [2, 3] の中核となる操作である。本稿は、その機構——クエリ・キー・バリューの射影、スケール済みドット積によるスコアリング、softmax 正規化、そしてバリューの集約——を明示し、これを高階関数として読む。すなわち、表現の集合に対する学習された重み付けによってパラメータ化された操作として、関数レベルプログラミングの意味で捉える [4, 5, 6]。注意は三つの再利用可能な部品——射影する、重み付けする、集約する——へと分解され、この構造こそが、注意をより大きな認知パイプラインの構成要素たらしめる。本稿は解説的な論述である。機構は Transformer のもの [1] であり、合成的な読み解きが本稿の寄与である。
1. はじめに
Transformer [1] は再帰を注意で置き換え、系列中のあらゆる位置を、単一かつ並列化可能な操作の中で他のあらゆる位置と直接に関係づけることを可能にした。これこそが、BERT [2] や GPT-3 [3] といったスケール化されたモデルが依拠する性質である。
注意は通常、固定的なアーキテクチャ部品として提示される。だが合成可能な認知を構築するうえでは、これを高階操作——表現の集合に対する学習されたスコアリング関数によってパラメータ化された関数——として読むほうが有用である。その読みこそ本稿の主題であり、機構を厳密に述べたうえで、合成される再利用可能な部品(射影・重み付け・集約)へと分解する。
2. 関連研究
自己注意と Transformer。 スケール済みドット積注意操作とマルチヘッド Transformer は Vaswani ら [1] による。事前学習された Transformer エンコーダ [2] とデコーダのみの言語モデル [3] は、単一の注意ベースのバックボーンが多数のタスクへ転移することを確立した。系列長に対して厳密な注意が時間・メモリともに二次であることは、この設計に記録された性質 [1] であり、姉妹稿 [7] で論じる効率化の文献を動機づけている。
関数合成。 注意を高階操作——表現の集合と学習されたスコアリング関数を受け取り、合成された表現を返す関数——として読むことは、関数レベルプログラミングの標準的な語彙を適用するものである。すなわち、von Neumann のボトルネックを介さないプログラミングに関する Backus の Turing 賞講演 [4]、高階関数と合成がなぜモジュール性に資するかを論じた Hughes [5]、そして関数を第一級の値として扱う Strachey の基礎的な論述 [6] である。注意はその自然な具体例をなす。重み付けのステップがバリューの集合に対する関数を生成し、集約のステップがそれを適用する。
3. 機構
各位置ごとの埋め込みを行とする行列 X として符号化された入力系列が与えられると、自己注意は学習された重み行列 W_Q、W_K、W_V により三つの線形射影を計算する。
Q = X · W_Q(クエリ)、K = X · W_K(キー)、V = X · W_V(バリュー)。
スコアはクエリとキーのスケール済みドット積であり、softmax によって行確率的な重み付けへと正規化され、それがバリューへ適用される。
Attention(Q, K, V) = softmax( (Q · Kᵀ) / √d_k ) · V
ここで d_k はキーの次元数であり、1/√d_k のスケーリングは softmax 前のドット積の分散を制御する [1]。このような層(および層ごとの複数ヘッド)を積み重ねることで、より深い層は、より浅い層が生成した表現の上で注意を行うことができる。
4. 高階操作としての注意(IPO による枠組み)
注意は、その Process(処理)ステップ自体が関数によってパラメータ化された Input–Process–Output(IPO)の記述に適合する。これこそが、関数型プログラミングの意味で注意を高階たらしめるものである [4, 5]。
- Input(入力)。 位置表現の集合
X(Xが担う位置符号化を超える順序の仮定はない)。 - Process(処理)。 (1)
X → Q, K, Vへ射影する。(2) 重み付けA = softmax(QKᵀ/√d_k)を形成する。(3)A · Vを集約する。ステップ (2) はバリューの集合に対する関数——重み付け——を生成し、ステップ (3) がそれを適用する。したがってこの操作は合成である。すなわちaggregate ∘ weight ∘ projectであり、[4, 6] の合成的なスタイルに則る。 - Output(出力)。 入力位置ごとに一つの文脈化された表現。
この枠組みは、より大きな認知パイプラインを構築する際に、注意のどの部分が再利用可能な合成部品——射影・重み付け・集約——であるかを明示する。これは高階関数が一般に提供するのと同じモジュール性である [5]。
5. 確立された性質
引用文献より、自己注意の以下の性質が合成的な読みへと引き継がれる。
- 直接的な長距離相互作用。 任意の二つの位置が単一の層で相互作用し、再帰に伴う経路長の増大 [1] を生じない。
- 並列性。 中核となる操作は、間に softmax を挟んだ二つの行列積であり、位置にわたって並列化可能である [1]。
- 転移可能なバックボーン。 単一の注意バックボーンが、事前学習と微調整を通じてタスク間で適応する [2, 3]。
証拠と範囲
本稿は解説的な論述である。機構は Transformer のもの [1] であり、厳密に述べた。寄与は、注意を射影・重み付け・集約へと分解し、それらを再利用可能な高階部品として名指す合成的な読みである [4, 5, 6]。本稿が引き継ぐ性質——直接的な長距離相互作用、並列性、転移可能なバックボーン——は、引用文献 [1, 2, 3] で確立されたものである。厳密な自己注意は系列長に対して時間・メモリともに二次である [1]。そのコストを近似なしに解消する IO 認識カーネル [7] は、姉妹稿で扱う。
参考文献
- Ashish Vaswani et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS). arXiv:1706.03762.
- Jacob Devlin et al. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. Proceedings of NAACL-HLT. arXiv:1810.04805.
- Tom B. Brown et al. (2020). Language Models are Few-Shot Learners. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2005.14165. [GPT-3]
- John Backus (1978). Can Programming Be Liberated from the von Neumann Style? A Functional Style and Its Algebra of Programs. Communications of the ACM. [1977 ACM Turing Award Lecture]
- John Hughes (1989). Why Functional Programming Matters. The Computer Journal.
- Christopher Strachey (2000). Fundamental Concepts in Programming Languages. Higher-Order and Symbolic Computation. [Reprint of 1967 lecture notes]
- Tri Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. Advances in Neural Information Processing Systems (NeurIPS). arXiv:2205.14135.