Linear（語彙数へ変換）（Linear Output Layer (LM Head)）

Level 1概要役割と処理の流れ

何をする機能か

Decoder の最終層が出した [B, T, D] を、語彙数ぶんのスコア [B, T, V] へ変換します。

重み行列は [D, V] の 1 枚だけで、非線形関数もありません。単純な行列積です。出てくる値を logits と呼びます。

「各位置について、次に来る語の候補 V 個すべてに点数を付ける」処理だと考えてください。

なぜ必要か

Decoder 内部を流れているのは D 次元（たとえば 512 次元）の連続ベクトルで、これは「意味の座標」であって語彙とは対応していません。

実際に出力するトークンを選ぶには、語彙表の各エントリに対する点数が必要です。その変換を行うのがこの層です。

入力埋め込みが ID → ベクトル だったのに対し、この層は ベクトル → 全IDのスコア という逆向きの対応を担います（そのため unembedding とも呼ばれます）。

全体の中でいつ呼ばれるか

Decoder スタックを抜けた直後、Softmax の直前です。図では出力ヘッドの枠の中、下から 2 番目にあたります。

Level 2内部処理Tensor形状・内部ステップ・設計上の意味

前段から受け取るもの

Decoder 最終層の出力 hidden_states [B, T, D]。 Pre-Norm 構成のモデルでは、この前に最終 LayerNorm が入ります。

後段へ渡すもの

logits [B, T, V]。Softmax へ渡されます。

形の変化とサイズ感

入力

hidden_states [B, T, D] = [1, 8, 512]

出力

logits [B, T, V] = [1, 8, 32000]

V は D より桁違いに大きいのが普通です（512 に対して 32,000 など）。そのため、この 1 層だけで次のような影響があります。

パラメータ数: $D \times V$ 。512 × 32,000 ≒ 1,640 万。モデルによっては全体の数割を占めます
メモリ: 出力 Tensor が $B \times T \times V$ 。T = 2048、V = 32,000、B = 8 なら float32 で約 2 GB になります。学習時にメモリを圧迫する代表的な箇所です

推論時には、次のトークンだけが必要なので最後の位置しか計算しません。 [B, 1, D] → [B, 1, V] とすることで、この負担を大きく減らせます。

例文での処理

共通例文

The keys on the table are lost.

概念上のトークン列（S = 8）: [The, keys, on, the, table, are, lost, .]

table までを入力し、次のトークンを予測する場面を考えます。 Decoder はこの位置のベクトル [1, 1, 512] を出し、Linear が語彙全体のスコアに変換します。

The
keys
on
the
table
are
lost
.

位置4（table）の出力ベクトルから、次のトークンのスコアを計算します。

候補トークン	logit（説明用の仮の値）
`are`	8.2
`is`	5.1
`were`	4.3
`on`	1.4
`table`	0.2
`banana`	-3.7

ここで重要なのは次の 2 点です。

値は確率ではありません。合計は 1 になりませんし、負の値も取ります
値の差に意味があります。全体に同じ数を足しても、Softmax 後の確率は変わりません

複数形の主語 keys を捉えられていれば are のスコアが is より高くなる、という形で文法的な判断が現れます。

EncoderとDecoderでの違い

この層は Decoder 側（正確には出力ヘッド）にのみあります。 Encoder-only モデル（BERT など）では、代わりに分類ヘッドなど別のヘッドが付きます。

Decoder-only の LLM でも、この層はまったく同じ形で存在します。出力ヘッドは原典 Transformer と Decoder-only で共通する部分です。

入力と出力

入力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`hidden_states`	Decoder 最終層の出力（各位置の表現）	float32 Tensor	`[B, T, D]`	B = バッチサイズ（同時に処理する文の数） / T = 出力系列長（Decoder側のトークン数） / D = モデル内部のベクトル次元（d_model）	[1, 8, 512]	重み [D, V] との行列積

出力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`logits`	語彙の各トークンに対するスコア（確率ではない）	float32 Tensor	`[B, T, V]`	B = バッチサイズ（同時に処理する文の数） / T = 出力系列長（Decoder側のトークン数） / V = 語彙数（Tokenizer が扱うトークンの総数）	[1, 8, 32000]	Softmax（確率へ変換）／損失計算

B: バッチサイズ（同時に処理する文の数）
T: 出力系列長（Decoder側のトークン数）
D: モデル内部のベクトル次元（d_model）
V: 語彙数（Tokenizer が扱うトークンの総数）

