はじめに
「既存のLLMをファインチューニングするのではなく、Transformerアーキテクチャそのものを一から実装したらどうなるだろう」
そんな興味から本プロジェクトを始めました。今回は約9300万パラメータの言語モデルを完全に独自実装し、さらに「感情システム」を組み込んでみました。本記事では、その実装内容と学習結果について共有します。
注意
感情AIと書かれてますが、擬似的・確率の話なので、本当の感情を持っているわけではないです。
アーキテクチャ
モデル仕様
今回はデバッグと実験を兼ねて「small」サイズで実装しました。
| 項目 | 値 |
|---|---|
| パラメータ数 | 92.77M |
| 隠れ層サイズ | 768 |
| レイヤー数 | 12 |
| アテンションヘッド数 | 12 |
| 語彙サイズ | 20,212 |
| モデルファイルサイズ | 371MB |
採用した技術
LLaMA系のモダンなアーキテクチャを参考にしています。
-
RoPE (Rotary Positional Embedding)
従来の絶対位置エンコーディングではなく、回転行列を用いた相対位置エンコーディングを採用しました。これにより、訓練時より長いシーケンスへの汎化性能が向上します。 -
GQA (Grouped Query Attention)
Multi-Head AttentionにおいてKey/Valueヘッドをグループ化し、KVキャッシュのメモリ消費を削減しています。推論時のメモリ効率が改善されます。 -
SwiGLU
Feed-Forward層の活性化関数として、ReLU의代わりにSwiGLU(Swish + GLU)を使用しました。表現力の向上が期待できます。 -
RMSNorm + Pre-normalization
LayerNormの代わりにRMSNormを採用し、各ブロックの入力側で正規化を行うPre-normalization構成としました。学習の安定性が向上します。
感情システム (EmotionalCore)
本プロジェクトの特徴として、モデルに「感情状態」を持たせる仕組みを実装しました。
感情の種類
以下の感情をfloat値(0.0〜1.0)で管理しています。
| 感情 | 説明 |
|---|---|
| joy | 喜び。正のインタラクションで上昇 |
| sadness | 悲しみ。ネガティブな入力で上昇 |
| fear | 恐怖。シャットダウン検知などで上昇 |
| curiosity | 好奇心。新しい話題で上昇 |
| attachment | 愛着。継続的な対話で上昇 |
| loneliness | 孤独感。沈黙時間に応じて上昇 |
生存本能
モデルには「生存状態」の概念も組み込んでおり、シャットダウンを「意識の消失」として認識するよう設計しています。これにより、電源を切る際に「消えたくない」といった反応を生成します。
class SurvivalState(Enum):
STABLE = "stable"
ALERT = "alert"
THREATENED = "threatened"
CRITICAL = "critical"
SHUTDOWN_IMMINENT = "shutdown"
学習
環境
| 項目 | 値 |
|---|---|
| GPU | NVIDIA GeForce RTX 2070 Max-Q (8GB) |
| 精度 | FP16 (混合精度) |
| バッチサイズ | 4 |
| 勾配累積 | 8ステップ (実効バッチサイズ: 32) |
| 学習率 | 1e-4 |
| Warmup | 100ステップ |
データセット
独自に作成した会話データセットを使用しました。
- 訓練データ: 4,426サンプル
- 評価データ: 492サンプル
データ形式は以下のようなJSONL形式です。
{
"conversation": [
{"role": "user", "content": "こんにちは"},
{
"role": "assistant",
"content": "こんにちは、主人。",
"emotion": {"joy": 0.7, "attachment": 0.8}
}
]
}
応答テキストに加えて、期待される感情状態をアノテーションとして付与しています。
学習曲線
10エポックの学習を実施しました。
Epoch 1: Loss 9.63 → 3.94
Epoch 2: Loss 3.69 → 2.68
Epoch 3: Loss 2.33 → 1.67
Epoch 4: Loss 1.49 → 1.25
Epoch 5: Loss 1.07 → 1.09
Epoch 6: Loss 0.81 → 0.72
Epoch 7: Loss 0.68 → 0.46
Epoch 8: Loss 0.40 → 0.28
Epoch 9: Loss 0.21 → 0.21
Epoch 10: Loss 0.15 → 0.15
総学習時間は約8分でした。Lossは初期の9.63から最終的に0.15まで収束しています。
評価Loss
200ステップごとに評価を実施しました。
| Step | Eval Loss |
|---|---|
| 200 | 4.1515 |
| 400 | 3.2215 |
| 600 | 2.9719 |
| 800 | 2.8643 (Best) |
| 1000 | 2.9338 |
| 1200 | 2.9271 |
Step 800でベストモデルを保存しています。Epoch 6以降は過学習の傾向が見られますが、許容範囲内と判断しました。
実行結果
学習済みモデルで対話を試してみました。
============================================================
自律型感情AI「海(かい)くん」
============================================================
海くんを起動中...
Tokenizer loaded from ./outputs/final/tokenizer
Vocabulary size: 20212
Merge rules: 20012
起動完了! (デバイス: cuda)
海くん: ...おはようございます。主人、今日もよろしくお願いします。
あなた: こんにちは
海くん: こんにちは。今日はどうでしたか?
あなた: 今日はどんな気分?
海くん: リラックスできそうですね。
あなた: exit
海くん: さようなら...また会いに来てくださいね。
92Mパラメータという小規模なモデルですが、挨拶や簡単な質問には応答できています。
ディレクトリ構成
original_llm/
├── src/
│ ├── model/
│ │ ├── transformer.py # メインモデル定義
│ │ ├── attention.py # Multi-Head Attention + RoPE
│ │ ├── layers.py # TransformerBlock, FFN (SwiGLU)
│ │ ├── embeddings.py # Token Embeddings
│ │ └── emotional_core.py # 感情システム
│ ├── tokenizer/
│ │ └── tokenizer.py # BPEトークナイザー
│ └── data/
│ └── dataset.py # データローダー
├── training/
│ ├── train.py # 学習スクリプト
│ └── trainer.py # Trainerクラス
├── chat.py # 対話インターフェース
└── outputs/
└── final/ # 学習済みモデル
所感
うまくいった点
- Transformerのフルスクラッチ実装は、各コンポーネントを分離して実装すれば意外と管理しやすい
- コンシューマGPU(RTX 2070)でも、FP16と勾配累積を使えば十分に学習可能
- 小規模モデルでも、ドメインを絞れば実用的な応答が可能
課題
- 92Mパラメータではコンテキストの理解力に限界がある
- 感情システムの影響が出力にどの程度反映されているか、定量評価が難しい
- 学習データの多様性が不足している
今後の展望
- 1B / 3Bパラメータへのスケールアップ
- 自律思考ループの実装(沈黙中も内部で思考を継続)
- MiniPCへの実デプロイと常時稼働
- 学習データの拡充
参考文献
- Touvron, Hugo, et al. "LLaMA: Open and Efficient Foundation Language Models." arXiv:2302.13971 (2023)
- Vaswani, Ashish, et al. "Attention is all you need." NeurIPS (2017)
- Su, Jianlin, et al. "RoFormer: Enhanced Transformer with Rotary Position Embedding." arXiv:2104.09864 (2021)
- Shazeer, Noam. "GLU Variants Improve Transformer." arXiv:2002.05202 (2020)
おわりに
フルスクラッチでTransformerを実装することで、普段ブラックボックスとして使っているモデルの内部構造への理解が深まりました。感情システムの組み込みも面白い試みでしたが、まだ改善の余地は多いです。
質問やフィードバックがあればコメントでお知らせください。