3 つの Transformer 形態

前回の primer のブロックは、別種のタスクに最適化された 3 つのアーキテクチャに組める。ブロックは変わらない —— 変わるのは注意マスク、トークンが他のトークンを見られるか。 5 つの短いトピック:理解用の Encoder-only(BERT);生成用のDecoder-only(GPT、現代の LLM はこれ);系列対系列のEncoder-Decoder(T5、BART、元の Transformer); BERT を GPT に変える規則 —— 因果マスク;そしてモデルがどの形態かを見分ける方法。

01

Encoder-only —— BERT

文全体を一度に、双方向に読む。理解のために作られ、生成はしない。

encoder-only モデルは、注意マスクを一切持たない Transformer ブロックの積み重ねだ:どの位置も他のどの位置も、左も右も自由に見られる。アクティブなトークンをドラッグして、それがこの行の他のトークンすべてに、両方向で届く様子を見てほしい:

トークン 2

ここでは向きがまったく関係ないことに注目してほしい。文の早い位置のトークンが後ろへ届くのも、遅い位置のトークンが前へ届くのも同じくらい容易だ。この対称性こそが encoder の定義そのもの —— どの位置の表現も、文全体から組み立てられ、決して前半分だけからではない。

BERT はまさにこれを利用して訓練される。15% のトークンを [MASK] 記号の後ろに隠し、両側の文脈から復元させる —— 一方向のモデルにはこの芸当ができない、まだ到達していない位置から答えを写せばいいだけになってしまうからだ。[MASK] スロットが使える情報を、1 段階ずつ開いていく:

非表示

予測が両側に届いてから初めて現れる様子を見てほしい:ステップ 1 ではこのスロットが左の「the cat」と右の「on the mat」の両方に注意しており、になって初めてsat に解決される。マスク言語モデリングは、双方向性を制約から訓練信号へと変える。

ここでは何も捨てられていない:クエリの行とキーの列からなる n × n グリッド全体が、文の長さに関わらず開いたままだ。n を変えて、どのセルも点いたままなのを見てほしい:

36 セル中 36 セルが開いている —— マスクは 0

どのセルも生き残るからこそ、どの出力トークンも入力全体を一度に条件としている。これが encoder の表現を分類・抽出・検索に強くしている理由だ —— 感情ラベル、固有表現のスパン、意味検索の埋め込み —— 系列をラベルへ、あるいはその一部へ写すタスク。長年、「BERT を fine-tune する」がこの種の問題への既定の第一手だった。

encoder にできないのは生成だ:どの位置もすでに文全体を見ている以上、「次のトークン」という概念自体が存在せず、 1 段階ずつ予測すべきものが何も残らない。同じ行を後方だけしか見られないよう制限し、すでに双方向に見える版と並べて比べてみよう:

トークン 2

下の行はもう前へ届かないので、生成器がそもそも持ってはいけない情報をちょうど捨てたことになる。この制限は欠陥ではない —— 次の形態がまさにその上に組み立てられる機構そのものだ。

02

Decoder-only —— GPT

左だけを見て次トークンを予測、追加、繰り返す。現代の LLM はすべてこの形態。

decoder-only モデルは同じブロックを保ちつつ、1 つの規則を加える:トークンは自分自身とそれより前のトークンにのみ注意でき、後に来るものには決して注意しない。アクティブなトークンを文の先へドラッグし、その届く範囲がどれだけ広がるかを見てほしい:

トークン 3

この範囲は広がるだけで、決して縮まないことに注目してほしい —— は が見えるものすべてに加え、その間のすべて、そして自分自身も見える。この一方向の規則こそが、この形態と前の形態との全ての違いであり、訓練目標をほとんど滑稽なほど単純にしている:各位置で、次のトークンを予測する。

この規則は任意ではない。マスクを外すと、同じ行は突然、本来予測すべき未来へ届いてしまう:

過去だけが見える

形状も損失も何も変わらないので、何もエラーにはならない —— モデルは単に、生成時には決して手が届かない位置から答えを写すことを学ぶだけだ。訓練損失は異常な速さで下がり、実際に生成させるとほぼランダムな続きを出す —— 学習した近道が訓練の外には存在しないからだ。訓練曲線が怪しいほど良く、生成が壊れているというこの落差こそが手がかりになる。

トークンは過去しか見えないので、生成は 1 トークンずつ行うしかない:予測し、新しいトークンを入力に追加し、また全体を実行する。 1 回分の生成ループを段階ごとに追ってみよう:

0 / 7 トークン生成済み

新しいトークンだけでなく、それ以前のすべてを毎回再処理していることに注目してほしい —— モデルは呼び出しの間に記憶を持たないので、すでに生成されたトークンは次のパスでも普通の入力として戻される。

