Skip to main content

Token to Logits Transformer Walkthrough

このページは、Aleksa Gordić の “Inside the Transformer: The Life of a Token” と 0xkato の “How LLMs Actually Work” の内容を、wiki 向けの学習メモとして日本語で要約し、再構成したものです。原文を逐語的に訳すのではなく、トークンが Transformer の中をどのように流れ、最終的に次トークンの確率分布になるのかを、助詞を省かず自然な日本語で説明します。

一枚で見る全体像​

Decoder-only LLM の forward pass は、概念的には次の流れです。

重要な点は、LLM が文字列そのものを直接読んでいるわけではないということです。まず文字列は token ID の列に変換され、その ID が embedding table の行を参照し、各 token は hidden size 次元の vector になります。その vector が Transformer block を何層も通り、最後に vocabulary 全体に対する logits へ写像されます。

Tokenization: 文字列を整数列にする​

Tokenizer は、入力文字列を固定 vocabulary の token ID 列へ変換します。現代の LLM では、token は単語全体ではなく、subword や byte-level piece であることが多いです。

  • Whole-word tokenizer は vocabulary が大きくなりすぎ、新語や未知語に弱くなります。
  • Character-level tokenizer は vocabulary が小さくなりますが、sequence が長くなりすぎ、単純な単語パターンまで model が学ぶ必要があります。
  • Subword tokenizer は、その中間を取り、よく出る部分列を一つの token にし、珍しい語はより小さな piece の組み合わせで表します。

この設計は、推論コストにも影響します。同じ文章でも、token 数が少ない tokenizer なら attention と FFN の計算回数が少なくなります。一方で、多言語対応や数字の扱いを雑にすると、特定の言語や数値推論で不利になります。

Training pipeline では、生の文字列を毎回その場で扱うとは限りません。大規模 pretraining では、事前に文書を token 化し、必要なら複数文書を一つの sequence に packing しておきます。そのとき、token ID だけでなく、位置情報に使う position index や、文書境界を越えて attention しないための segmentation mask も用意されます。

Special token と FIM​

通常の文章から自然に出てくる token とは別に、training や inference の意図を model に伝えるための special token があります。たとえば、コード補完では Fill-in-the-Middle (FIM) の special token が使われます。

<FIM_PRE> prefix <FIM_SUF> suffix <FIM_MID> middle

この形式で学習しておくと、IDE で cursor の前後にあるコードから、cursor 位置に入るべき middle を生成しやすくなります。つまり special token は、単なる飾りではなく、「今から何をしてほしいのか」を model に伝える制御信号として働きます。

Embedding: ID を意味のある vector に変える​

Token ID そのものは、vocabulary の何番目かを示す index でしかありません。その ID に意味を持たせるために、LLM は embedding matrix を使います。

X0=E[token_ids]X_0 = E[\text{token\_ids}]

ここで EE は vocabulary size × hidden size の行列です。各 token ID は、対応する行 vector に置き換えられます。7B 級の model では hidden size が 4096 前後になることが多く、一つの token が数千次元の vector として表されます。

Embedding は training によって学習されます。そのため、意味的に近い token は、vector 空間でも近い位置に配置されやすくなります。ただし、この段階の embedding だけでは、token が sequence のどこにあるかはわかりません。同じ token であれば、文頭にあっても文末にあっても同じ embedding が lookup されるからです。そこで positional information が必要になります。

Positional information: 順序を model に渡す​

Self-attention は、入力 token の集合を比較する仕組みなので、それだけでは token の順序を直接知ることができません。たとえば、「犬が人を噛む」と「人が犬を噛む」は token の集合が似ていても意味が大きく違います。したがって、model には位置に関する情報を与える必要があります。

Original Transformer では、sinusoidal positional encoding を embedding に足していました。現代の decoder-only LLM では、RoPE (Rotary Position Embedding) が広く使われます。RoPE は位置 vector を embedding に足すのではなく、attention の Query と Key を、position に応じた角度で回転させます。

RoPE の直感は次の通りです。

  • 各 token の Query / Key は、その token の位置に応じて回転します。
  • Attention score を計算するとき、Query と Key の相対的な回転差が効きます。
  • そのため、絶対位置だけでなく、token 同士の相対距離を attention に反映できます。

長い context に対応するためには、RoPE の周波数設計を拡張する必要があります。YaRN (Yet another RoPE extensioN) は、RoPE を長文 context に外挿しやすくするための方法で、長い sequence でも相対位置情報が破綻しにくいように周波数を調整します。

Transformer block: token vector を深く更新する​

現代的な decoder-only Transformer block は、おおまかに次の構造を持ちます。

ここで注意したいことは、attention 以外の多くの submodule は、各 token vector に独立に作用するということです。RMSNorm も FFN も、基本的には token ごとに同じ処理を適用します。Token 間で情報を混ぜる主な場所は self-attention です。

また、深い Transformer では residual connection が非常に重要です。各 block の出力が前の表現を完全に置き換えるのではなく、前の表現に新しい情報を足し込むことで、情報と gradient が深い層まで流れやすくなります。この加算され続ける流れを、解釈性の文脈では residual stream と呼ぶことがあります。

