B-1：TransformerブロックをPyTorchで実装する

A-7では、Transformerを設計・実装・評価を繰り返す対象として捉え、まず小さな基準構成を正しく動かす重要性を確認しました。実装編の最初に必要なのは、いきなり大規模な学習環境を用意することではありません。Attention、FFN、残差接続、正規化が、期待する形状と順序で計算されていることを確かめることです。

今回は、固定長のトークン列を対象に、最小のDecoder-only TransformerをPyTorchで実装します。Causal Self-Attention、Pre-Norm、GELU FFN、学習可能な絶対位置Embeddingを組み合わせ、次トークン予測のlogitsを返すところまでを扱います。

ここで作るモデルは、すぐに実サービスへ使うための完成品ではありません。可変長バッチ、RoPE、最適化、混合精度、分散学習は後続の記事で段階的に加えます。まずは、各部品を自分で説明でき、異常が起きたときに切り分けられる最小構成を作りましょう。

1. 今回、どこまで実装するのか？
2. 実装前に、テンソルの形状をどう決めるのか？
3. Causal Self-Attentionでは何を確認するのか？
4. 最小のTransformerを実装する
5. 出力の形状と次トークン予測を確認する
6. 小さな過学習テストが、なぜ有効なのか？
7. この段階で、何を実装しないのか？
8. よくある実装ミスは何か？
9. まとめ：最小実装を、次の拡張の基準にする
10. 今回のブログの考察：B-2への橋渡し
参考文献
このシリーズの案内

1. 今回、どこまで実装するのか？

Decoder-only Transformerは、入力トークンから次のトークンの確率分布を作るモデルです。今回の実装では、次の処理を順に行います。

処理	今回の実装	役割
Token Embedding	`nn.Embedding`	トークンIDを隠れベクトルへ変換する
位置情報	学習可能な絶対位置Embedding	固定長の入力で順序を区別できるようにする
Causal Self-Attention	`scaled_dot_product_attention`	過去と現在のトークンだけから文脈を集める
FFN	Linear → GELU → Linear	各トークンの表現を非線形に変換する
Pre-Normと残差接続	LayerNorm + 加算	深い層へ拡張しやすい基本ブロックにする
LM Head	Linear	各位置の次トークン候補のlogitsを出す

今回の範囲を意図的に絞る理由は、問題を分離するためです。paddingを含む可変長バッチとRoPEはB-2、勾配計算と最適化はB-3、混合精度はB-4で扱います。最初から全てを入れると、出力がおかしいときにAttention、マスク、位置情報、データ処理のどこを確認すべきか分からなくなります。

2. 実装前に、テンソルの形状をどう決めるのか？

実装の混乱の多くは、数式よりテンソルの形状から起きます。今回のコードでは、次の記号を使います。

記号	意味	例
`B`	バッチサイズ	2
`T`	系列長	16
`C`	隠れ次元（`d_model`）	128
`H`	Attentionヘッド数	4
`C / H`	1ヘッドあたりの次元	32
`V`	語彙サイズ	1,000

モデルの入力は(B, T)の整数トークンID、Embedding後の表現は(B, T, C)です。Attention内部ではQ・K・Vを(B, H, T, C / H)へ分割します。最終出力のlogitsは(B, T, V)です。

この形状をコードのコメントとassertで明示すると、ヘッド数と隠れ次元の不整合、誤った転置、マスクの向きを早く見つけられます。特にC % H == 0は、Multi-Head Attentionの基本的な前提です。

3. Causal Self-Attentionでは何を確認するのか？

言語モデルのSelf-Attentionでは、位置tのトークンが位置t+1以降を見てはいけません。学習時に正解となる未来のトークンを参照できてしまうと、Lossが下がっていても、生成時には使えない情報を使っていることになります。

この制約がCausal Maskです。PyTorchのtorch.nn.functional.scaled_dot_product_attentionでは、is_causal=Trueを指定すると、Queryの位置より後ろにあるKeyを参照しないためのマスクを適用できます。

is_causal=Trueは、paddingを含まない固定長の入力には分かりやすい選択です。ただし、可変長バッチでpaddingも無視したい場合は、Causal Maskに加えてpadding用のマスクが必要になります。この扱いはB-2で実装します。

4. 最小のTransformerを実装する

以下のコードは、今回の構成を一つにまとめた最小例です。TinyLanguageModelはトークンIDを受け取り、各位置について語彙全体のlogitsを返します。Softmaxは損失計算時にcross_entropyへ任せるため、モデルの最後には入れません。

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

@dataclass
class TransformerConfig:
    vocab_size: int = 1_000
    block_size: int = 64
    d_model: int = 128
    n_heads: int = 4
    n_layers: int = 2
    d_ff: int = 512
    dropout: float = 0.1

