精读笔记

Problem Setting

《Learning to Focus: Prioritizing Informative Histories with Structured Attention Mechanisms in Partially Observable Reinforcement Learning》（Awesome World Models / 2025）关注的是 Transformer-based world model 在部分可观测 RL 中如何选择历史，而不是泛泛地提升 Transformer 表达能力。

更具体地说，UniZero 这类模型把 latent state-action sequence 交给 causal Transformer dynamics head，用 self-attention 聚合过去信息来预测 next latent state 和 reward。问题在于：POMDP 下确实需要 history，但不是所有 history 都有用；Atari 轨迹里大部分 transition 对当前控制和 reward prediction 没有信息量。标准 attention 的参数化能力足够，但在 100k interaction 的低数据 regime 下，它要同时学习表示、动力学、价值、策略以及“哪些历史重要”，这个 credit assignment 太分散。

以前方法卡住的地方不是没有 memory，而是 memory selection 太弱。RNN / RSSM 把历史压进隐状态，可能丢长依赖；Transformer 保留全历史，但把选择问题交给数据驱动 attention。二者都没有明确利用 RL 轨迹的一个事实：有效依赖通常稀疏、局部但不严格局部，并且和 reward / transition prediction 相关，而不是和语言 token 共现类似。

关键矛盾是：POMDP 需要足够长的历史来 disambiguate hidden state，但低样本 RL 又承受不起对所有历史 token 做无偏学习。论文试图在“保留可访问历史”和“限制无效历史干扰”之间加一个结构化折中。

Motivation

作者的核心观察是，Transformer 在 NLP 中成功依赖的是大规模、重复、密集监督的序列数据；而 RL trajectory 是 on-policy / replay-buffer 相关数据，奖励稀疏、分布偏斜、有效事件少。把标准 self-attention 直接用于 dynamics modeling，等价于假设模型可以从少量 noisy rollout 中自己发现 temporal relevance，这在 Atari-100k 这种 setting 下并不现实。

已有路线缺的不是更大的 backbone，而是对 attention search space 的约束。UniZero 已经证明 Transformer dynamics 比 recurrent unroll 更适合处理历史，但它仍然继承了“所有 past token 初始等价”的建模方式。对于 reward-driven trajectories，这个先验很差：模型早期会把容量浪费在大量平庸 transition 上，reward-relevant dependencies 反而被稀释。

因此作者想到直接在 attention logits 里注入 temporal prior：不是让模型先学会全部 pairwise relevance，再从中筛选；而是先假设某些时间距离更可能有用，让 content attention 在这个结构上微调。关键缺口就是：RL world model 中 attention prior 目前大多服务于位置编码或计算效率，而不是服务于 task-relevant history selection。

Core Idea

论文真正的核心思想是：把 dynamics head 的历史聚合从“无结构的 content-based attention”改成“content attention + 可学习时间先验”。这不是改变 planner，也不是改变 MuZero/UniZero 的训练目标，而是改变 world model 在预测 transition/reward 时接收历史信息的方式。它给每个 attention head 一个时间尺度偏好，让 head 不只是学习语义子空间，也学习 temporal filter。

其中最有效的版本是 Gaussian positional prior。它假设对当前 timestep 有用的信息往往集中在某个相对时间 offset 附近，但这个集中不是硬边界，而是平滑衰减。相比固定窗口或 adaptive span，这个假设更符合 RL：很多依赖有大致时延，但边界不稳定；同一游戏中不同状态的有效 history length 也会波动。Gaussian bias 允许 attention 在局部性和长尾之间连续调节，而不是做离散截断。

和 prior 的本质区别在于，它不是用 Transformer 获得更长上下文，而是让 Transformer 更早、更稳定地忽略大部分无效上下文。换句话说，它提升的不是理论表达上限，而是样本效率和优化路径。这个区别很重要：标准 Transformer 已经能表示这些依赖，论文的贡献是把有效依赖变得更容易学到。

Method

方法上可以抽象成三种 attention logit bias。

第一，memory-length prior / Adaptive Attention：每个 head 学一个 look-back span，span 外 token 被 hard mask 掉。它解决的是历史冗余和二次 attention 中无效 token 干扰的问题。必要性来自 POMDP 中有限有效记忆的假设；核心变化是把每个 head 的 receptive field 从完整 causal history 变成可学习窗口。问题也在这里：hard cutoff 假设有效依赖有清楚边界，而 RL 中这个边界通常不存在或随状态变化。