RMSNorm: vector の scale を安定させる​

Layer normalization は、深い network の activations が大きくなりすぎたり小さくなりすぎたりすることを防ぎます。多くの現代 LLM では、LayerNorm の簡略版である RMSNorm が使われます。

RMSNorm は、平均を引く処理を省き、root mean square によって vector の大きさを揃えます。概念的には次のように書けます。

RMSNorm(x)=g⊙x1d∑i=1dxi2+ϵ\mathrm{RMSNorm}(x) = g \odot \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2 + \epsilon}}

ここで gg は学習可能な scale parameter です。RMSNorm は各 token vector に独立に適用され、出力 shape は入力と同じです。実用上は、pre-norm 構成、つまり attention や FFN の前に normalization を置く構成が、深い Transformer を安定して学習するためによく使われます。

FFN / MLP: token ごとの非線形な処理​

Attention が token 間の情報交換を担うのに対して、FFN は各 token vector を個別に変換します。典型的には、いったん hidden dimension を大きく広げ、非線形関数を通し、元の dimension に戻します。

Original Transformer では ReLU が使われましたが、GPT / BERT 系では GELU が広く使われ、現代の Llama / Mistral / PaLM 系では SwiGLU や GeGLU のような gated MLP がよく使われます。

GeGLU は、概念的には次のように書けます。

GeGLU(x)=W2(GELU(W0x)⊙W1x)\mathrm{GeGLU}(x) = W_2\bigl(\mathrm{GELU}(W_0x) \odot W_1x\bigr)

ここで、⊙\odot は要素ごとの積です。Gating は、feature を単に通すか止めるだけではありません。GELU の場合、値が負になることもあるため、feature の符号を反転させるような働きもできます。

Dense Transformer では、parameter の多くが attention ではなく FFN にあります。そのため、FFN は単なる補助部品ではなく、model が事実、概念、文法的パターンなどを蓄える重要な場所だと考えられています。MoE (Mixture of Experts) は、この FFN 部分を複数の expert に分け、token ごとに一部の expert だけを使うことで、総 parameter 数を増やしながら token あたりの計算量を抑える設計です。

Self-Attention: token 同士が情報を渡し合う​

Self-attention では、各 token vector から Query、Key、Value を作ります。

Q=XWQ,K=XWK,V=XWVQ = XW_Q, \quad K = XW_K, \quad V = XW_V

直感的には、次のように考えるとわかりやすいです。

要素直感役割
Query自分が探している条件どの token から情報を受け取りたいかを問い合わせる
Key自分が持っている見出し他の token の Query と照合される
Value実際に渡す情報attention weight によって混ぜ込まれる

Attention の基本式は次の通りです。

Attention(Q,K,V)=softmax(QK⊤dk+M)V\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V

QK⊤QK^\top は token 間の照合 score です。dk\sqrt{d_k} で割るのは、内積の scale が大きくなりすぎて softmax が極端に尖ることを防ぐためです。MM は mask であり、decoder-only LLM では未来 token を見ないための causal mask が入ります。

Causal mask と segmentation mask​

Causal mask は、position ii の token が position i+1i+1 以降を参照できないようにします。これにより、model は「未来を見て」次 token を予測することができません。

Document packing を使う場合には、別の文書同士が同じ sequence に入ることがあります。その場合、document 1 の token が document 2 の token を参照してしまうと、学習信号が壊れます。そこで segmentation mask を使い、同じ packed sequence 内でも文書境界を越えた attention を禁止します。

Multi-Head Attention と GQA​

一つの attention head だけでは、token 間の関係を一種類の見方でしか捉えられません。Multi-head attention では、複数の head がそれぞれ独立した projection を持ち、異なる subspace で attention を計算します。

重要なのは、各 head が元の vector の固定された slice を見るわけではないということです。各 head は、full hidden vector から学習済みの projection によって自分用の Q / K / V を作ります。したがって、head は「元 vector の別々の切れ端」ではなく、「同じ token に対する異なる learned view」です。

推論時には、過去 token の Key / Value を保存する KV cache が必要です。Full multi-head attention では、head ごとに K / V を持つため、長い context では memory cost が大きくなります。Grouped-Query Attention (GQA) では、多数の Query head が、より少数の Key / Value head を共有します。これにより、多くの query の見方を保ちながら、KV cache の memory pressure と inference cost を下げることができます。

Long context: global attention と block-local attention​

Full attention は、各 token が見られるすべての token と照合するため、sequence length に対して概ね二乗で重くなります。Context length が長くなると、attention の計算量と memory 消費が大きな問題になります。

その対策の一つが、global attention と block-local attention を混ぜる hybrid attention です。

  • Global attention layer では、token が過去の広い範囲を参照できます。
  • Block-local attention layer では、token は近い block 内の token だけを参照します。
  • 多くの layer を block-local にし、一部の layer を global にすることで、長距離依存と効率の trade-off を取ります。

