Attention 再入門 is all you need (は言い過ぎ

Transcript

1. 許諾なく撮影や第三者へ の開示を禁止します Attention 再入門 is all you need (は言い過ぎ) 2025/12/18
2. 2 1. Attention is all you needの功罪 2. 効率化との戦い 3. KVcacheと俺たちのvLLM 4. Test time scalingの立役者である長文処理 6. まとめ 2 Agenda
3. “Attention is all you need”でお馴染みの Attention機構．デコーダ型 transformerの約1/4を程度のパラメータを占めるコア技術． トークンと呼ばれる単位要素の関係性を測るための機構．現在の AI技術の様々なブレイクスルーの起源的な存在． 導入：Attention機構とは？ は X 2 春 X 1 曙 X 3 F query F key F value Q 1 Q 2 Q 3 K 1 K 2 K 3 V 1 V 2 V 3 内積＋正規化 [ Q 3 K 1 T Q 3 K 2 T Q 3 K 3 T ] “Attention Weight” ※Q 3 K 1 T・・・「曙」から「春」への重み ※Q 3 K 2 T・・・「曙」から「の」への重み ※Q 3 K 3 T・・・「曙」から「曙」への重み 加重和 Q 3 K 1 TV 1 + Q 3 K 2 TV 2 + Q 3 K 3 T V 3 Attention Weightを加味した「曙」ベクトル →「春」と「は」との関係性情報を持ったベクトル
4. QKVを得るための線形変換層を増やしたバージョンの Attention機構． 現行のLLMのAttention層はほぼマルチヘッドになっている．たとえば Qwen3の32Bのモデルだとヘッドが 32個ある． 導入：Multi head Attention機構とは？ は X 2 春 X 1 曙 X 3 内積 ＋正規化 ＋加重和 Q 1,1 Q 1,2 Q 1,3 Q 2,1 Q 2,2 Q 2,3 Q 3,1 Q 3,2 Q 3,3 F query F1 query F1 query F key F key F1 key F query F1 query F1 value K 1,1 K 1,2 K 1,3 K 2,1 K 2,2 K 2,3 K 3,1 K 3,2 K 3,3 V 1,1 V 1,2 V 1,3 V 2,1 V 2,2 V 2,3 V 3,1 V 3,2 V 3,3 Q 3,1 K 1,1 TV 1,1 +Q 3,1 K 2,1 TV 2,1 +Q 3,1 K 3,1 T V 3,1 Q 3,2 K 1,2 TV 1,2 +Q 3,2 K 2,2 TV 2,2 +Q 3,2 K 3,2 T V 3,2 Q 3,3 K 1,3 TV 1,3 +Q 3,3 K 2,3 TV 2,3 +Q 3,3 K 3,3 T V 3,3 線形変換 Multi head Attention = Z 曙
5. Attention機構はtoken間の関係性を定量化することによって，強力な言語モデルを作り出す源泉となっている． が，計算がめちゃ大変．これを如何にして効率化していくかという戦いを LLM(transformer)の研究者たちは行ってきた． 解決方法を分類してみると，対処対象とアプローチによって分類することができる． (と思う) Attention機構の進化の方向性 対応対象 アルゴリズム的なアプローチ システム的なアプローチ 効率化 計算量削減 　　　　　　Sparse Attention 　　　　　 Liner Attention Mamba (Granite 4.0) FlashAttention KVcache 処理 　Grouped Query 　　　　　　 Attention 　Multihead Latent 　　　　　　 Attention 　　　　　 PagedAttention 長文処理 　　　　　　Rotary Positional 　　　 Ebeddings YaRN 　　　　 　Ring Attention 　　　　 　Ulysses
6. Attention機構はtoken間の関係性を使って，強力な言語モデルを作り出す源泉となっている． が，計算が死ぬほど大変．これを如何にして効率化していくかという戦いを LLM(transformer)の研究者たちは行ってきた． 解決方法を分類してみると，対処対象とアプローチによって分類することができる． (と思う) 効率化との戦い 対応対象 アルゴリズム的なアプローチ システム的なアプローチ 計算量 　　　　　　Sparse Attention 　　　　　 Liner Attention Mamba (Granite 4.0) FlashAttention KVcache 　Grouped Query 　　　　　　 Attention 　 Multihead Latent 　　　　　　 Attention 　　　　　 PagedAttention 長文処理 　　　　　　Rotary Positional 　　　 Ebeddings YaRN 　　　　 　Ring Attention 　　　　 　Ulysses 研究者 が知っておかないといけない方
7. 1. LLMの進化そのものである Attention機構の起源と進化を 復習・理解して，実装や LLM学習時に活かす 2. Attentionの進化の歴史と方向性を構造的に理解して，課題とそ の解決を学び， アイディエーションに使ってもらう． 小言：サーベイ論文も見つからなかったので 構造化にチャレンジしてみる 本日の発表の意義
8. Attention機構はtoken間の関係性を使って，強力な言語モデルを作り出す源泉となっている． が，計算が死ぬほど大変．これを如何にして効率化していくかという戦いを LLM(transformer)の研究者たちは行ってきた． 解決方法を分類してみると，対処対象とアプローチによって分類することができる． (と思う) 効率化との戦い 対応対象 アルゴリズム的なアプローチ システム的なアプローチ 計算量 　　　　　　Sparse Attention 　　　　　 Liner Attention Mamba (Granite 4.0) FlashAttention KVcache 　Grouped Query 　　　　　　 Attention 　Multihead Latent 　　　　　　 Attention 　　　　　 PagedAttention 長文処理 　　　　　　Rotary Positional 　　　 Ebeddings YaRN 　　　　 　Ring Attention 　　　　 　Ulysses エンジニア が知っておかないといけない方
9. 効率化との戦い 対応対象 アルゴリズム的なアプローチ システム的なアプローチ 計算量 　　　　　　Sparse Attention 　　　　　 Liner Attention Mamba (Granite 4.0) FlashAttention KVcache 　Grouped Query 　　　　　　 Attention 　Multihead Latent 　　　　　　 Attention 　　　　　 PagedAttention 長文処理 　　　　　　Rotary Positional 　　　 Ebeddings YaRN 　　　　 　Ring Attention 　　　　 　Ulysses データサイエンティスト が知っておかないといけない (両)方
10. 効率化との戦い 対応対象 アルゴリズム的なアプローチ システム的なアプローチ 計算量 　　　　　　Sparse Attention 　　　　　 Liner Attention Mamba (Granite 4.0) FlashAttention KVcache 　Grouped Query 　　　　　　 Attention 　Multihead Latent 　　　　　　 Attention 　　　　　 PagedAttention 長文処理 　　　　　　Rotary Positional 　　　 Ebeddings YaRN 　　　　 　Ring Attention 　　　　 　Ulysses ・太字_：ぜひ知っておきたいやーつ ・太字：知っているとベターなやーつ ・無協調：んーその都度必要になればキャッチアップで余裕かも
11. よくtransformerは「計算量が膨大」などの「重たい」指摘がある． これはQKTを計算するときに発生する膨大な計算量が根本の原因で，これが全体トークン量 nが伸びれば伸びるほど O(n2)で増える． 計算量が膨大になってしまう原因 は X 2 春 X 1 曙 X 3 ・ ・ ・ F query F1 query F1 query F key F key F1 key F query F1 query F1 value K 1,h K 2,h ・・・ K n,h 0.20 0.30 ・ ・ ・ 0.10 d次元 ※Qwenとかはd=4096 N：センテンス内の全トークン数，　 H：ヘッド数，　 h：h番目のヘッド，　 d：1つのQ・Kベクトルが持つ次元数 ,　 Q 1,h Q 2,h ・・・ Q n,h 0.50 0.00 ・ ・ ・ 0.20 d次元 QKTはn×n次元 　→　総計算量は n×n×d → O(n2) 1k tokenが 2k tokenになると計算量は 4倍 8k tokenになると計算量は 64倍 メモリ量も 1k tokenで2MBくらい 2k tokenになるとメモリ量は 4MB 32k tokenになるとメモリ量は 2GB 1m tokenになるとメモリ量は 2TB
12. Attention機構の構造上， Z 曙 ベクトルの計算時には Z 春 とZ は は計算されているはず． そのため，その際に既に計算されていてかつ，使いまわすことができる KVベクトルはキャッシュしておくことが一般的． KVcacheとは は X 2 春 X 1 曙 X 3 内積 ＋正規化 ＋加重和 Q 1,1 Q 1,2 Q 1,3 Q 2,1 Q 2,2 Q 2,3 Q 3,1 Q 3,2 Q 3,3 F query F1 query F1 query F key F key F1 key F query F1 query F1 value K 1,1 K 1,2 K 1,3 K 2,1 K 2,2 K 2,3 K 3,1 K 3,2 K 3,3 V 1,1 V 1,2 V 1,3 V 2,1 V 2,2 V 2,3 V 3,1 V 3,2 V 3,3 Q 3,1 K 1,1 TV 1,1 +Q 3,1 K 2,1 TV 2,1 +Q 3,1 K 3,1 T V 3,1 Q 3,2 K 1,2 TV 1,2 +Q 3,2 K 2,2 TV 2,2 +Q 3,2 K 3,2 T V 3,2 Q 3,3 K 1,3 TV 1,3 +Q 3,3 K 2,3 TV 2,3 +Q 3,3 K 3,3 T V 3,3 青色： Z 春 ベクトルを計算するときに既に計算している場所 赤色： Z は ベクトルを計算するときに既に計算している場所 Z 曙 ベクトルを計算するときに 以前計算しているところは 再度計算しなくていいように キャッシュしておく
13. 計算量を直接減らす代表的な方法として Sparse AttentionとLiner Attentionがある． これらが発展し Sliding windowとGlobal Attentionを活用した Longformerなどが生まれ，長文処理への扉が開かれた． 計算量をなんとかする Sparse Attention Liner Attention Attention Matrix 吾輩 は 猫 で ある 吾輩 は で 猫 ある Zベクトルの大きさを関係性の 強さと考えると ←マトリクスでトークン間の関 係性を理解しやすくなる input N×d Q N×d K N×d V N×d QKT N×N Z N×d 各ベクトルの大きさだけ見ると input N×d Q N×d K N×d V N×d KTV m×d Z N×d Q’ N×m K’ N×m 先にKTVが計算できれば N×Nの計算は避けられる
14. Flash Attention 計算量をアルゴリズムの観点からではなく， GPUの制御によって減らそうという思想から始まったのが Flash Attention． Attention計算をできるだけ SRAMで完結させることでデータ転送を最小化することを目指している． 計算量をなんとかする (物理)(物理ではないか GPUの中身 プロセッサ VRAM・HBM (ビデオメモリ プロセッサ レジスタ L1 cache L2cache SRAM ※所謂A100やH100など の業務用 GPUが高いの は このHBMが高い GPU Core SRAM GPU Core SRAM QKT QKT VRAM VRAM ・早く計算できる場所で計算を完結させるという考え方 →SRAM内で計算できる大きさに QKTの計算項を分割する． 　SRAM内で計算し，それをVRAMに書き出さない． w/o Flash Attention w/ Flash Attention 基本この１行で劇的に早くなるが， バージョン合わせなどに毎回躓く
15. KVcacheを節約する 2大手法が， Grouped Query AttentionとMulti head latent Attention． GQAはKとVのみヘッド数を減らす手法． MLAはキャッシュを圧縮しメモリ圧迫を減らす手法． KVcacheをなんとかする Grouped Query Attention Multi head latent Attention Multi-head attentionのQKVの各 ヘッドは同じ数でなくてもいい． 3つのQに対して 1つのKVが対応し ていても原理的に問題ない． input Q‘ K‘ V‘ QKT Z W Q W KV 次元 圧縮 復元 キャッシュ
16. w/o Paged Attention w/ Paged Attention 従来予測できない全トークン長への対応のために大量のメモリブロックを 1シーケンスのために確保していたが， Paged Attentionでは細かなメモリブロックを必要なだけ割り当てていくことで無駄をなくし推論を高速化する． KVcacheをなんとかする (物理)(物理ではないか 文章➀用 文章②用 文章③用 メモリ 内部断片化 Internal Fragmentation メモリ ・・・・・・・・・ 文章➀用 文章②用 import time from transformers import AutoModelForCausalLM, AutoTokenizer from config import MODEL_NAME, MAX_TOKENS, BATCH_MULTIPLIER from prompts import PROMPTS def main(): prompts = PROMPTS * BATCH_MULTIPLIER tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) tokenizer.pad_token = tokenizer.eos_token 　 inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device) outputs = model.generate(**inputs, max_new_tokens=MAX_TOKENS) results = tokenizer.batch_decode(outputs, skip_special_tokens=True) import time from vllm import LLM, SamplingParams from config import MODEL_NAME, MAX_TOKENS, BATCH_MULTIPLIER from prompts import PROMPTS def main(): prompts = PROMPTS * BATCH_MULTIPLIER llm = LLM(model=MODEL_NAME) sampling_params = SamplingParams(max_tokens=MAX_TOKENS) 　 outputs = llm.generate(prompts, sampling_params) results = [o.outputs[0].text for o in outputs] 10プロンプトでロード抜きで約 300秒くらい 10プロンプトでロード抜きで約 20秒くらい
17. 正弦波 Encoding/Embedding YaRN (RoPEの上位) 最近のLLMのほとんどのベンチマークは Reasoningモデルが SOTA． そのためロングコンテキストを取り扱うことができる Attention機構が必要． RoPEやYaRNを使うことで対応ウインドウを伸ばせる Test time scalingを成立させる技術 Transformerは元来位置関係を理解できない →位置関係を理解するための情報を入れる X 春 = Word Embedding 春 + Positional Embedding 春 X は = Word Embedding は + Positional Embedding は X 曙 = Word Embedding 曙 + Positional Embedding 曙 かなり遠いとほぼ 0(or1)に近づいて 位置関係の情報が消失してしまうという課題 ABEJAの服部さんの神記事 config = AutoConfig.from_pretrained(model_name) config.rope_scaling = { "type": "yarn", "factor": 2.0, # コンテキスト長を何倍に拡張するか "original_max_position_embeddings": 4096, } # モデルを読み込み model = AutoModelForCausalLM.from_pretrained( model_name, config=config) これだけで OK
18. まとめ 対応対象 アルゴリズム的なアプローチ システム的なアプローチ 計算量 　　　　　　Sparse Attention 　　　　　 Liner Attention Mamba (Granite 4.0) FlashAttention 　Grouped Query 　　　　　　 Attention 　Multihead Latent 　　　　　　 Attention 　　　　　 PagedAttention 長文処理 　　　　　　Rotary Positional 　　　 Ebeddings YaRN 　　　　 　Ring Attention 　　　　 　Ulysses 現在，テックブログ執筆中です．