class CausalSelfAttention(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        assert config.d_model % config.n_heads == 0

        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model)
        self.out_proj = nn.Linear(config.d_model, config.d_model)
        self.dropout = config.dropout

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, d_model = x.shape

        qkv = self.qkv(x)
        qkv = qkv.view(
            batch_size, seq_len, 3, self.n_heads, self.head_dim
        ).permute(2, 0, 3, 1, 4)
        query, key, value = qkv.unbind(dim=0)

        attended = F.scaled_dot_product_attention(
            query,
            key,
            value,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=True,
        )
        attended = attended.transpose(1, 2).contiguous()
        attended = attended.view(batch_size, seq_len, d_model)
        return self.out_proj(attended)

class FeedForward(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.d_model, config.d_ff),
            nn.GELU(),
            nn.Linear(config.d_ff, config.d_model),
            nn.Dropout(config.dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

class TransformerBlock(nn.Module):
    """Pre-Norm型のAttentionとFFNからなる1ブロック。"""

    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.norm1 = nn.LayerNorm(config.d_model)
        self.attention = CausalSelfAttention(config)
        self.norm2 = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attention(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

class TinyLanguageModel(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.position_embedding = nn.Embedding(config.block_size, config.d_model)
        self.blocks = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.n_layers)]
        )
        self.final_norm = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len = input_ids.shape
        if seq_len > self.config.block_size:
            raise ValueError("入力長が block_size を超えています")

        positions = torch.arange(seq_len, device=input_ids.device)
        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        for block in self.blocks:
            x = block(x)

        return self.lm_head(self.final_norm(x))

この実装で、A-3のPre-Normはx + attention(norm1(x))とx + ffn(norm2(x))に対応します。A-2のQ・K・Vは一つの線形層でまとめて作り、viewとpermuteでヘッドごとの形状に分けています。A-5で扱った位置情報は、まず学習可能な絶対位置Embeddingとして加えています。B-2では、この部分をRoPEと可変長バッチに拡張します。

5. 出力の形状と次トークン予測を確認する

実装したら、学習の前にforwardが期待どおりの形状を返すか確認します。次の例では、ランダムなトークンIDを入力し、logitsと次トークン予測用のLossを作ります。

config = TransformerConfig()
model = TinyLanguageModel(config)

input_ids = torch.randint(0, config.vocab_size, (2, 16))
logits = model(input_ids)
assert logits.shape == (2, 16, config.vocab_size)

# 位置0〜14の出力で、位置1〜15のトークンを予測する
targets = input_ids[:, 1:]
next_token_logits = logits[:, :-1, :]
loss = F.cross_entropy(
    next_token_logits.reshape(-1, config.vocab_size),
    targets.reshape(-1),
)

print(f"logits: {logits.shape}")
print(f"loss: {loss.item():.4f}")

ランダムに初期化したモデルのLossは、まだ意味のある品質を示しません。ここで確認したいのは、(B, T, V)という出力形状、ラベルを一つ左へずらす処理、Lossが有限値になることです。Softmax済みの確率をcross_entropyへ渡す必要はありません。F.cross_entropyは内部でlogitsに対する対数Softmaxを計算します。

6. 小さな過学習テストが、なぜ有効なのか？

モデルを正しく書けたように見えても、Causal Maskの向き、ラベルのずれ、損失計算の形状が一つ違うだけで、学習は期待どおりに進みません。そこで、大規模データで学習する前に、少数の固定サンプルへ意図的に過学習できるかを確認します。

結果	まず考えること
少数データでLossが十分に下がる	forward・ラベル・更新処理の基本的な接続は確認できる
Lossがほとんど下がらない	ラベルのシフト、学習率、パラメータ更新、Causal Maskを確認する
LossがNaNやInfになる	学習率、入力IDの範囲、勾配、数値精度を確認する
学習時だけ良く、生成が崩れる	`eval()`への切り替え、生成時の入力更新、マスクを確認する

過学習テストは、モデルの汎化性能を示すものではありません。「実装した計算グラフを最適化できるか」を確認するための単体テストに近い位置付けです。実際のパラメータ更新、勾配クリッピング、学習率スケジュールはB-3で扱います。

7. この段階で、何を実装しないのか？

最小モデルを動かす段階では、機能を足すことより、前提を明確にすることが大切です。今回のコードには、次の機能をあえて入れていません。

後回しにする機能	後続で扱う回	今回分ける理由
paddingと可変長バッチ	B-2	Causal Maskとは別のマスク処理が必要になる
RoPE	B-2	Q・Kへ位置を適用する処理を独立して検証する
Optimizerとbackward	B-3	forwardの正しさと勾配更新の問題を分ける
混合精度	B-4	数値精度の影響を単精度の基準構成と比較する
FlashAttention	B-5	Attentionの意味を変えず、実装効率を改善する手法として扱う
分散学習とメモリ最適化	B-6・B-7	単一GPUの動作確認後に、計算資源の制約へ対応する

