JA-Engram-117M-Experiment

約33Mの日本語Transformerに、約84Mのn-gram記憶表と選択式の判断ヘッドを加えた実験モデルです。事前学習、日本語SFT、選択課題との共同学習、量子化、評価まで実施しました。重みと実装を、結果が十分でなかった点も含めて公開します。

性能は今回比較した同規模の既存SLMを下回りました。とくに自由生成と会話の品質には大きな課題があります。用途は小型モデルの構成、選択式推論、記憶表の量子化を調べるための実験です。実用的な日本語チャットモデルとしては推奨しません。

主候補は検証データで選んだ joint1000.pt です。テスト結果を見て候補を選び直していません。事前学習のみのモデル、末尾のチェックポイント、圧縮版も併せて公開しています。

構成と学習

構成要素	パラメータ数
Transformer本体	32,833,552
Engram記憶表	83,886,080
記憶と本体をつなぐ融合部	493,824
選択式判断ヘッド	147,456
合計	117,360,912

本体は8層、幅384、6 attention heads / 2 KV headsのGQA、SwiGLU、RoPEで構成しています。語彙は32,768、文脈長は2,048 tokenです。入力embeddingとLM headは重みを共有します。Engramは2-gramと3-gramを各2表へハッシュし、2層目と5層目で本体に加えます。判断ヘッドは問題と選択肢を読み、各選択肢のスコアを返します。

これは独自のPyTorch実装です。DeepSeek Engramなどの考え方を参考にしていますが、公式モデルの再現やTypeSafe Jevの公開実装を意味しません。コード内のEJ-127やnanochat-26mに関する説明には、先行する英語版実験の履歴が含まれます。この公開モデルの構成は上表とチェックポイント内の model_config を参照してください。

日本語の事前学習はFineWeb-2-Edu Japanese由来の約2B tokenです。その後、Swallow-Gemma-MagpieとAyaから作成した94,950例でSFTを行いました。SFTの教師信号は回答と終了tokenの約37.27M tokenです。判断ヘッドのみの2,000更新に続き、LM replay、JCommonsenseQA、算術選択課題を混ぜて5,000更新しました。主候補はそのうちjoint 1,000更新時点です。全学習を複数seedで繰り返す検証は行っていません。

日本語の主評価

公開重み	LM BPB ↓	JCommonsenseQA正解率 ↑	JNLI 3-shot正解率 ↑	JNLI 3-shot macro-F1 ↑	JNLI 0-shot正解率 ↑
joint1000（主候補）	1.081391	52.81%	45.06%	0.3470	34.69%
joint5000（末尾）	1.069406	56.60%	35.65%	0.2433	35.49%
joint1000・Engram 4bit	1.081523	52.95%	45.10%	0.3468	34.89%
joint1000・Engram 2bit	1.093712	52.75%	46.81%	0.3553	36.36%
Engram 117M・事前学習のみ	1.024941	未測定	50.48%	0.2950	32.10%
Backbone 33M・事前学習のみ	1.041784	未測定	39.51%	0.2601	30.98%

JCommonsenseQAは課題学習済みの5択評価です。独自に固定したtest 1,118行、質問単位で1,115件を使い、選択肢順を5通りに変えています。5回の並べ替えを別々の独立した問題とは数えていません。

JNLIは追加学習なしの回答尤度評価です。trainから各ラベル1件ずつを例示した3-shotを主評価とし、0-shotも併記しました。testは2,508行、重複をまとめると2,497組です。多数派ラベルだけを返す基準は正解率54.43%、macro-F1 0.2350でした。主候補はmacro-F1では上回るものの、正解率ではこの基準を下回ります。回答ラベルのtoken長が異なるため、尤度和による長さの偏りも残っています。

この表のBPBは固定した検証256窓をtoken片のbyte数で換算した旧評価です。以下の既存SLM比較やトークナイザー比較とは計算条件が異なり、表をまたいで数値を比較できません。SFTや共同学習後には、事前学習のみの状態よりLM BPBが悪化しました。

集計値と評価条件は evaluation-small-ja.json に収録しています。

本体規模と総規模で分けた既存SLM比較

事前学習のみのモデルについて、同じ日本語220文書、619,479 UTF-8 bytesを使って比較しました。BPBは本文の負の対数尤度を実byte数で割った値で、低いほど良好です。