この設計では、すべての layer がすべての過去 token を見るわけではありません。それでも、少数の global layer が長距離の情報経路を提供するため、長い context を扱いやすくなります。

Logits と next-token prediction​

Transformer block をすべて通過した後、model は各 position に対する最終 hidden vector を持っています。生成時に次 token を決めるときには、通常、最後の token の hidden vector だけを取り出し、LM head によって vocabulary size 次元の logits に変換します。

z=hTWUz = h_T W_U

Logits は、まだ確率ではありません。Softmax を通すことで、次 token に対する確率分布になります。

p(xT+1∣x≤T)=softmax(z)p(x_{T+1}\mid x_{\le T}) = \mathrm{softmax}(z)

Training では、この分布と正解 token の cross-entropy loss を最小化します。Inference では、この分布から token を選びます。常に最大確率の token を選ぶ greedy decoding もありますが、実際には temperature、top-k、top-p などを使い、出力の多様性と安定性を調整します。

生成 loop は単純です。

  1. Prompt を token ID に変換します。
  2. Model が次 token の確率分布を出します。
  3. その分布から token を一つ選びます。
  4. 選んだ token を sequence の末尾に追加します。
  5. 終了 token が出るか、長さ上限に達するまで繰り返します。

Base LLM の中心的な training signal は、この next-token prediction です。会話の上手さ、指示追従、安全性、好ましい文体などは、SFT や RLHF / DPO などの post-training によって後から調整されます。

KV Cache: autoregressive generation を現実的にする​

Autoregressive generation では、token を一つずつ追加します。もし毎 step で prompt 全体の Key / Value を再計算すると、生成は非常に遅くなります。しかし causal Transformer では、過去 token の Key / Value は、新しい token が追加されても変わりません。

そのため、inference では過去 token の Key / Value を KV cache に保存します。新しい token を生成するときには、新しい token の Q / K / V だけを計算し、過去の K / V は cache から再利用します。

KV cache は速度を上げる一方で、long context や大きな batch では memory の主要な消費源になります。GQA、MQA、PagedAttention、quantization などの推論最適化は、この memory 問題と深く関係しています。

Transformer の mental math​

大まかな parameter 数を見積もるときには、細かな normalization parameter などは無視し、各 layer の大きな行列だけを見ると理解しやすいです。

  • Attention では、主に WQ,WK,WV,WOW_Q, W_K, W_V, W_O の 4 種類の行列を数えます。
  • Gated MLP では、主に W0,W1,W2W_0, W_1, W_2 の 3 種類の行列を数えます。
  • Embedding / unembedding は vocabulary size と hidden size に比例します。

Training compute の概算では、parameter 数を NN としたとき、一 token あたりの FLOPs が roughly 6N6N になるという近似がよく使われます。ただし、この近似は sequence length が model の inner dimension に比べて十分小さく、attention の二乗項が支配的でない状況で使いやすいものです。Long context では attention の計算量が無視できなくなるため、より丁寧な見積もりが必要になります。

このような計算は、単なる豆知識ではありません。Pretraining の cluster sizing、GPU 数、training duration、必要 FLOPs、data token 数を見積もるときの基礎になります。

Architecture と trained weights の違い​

GPT、Claude、Gemini、Llama、Qwen などの model は、公開情報の範囲では細部が異なります。しかし、多くの現代的な Transformer-based LLM は、次の骨格を共有しています。

  • tokenization
  • embedding
  • positional information
  • stacked Transformer blocks
  • multi-head attention / GQA
  • FFN / gated MLP / MoE
  • residual stream
  • LayerNorm / RMSNorm
  • next-token prediction

一方で、model 間の違いは、次のような部分に現れます。

違い具体例
Trained weightsどの data で、どの規模で、どの objective で学習したか
Configurationlayer 数、hidden size、head 数、vocabulary size、context length、MoE の有無
Positional methodRoPE、YaRN、ALiBi、long-context extension など
Attention patternfull、sliding-window、block-local、global + local など
Post-trainingSFT、RLHF、DPO、安全性 tuning、reasoning RL など
Serving optimizationKV cache、GQA、speculative decoding、quantization など

つまり、architecture は「計算の形」を決め、trained weights は「その形の中に何が学習されたか」を決めます。LLM を理解するときには、この二つを分けて考えると混乱しにくくなります。

読むときの要点​

  • LLM は文字列を直接読まず、token ID と embedding vector を処理します。
  • Position 情報がなければ、attention は語順を直接扱えません。
  • Attention は token 間の情報交換であり、FFN は token ごとの非線形変換です。
  • Residual connection と normalization があるからこそ、深い Transformer を安定して学習できます。
  • GQA と KV cache は、特に inference cost と long context の理解に重要です。
  • YaRN や hybrid attention は、長い context を扱うための実用的な改良です。
  • Logits から softmax、sampling、次 token 追加という loop が、文章生成の基本です。
  • 現代 LLM の多くは似た骨格を持ちますが、data、scale、configuration、post-training、serving 実装が性能と性格を大きく変えます。

関連ページ​

主なソース​