ここで「機能が少なすぎる」と感じるかもしれません。しかし、最初の実装で問題の切り分けができることは、最終的な実装の速さにつながります。まずは今回のモデルで、入力からlogitsまでの経路を確実に理解しましょう。

8. よくある実装ミスは何か？

8.1 Causal Maskを適用しない

訓練時に未来のトークンが見えると、Lossは不自然に下がる可能性があります。しかし、自己回帰生成では未来のトークンが存在しないため、同じ性能は再現できません。is_causal=Trueが有効になっていることを確認します。

8.2 d_modelとn_headsが割り切れない

Multi-Head Attentionでは、隠れ次元をヘッドへ均等に分けます。たとえばd_model=128、n_heads=4なら1ヘッドは32次元です。割り切れない設定ではviewの形状を作れないため、初期化時のassertで早く失敗させることが重要です。

8.3 logitsへ先にSoftmaxを適用する

PyTorchのF.cross_entropyは、生のlogitsを受け取る関数です。先にSoftmaxを適用すると、数値的に不利になり、意図した損失計算になりません。確率が必要なのは、評価や生成の段階です。

8.4 ラベルをずらさず、同じトークンを正解にする

次トークン予測では、位置tの出力に対する正解は位置t+1のトークンです。入力とラベルを同じ位置で比べると、タスク自体が変わってしまいます。入力、logits、labelsの系列長を紙に書き、[:, :-1]と[:, 1:]が対応しているか確認しましょう。

9. まとめ：最小実装を、次の拡張の基準にする

今回の要点を整理します。

最初の実装は、固定長・単一GPU・最小構成に絞ると問題を分離しやすい
Decoder-only Transformerは、Embedding、Causal Self-Attention、FFN、Pre-Norm、LM Headで構成できる
Causal Maskは、学習時に未来のトークンを参照させないために必須である
d_model、ヘッド数、Q・K・Vの形状を明示すると、Attention実装の不整合を見つけやすい
F.cross_entropyにはSoftmax前のlogitsを渡し、次トークン予測ではラベルを一つずらす
少数データへの過学習テストは、実装・損失・更新処理の接続を確認する有効な手段である

今回のモデルは、実装の完成形ではなく、比較のための基準構成です。この基準があるから、次に位置表現をRoPEへ変えたときや、paddingを導入したときに、何が変わったかを確認できます。

10. 今回のブログの考察：B-2への橋渡し

B-1では、同じ長さのトークン列を入力として、最小のTransformerブロックを実装しました。しかし、実際のテキストは同じ長さではありません。短い文と長い文を同じバッチへ入れるにはpaddingとマスクが必要であり、長い文脈を扱うには位置情報の設計も重要になります。

次回B-2では、可変長テキストを効率よくバッチ化する方法と、A-5で扱ったRoPEをQ・Kへ適用する実装を扱います。B-1のAttentionブロックを土台に、現実のデータを扱えるモデルへ一歩進めましょう。

参考文献

Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). “Attention Is All You Need.” NeurIPS 2017. https://arxiv.org/abs/1706.03762
PyTorch. “torch.nn.functional.scaled_dot_product_attention.” https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
PyTorch. “CrossEntropyLoss.” https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
Xiong, R., Yang, Y., He, D., et al. (2020). “On Layer Normalization in the Transformer Architecture.” ICML 2020. https://arxiv.org/abs/2002.04745

このシリーズの案内

ブログB「Transformer実装」では、Aシリーズで整理した設計を、PyTorchの実装・最適化・運用上の工夫へ段階的に落とし込みます。

#	テーマ	内容
B-1	Transformerブロック実装	Causal Attention、FFN、Pre-Norm、次トークン予測（今回）
B-2	バッチ処理と位置エンコーディング	padding、DataLoader、RoPE
B-3	自動微分と最適化	backward、AdamW、勾配クリッピング
B-4	混合精度演算	AMP、損失スケーリング、数値安定性
B-5	FlashAttention	Attentionの高速化とメモリ効率
B-6	分散学習とDDP	複数GPUでの学習
B-7	メモリ最適化	活性化・KVキャッシュ・チェックポイント

A-7：理論を実装へつなぐ

次に読む

B-2：バッチ処理と位置エンコーディングの実装

実装を深掘りしたい方へ

B-3：自動微分とバックプロパゲーション最適化
B-5：FlashAttention実装と速度検証

LLM【事前学習Ⅱ：基礎理論 B-1】