mlab-LLM MoE 900M
日本語QAを対象にした、実験用の小規模MoEモデルです。 Dense 300Mモデルを4 experts / top-2のMoEに変換し、層別の負荷分散補助lossを使用して 全パラメータを3 epoch更新しました。Baseモデルではなく、チャット形式のSFTモデルです。 研究室の公式見解・最新情報を保証する案内システムではありません。
Architecture
- Native Transformers
MixtralForCausalLM; custom code / trust_remote_code不要 - Total parameters: 906,186,624 (約906M。900Mは丸めたモデル名)
- 28 layers, hidden size 896, intermediate size 2,688
- 4 FFN experts per layer, top-2 token routing
- 14 attention heads / 7 KV heads; shared attention
- 32,768 vocabulary; tied embeddings; context configuration 4,096 tokens
- SFT sequence length: 2,048 tokens; BF16 weights
- この最終学習ではattention・埋め込み・正規化も更新しています。
DenseからMoEへの変換
元のDenseモデルの各層にあるFFN(トークンの特徴を変換するフィードフォワード部分)を、 同じ重みを持つ4個のexpertに複製しました。attention・埋め込み・正規化などの共有部分は 元の重みを引き継ぎ、各層のrouterだけをランダム初期化しています。 異なる4つの事前学習モデルを合体させた構成ではなく、1つのDenseモデルからのupcyclingです。
routerは各tokenの隠れ状態からexpertのスコアを計算し、上位2個の出力を重み付きで合成します。 選択は層ごと・tokenごとに行われます。「論文担当」「ニュース担当」などの役割を 人手で割り当てたり、質問全体に対して1個のexpertを選んだりする構成ではありません。 expertの意味的な専門分化が成立したかは、独立に検証していません。
総重みは約906Mですが、1 tokenで使う重みの概算は共有部分を含め約502Mです。 これは速度やGPUメモリが同規模Denseと等しくなるという意味ではありません。 読み込みには全expertの重みが必要で、routingにも処理コストがあります。
負荷分散と更新対象
一部expertに選択が集中するのを抑えるため、各層で選択頻度とrouter確率に基づく 負荷分散補助lossを計算し、28層の平均に0.01を掛けて言語モデルlossへ加えています。 paddingは集計から除外します。全層をまとめて集計する方式と異なり、ある層の偏りが 別の層の逆方向の偏りで見かけ上相殺されるのを避ける目的です。 均等なexpert利用そのものが、正確な回答や知識の獲得を保証するわけではありません。
初期の比較実験では共有層を固定しましたが、公開モデルは固定せず、全906Mパラメータを更新しています。 共有層固定版の学習済み重みから継続したのではなく、同じ初期MoE重みから学習し直したものです。
Training History and Data
- 299,157,376-parameter Denseモデルをscratchから事前学習。
hotchpotch/fineweb-2-edu-japaneseのsample_10BTを使用し、約10.642B tokensで学習。 - 一般Instruction SFT:
llm-jp/magpie-sft-v1.0とkunishou/oasst1-chat-44k-ja。準備済みtrainは計23,809例。 - 研究室の公開ページ・ニュース・論文由来の合成QAでSFTし、研究室の30の基本事項を 集中学習したDense checkpointからMoEを初期化。 基本事項QAは一般常識QAではなく、研究室固有の知識です。
- DenseのFFNを4個に複製し、routerを初期化。研究室QA12,243例と基本事項の replay3,061例、計15,304例でMoEを全層更新。 研究室QAには資料なし・資料付き、資料横断・個別論文などを含みます。 合成QAはローカルのQwen teacherを用いた生成パイプラインで作成しました。
本リポジトリには元の論文PDF・抽出本文・学習会話データを含めていません。 学習例数には同一会話からの派生例やreplayを含み、独立した論文数や会話数ではありません。
合成QAの生成方法
公開ページ・ニュース・論文PDF由来のテキストを、出典ID付きの断片(chunk)に整理しました。
生成runの記録では資料572件を扱っています。これはすべてが論文という意味ではありません。
teacherはLM StudioのOpenAI互換APIで呼び出したQwenで、ローカルのモデル識別子は
qwen/qwen3.8-27bでした。この識別子は実行ログ上の値であり、配布元の重みrevisionや
量子化方式まで本モデルカードで特定・保証するものではありません。
同じteacherを、別々のプロンプトによる「質問設計役」と「回答作成役」に分けて呼び出しました。 別々の重みを持つ2モデルを訓練した、という意味ではありません。
| 質問の系統 | 質問役が参照する情報 | 採用された会話数 |
|---|---|---|
| blind | 資料一覧を与えず、研究室への一般的な質問を設計 | 998 |
| catalog | 公開資料一覧を見て、関連研究の探索や資料横断の質問を設計 | 1,793 |
| paper | 個別論文の資料を見て、その内容を掘り下げる質問を設計 | 3,933 |
| 合計 | 当初の7,000会話目標に対して採用済みの範囲を使用 | 6,724 |
質問役には、学部生・研究者・企業担当者・一般来訪者などのペルソナと、口語・丁寧語・ 比較・具体例の要求などの言い方をプログラム側から割り当てました。複数turnの会話では 関連研究を尋ねてから詳細を掘り下げる流れも作り、途中で質問者の設定を変更する仕組みも使いました。 最終runの交代確率設定は0.7です。既生成の質問も文字数上限付きで提示し、重複を減らしました。 ただし、意味的な重複を完全に除去した保証はありません。
回答役には、質問に応じて選択した資料chunkを渡し、提示資料だけに基づく短い日本語回答、 資料で確認できない場合の回答拒否、Web検索・外部URL参照の禁止を指示しました。 回答役が資料全文を常に読んでいる構成ではなく、chunk選択や文字数上限による情報欠落があり得ます。
構造化出力と学習用テキストの分離
teacherの出力にはJSON Schemaを使い、回答文、回答可能性のフラグ、引用chunk ID、 根拠の原文引用を別フィールドに保存しました。構造・ID・turn数などを検査し、不正な出力は 修正指示を添えて再試行しました。引用照合の情報も保存しましたが、人手で全会話の 事実性・引用の妥当性を確認したデータセットではありません。 teacherの推論過程や根拠管理用のJSONを、そのままassistantの回答として学習させてはいません。
学習用には、質問・回答だけを含む資料なしのclosed_book形式と、
引用根拠のテキストをuser側へ添えたgrounded形式へ展開しました。
資料付き形式の根拠はteacherの引用から構成したもので、実運用の検索精度を測る評価ではありません。
assistantの回答tokenだけを学習lossの対象とし、根拠ID・引用などは別途provenanceとして保持しています。
最終MoE学習では研究室例12,243件(catalog 3,391/paper 7,434/blind 1,418)に、 集中学習した30の基本事項のtrain例を再利用するreplayを3,061件混ぜ、約80:20としました。 この最終段階のreplayは一般Instructionデータではありません。 同じ会話threadがtrainとvalidationに跨がない分割を引き継ぎ、長すぎる例を除外していますが、 同じ論文・同じ事実まで完全に分離する評価設計ではありません。
Final Run
- Seed: 20260930; 2,870 optimizer steps (約3 epochs)
- Microbatch 2, gradient accumulation 8, effective batch 16
- LR: 1e-5 for non-router parameters, 1e-4 for routers
- Warmup 30 steps, then constant LR
- Auxiliary loss: 0.01 * mean(per-layer load-balancing losses), excluding padding
- Final checkpoint: 2870; training completed 2026-10-02 (JST)
固定の研究室validation128例、および基本事項validation120例でのassistant-token loss:
| Model | Lab validation | Core QA validation |
|---|---|---|
| Initial MoE | 3.535121 | 0.067417 |
| Shared layers frozen, layerwise loss, 3 epochs | 2.381652 | 0.245047 |
| This model: full update, layerwise loss, 3 epochs | 2.358875 | 0.242095 |
補助lossは検証値に含みません。基本事項validationは既知の事実への異なる質問表現であり、 未知の知識への汎化指標ではありません。loss改善は回答の正確性改善を保証しません。
Usage
検証環境: Python / PyTorch 2.12.1 / Transformers 5.15.0。 古いTransformersや別の推論エンジンとの互換性は未検証です。
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "hmwri/mlab-LLM-MoE-900M"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
repo, dtype=torch.bfloat16, attn_implementation="sdpa"
).to("cuda").eval()
with open(hf_hub_download(repo, "system_prompt.txt"), encoding="utf-8") as f:
system = f.read().strip()
messages = [
{"role": "system", "content": system},
{"role": "user", "content": "この研究室はどのような研究をしていますか?"},
]
inputs = tokenizer.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True,
return_dict=True, return_tensors="pt",
).to(model.device)
with torch.inference_mode():
result = model.generate(**inputs, max_new_tokens=128, do_sample=False, use_cache=True)
print(tokenizer.decode(result[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
必要なら質問に参考資料を添えてください。モデル自身に検索機能はありません。
Generation Review
公開した最終重みから2026-10-02にローカルで6候補の回答を生成し、内容を確認しました。
付属system_prompt.txt、BF16、SDPA、greedy (do_sample=False)、最大128 tokens、
repetition_penalty=1.0で、各質問を独立した新規会話として実行しました。
研究室概要の1例は根拠資料と整合しましたが、他の候補には未確認の数値・日時の生成や
反復が見られました。少数の手動確認であり、モデル全体の正答率を示す評価ではありません。
研究室概要は集中学習した内容であり、未知の知識への性能実証ではありません。
対象の固有名を説明に掲載しない方針のため、質問・回答の原文は掲載していません。
Limitations and Licensing
- 反復、不自然な日本語、資料にない数値・名称・日時の生成が起こり得ます。
- 未確認情報への適切な回答拒否、算数、一般的な指示追従を保証しません。
- 小規模な研究実験用です。重要な判断には原資料の確認が必要です。
- モデル重みの利用ライセンスは現時点で未指定です。公開ダウンロード可能であることは、 無制限の再利用・商用利用の許諾を意味しません。
- 学習に使った外部データ・論文にはそれぞれの権利・利用条件があります。 包括的な権利確認を済ませた、という主張はしていません。
公開用にはoptimizer状態を除外し、generation_configのuse_cacheをtrueに変更しています。 学習済みモデル重みは最終checkpointのままです。