02 方法Method
LoopWM 由四个模块组成：观测编码器 $\mathcal{E}_\phi$、动作嵌入器 $\mathcal{A}_\psi$、作为核心贡献的循环动力学核（Looped Dynamics Core） $\mathcal{L}_\theta$，以及轻量预测头 $\mathcal{D}_\xi$（解码观测重建、奖励、终止信号）。整个前向过程为 $e_k=\mathcal{E}_\phi(o_k),\ u_k=\mathcal{A}_\psi(a_k),\ h_k=\mathcal{L}_\theta(h_{k-1},e_k,u_k),\ (\hat o_{k+1},\hat r_k,\hat c_k)=\mathcal{D}_\xi(h_k)$。LoopWM comprises four modules: the observation encoder $\mathcal{E}_\phi$, the action embedder $\mathcal{A}_\psi$, the Looped Dynamics Core $\mathcal{L}_\theta$ that constitutes the core contribution, and a lightweight prediction head $\mathcal{D}_\xi$ (decoding observation reconstruction, reward and termination signals). The whole forward pass is $e_k=\mathcal{E}_\phi(o_k),\ u_k=\mathcal{A}_\psi(a_k),\ h_k=\mathcal{L}_\theta(h_{k-1},e_k,u_k),\ (\hat o_{k+1},\hat r_k,\hat c_k)=\mathcal{D}_\xi(h_k)$.
循环动力学核：Prelude → 共享 Recurrent Block → CodaLooped Dynamics Core: Prelude → shared Recurrent Block → Coda
核心把动力学核拆成三段：Prelude（少量层，融合 $[h_{k-1};e_k;u_k]$ 并做 LayerNorm 得到条件信号 $e$）；Recurrent Block（参数共享的一叠 transformer 层，循环 $T$ 次，更新规则为 $h^{(t+1)}=\bar A h^{(t)}+\bar B e+\bar{\mathcal{R}}(h^{(t)},e)$，其中 $\bar{\mathcal{R}}$ 是 transformer 的非线性部分）；Coda（独立参数的收尾层，将终态 $h^{(T)}$ 投影为下一步的隐状态 $h_k$）。关键点是"$\mathcal{R}$ 的参数在全部 $T$ 次迭代间共享"，因此计算深度与参数量解耦。终态 $h^{(T)}$ 还可以作为下一环境步的初始隐状态 $h^{(0)}$，形成"内层迭代 refine 单步转移、外层跨时间步传播"的双层循环结构。The core splits the dynamics core into three stages: Prelude (a few layers that fuse $[h_{k-1};e_k;u_k]$ and apply LayerNorm to obtain the conditioning signal $e$); Recurrent Block (a parameter-shared stack of transformer layers, looped $T$ times, with update rule $h^{(t+1)}=\bar A h^{(t)}+\bar B e+\bar{\mathcal{R}}(h^{(t)},e)$, where $\bar{\mathcal{R}}$ is the nonlinear part of the transformer); Coda (a closing layer with its own parameters that projects the final state $h^{(T)}$ into the next-step hidden state $h_k$). The key point is that "the parameters of $\mathcal{R}$ are shared across all $T$ iterations", so computational depth is decoupled from parameter count. The final state $h^{(T)}$ can further serve as the initial hidden state $h^{(0)}$ of the next environment step, forming a two-level loop in which the inner iterations refine a single-step transition while the outer loop propagates across time steps.
谱稳定性约束（Spectral Stability Constraint）Spectral Stability Constraint
为保证隐状态不随循环次数 $T$ 增大而爆炸，作者把状态保留矩阵 $\bar A$ 参数化为连续时间负对角矩阵的离散化：$A:=\mathrm{diag}(-\exp(\mathbf a))$，再零阶保持离散化 $\bar A=\exp(\Delta A)$。因为 $A$ 的对角元恒为负，指数映射后 $\bar A$ 的所有特征值都落在 $(0,1)$ 区间，"guaranteeing $\rho(\bar A) < 1$... This constraint holds by construction throughout training; no gradient clipping, post-hoc normalisation, or sensitive hyperparameter tuning is required."To keep the hidden state from blowing up as the number of loops $T$ grows, the authors parameterise the state-retention matrix $\bar A$ as the discretisation of a continuous-time negative diagonal matrix, $A:=\mathrm{diag}(-\exp(\mathbf a))$, followed by zero-order-hold discretisation $\bar A=\exp(\Delta A)$. Because the diagonal entries of $A$ are always negative, after the exponential map every eigenvalue of $\bar A$ falls in the interval $(0,1)$, "guaranteeing $\rho(\bar A) < 1$... This constraint holds by construction throughout training; no gradient clipping, post-hoc normalisation, or sensitive hyperparameter tuning is required."
训练目标：随机深度 + 熵正则Training objective: stochastic depth + entropy regularisation
训练时每个样本独立地从 $T\sim\mathrm{Poisson}(\mu_{\mathrm{rec}})$ 采样循环次数（而非按 micro-batch 采样，作者称这能"reduces variance in the training objective and empirically eliminates most loss spikes"），总损失为观测/奖励/终止预测损失之和；若启用自适应早退，再加一个熵正则项防止退出门坍缩到"总在第一次退出"或"永不退出"两个极端。反向传播只回传到 $\mu_{\mathrm{bwd}}=\lceil\mu_{\mathrm{rec}}/2\rceil$ 步以控制显存。During training each sample independently draws its loop count from $T\sim\mathrm{Poisson}(\mu_{\mathrm{rec}})$ (rather than sampling per micro-batch; the authors say this "reduces variance in the training objective and empirically eliminates most loss spikes"), and the total loss is the sum of the observation, reward and termination prediction losses; when adaptive early exit is enabled, an entropy regularisation term is added to stop the exit gate collapsing to either extreme of "always exiting on the first iteration" or "never exiting". Backpropagation is carried back only $\mu_{\mathrm{bwd}}=\lceil\mu_{\mathrm{rec}}/2\rceil$ steps in order to bound memory.
推理时自适应早退（Adaptive Early Exit）Adaptive Early Exit at inference time
推理阶段用一个单层 MLP + sigmoid 的退出门 $g^{(t)}=\sigma(\mathbf w_g^\top h^{(t)}+b_g)$，一旦 $g^{(t)}$ 超过阈值 $\tau$ 就提前终止内层循环。论文给出一个量级估计：对 100 层的定深基线，简单的自由飞行片段若 LoopWM 只需 1 次循环（如 4 层）即可退出，"reducing inference FLOPs by a factor of ∼25× for that step"；把简单/复杂转移混合在一条长 rollout 里，"the aggregate FLOPs reduction can reach up to two orders of magnitude"。推理时的最大循环数 $T_{\max}$ 还可以超过训练时的均值 $\mu_{\mathrm{rec}}$，从而支持 test-time compute scaling。At inference an exit gate made of a single-layer MLP plus sigmoid, $g^{(t)}=\sigma(\mathbf w_g^\top h^{(t)}+b_g)$, terminates the inner loop early as soon as $g^{(t)}$ exceeds the threshold $\tau$. The paper offers an order-of-magnitude estimate: against a fixed-depth baseline of 100 layers, if LoopWM can exit on an easy free-flight segment after only 1 loop (say 4 layers), this is "reducing inference FLOPs by a factor of ∼25× for that step"; mixing easy and hard transitions along one long rollout, "the aggregate FLOPs reduction can reach up to two orders of magnitude". The maximum loop count $T_{\max}$ at inference may also exceed the training-time mean $\mu_{\mathrm{rec}}$, which is what enables test-time compute scaling.
Deferred Decoding：只在终点解码Deferred Decoding: decode only at the end point
标准世界模型在每个环境步都调用一次预测头。LoopWM 提出"encode, think, then decode"式的 Deferred Decoding：给定 $K$ 步动作序列，模型只在隐空间里连续推进 $K$ 次（每步内部仍是 $T$ 次内层循环），中间不产生任何观测/奖励/终止预测，只在第 $K$ 步才调用一次解码器。训练上用终态预测损失 + 隐一致性正则（用冻结编码器对中间真实观测编码作为软监督目标）+ 谱收缩预算（限制 $\sum_k\|h_{k+1}-h_k\|_2$ 的累计量），并配合"K 从 1 逐步增大"的课程学习，因为直接用大 $K$ 训练要反传 $K\times T$ 次共享参数应用，容易不稳定。规划模式下，一次 decode-free rollout 只需调用解码器 1 次而不是 $K$ 次，"saving approximately $(K-1)\times\text{cost}(\mathcal{D}_\xi)$ FLOPs per candidate sequence"。A standard world model calls the prediction head at every environment step. LoopWM proposes an "encode, think, then decode" style of Deferred Decoding: given a $K$-step action sequence, the model advances $K$ times consecutively inside the latent space (each step still running $T$ inner loops), produces no observation, reward or termination prediction in between, and calls the decoder only once, at step $K$. Training uses a terminal-state prediction loss + a latent consistency regulariser (a frozen encoder encodes the intermediate ground-truth observations as a soft supervision target) + a spectral contraction budget (bounding the cumulative $\sum_k\|h_{k+1}-h_k\|_2$), together with a curriculum that grows $K$ from 1 upwards, because training directly with a large $K$ means backpropagating through $K\times T$ applications of the shared parameters and is prone to instability. In planning mode, one decode-free rollout calls the decoder once rather than $K$ times, "saving approximately $(K-1)\times\text{cost}(\mathcal{D}_\xi)$ FLOPs per candidate sequence".