この再処理にはコストがかかる:毎ステップ、前半分全体で注意を再計算するコストはステップ n で O(n²)、N ステップ分を合計すると O(N³) になる —— 注意 1 回分の広さを、何度も何度も計算し直しているようなものだ。昨日の key と value を取っておくキャッシュなら、毎ステップ O(n) の新規作業だけで済み、合計は O(N²) になる。ステップ数が増えるにつれて差が開くのを見てほしい:

キャッシュ 10 · 再計算 30

同じ 2 つの式を N = 1,000 まで計算すると、比はちょうど 667 倍になる —— 大まかな桁数ではなく、上の計算がそのまま出す数字であり、大きさが棒グラフに収まらないだけだ。だからこそ、どの推論サーバーもKV キャッシュを持っている:キャッシュは 1 ステップを安くするのではなく、新しい 1 ステップのたびにそれまでの N − 1 ステップをやり直すことをやめさせているのだ。

とはいえキャッシュもどこかに置かねばならず、その大きさもタダではない。各トークンは各層で自分のkey と valueを確保するので、キャッシュはまだ開いている文脈の長さに比例して増える。文脈長をドラッグして、それが埋まっていく様子を見てほしい:

2.00 MiB キャッシュ済み

Llama-2-7B の公開されている形状 —— 32 層、隠れ次元 4,096、fp16 で保存 —— では、トークン 1 つあたり、全層を合わせて 512 KiB になる。つまりは、モデルが新しい単語を 1 つも生成する前から、ちょうど 2 GiB のメモリを確保している。だからこそ文脈長は計算の問題であると同時に、メモリ予算の問題でもある。

03

Encoder-Decoder —— T5、BART

入力を全部読み、それから出力を生成 —— cross-attention が橋渡し。

encoder-decoder 形態は 2017 年の元祖 Transformer そのもので、文字通り前の 2 つの形態を組み合わせたものだ:双方向のencoder スタックが入力全体を読み込み、因果的なdecoder スタックが出力を 1 トークンずつ生成する。両方のスタックの層を並べて、1 層ずつ追ってみよう:

層 0

ここまでは 2 つの列を何もつないでいないことに注目してほしい —— どちらのスタックも、単独だったときとまったく同じ形態のままだ。唯一新しいのは、両者の間の橋だけだ。

その橋がcross-attentionだ:各 decoder 層で、生成中のトークンは単一の固定された入力の要約 —— encoder が仕上げた出力 —— へ振り返って手を伸ばす。decoder 自身が生成した何かへではない。decoder の層を追いながら、橋が毎回どこに着地するかを見てほしい:

層 0

を見てもを見ても弧が同じmemoryに着地する以上、ここは正確に言っておく価値がある:cross-attention のクエリは decoder 由来だが、キーと値は encoder 由来だ —— 1 度だけ計算され、すべての層で使い回され、生成が進んでも再計算されることはない。

decoder ブロックは実は 1 つではなく 2 つの注意サブ層を続けて実行しており、それぞれがまったく異なる情報源を使う。1 つの decoder 位置について、両者を並べて比較してみよう:

位置 1

上のパネルと下のパネルは同じクエリを共有しているので、変わったのはキーと値がどこから来たかだけだ:上は自分自身の過去、下はencoder の memory。どちらも各層で実行され、片方がもう片方の代わりになることはない。

そしてcross-attentionには三角形が一切ない —— 生成が始まる前から入力は完全に利用可能で、decoder の位置から隠すべきものが何もないからだ。decoder の行を追いながら、encoder のどの列も開いたままなのを見てほしい:

行 0

その完全な長方形こそ、対訳ペアが必要としているものだ:ソースとターゲットは境界のはっきりした 2 つの別系列で、ソース全体が最初から分かっている。翻訳、要約、音声認識はどれもこの形をしており、だからこそ T5 や BART は今でもこれらに対する標準的な選択肢であり続けている。

とはいえ、素の decoder でも翻訳はできる —— 「French: … English: …」を与え、因果マスクに 2 つのスタックではなく 1 つのスタックの中で両半分を処理させればいい。2 つのレイアウトを切り替えてみよう:

2 スタック(T5)

2 スタック版は、対応する encoder と decoder、固定された入出力の分割、そして独自のcross-attention配線を必要とする。一方 1 スタック版はそのどれも要らず、必要なのは長めのプロンプトだけだ。decoder が両半分を文脈内に収められるほど大きくなると、分野は概ねより単純な形態へと集約した。 T5 系のモデルは焦点を絞った系列対系列の仕事では依然強いが、最前線は decoder-only へ移った。

04

因果マスク

対角線に沿った 1 本の −∞ の規則が、BERT と GPT の全ての違い。

