evex-2: 正規化記号を損失から外して記号だけの返答を 38.5% → 0%

README.md

@@ -0,0 +1,182 @@
+---
+language:
+  - ja
+license: mit
+library_name: pytorch
+pipeline_tag: text-generation
+tags:
+  - causal-lm
+  - japanese
+  - from-scratch
+  - tiny
+---
+
+# evex-2
+
+ある Discord サーバーの過去ログ **だけ** をゼロから学習した 5.87M パラメータの言語モデル。
+既存モデルからの継続学習ではなく、tokenizer も含めて全部そのログから作った。
+**このモデルにとって世界はそのログだけで、一般知識は一切持っていない。**
+
+[evex-1](https://huggingface.co/tako080614/evex-1) の第2世代。コーパスも語彙も構成も
+同じで、変えたのは2つだけ。
+
+## evex-1 から変えたこと
+
+### 1. 学習率 3e-4 → 1e-3
+
+evex-1 は 10 epoch 通して val loss が単調に下がり続けていた = **収束していなかった**。
+上げて 12 epoch 回したら val 4.2404 → 4.0384 (同じ尺度で測った値)。
+
+### 2. 正規化記号を損失から外した
+
+evex-1 の一番大きい問題は、**返答の 38.5% が `<url>` や `<file>` だけ**だったこと。
+正規化で作った記号なので発言ではなく、bot は画像もリンクも実際には出せない。
+
+原因は密度ではなく1トークンであること。添付だけの発言 (コーパスの 3.6%) が `<file>`
+1個に潰れていたので、発言の先頭で確率が1点に集中していた。
+
+そこで `<file>` `<url>` `<mention>` `<channel>` `<time>` を**損失から外して**学習した。
+文脈からは消していない — 「誰かが画像を貼って、他の人が反応する」流れは学ばせたいので、
+入力には残したまま、その位置の予測だけ評価しない。
+
+`<nl>` と `<code>` は外していない。改行もコードブロックもモデルが書いて良いもの。
+
+## 効果
+
+同じチェックポイント同士を**同じ尺度で測り直した**結果 (12 epoch / 96 生成 /
+プロンプトと乱数を固定):
+
+| | 素 val | マスク val | 記号だけの返答 |
+|---|---|---|---|
+| マスクなし | **4.0384** | 4.0465 | **38.5%** |
+| マスクあり (これ) | 4.0970 | **4.0330** | **0.0%** |
+
+- **素 val** … `<url>` `<file>` も予測対象に含める尺度
+- **マスク val** … 記号を損失から外す尺度 = 実際に読まれる語だけの尺度
+
+素の尺度で 0.059 負けているのは、**その尺度が `<url>` を当てることを点数にしている**
+から。出してほしくないトークンなので、そこで負けるのは払って良いコスト。実際に読まれる
+語だけで測れば勝っていて、記号だけの返答は消えた。
+
+**推論時に記号を禁止する必要がなくなった。** evex-1 では `ban_ids` で潰して 38% → 12%
+に抑えるしかなかったが、あれは高確率のトークンを削って再正規化するので出てくる第二候補が
+歪む。損失から外せば確率の質量が最初から実際の語に乗る。
+
+```
+evex-1  '<mention>' / '<file>' / 'AGPLだったら'
+evex-2  'Cloudflare Codexがバグってる?' / 'おぉ'
+        'RadeonはCloudflare Cloudflare Pagesですらにゃい'
+```
+
+## 使い方
+
+`transformers` の Auto クラスには当てはまらない構成なので、同梱の `modeling_evex.py` を使う。
+**tokenizer は evex-1 と同一** (`tok.model` の md5 が一致) なので、既に evex-1 を持って
+いるならそちらを使い回せる。
+
+```bash
+pip install torch safetensors sentencepiece huggingface_hub
+hf download tako080614/evex-2 --local-dir evex-2
+```
+
+`git clone` で取るなら git-lfs が必要。入れずに clone すると `model.safetensors` が
+133 バイトのポインタになり、読み込もうとしても壊れる。`ls -la` して 22MB あるか確かめる。
+
+```python
+import json, sys, torch, sentencepiece as spm
+from safetensors.torch import load_file
+
+sys.path.insert(0, "evex-2")
+from modeling_evex import Config, MicroLM
+
+cfg_json = json.load(open("evex-2/config.json"))
+cfg = Config(
+    vocab_size=cfg_json["vocab_size"], n_layers=cfg_json["n_layers"],
+    d_model=cfg_json["d_model"], n_heads=cfg_json["n_heads"],
+    context=cfg_json["context"], dropout=0.0, attn_dropout=0.0,
+)
+
+model = MicroLM(cfg)
+state = load_file("evex-2/model.safetensors")
+state["head.weight"] = state["embed.weight"]      # weight tying を結び直す
+model.load_state_dict(state)
+model.eval()
+
+sp = spm.SentencePieceProcessor(model_file="evex-2/tok.model")
+end_id = sp.piece_to_id("<|end|>")
+
+prompt = "<|conv|><|s3|>これバグってる？<|other|>"
+ids = torch.tensor([sp.encode(prompt, out_type=int)])
+out = model.generate(ids, max_new_tokens=60, temperature=0.9, top_k=40, stop_id=end_id)
+print(sp.decode(out[0].tolist()))
+```
+
+`head.weight` は入っていない。weight tying で `embed.weight` と同じテンソルを指しており、
+safetensors はストレージを共有したテンソルを保存できないので落としてある。上のよ��に結び直す。
+
+`ban_ids` は要らない。渡しても害は無いが、記号だけの返答は既に出ない。
+
+### プロンプトの形
+
+学習データと同じ直列化でないと、モデルは一度も見ていない形を受け取って崩れる。
+
+```
+<|conv|><|s3|>今日ひま？<|s7|><|re|>ひま<|end|>
+```
+
+| トークン | 意味 |
+|---|---|
+| `<\|conv\|>` | 会話の開始 |
+| `<\|end\|>` | 会話の終了 |
+| `<\|s0\|>`〜`<\|s47\|>` | 話者。発言数の多い上位48人 (人間の発言の 85.3% を被覆) |
+| `<\|other\|>` | それ以外の 2,599 人 |
+| `<\|re\|>` | 直前の誰かへの返信 |
+| `<nl>` | 発言内の改行 |
+| `<url>` `<mention>` `<channel>` `<time>` `<file>` | 正規化した URL / メンション / チャンネル / 時刻 / 添付。**このモデルはこれらを出さない** |
+| `<code>` `</code>` | コードブロック |
+
+末尾に話者トークンを置くと、その話者として続きを書く。`speakers.json` に各話者の
+発言数が入っている (実アカウントとの対応は公開していない)。
+
+**同じ話者トークンが続けて出ることがある。** 学習データでは話者の塊のうち 27.4% が
+2連続以上なので、モデルもそう書く。1発言だけ欲しいなら最初の話者トークンで切る。
+
+絵文字・`草`・`www`・顔文字は正規化せず残してあるので、そのまま出る。
+
+## 数字
+
+| | evex-2 | evex-1 |
+|---|---|---|
+| パラメータ | 5,868,800 | 5,868,800 |
+| 学習トークン | 6,685,152 | 6,685,152 |
+| 語彙 | 4,096 (SentencePiece BPE / byte fallback) | 同じ |
+| context | 512 | 512 |
+| 構成 | decoder-only / 6 層 / d_model 256 / 4 head / d_ff 704<br>RoPE + RMSNorm + SwiGLU + weight tying | 同じ |
+| 学習 | 12 epoch / lr 1e-3 / AdamW / cosine / **CPU のみ** | 10 epoch / lr 3e-4 |
+| 損失マスク | `<file>` `<url>` `<mention>` `<channel>` `<time>` | なし |
+| train / val loss | 3.5535 / 4.0559 (マスク尺度) | 3.8685 / 4.2404 (素) |
+| 記号だけの返答 | **0.0%** | 38.5% |
+
+**データは足りていない。** Chinchilla 最適 (20 トークン/パラメータ) はこの規模だと
+33万パラメータで、5.87M は最適の 6%。val は 12 epoch でもまだ下がり続けているので、
+学習の余地は残っている。
+
+## できること / できないこと
+
+**できる**: チャットの口調、短い応答、ネットスラング、そのサーバー特有の語彙と話題、
+話者ごとの癖 (ある話者は「〜にゃい」と書くが、それを再現する)
+
+**できない**: 一般常識、推論、数学、コード生成、長い整合性、知らない話題への応答
+
+固有名詞は形だけ真似て中身が合わない。`Cloudflare Codex` `AGP Core MT6589` のように、
+見たことのある語を組み替えたものが出る。**事実として読むものではない。**
+
+## 出どころと制限
+
+学習データは**同意を明示的に取っていない実在の人物の会話**。
+逐語での再生は 20 文字以上の完全一致で 0 箇所だったが、669万トークンを12周している
+ので**実際の発言に近いものが出る可能性は残る**。
+
+- 生成物を事実として扱ってはいけない
+- 特定の人物の発言として扱ってはいけない
+- 学習に使ったログそのもの、話者と実アカウントの対応表は公開していない

config.json

@@ -0,0 +1,15 @@
+{
+  "vocab_size": 4096,
+  "n_layers": 6,
+  "d_model": 256,
+  "n_heads": 4,
+  "context": 512,
+  "dropout": 0.1,
+  "attn_dropout": 0.0,
+  "architecture": "decoder-only transformer (RoPE + RMSNorm + SwiGLU)",
+  "tie_word_embeddings": true,
+  "trained_epoch": 12,
+  "train_loss": 3.5534612417221068,
+  "val_loss": 4.055904150009155,
+  "params": 5868800
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57585ec95d58e7f57b451f5a22f915e055287fe2f6c75e1f625201f3cdf34f78
+size 23479240

modeling_evex.py

@@ -0,0 +1,222 @@
+"""evex-1 のモデル定義 (単体で動く)。
+
+    from modeling_evex import Config, MicroLM
+
+学習に使ったものと同じコード。HF のどの Auto クラスにも当てはまらない構成なので、
+transformers ではなくこれを直接使う。
+"""
+
+"""Decoder-only Transformer。RoPE + RMSNorm + SwiGLU + weight tying。
+
+669万トークンしか無いので、パラメータは意図的に小さく取る (Chinchilla 最適は 33万)。
+d_model / n_layers は環境変数で振れるようにしてある。
+
+    .venv-llm/bin/python scripts/llm/model.py     # パラメータ数と過学習テスト
+"""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+@dataclass
+class Config:
+    vocab_size: int = 4096
+    n_layers: int = 6
+    d_model: int = 256
+    n_heads: int = 4
+    context: int = 512
+    dropout: float = 0.0
+    # アテンション内の dropout は既定で切る。
+    #
+    # dropout_p > 0 を渡すと scaled_dot_product_attention は融合カーネルを使えず、
+    # B×H×T×T のアテンション行列を実体化する math 経路に落ちる
+    # (24×4×512×512 で 1 層あたり 100MB。6 層ぶんの往復でメモリ帯域を食い潰す)。
+    # 正則化は残差側の dropout で足りるので、ここは 0 にして融合経路に乗せる。
+    attn_dropout: float = 0.0
+
+    @property
+    def d_ff(self):
+        # SwiGLU は行列が 3 つなので、4*d_model 相当に合わせて 2/3 に縮める。
+        # 64 の倍数に丸めて行列積を素直にする。
+        raw = int(self.d_model * 4 * 2 / 3)
+        return (raw + 63) // 64 * 64
+
+    @property
+    def d_head(self):
+        return self.d_model // self.n_heads
+
+
+class RMSNorm(nn.Module):
+    """LayerNorm から平均を引く処理を落としたもの。小さいモデルでは差が出ないが安い。"""
+
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x.float() * norm).type_as(x) * self.weight
+
+
+def rope_cache(context, d_head, device, base=10000.0):
+    """RoPE の cos/sin を先に作っておく。学習中は使い回すだけ。"""
+    inv = 1.0 / (base ** (torch.arange(0, d_head, 2, device=device).float() / d_head))
+    pos = torch.arange(context, device=device).float()
+    freqs = torch.outer(pos, inv)
+    return freqs.cos(), freqs.sin()
+
+
+def apply_rope(x, cos, sin):
+    # x: (B, heads, T, d_head)
+    t = x.shape[2]
+    cos = cos[:t].view(1, 1, t, -1)
+    sin = sin[:t].view(1, 1, t, -1)
+
+    even, odd = x[..., 0::2], x[..., 1::2]
+    rotated = torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1)
+    return rotated.flatten(-2)
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.qkv = nn.Linear(cfg.d_model, cfg.d_model * 3, bias=False)
+        self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.dropout = cfg.attn_dropout
+
+    def forward(self, x, cos, sin):
+        b, t, _ = x.shape
+        h, dh = self.cfg.n_heads, self.cfg.d_head
+
+        q, k, v = self.qkv(x).split(self.cfg.d_model, dim=2)
+        q = q.view(b, t, h, dh).transpose(1, 2)
+        k = k.view(b, t, h, dh).transpose(1, 2)
+        v = v.view(b, t, h, dh).transpose(1, 2)
+
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+
+        # is_causal で三角マスクは自前で持たない (CPU でも flash 経路に乗る)
+        out = F.scaled_dot_product_attention(
+            q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0.0
+        )
+        return self.proj(out.transpose(1, 2).contiguous().view(b, t, self.cfg.d_model))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.gate = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
+        self.up = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
+        self.down = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Block(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.n1 = RMSNorm(cfg.d_model)
+        self.attn = Attention(cfg)
+        self.n2 = RMSNorm(cfg.d_model)
+        self.ff = SwiGLU(cfg)
+        self.drop = nn.Dropout(cfg.dropout)
+
+    def forward(self, x, cos, sin):
+        x = x + self.drop(self.attn(self.n1(x), cos, sin))
+        return x + self.drop(self.ff(self.n2(x)))
+
+
+class MicroLM(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop = nn.Dropout(cfg.dropout)
+        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
+        self.norm = RMSNorm(cfg.d_model)
+        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+
+        # weight tying。669万トークンで語彙 4096 ぶんの出力行列を別に学ぶ余裕はない
+        self.head.weight = self.embed.weight
+
+        cos, sin = rope_cache(cfg.context, cfg.d_head, torch.device("cpu"))
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+
+        self.apply(self._init)
+        # 残差の出口だけ層数でスケールを落とす (深くしたときに発散させない)
+        for name, param in self.named_parameters():
+            if name.endswith("proj.weight") or name.endswith("down.weight"):
+                nn.init.normal_(param, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layers))
+
+    @staticmethod
+    def _init(module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        x = self.drop(self.embed(idx))
+        for block in self.blocks:
+            x = block(x, self.cos, self.sin)
+        logits = self.head(self.norm(x))
+
+        if targets is None:
+            return logits, None
+
+        loss = F.cross_entropy(
+            logits.view(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-1
+        )
+        return logits, loss
+
+    def parameter_count(self):
+        # tying しているので head は数えない (embed と同じテンソル)
+        seen = set()
+        total = 0
+        for param in self.parameters():
+            if id(param) in seen:
+                continue
+            seen.add(id(param))
+            total += param.numel()
+        return total
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=0.9, top_k=40, stop_id=None,
+                 ban_ids=None, min_new_tokens=0):
+        """ban_ids: 絶対に出させないトークン。min_new_tokens: それまでは stop_id も出させない。
+
+        チャットに使うと `<url>` や `<file>` だけを吐いて終わることが多い
+        (実測 38%)。あれは正規化が作った記号で発言ではないので、
+        呼び出し側から外せるようにしてある。
+        """
+        self.eval()
+        for step in range(max_new_tokens):
+            window = idx[:, -self.cfg.context:]
+            logits, _ = self(window)
+            logits = logits[:, -1, :] / max(temperature, 1e-5)
+
+            if ban_ids:
+                logits[:, ban_ids] = float("-inf")
+            # 何か言う前に終わらせない
+            if stop_id is not None and step < min_new_tokens:
+                logits[:, stop_id] = float("-inf")
+
+            if top_k:
+                kth = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:]
+                logits = logits.masked_fill(logits < kth, float("-inf"))
+
+            nxt = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+            idx = torch.cat((idx, nxt), dim=1)
+
+            if stop_id is not None and int(nxt) == stop_id:
+                break
+        return idx

speakers.json

@@ -0,0 +1,242 @@
+[
+  {
+    "token": "<|s0|>",
+    "rank": 0,
+    "messages": 55316
+  },
+  {
+    "token": "<|s1|>",
+    "rank": 1,
+    "messages": 38384
+  },
+  {
+    "token": "<|s2|>",
+    "rank": 2,
+    "messages": 30599
+  },
+  {
+    "token": "<|s3|>",
+    "rank": 3,
+    "messages": 26538
+  },
+  {
+    "token": "<|s4|>",
+    "rank": 4,
+    "messages": 26056
+  },
+  {
+    "token": "<|s5|>",
+    "rank": 5,
+    "messages": 24511
+  },
+  {
+    "token": "<|s6|>",
+    "rank": 6,
+    "messages": 23817
+  },
+  {
+    "token": "<|s7|>",
+    "rank": 7,
+    "messages": 22628
+  },
+  {
+    "token": "<|s8|>",
+    "rank": 8,
+    "messages": 21982
+  },
+  {
+    "token": "<|s9|>",
+    "rank": 9,
+    "messages": 20735
+  },
+  {
+    "token": "<|s10|>",
+    "rank": 10,
+    "messages": 17102
+  },
+  {
+    "token": "<|s11|>",
+    "rank": 11,
+    "messages": 16132
+  },
+  {
+    "token": "<|s12|>",
+    "rank": 12,
+    "messages": 13383
+  },
+  {
+    "token": "<|s13|>",
+    "rank": 13,
+    "messages": 12956
+  },
+  {
+    "token": "<|s14|>",
+    "rank": 14,
+    "messages": 10498
+  },
+  {
+    "token": "<|s15|>",
+    "rank": 15,
+    "messages": 8806
+  },
+  {
+    "token": "<|s16|>",
+    "rank": 16,
+    "messages": 8699
+  },
+  {
+    "token": "<|s17|>",
+    "rank": 17,
+    "messages": 8459
+  },
+  {
+    "token": "<|s18|>",
+    "rank": 18,
+    "messages": 8369
+  },
+  {
+    "token": "<|s19|>",
+    "rank": 19,
+    "messages": 7698
+  },
+  {
+    "token": "<|s20|>",
+    "rank": 20,
+    "messages": 7586
+  },
+  {
+    "token": "<|s21|>",
+    "rank": 21,
+    "messages": 7184
+  },
+  {
+    "token": "<|s22|>",
+    "rank": 22,
+    "messages": 6746
+  },
+  {
+    "token": "<|s23|>",
+    "rank": 23,
+    "messages": 6308
+  },
+  {
+    "token": "<|s24|>",
+    "rank": 24,
+    "messages": 5715
+  },
+  {
+    "token": "<|s25|>",
+    "rank": 25,
+    "messages": 5308
+  },
+  {
+    "token": "<|s26|>",
+    "rank": 26,
+    "messages": 4541
+  },
+  {
+    "token": "<|s27|>",
+    "rank": 27,
+    "messages": 4419
+  },
+  {
+    "token": "<|s28|>",
+    "rank": 28,
+    "messages": 3834
+  },
+  {
+    "token": "<|s29|>",
+    "rank": 29,
+    "messages": 3649
+  },
+  {
+    "token": "<|s30|>",
+    "rank": 30,
+    "messages": 3628
+  },
+  {
+    "token": "<|s31|>",
+    "rank": 31,
+    "messages": 3138
+  },
+  {
+    "token": "<|s32|>",
+    "rank": 32,
+    "messages": 2834
+  },
+  {
+    "token": "<|s33|>",
+    "rank": 33,
+    "messages": 2650
+  },
+  {
+    "token": "<|s34|>",
+    "rank": 34,
+    "messages": 2527
+  },
+  {
+    "token": "<|s35|>",
+    "rank": 35,
+    "messages": 2504
+  },
+  {
+    "token": "<|s36|>",
+    "rank": 36,
+    "messages": 2302
+  },
+  {
+    "token": "<|s37|>",
+    "rank": 37,
+    "messages": 2280
+  },
+  {
+    "token": "<|s38|>",
+    "rank": 38,
+    "messages": 2243
+  },
+  {
+    "token": "<|s39|>",
+    "rank": 39,
+    "messages": 2229
+  },
+  {
+    "token": "<|s40|>",
+    "rank": 40,
+    "messages": 2049
+  },
+  {
+    "token": "<|s41|>",
+    "rank": 41,
+    "messages": 2047
+  },
+  {
+    "token": "<|s42|>",
+    "rank": 42,
+    "messages": 1936
+  },
+  {
+    "token": "<|s43|>",
+    "rank": 43,
+    "messages": 1932
+  },
+  {
+    "token": "<|s44|>",
+    "rank": 44,
+    "messages": 1925
+  },
+  {
+    "token": "<|s45|>",
+    "rank": 45,
+    "messages": 1896
+  },
+  {
+    "token": "<|s46|>",
+    "rank": 46,
+    "messages": 1885
+  },
+  {
+    "token": "<|s47|>",
+    "rank": 47,
+    "messages": 1880
+  }
+]

tok.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b420040c24f8b3ef739d8f2230820300e8a17d8371333fd731cc19444b0f2f3
+size 58498