arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2512.03400cs.LGcs.AI
收藏

更好的世界模型可以导致更好的训练后性能

Better World Models Can Lead to Better Post-Training Performance

University of Michigan(密歇根大学);Princeton University(普林斯顿大学);Harvard University(哈佛大学)

Prakhar Gupta, Henry Conklin, Sarah-Jane Leslie, Andrew Lee

AI总结:

本研究通过比较不同世界建模策略,发现显式建模能提升Transformer的状态表示质量,从而增强强化学习后训练的效果。

摘要:

在本工作中,我们研究了显式世界建模目标如何影响Transformer在不同训练阶段的内部表示和下游能力。我们使用一个受控的2x2x2魔方,并提出问题:(1)显式预训练世界模型如何影响模型的潜在表示?(2)世界模型质量如何影响模型在强化学习后训练中的性能?我们比较了标准的下一个标记预测与两种显式世界建模策略——(i)状态预测预训练和(ii)联合状态预测+下一个标记目标,并评估在应用组相对策略优化(GRPO)作为后训练后的任务性能。我们通过线性探针和因果干预评估表示质量。我们发现,显式世界建模会产生更具线性可解码性和因果可操控性的状态表示。更重要的是,我们发现改进的状态表示会为GRPO带来更高的收益,尤其是在更难的魔方状态上。我们的结果表明,细化状态表示可以提高序列规划任务的后训练效果。

英文摘要:

We study how explicit world-modeling objectives affect the internal representations and downstream capability of Transformers, using Rubik's Cubes as our training domain. We ask: (1) how does explicitly pretraining a world model affect a model's latent representations, (2) how does world-model quality affect post-training performance, and (3) how should a finite data budget be split between pretraining and task-specific fine-tuning? We compare standard action-prediction fine-tuning with two strategies that add explicit state-prediction supervision: state pretraining followed by fine-tuning, and joint action and state training. We measure task accuracy after Group Relative Policy Optimization (GRPO). We further find that explicit world-modeling yields better representations in terms of higher probing accuracy and steerability of the model, and that better representations yield larger gains from GRPO, especially on harder cube states. Finally, when the fine-tuning budget is held fixed, probe accuracy strongly predicts the GRPO improvement. Under a fixed total data budget, accuracy is maximized by allocating only a small fraction to pretraining.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