比較する規模	モデル	総パラメータ	BPB ↓
本体規模	JA-Backbone-33M	32.83M	1.029758
本体規模	rinna/japanese-gpt2-xsmall	37.40M	0.916555
総規模	JA-Engram-117M	117.36M	1.014094
総規模	rinna/japanese-gpt2-small	110.42M	0.806208
総規模	llm-jp/llm-jp-3-150m	152.32M	0.844785

本体33Mはrinna 37MよりBPBが約12.35%高く、Engram 117Mはrinna 110Mより約25.79%、LLM-jp 152Mより約20.04%高い結果でした。これらは損失の相対差であり、能力全体の差を百分率で示したものではありません。Engram付きは本体のみよりBPBが約1.52%低いものの、記憶表を加えたことの効果を確定するには追加検証が必要です。本体のみのモデルは独立に学習した比較対象で、Engramを取り外したモデルではありません。

比較では共通のNFKC正規化、小文字化、空白整理を行い、全tokenizerで未知語がなく復号が一致する文書を採用しました。元の256文書から36文書を除外しています。文脈長1,024 token、stride 512、batch 1、FP32重みとBF16 autocastを使い、本文tokenを一度ずつ採点しました。

これは小規模な独自診断です。モデルごとに学習量、データ、語彙、構造が異なり、学習条件を揃えた性能順位ではありません。外部モデルの学習データとの重複も未監査です。正規化前の文章、会話能力、推論能力まで評価した結果ではありません。

結果は comparison-bases.json、条件は comparison-protocol.json にあります。

純正LLM-jp v3トークナイザーでの短期比較

別実験として、同じ92,838文書、396,250,996 UTF-8 bytesを1 epoch学習しました。本体とEngramの初期値を共通にし、語彙とembedding容量を変更しています。どちらもRTX 3090上で実行しました。LLM-jpの学習済みモデル重みは使用していません。

学習進捗	現行32K・総117.36M	純正v3・総143.02M
0%	3.037758	3.137269
25%	1.541448	1.525941
50%	1.403973	1.381693
75%	1.349743	1.328984
100%	1.328392	1.308022

値は同じ検証256文書の本文NLLから求めたBPBです。最終値はv3側が約1.53%低くなりました。v3側では本体が32.83Mから58.49Mへ増えており、トークナイザー単独の改善とは言えません。1 seed・短期学習の結果で、近似重複の監査も未実施です。この差だけで優劣や「トークナイザーは関係ない」といった結論は出していません。

この短期比較ではFP32の学習用重みを使っています。後述する旧事前学習のBF16更新の懸念とは分けて扱います。曲線と重みの識別情報は tokenizer-pilot.json に収録しました。143Mの学習重みはこのリポジトリの配布対象に含めていません。

圧縮、速度、未達事項

Engramの記憶表を4bit・2bitで実際に詰めて保存し、参照した行だけをBF16へ戻す実装です。本体の行列演算は浮動小数のままです。

joint1000の形式	配布ファイル容量
浮動小数	301,694,475 bytes
Engram 4bit	177,964,479 bytes
Engram 2bit	156,992,959 bytes

4bit版は今回測ったスコアをほぼ維持しました。ただし、事前に許容幅を定めた同等性検定は行っていません。モデル全体の4bit化や、全体で100MiB以下という目標の達成でもありません。

RTX 3090、batch 1、BF16、KV cacheなしの測定では、主候補のpointer推論は中央値14.881ms、95パーセンタイル22.772msでした。32問を繰り返し計測した結果で、正解は15問です。自由本文生成の中央値は1,976.613msでしたが、正答が0だったため、同じ品質での速度比は算出できません。evaluation-latency.json に結果を残しています。

会話の固定12問では客観採点の正答が0で、反復も見られました。算術の自由生成は最大16 tokenなどの固定条件で厳密一致0でした。一方、選択式2桁算術は主候補57.54%、末尾88.33%です。選択肢から答えを選ぶ性能と、答えを自力で生成する性能には差があります。

当初のGo/No-Go条件は全面達成していません。主候補のBPBは本体のみより悪く、算術自由生成の改善も確認できませんでした。校正はBrier scoreを改善した一方、比較した生成head分類よりECEが悪く、両方改善という条件を満たしていません。複数seedでの検証、人間によるblind比較、同品質での速度優位の確認も残っています。