このサイトでは Tensor 形状を PyTorch の `batch_first=True` 相当、つまり `[B, S, D]`（バッチ次元が先頭）で統一しています。実装によっては `[S, B, D]` を採用するものもあります。

よくある誤解

Level 3数式・Python数式・疑似実装・計算量・実装差

数式

\mathrm{logits} = H W_{out} + b, \qquad H \in \mathbb{R}^{B \times T \times D},\; W_{out} \in \mathbb{R}^{D \times V}

weight tying を使う場合は、入力側の埋め込み行列 $E \in \mathbb{R}^{V \times D}$ を転置して使います。

\mathrm{logits} = H E^{\top}

この形は、「出力ベクトルと各トークンの埋め込みベクトルの内積を取っている」と読めます。つまり、Decoder が出した「次に来るべき意味の方向」に近い埋め込みを持つトークンほど、高いスコアになります。

Pythonで表したインターフェース

from typing import Optional
import torch
from torch import Tensor, nn
 
 
class LMHead(nn.Module):
    """[B, T, D] を語彙ごとの logits [B, T, V] へ変換する。"""
 
    def __init__(
        self,
        d_model: int,
        vocab_size: int,
        tied_embedding: Optional[nn.Embedding] = None,
    ) -> None:
        super().__init__()
        self.projection = nn.Linear(d_model, vocab_size, bias=False)  # [D, V]
        if tied_embedding is not None:
            # 入力側 Embedding [V, D] と重みを共有する（weight tying）
            self.projection.weight = tied_embedding.weight
 
    def forward(self, hidden_states: Tensor) -> Tensor:
        """
        Args:
            hidden_states: [B, T, D]。Decoder 最終層の出力。
 
        Returns:
            [B, T, V]。確率ではなく logits（実数、合計は 1 にならない）。
        """
        return self.projection(hidden_states)
 
 
def project_to_vocabulary(hidden_states: Tensor, head: LMHead) -> Tensor:
    """[B, T, D] を語彙ごとの logits [B, T, V] へ変換する。"""
    return head(hidden_states)
 
 
def project_last_position(hidden_states: Tensor, head: LMHead) -> Tensor:
    """推論用。最後の位置だけを変換してメモリを節約する。
 
    Returns:
        [B, 1, V]
    """
    return head(hidden_states[:, -1:, :])
 
 
# 最小使用例
B, T, D, V = 1, 8, 512, 32_000
head = LMHead(d_model=D, vocab_size=V)
 
hidden_states = torch.randn(B, T, D)             # [B, T, D]
logits = project_to_vocabulary(hidden_states, head)
 
print(logits.shape)                               # torch.Size([1, 8, 32000]) == [B, T, V]
print(logits[0, -1].sum())                        # 1 にはならない（確率ではない）
print(project_last_position(hidden_states, head).shape)  # torch.Size([1, 1, 32000])

この例で単純化している点

weight tying の有無を引数で切り替えられるようにしています（実際はモデル定義で固定）。
推論時に最後の位置だけを計算する最適化は、別関数として示しています。
損失計算は shifted-output のページを参照してください。

型と Tensor 形状の見え方を掴むための擬似実装です。そのまま学習・推論に使えるものではありません。

計算量とメモリ

項目	式	T=2048, D=4096, V=32000, B=1 の目安
計算量	$O(B \cdot T \cdot D \cdot V)$	1 層だけで Transformer 数層ぶんに匹敵
パラメータ	$D \times V$	約 1.3 億
出力メモリ	$B \cdot T \cdot V$ × 4 byte	約 260 MB（float32）

学習時にこの出力がメモリを圧迫するため、チャンクに分けて損失を計算する実装や、 logits を float32 に上げつつ分割処理する工夫が使われます。

実装差と派生

weight tying: 原典 Transformer および多くの LLM が採用。パラメータを $D \times V$ ぶん節約できます。
bias の有無: 多くの実装で bias なしです。
Adaptive Softmax / 階層 Softmax: 語彙を頻度で階層化して計算を減らす手法。現在の GPU 環境ではあまり使われません。
logit soft-capping: logits の絶対値が大きくなりすぎないよう tanh で抑える実装もあります（Gemma など）。

Linear（語彙数へ変換）Linear Output Layer (LM Head)