上の 3 形態のどのマスクも、突き詰めれば同じ問いに 2 通りで答えているにすぎない:クエリは対角線からどこまで先を見てよいか?この窓をドラッグして、各行の重みが、残った開いているキーの上でどう再分配されるかを見てほしい:

k = 0 — どの行も合計は 1.00

どの行の取り分も、開いているキーの数に関わらず必ず合計 1 になることに注目してほしい —— 開いているキーが少ないほど、 1 つあたりの取り分が増えるだけだ。今度は窓をちょうど因果的な位置を超えて までドラッグしてみよう:行 0 は自分自身すら失い、この不変量にはもう分配するものが何も残らない。実際の実装は窓を決してそこまで行かせてはならない、だからこそ対角線は常に含まれる。

この再分配は口先だけの話ではない —— 実際のスコアに対する本物の softmax だ。ここに小さな因果マスク 1 つ分の生スコアが 4 つある。行を 1 つずつ追って、その行が実際に落ち着く重みを読んでほしい:

i = 2

行 0 に注目してほしい:許されたキーがちょうど 1 つなので、自分のスコアが何であれ重みは 1.00 になる —— 最初の位置には選択の余地がない。 になると全キーが開いており、重みは 4 つに分かれなければならない。

許されないキーをちょうどゼロに保つ機構は、softmax の前にそのスコアを −∞ にすることだ、exp(−∞) = 0 だから。典型的なバグはそれを有限の代用品に取り替える —— 代わりに 0でマスクする —— これは無害に見えるが、許されないキーがたまたま大きな生スコアを持っていると話は別だ。規則を切り替えて、マスクされた未来の位置がゼロでなくなる様子を見てほしい:

マスクされた位置が今、注意の 0.0% を占めている

正しいマスクなら、その位置は毎回この行のちょうど 0% を占める。バグのある方では 64.4% を占める —— 正しいマスクの下でこの行の最良の正当なキーが得る量とほぼ同じだ —— exp(0) = 1 が依然として合計に完全に参加するからだ。例外は何も投げられず、形状もおかしくない。モデルはただ静かに、本来見てはいけない答えを読むことを学び、訓練損失はほんの少し速く下がりすぎて、誰も怪しまない程度に収まる。

この primer のどの形態も、同じ三角行列が、あるかないか、1 か所か 2 か所かの違いにすぎない:マスクなしならencoderになり、因果マスクならdecoderになり、因果マスクに cross-attention を足せばencoder-decoderになる。同じ 1 行の上で 3 つを切り替えてみよう:

encoder

ブロックも、注意の数学も、feed-forward 層も何一つ変わらないので、アーキテクチャ上の決定はすべてこの1 つのダイヤルに帰着する。トークン化、埋め込み、 N 個のマスク付きブロックを通し、最終 norm を取り、次トークンの logits を読み、サンプルし、ループする —— たった今組み上げたのがこの 3 台のうちどれかを決めたのは、マスクだった。

05

どの形態か見分ける

3 つの問いで決まる:そのモデルは両方向を見るか、後ろだけを見るか、 2 つ目のスタックが 1 つ目の出力を読むか。

すでに知っている名前から始めよう。モデルを選ぶと、同じマスクのパターンがこの行に点灯し、それに付随する数字も出てくる:

BERT-base — encoder-only · 12 層 · 1.1 億パラメータ

BERT と GPT はブロック数もパラメータ予算もほぼ同じ —— 12 層、およそ 1.1〜1.17 億パラメータ —— なのに一方はencoder、もう一方はdecoderという正反対のマスクに落ち着くことに注目してほしい。マスクは常にアーキテクチャ上の選択であり、規模の結果ではないからだ。 T5 は同じ層あたりの幅で、encoder-decoderとしての 2 つ目のスタックの分だけおよそ 2 倍のパラメータを払っている。

実践では、形態を見分ける一番早い方法はパラメータ数ではなく出力を見ることだ。出力の見た目を選んで、どの形態が答えるかを見てほしい:

encoder-only

ラベルやスパンは 1 トークンずつ生成する必要が一切ないので、因果マスクもまったく要らない ——encoder 単体で十分で、しばしば実行コストも安い。出力を 1 トークンずつ順番に作らなければならなくなった瞬間、2 つ目のスタックがあろうとなかろうと、何らかの因果マスクが登場する。

ここまでで、この primer と前回の primer のあらゆる部品が、実際に配備されているどのモデルも走らせている、同じ 7 段階のパイプラインへと組み上がる。最初から最後まで、 1 段階ずつ追ってみよう:

段階 0:トークン化

たった 1 段階 —— —— だけが今日の決定の宿る場所だと注目してほしい:そのマスクを差し替えれば、他を変えずに同じ 7 段階が BERT 分類器にも、GPT 式の生成器にも、T5 式の翻訳器にもなる。難しいのはブロックでなく、マスクだった。