性能不足について分かったことと改善候補

性能差の原因を一つに特定する実験は行っていません。今回の記録から、次に確かめる候補は以下です。いずれも改善効果を確認したという意味ではありません。

1. 学習時の数値精度。旧事前学習の実装ではembeddingとEngram表がBF16で、optimizerのモーメントはFP32ですが、更新先の重みはBF16のままです。小さい更新が丸めで消える可能性があります。実際の学習で消えた更新の割合や、性能への寄与は未測定です。FP32の学習用重みを保持する方式との対照実験が候補になります。
2. 検証と途中再開。旧 train_lm.py のログ上の val_loss は学習データの固定サンプルを使っています。今回掲載した別実行のheld-out評価とは区別が必要です。旧事前学習ではoptimizer・乱数状態を保存せず再開した履歴もあります。独立した検証データと完全な再開状態を、事前学習から揃える余地があります。
3. データ量と品質。約2B tokenで十分に学習したとは判断していません。品質、分野の偏り、重複を点検したうえで、例えば2B、4B、8Bで検証曲線を比べる方法が考えられます。同じデータの反復と、新しい文書の追加は分けて評価する必要があります。ここに挙げた追加学習は未実施です。
4. 本体と記憶の配分。総117MのうちTransformer本体は約33Mです。記憶表の利用率、ハッシュ衝突、融合ゲート、本体へ配分する容量を調べる余地があります。現状の結果からEngram全般の有効性を否定したり、容量を増やせば解決すると断定したりはできません。
5. 追加学習と評価。SFT・共同学習でLM損失が悪化し、会話品質も不足しました。データの品質、LM replayの比率、学習率を検証データで調べることが候補です。合成データの自己採点だけを正しさの保証にせず、複数seedと評価データへの混入確認も必要です。

今回は学習済み成果を実験記録として公開し、上記の改善は今後の候補として残します。

使い方と配布物

独自のPyTorch形式です。AutoModelForCausalLM.from_pretrained やGGUF用ランタイムへの直接読み込みには対応していません。REPRODUCE.md に取得、推論、形式確認、再現できる範囲を記載しています。

python -m pip install -r requirements.txt
python -m zipfile -e source.zip .
python -m ejv.decide --checkpoint joint1000.pt --tokenizer ja32k.model --calibration joint1000.calibration.json --task arith --question "12 + 34 は？" --option "46" --option "45" --option "35" --option "56"

このコマンドは選択結果、各選択肢の確率、生スコアを返します。確率は渡した選択肢内の分布であり、正しさの保証ではありません。圧縮版を使う場合はcheckpointと校正ファイルを同じ名前の組に替えます。

ファイル	内容
joint1000.pt	validationで固定した主候補
joint5000.pt	共同学習の末尾
joint1000-engram-int4.pt / joint1000-engram-int2.pt	主候補の記憶表を実圧縮した推論専用版
engram117-base.pt	評価に使った事前学習のみのモデル
backbone33-base.pt	独立に学習した本体のみの比較対象
*.calibration.json	各joint重みに対応する検証時の温度校正
ja32k.model	現行SentencePiece tokenizer
source.zip	推論・学習・評価コードと構成設定
SHA256SUMS / weights.json	配布物のhashと重みの対応

データの出典とライセンス

事前学習データは FineWeb-2-Edu Japanese、SFTは Swallow-Gemma-Magpie と Aya Dataset を使用しました。JCommonsenseQAとJNLIは JGLUE に由来します。データ本文や他組織のモデル重みは同梱していません。

元データの表示ライセンスは順にODC-BY、Gemma、Apache-2.0、JGLUEのCC-BY-SA-4.0です。独自コードと事前学習のみの配布重みにはApache-2.0を適用します。Gemma生成データで追加学習したjoint系列とその圧縮版には、Gemma Terms of Useと利用制限も適用します。全ファイルが一律にApache-2.0という配布ではありません。LICENSE、NOTICE、同梱のGemma規約を参照してください。

設計上の参考は nanochat、DeepSeek Engram、Low-Bit N-Gram です。各プロジェクトがこの実験や結果を承認したことを示すものではありません。