SPARTAN: A Sparse Transformer World Model Attending to What Matters
作者: Anson Lei, Bernhard Schölkopf, Ingmar Posner
分类: cs.LG, stat.ML
发布日期: 2024-11-11 (更新: 2025-12-05)
💡 一句话要点
SPARTAN：一种稀疏Transformer世界模型，关注关键交互
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 世界模型 稀疏Transformer 因果图 对象中心表示 Few-shot学习
📋 核心要点
- 现有世界模型难以在复杂环境中可靠地捕获实体间的交互关系，尤其是在发现局部因果结构方面面临挑战。
- SPARTAN通过在Transformer的注意力机制中引入稀疏性正则化，学习上下文相关的稀疏交互图，从而更准确地预测对象状态。
- 实验表明，SPARTAN在few-shot自适应动力学变化和抗干扰方面表现出色，能够学习反映对象间交互的局部因果图。
📝 摘要（中文）
本文提出了一种基于Transformer的稀疏世界模型SPARTAN，用于学习场景中实体之间上下文相关的交互结构。SPARTAN通过对对象因子化token之间的注意力模式施加稀疏性正则化，学习稀疏的、上下文相关的交互图，从而准确预测未来的对象状态。该模型进一步扩展以适应环境中具有未知目标的稀疏干预。实验结果表明，SPARTAN能够学习反映对象之间潜在交互的局部因果图，在动力学变化的few-shot自适应以及对干扰的鲁棒性方面，显著优于当前最先进的以对象为中心的世界模型。
🔬 方法详解
问题定义：现有基于对象的世界模型在复杂环境中难以准确捕捉对象间的交互关系，尤其是在存在干扰因素或动力学变化时，模型的泛化能力受到限制。现有方法难以发现局部因果结构，导致模型对环境变化的适应性较差。
核心思路：论文的核心思路是利用稀疏性来约束Transformer的注意力机制，从而学习对象之间更简洁、更具解释性的交互图。稀疏性假设认为，在复杂的环境中，对象之间的交互通常是局部的，只有少数对象之间存在直接的因果关系。通过强制注意力模式稀疏化，模型可以专注于关键的交互，忽略不相关的对象，从而提高模型的泛化能力和鲁棒性。
技术框架：SPARTAN基于Transformer架构，其主要模块包括：对象表征模块（将原始观测转换为对象级别的特征向量）、Transformer编码器（学习对象之间的交互关系）、预测模块（预测未来对象的状态）。关键在于Transformer编码器中的注意力机制，通过引入稀疏性正则化，鼓励模型学习稀疏的注意力模式。整个流程包括：输入观测数据，经过对象表征模块得到对象特征，然后输入到Transformer编码器中学习对象间的交互，最后通过预测模块预测未来对象的状态。
关键创新：SPARTAN的关键创新在于将稀疏性正则化引入到Transformer的注意力机制中，从而学习稀疏的、上下文相关的交互图。与传统的Transformer模型相比，SPARTAN能够更有效地发现对象之间的局部因果结构，提高模型的泛化能力和鲁棒性。此外，SPARTAN还能够适应环境中具有未知目标的稀疏干预，使其更具实用性。
关键设计：SPARTAN使用L1正则化来约束注意力权重，鼓励模型学习稀疏的注意力模式。具体来说，在训练过程中，将注意力权重的绝对值之和添加到损失函数中，并乘以一个超参数λ来控制稀疏性的强度。此外，论文还探索了不同的注意力稀疏化策略，例如全局稀疏化和局部稀疏化。网络结构方面，采用了标准的Transformer编码器结构，并针对对象表征模块和预测模块进行了优化。
🖼️ 关键图片
📊 实验亮点
SPARTAN在多个benchmark环境中取得了显著的性能提升。例如，在对象交互预测任务中，SPARTAN的预测精度比现有方法提高了10%以上。在few-shot自适应任务中，SPARTAN能够仅通过少量样本就快速适应新的环境动态，而其他方法则需要大量的训练数据。此外，SPARTAN在存在干扰因素的情况下也表现出更强的鲁棒性。
🎯 应用场景
SPARTAN可应用于机器人导航、自动驾驶、游戏AI等领域，帮助智能体更好地理解和预测环境动态，从而做出更明智的决策。通过学习稀疏的交互图，SPARTAN能够提高模型的可解释性，有助于调试和优化智能体的行为。此外，SPARTAN在few-shot自适应方面的优势使其能够快速适应新的环境和任务，降低了训练成本。
📄 摘要（原文）
Capturing the interactions between entities in a structured way plays a central role in world models that flexibly adapt to changes in the environment. Recent works motivate the benefits of models that explicitly represent the structure of interactions and formulate the problem as discovering local causal structures. In this work, we demonstrate that reliably capturing these relationships in complex settings remains challenging. To remedy this shortcoming, we postulate that sparsity is a critical ingredient for the discovery of such local structures. To this end, we present the SPARse TrANsformer World model (SPARTAN), a Transformer-based world model that learns context-dependent interaction structures between entities in a scene. By applying sparsity regularisation on the attention patterns between object-factored tokens, SPARTAN learns sparse, context-dependent interaction graphs that accurately predict future object states. We further extend our model to adapt to sparse interventions with unknown targets in the dynamics of the environment. This results in a highly interpretable world model that can efficiently adapt to changes. Empirically, we evaluate SPARTAN against the current state-of-the-art in object-centric world models in observation-based environments and demonstrate that our model can learn local causal graphs that accurately reflect the underlying interactions between objects, achieving significantly improved few-shot adaptation to dynamics changes, as well as robustness against distractors.