第二，distributional prior / Gaussian Attention：每个 head 学一个相对时间中心 μ 和宽度 σ，把 Gaussian bias 加到 attention logits。它解决的是“哪些过去时间距离先验更可能有用”的问题。相比 adaptive span，它不是删除 token，而是重加权 token；因此即使主要关注近期或中程依赖，也保留远处 token 通过 content attention 被选中的可能。这是本文真正值得保留的机制。

第三，Gaussian Adaptive Attention：把 Gaussian bias 和 hard span mask 相加。直觉上想结合软 saliency 与有限 horizon，但实验显示组合不成立。机制原因很直接：Gaussian 的优势来自平滑 tail 和可连续调节，hard span 会破坏这个性质。这个结果反而强化了论文的主结论：在不确定依赖长度的 RL 中，软先验比硬约束更稳。

这些机制都只作用在 dynamics head 的 self-attention 中，prediction head、MCTS、UniZero 主体基本保持不变。因此它是一种 attention-level inductive bias，而不是完整的新 model-based RL 算法。

Key Insight / Why It Works

最重要的 insight 是：在低数据 POMDP world modeling 中，attention 的主要瓶颈不是容量，而是 relevance discovery 的样本复杂度。Gaussian prior 有效，是因为它把“从所有历史 token 中发现少数有用 token”的问题，改写成“在若干平滑时间尺度附近微调 content relevance”。这显著降低了优化难度。

这更像 better inductive bias，而不是 scaling、test-time compute 或 planning 改进。MCTS 没变，数据量没变，模型参数和 FLOPs 几乎没变；性能增益大概率来自 dynamics representation 更稳定，尤其是 reward / next latent prediction 对短中期历史的利用更好。它不是 retrieval，因为没有外部 memory 或 explicit nearest-neighbor；也不是 curriculum；也不是显式 belief-state construction。它是 temporal filtering prior。

Gaussian 比 Adaptive span 好的原因也很清楚：hard window 把“历史是否有用”建模成二值决策，但 RL 中很多 dependency 是 soft、state-dependent、non-stationary 的。一个事件可能在某些状态下 3 步前有用，另一些状态下 8 步前有用；hard span 一旦学错，就要么截断 delayed signal，要么保留太多噪声。Gaussian 允许错误是连续的：μ 不准、σ 稍宽时仍能覆盖有效 token，优化更宽容。

真正核心贡献是 Gaussian logit bias 这个很小的改动在 RL world model 中的归因实验，而不是 Adaptive span 本身。Adaptive span 是已有思想迁移，且在这里暴露出局限；Gaussian Adaptive 的失败也说明“组合更多 prior”不等于更好。

需要警惕的是，论文把问题表述为 partial observability 下 prioritizing informative histories，但实验 context length 只有 10、inference context length 更短，且 Atari 很多依赖本身偏短期视觉动态。Gaussian prior 的收益可能主要来自一个更好的 recency / short-mid horizon bias，而不是证明它解决了真正长程 POMDP belief inference。所谓 long-term dependency 的 claim 有些过强。

Relation To Prior Work

这篇最接近三条线：MuZero/UniZero 式 latent world model + planning，Transformer world models for RL，以及 adaptive attention span / positional bias。

相对 MuZero，它没有重新设计 latent dynamics 或 planning，只是在 UniZero 的 Transformer dynamics 上做 attention bias。相对 UniZero，它的本质差异是从无结构 self-attention 转向带 temporal inductive bias 的 self-attention；不是让模型看得更长，而是让模型更有偏地看历史。

相对 Adaptive Span Transformer，这篇的新意不在 adaptive span，而在说明 hard span 在 RL world modeling 中并不可靠。这个负结果有价值：NLP 中为效率设计的可学习上下文长度，迁移到 reward-driven trajectories 时可能会切掉关键信号。

相对 IRIS、TransDreamer、Decision Transformer 等 Transformer-in-RL 工作，这篇不把 Transformer 当通用序列建模器，而是明确针对 dynamics head 的 temporal relevance 加先验。这个定位更窄，但也更干净：它研究的是 world model 内部的 history aggregation，而不是把 control 整体改写成 sequence modeling。

