用线性探针提升世界模型的编码能力,让小模型也能表现更好。
Improving World Models using Deep Supervision with Linear Probes
- 在损失函数中加入线性探针,引导网络编码真实世界特征。
- 训练和测试性能提升,隐藏状态更易解码,即使未显式训练的特征也有效。
- 显著降低分布漂移,适合计算资源有限时实现高性能小模型。
构建高效的世界模型对实现能推理与导航复杂环境的人工智能体至关重要。本文研究一种深度监督技术,用于在端到端预测下一观察值的网络中促进世界模型的发展。尽管深度监督广泛应用于特定任务学习,本文聚焦于改进世界模型。基于Flappy Bird游戏的实验环境,代理仅接收激光雷达(LIDAR)测量作为观测,探索在损失函数中加入线性探针组件的效果。该附加项促使网络将部分真实世界特征编码至其隐藏状态。实验表明,该监督技术提升了训练与测试性能,增强了训练稳定性,并使世界特征更易解码——即使这些特征未被明确包含在训练中。此外,采用线性探针的网络在游戏高变异性阶段(如连续管道间飞行)表现出更小的分布漂移。引入世界特征损失大致相当于模型规模翻倍,表明该技术在计算资源受限或追求小模型最优性能时尤为有益。这些发现有助于理解如何构建更鲁棒、更复杂的智能体世界模型,为该领域进一步发展奠定基础。
原文摘要 · Abstract (English)
Developing effective world models is crucial for creating artificial agents that can reason about and navigate complex environments. In this paper, we investigate a deep supervision technique for encouraging the development of a world model in a network trained end-to-end to predict the next observation. While deep supervision has been widely applied for task-specific learning, our focus is on improving the world models. Using an experimental environment based on the Flappy Bird game, where the agent receives only LIDAR measurements as observations, we explore the effect of adding a linear probe component to the network's loss function. This additional term encourages the network to encode a subset of the true underlying world features into its hidden state. Our experiments demonstrate that this supervision technique improves both training and test performance, enhances training stability, and results in more easily decodable world features -- even for those world features which were not included in the training. Furthermore, we observe a reduced distribution drift in networks trained with the linear probe, particularly during high-variability phases of the game (flying between successive pipe encounters). Including the world features loss component roughly corresponded to doubling the model size, suggesting that the linear probe technique is particularly beneficial in compute-limited settings or when aiming to achieve the best performance with smaller models. These findings contribute to our understanding of how to develop more robust and sophisticated world models in artificial agents, paving the way for further advancements in this field.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。