bingreeky/MemGen

📚 MemGen とは？

MemGen は、ICLR-2026 論文で紹介された MemGen モデルを実装した研究用コードベースです。中心となるアイデアは、外部データベースではなく、モデル自身の推論ストリーム内に存在する 潜在記憶 を大規模言語モデルに与えることです。2 つのコンポーネントが連携します：

メモリートリガー – 生成中に、モデルが潜在記憶から引き出すべきかどうかを決定します。
メモリーウィーバー – 過去の経験を要約し、現在のコンテキストに織り込むコンパクトな「メモリートークン」を作成します。これらが連携して、エージェントは基礎となるモデル重みを変更せずに、相互作用を通じて行動を進化させることができます。

⚙️ 始め方

環境を作成 – requirements.txt を pip でインストールするか、提供されている memgen.yml conda ファイルを使用します。
（オプション）リトリーバーを設定 – リポジトリは外部検索用に Search‑R1 環境に接続できますが、コアの MemGen 実験には必須ではありません。
モデルとデータセットを選択 – README には、KodCode、GSM8K、TriviaQA 用のすぐに使えるチェックポイント（Qwen2.5‑1.5B‑Instruct および SmolLM3‑3B）がリストされています。
スクリプトを実行 – トレーニングと評価は、scripts/weaver_sft.sh（Weaver モジュール用）や scripts/trigger_train.sh（Trigger モジュール用）などの単純な bash ファイルにラップされています。scripts/eval.sh の LOAD_MODEL_PATH をテストしたいチェックポイントに調整します。

🚀 何ができますか？

論文の結果を再現 単一ターンおよび複数ターンの推論タスクで。
潜在記憶を実験 Weaver または Trigger コンポーネントを交換するか、configs/latent_memory/*.yaml ファイルを編集します。
バニラベースラインと比較 – リポジトリには標準の generate 関数のコメントアウトされたバージョンが含まれています。それを切り替えると、潜在記憶がパフォーマンスをどれだけ向上させるかを確認できます。

📦 中身は？

フォルダ / ファイル	目的
`memgen/model/modeling_memgen.py`	メモリートークンを処理するカスタム `generate` を持つコアモデルクラス。
`scripts/`	トレーニング（`_sft.sh`、`_grpo.sh`、`trigger_train.sh`）および評価（`*_eval.sh`）用の Bash ラッパー。
`configs/latent_memory/`	各データセットの YAML 設定。アテンション実装や高速 SFT モードなどのオプションを制御します。
`requirements.txt` / `memgen.yml`	依存関係リスト（PyTorch、FlashAttention‑2、HuggingFace Transformers など）。
`assets/memgen.png`	Trigger‑Weaver パイプラインの図。

📊 評価とチェックポイント

README の表には、各モデル‑データセットペアの HuggingFace リンクがリストされています。scripts/eval/qwen2_5_gsm8k_sft_trigger.sh などのスクリプトは、チェックポイントをダウンロードし、テスト分割で推論を実行し、標準メトリクス（例：GSM8K の精度、TriviaQA の完全一致）を報告します。

🛠️ 既知の制限

マルチ GPU トレーニング は、分散データ並列（DDP）では機能しますが、完全シャーディングデータ並列（FSDP）では機能しません。著者は FAQ でこの点に言及しています。
マルチターン GRPO スクリプトとチェックポイントはまだ保留中です。現在はシングルターン SFT/GRPO のみがリリースされています。
小さなフォーマットの問題（例：プロンプトが \boxed{} の後にピリオドまたは改行で終わるかどうか）は、1.5B モデルのパフォーマンスに影響を与える可能性があります – 最新のコードは一貫した ChatML テンプレートを強制することでこれを修正しています。

📖 引用

研究で MemGen を使用する場合は、ICLR‑2026 論文を引用してください：

@article{zhang2025memgen,
  title={MemGen: Weaving Generative Latent Memory for Self-Evolving Agents},
  author={Zhang, Guibin and Fu, Muxin and Yan, Shuicheng},
  journal={arXiv preprint arXiv:2509.24704},
  year={2025}
}

🙏 謝辞

このプロジェクトは、Search‑R1 検索環境、以前の潜在推論研究（LatentSeek、SoftCoT、R3Mem）、および LAVIS コードフレームワークに基づいています。

TL;DR： MemGen は、Trigger‑Weaver メカニズムを介して LLM に潜在的な内部生成記憶を追加する研究コードベースであり、エージェントが重み更新なしで時間とともに改善できるようにします。複数のデータセットと小規模から中規模のモデル向けのすぐに実行できるトレーニング/評価スクリプトが付属しています。