今メインで使っている 512x8 ネットワークの構造

現在のメイン会話モデルは checkpoints/dazai-friend-peers-512x8-v1/best.pt を優先して使う構成で、 実体は decoder-only Transformer です。ここでは 512x8 が何を意味するか、1 ブロックの中で何が起きているか、入出力の形がどう流れるかを、 実装とチェックポイントのメタ情報に沿って整理します。

モデル名
dazai-friend-peers-512x8-v1
総パラメータ数
27,534,848
語彙サイズ
4,265 文字
文脈長
256 tokens
8 層の Transformer Block 隠れ次元 512 8 heads x 64 次元 FFN 2048 dropout 0.1

1. まず 512x8 の意味

このリポジトリでの 512x8 は、基本的に 幅 512 の表現を、8 段の Transformer Block で処理する という意味です。 つまり 512 は各 token が内部で持つベクトルの太さ、x8 はその処理ブロックの段数です。

512 = 幅

1 token は埋め込み後に 512 次元のベクトルになります。 Attention でも FFN でも、基本の通り道はこの 512 次元です。

x8 = 深さ

同じ形の Transformer Block を 8 回 連続で通します。 深くするほど表現力は上がりますが、そのぶん計算量も増えます。

補足: ここでの x8 は「8 heads」ではありません。 このモデルは結果として n_head=8 でもありますが、名前の主眼は d_model=512 と n_layer=8 にあります。

2. 全体フロー

入力は文字 tokenizer で token ID 列に変換されます。そこに token embedding と position embedding を足し、 8 層の block を通し、最後に各位置ごとの次 token 候補を出します。

入力 token IDs 形: [B, T] / T は最大 256
↓
Token Embedding + Position Embedding [B, T] → [B, T, 512]
↓
Dropout 埋め込み和に dropout 0.1
↓
Transformer Block x 8 各 block: LayerNorm → Causal Self-Attention → Residual → LayerNorm → FFN → Residual
↓
Final LayerNorm [B, T, 512]
↓
LM Head [B, T, 512] → [B, T, 4265]
↓
各位置の次 token の logits 語彙 4265 個へのスコア

3. 1 ブロックの中で何が起きるか

8 層とも中身は同じです。違うのは重みだけで、処理の型は共通です。

A. Causal Self-Attention

  • 入力: [B, T, 512]
  • qkv_proj で [B, T, 1536] を作る
  • Q / K / V に 3 分割し、さらに 8 heads に分ける
  • 各 head の次元は 64 = 512 / 8
  • 未来 token を見ないように causal mask をかける
  • 最後に head を結合して [B, T, 512] に戻す

B. Feed Forward Network

  • 入力: [B, T, 512]
  • まず 512 → 2048 に拡張
  • GELU で非線形変換
  • 次に 2048 → 512 に圧縮
  • 最後に dropout 0.1
入力 x [B, T, 512]
↓
LayerNorm attention の前で正規化
↓
Causal Self-Attention 未来を見ずに、過去と現在の token から文脈を集める
↓
Residual Add 元の x に attention 出力を足す
↓
LayerNorm FFN の前で再度正規化
↓
FFN: 512 → 2048 → 512 各 token を位置ごとに太く計算する
↓
Residual Add attention 後の表現に FFN 出力を足す

4. 形を数で見る

項目 値 意味
vocab_size 4265 文字 tokenizer の語彙数
context_length 256 一度に見られる最大 token 数
d_model 512 各 token の内部表現の幅
n_layer 8 Transformer Block の段数
n_head 8 Attention の分割数
head_dim 64 各 head の次元。512 / 8
ffn_hidden 2048 FFN の中間幅。4 x d_model
dropout 0.1 学習時の過学習抑制

5. パラメータ数の内訳

このモデルの総パラメータ数は 27,534,848 です。 大部分は 8 個の Transformer Block に入っています。

部品 個数 備考
Token Embedding 2,183,680 4265 x 512
Position Embedding 131,072 256 x 512
1 block あたり 3,152,384 Attention + FFN + 2 つの LayerNorm
8 blocks 合計 25,219,072 全体の大半を占める
Final LayerNorm 1,024 出力直前の正規化
LM Head 追加なし token embedding と重み共有
重み共有: lm_head.weight = token_embedding.weight なので、 出力層のために別の巨大な語彙行列をもう 1 枚持っていません。

6. このモデルが今のメインである理由

CLI 側では会話用の既定チェックポイントとして dazai-friend-peers-512x8* を最優先で探すようになっています。 つまり「今メインで使っているモデル」は、運用上もこの 512x8 系です。

なお dazai-friend-peers-512x8-v1 は dazai-peers-512x8-base-mix/best.pt から resume して会話向けに微調整されていますが、 ネットワーク構造そのものは同じ 512x8 のまま です。

7. 実装上の根拠

  • src/original_llm/model.py: decoder-only Transformer 本体
  • src/original_llm/config.py: ModelConfig の定義
  • src/original_llm/cli.py: 会話用既定 checkpoint の優先順
  • checkpoints/dazai-friend-peers-512x8-v1/best.pt: 実際の model_config

このページの数値は、上の実装と checkpoint metadata を読み取って整理したものです。