看似新的地方有一部分是已有思想重组：Gaussian positional bias、adaptive span、causal attention 都不是新概念。实质创新在于把这些 prior 放进 model-based RL dynamics head，并通过同一 UniZero 框架比较软 prior、硬 prior 和组合 prior，得出“软时间分布先验比硬记忆窗口更适合低样本 POMDP dynamics modeling”的经验结论。

Dataset / Evaluation

评估集中在 Atari-100k，single-task，每个游戏单独训练。这个 benchmark 适合检验 sample efficiency，也适合观察不同游戏中 temporal dependency 和 reward sparsity 的影响。因此它能够支持一个有限 claim：在低交互预算 Atari world model 中，Gaussian attention prior 比 vanilla UniZero attention 更稳。

但它没有验证更强 claim：跨任务泛化没有做，连续控制没有做，真实机器人或 embodied setting 没有做，更长 horizon POMDP 没有做。Atari 虽然是部分可观测的，但许多游戏的有效控制依赖主要来自最近几帧或短中期运动趋势；这和真实 POMDP 中需要长期 memory / belief tracking 的情形差别很大。

实验设计的优点是 ablation 比较直接：Adaptive、Gaussian、Gaussian Adaptive 放在同一个 UniZero 框架中，参数和 FLOPs 基本不变，说明提升不太可能来自模型规模。初始化敏感性也显示 σ 比 μ / span 更关键，这支持“窄而平滑的时间先验”这一解释。

不足是归因仍不够细。文中没有充分拆分 reward prediction loss、latent dynamics loss、policy/value improvement 与 final return 的关系；也没有展示 attention map 是否真的对齐 reward-relevant events。最终 score 提升可能来自更稳定的短期 dynamics，而不一定是更好的 informative history discovery。

Limitation

第一，方法成立依赖一个强前提：有效历史可以被相对时间距离上的平滑单峰或少数 head 的单峰分布近似。如果任务依赖的是事件类型而非时间距离，例如“上一次看到钥匙的位置”“若干秒前隐藏变量切换”，Gaussian offset prior 可能不够，甚至会误导 attention。

第二，Gaussian prior 是 position-conditioned，不是 state/event-conditioned。μ 和 σ 是 head-level 参数，不随当前 observation 或 belief 动态变化。它能表达“这个 head 通常看 6 步前”，但不能表达“当屏幕出现某个 cue 时看 30 步前，否则看 2 步前”。这限制了它在真正非平稳 POMDP 中的上限。

第三，context length 很短，使得论文关于 long-term dependency 的论证打折。训练 context length 为 10，max adaptive span 20，inference context 也有限；这更像 short-context temporal bias，而不是解决长程 memory。若扩展到数百步 history，Gaussian bias 的计算和表示是否仍然有效，文中未充分说明。

第四，Adaptive span 需要正则避免 collapse，Gaussian width 初始化也敏感。尤其 σ=1 明显好于 σ=3，说明性能依赖较强的先验设定。它不是完全从数据中发现结构，而是把一个合理的 Atari 时间尺度写进模型。

第五，增益不等于 planner 更会推理。MCTS 和训练目标未变，planning 能力的改善大概率来自 one-step / short rollout model 更准。若 world model 仍然只是在常见局部 transition 上更准，那么所谓部分可观测下的 history focusing 可能只是 improved local dynamics fitting。

第六，泛化范围尚未证明。Atari-100k 的游戏集合足以说明经验有效，但不足以说明该 prior 是通用 world model 组件。跨 domain、continuous control、多任务共享、offline RL、真实 embodied 数据都可能暴露不同 failure mode。

Takeaway

• 1. 对 RL Transformer world model 来说，attention prior 不是纯工程细节；在低数据 regime 下，它直接决定 relevance discovery 的样本复杂度。
• 2. 软时间先验比硬窗口更值得迁移。
• 很多 RL 依赖没有清晰边界，hard memory span 容易把问题从“学 relevance”变成“学对 cutoff”，后者在 sparse reward 下更脆弱。
• 3. 这篇推动的是一个很具体的方向：从扩大 context / backbone，转向设计 task-aware temporal bias。

一句话总结

这篇论文是 UniZero-style Transformer world model 上的一次小而有判断力的 inductive-bias 修正：它证明在低样本部分可观测 RL 中，给 dynamics attention 加平滑时间分布先验，比让模型无结构地从历史中自学 relevance 或用硬窗口截断历史更可靠。