MDAgent2: LLM for Code Generation and Q&A in Molecular Dynamics

MDAgent2: LLM for Code Generation and Knowledge Q&A in Molecular Dynamics

LAMMPS特化の3段階LLM訓練（CPT+SFT+MD-GRPO）＋マルチエージェント自己修正ループ（arXiv:2601.02075, Shi & Jing et al., 2026）

🎯 自然言語→LAMMPS実行可能コードの完全自動化。実行フィードバックRLでExecSucc@3を約2.7倍に向上

① 背景と課題

LAMMPSスクリプト作成には高度な専門知識と多くの手作業が必要。汎用LLMはMDドメイン知識が不足し、1ショット生成では実行不能コードを多産する。先行MDAgentはSFTのみで実行フィードバックRLを欠いており、ExecSucc@3が14.23%に留まっていた。

MDドメイン専用の大規模高品質データセットとベンチマークが存在しない

実行ベースの閉ループ最適化がなく、物理的正確性のreward設計が課題

→ CPT+SFT+MD-GRPO の3段階訓練 + MDAgent2-RUNTIMEマルチエージェントシステムで解決

② 3段階訓練パイプライン

① CPT（継続事前学習）
MD-Knowledge corpus: MD論文・教科書・LAMMPSマニュアル
MinHash/LSH重複排除 + LLM品質フィルタ
↓
② SFT（指示チューニング）
MD-InstructQA: 多形式QA（単択/多択/穴埋め/記述）
↓
③ MD-GRPO（実行フィードバックRL）
MD-CodeGen: LAMMPS実行結果をrewardに
低報酬軌跡リサイクル機構

③ MDAgent2-RUNTIME

自然言語タスク記述
↓ CodeGenerator（MD-Code-8B）
↓ CodeRunner（LAMMPS実行）
↓ ResultEvaluator（物理量検証）
↓ 自己修正ループ（max_retry=3）
↓ 実行可能 LAMMPS スクリプト

LAMMPS専用ツールが実行エラーと物理量収束を評価し、自己修正のフィードバック信号を生成

④ QA評価結果（MD-EvalBench）

モデル	サイズ	Overall Avg
Qwen3-max	Large/閉	82.49
Qwen3-32b	32B	77.34
MD-Instruct-8B	8B	74.67
Qwen-flash	Large/閉	73.47
Qwen3-14b	14B	72.91
Qwen3-8b (base)	8B	70.50

8Bで14Bとflashを超過：CPT+SFTの効果

④ コード生成評価結果

設定	ExecSucc@3	Code Score
Direct Prompt (MD-Code-8B)	14.23%	9.29
MDAgent framework	〜20%	〜9.20
MDAgent2-RUNTIME	37.95%	9.32

2.7×

RUNTIME導入によるExecSucc@3の改善倍率

⑤ MD-GRPO の仕組み

GRPOベース：K並列ロールアウトの相対報酬でベースライン分散を低減
Reward = LAMMPS実行成功 + 物理量正確性スコア
低報酬軌跡リサイクル：失敗例を再利用して継続改善
SFTのみに比べ、コード実行可能性と物理的正確性を大幅向上

⑥ MD-EvalBench（初のLAMMPS統合ベンチマーク）

サブセット	問題数	評価観点
MD-KnowledgeEval	336問	MD基礎理論・熱力学
LAMMPS-SyntaxEval	333問	コマンド・文法理解
LAMMPS-CodeGenEval	-	コード生成品質

単択・多択・穴埋め・記述の4形式。ExecSucc@k + Code-Score-Human（0-10）で評価

⑦ テイクホームメッセージ

🔬 LAMMPS特化LLMの初の総合フレームワーク
データ構築・3段階訓練・マルチエージェントRUNTIME・ベンチマークを一体提供。

⚡ 実行フィードバックRLが鍵
MD-GRPOで実際の実行成功/失敗をrewardに。SFT単独に比べコード品質が大幅向上。

🤖 自己修正ループで2.7倍改善
MDAgent2-RUNTIMEがExecSucc@3を14%→38%へ。LAMMPS専用ツールがフィードバック精度を支える。

📊 8Bで14Bを超えるQA性能
MD-Instruct-8BがQwen3-14B・Qwen-flashを上回る。ドメイン適応の効率性を実証。

ケムインフォマティクスへの応用

適用先	ユースケース
lib/md	GROMACS入力ファイル自動生成 + 自己修正ループ
lib/md	MD-EvalBenchをGROMACS向けに転用したベンチマーク構築
lib/fep	Alchemd入力自動生成にMD-GRPO閉ループ訓練を適用

LAMMPS→GROMACSの差分はCodeRunnerの出力パーサ入替えのみで多くが再利用可能

本研究のインパクト

LLMによるMDシミュレーション自動化の総合フレームワークを初公開
実行フィードバックRL（MD-GRPO）がコード生成品質向上の鍵を実証
GROMACSへの転用でlib/mdのセットアップ工数を大幅削減できる可能性