[close ]
k -WM: Predicting k Steps at Once for Long-Horizon Planning with World ModelsThe usual JEPA-style planning world model rolls out autoregressively, and that brings the standard problems of an RNN-like loop: errors accumulate along the rollout, and in gradient-based planning every step’s Jacobian enters one long product, leaving a badly conditioned optimization problem. k -WM avoids the loop. Given a few context frames and $k$ actions, the world model predicts all $k$ future states at once; $k=1$ recovers the usual next-step world model. This change improves planning, and, more surprisingly, it also improves the latent geometry without any human prior. The last plot below is one example; the project page has more analysis.
k -WM fits any world model with a transformer predictor: the loss target becomes the next $k$ latents instead of one, and at inference one forward pass rolls out $k$ steps. The second tab shows the next-step model for comparison.
[k -WMnext-step ] [replay ]
k -WM training
o1:c+k
Obs.
Encoder
Eθ
reg.
z1:c
a1:c
c = 3
zc
ac+1
zc
ac+2
zc
ac+3
zc
ac+k−1
⋯
k steps (repeat last latent) Causal transformer predictor
Tθ k-step
ẑc+1
ẑc+2
ẑc+3
ẑc+k−1
ẑc+k
⋯
zc+1:c+k
ℒ
k targets
k -WM inference
o1:c
Obs.
Encoder
Eθ
z1:c
a1:c
c = 3
zc
ac+1
zc
ac+2
zc
ac+3
zc
ac+k−1
⋯
k steps (repeat last latent) Causal transformer predictor
Tθ k-step
ẑc+1
ẑc+2
ẑc+3
ẑc+k−1
ẑc+k
⋯
1 call
calls: 1
Next-step WM training
o1:c+1
Obs.
Encoder
Eθ
reg.
z1
a1
z2
a2
zc
ac
c = 3
Causal transformer predictor
Tθ 1-step
ẑ2
ẑc
ẑc+1
z2:c+1
ℒ
1-step targets
Next-step WM inference
o1:c
Obs.
Encoder
Eθ
z1:c
a1:c
c = 3
Causal transformer predictor
Tθ 1-step
ẑc+1
calls: 1
ẑc+1
ac+1
2
ẑc+2
calls: 2
ẑc+2
ac+2
3
ẑc+3
calls: 3
ẑc+k−1
ac+k−1
k
⋯
⋯
ẑc+k
calls: k k steps
Training above, inference below. The c context frames are teacher forced; “reg.” is whatever anti-collapse regularizer the base model uses. A next-step model predicts one latent per call and feeds it back in, so k steps take k calls. k -WM fills the future positions with copies of the last context latent and the future actions, and one call returns all k latents. The loss is averaged over the k steps, with no curriculum.
We use $k=8$ and apply it to three families of world models: DINO-WM plans on frozen DINOv2 features, DINO FT fine-tunes DINOv2 end to end, and LeWM trains a small ViT from scratch with SIGReg. Each is compared with its own next-step version, trained identically otherwise.
[DINO FT ] [LeWM ] [DINO-WM ]
0
25
50
75
100
CEM
success rate (%)
Push-T
Maze
Cube
Reacher
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
0
25
50
75
100
GD
success rate (%)
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
D
I
N
O
F
T
w
i
t
h
-
W
M
k
DINO FT, next-step
0
25
50
75
100
CEM
success rate (%)
Push-T
Maze
Cube
Reacher
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
0
25
50
75
100
GD
success rate (%)
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
L
e
W
M
w
i
t
h
-
W
M
k
LeWM, next-step
0
25
50
75
100
CEM
success rate (%)
Push-T
Maze
Cube
Reacher
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
0
25
50
75
100
GD
success rate (%)
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
5
6
7
8
9
10
p
l
a
n
n
i
n
g
h
o
r
i
z
o
n
H
D
I
N
O
-
W
M
w
i
t
h
-
W
M
k
DINO-WM, next-step
Success rate against the planning horizon on four tasks, with sampling-based planning (CEM) above and gradient-based planning (GD) below. Solid is the model with k -WM, dashed its next-step version. Mean and standard deviation over three seeds.
My favorite result, though, is not about planning success but about representation learning. As $k$ grows, the encoder learns a latent that is more and more reflective of the features that the task requires. The plot below uses DINO FT on the Maze environment. We pick an anchor cell (the star) and color every other cell by how far its latent is from the anchor’s. With next-step training, $k=1$, the latent mostly encodes local geometry and the map has a grid-like pattern. As $k$ grows, the latent increasingly groups cells by their distance along the corridors, and by $k=16$ the map looks like the corridor distance in the first column. A long-horizon prediction target induces a latent geometry that mirrors the task-relevant geometry, with no human prior about what matters.
corridor distance
k
=
1
k
=
4
k
=
8
k
=
1
6
0
4
8
corridor distance (blocks)
0.0
0.6
1.2
latent sq. distance (1 = random pair)
DINO FT on Maze, one anchor per row. Left: distance along the corridors. Then: latent distance for models trained with k = 1, 4, 8 and 16, normalized so that two random frames are at distance 1.
[close ]
Modular Hamiltonians for Future-Perturbed States
O1 O2 O3 ξ+ ξ−
O1 O2 O3 ξ+ ξ−
O1 O2 O3 ξ+ ξ−
O1 O2 O3 ξ+ ξ−
A fun diagram from the paper; for the context, read on. Light travels along the diagonals, and the blue curves are hyperbolas. In these four configurations infinities appear in our calculation. Taming them lets us extract further results.
While the name of this project seems technical and very specific (it is), we are actually studying entanglement (again), but in a slightly different context!
Let’s start with the computation of entropy. Suppose you are handed a probability distribution $p$. We recognize that the entropy is defined as
\[S = -\sum p \log p = \mathbb{E}(-\log p) = \langle-\log p\rangle,\]
where the last two equality signs are just juggling notations, physicists like to use $\langle\,\cdot\,\rangle$ to denote expectation value. Here, the interesting object of study is $\log p$.
In quantum mechanics, the object that is similar to the classical probability distribution is the density matrix $\rho$, and similar to how the probability distribution sums to one, $\sum p = 1$, the density matrix (a positive semi-definite matrix) traces to one, $\operatorname{tr}\rho=1$. Similar to classical probability, we define the entropy quantum mechanically via
\[S = \langle-\log \rho \rangle = -\operatorname{tr}(\rho \log\rho),\]
and we call $H_\text{mod}=-\log\rho$ the modular Hamiltonian, we have $S = \langle H_\text{mod}\rangle$.
However, it turns out that the $S$ defined above can be rather boring, depending on whether the density matrix is boring, and for a pure state, which is a state that we know everything about, the density matrix is boring: it has rank one and zero entropy, just like a deterministic distribution (e.g. Kronecker delta).
What is interesting is entanglement, namely how much information about a subsystem $A$ we lose if we throw away the rest of the system, $B$. Throwing away $B$ is the partial trace $\operatorname{tr}_B$, the matrix version of marginalizing. The interesting quantity is therefore
\[S_A = -\operatorname{tr}(\rho_A\log\rho_A),\quad\text{where}\quad \rho_A = \operatorname{tr}_B\rho.\]
Perhaps surprisingly, while a deterministic joint distribution has deterministic marginals, a pure quantum state can have a mixed $\rho_A$, meaning $S_A$ can be nonzero even when $S=0$! We interpret $S_A$ as measuring exactly how entangled the two subsystems $A$ and $B$ are.
The object this paper is about is the modular Hamiltonian of a subsystem, $H_A = -\log\rho_A$. It carries everything about the entanglement, starting with $S_A = \langle H_A\rangle$. For a few special subsystems it is known exactly. We want to understand it more generally, so we reach into physicists’ favorite toolbox — perturbation theory.
The above quantum mechanical picture can be promoted to quantum field theory, if one is happy to let some infinities slip. The subsystems $A$ and $B$ are now regions of spacetime, and we are studying entanglement across these regions. In field theory, correlations between two points become singular when a light ray connects them. These light-cone singularities are how causality shows up in the calculation. The diagram at the top is the catalogue of where they occur. Taming those singularities lets us progress in understanding the modular Hamiltonian. For the rest, the paper has it all, check it out yourself!
[close ]
Entanglement Groups
Everyone has heard of entanglement, from your undergraduate physics friends to your coworker obsessed with popular science: “spooky action at a distance”! The fame is well deserved: entanglement is fundamental to our understanding of quantum phenomena. It also plays a central role in recent advances in quantum gravity, which we (I) care a lot about.
In this paper, we aim to investigate entanglement and propose a different way of thinking about it. A familiar way to characterize entanglement goes as follows. Pick your favorite pure state (vector) $\lvert\psi\rangle \in \mathcal{H}$, where $\mathcal{H}$ is the space of possible states (it’s a vector space!). We divide the system into two parts, $A$ and $B$, so that $\mathcal{H} = \mathcal{H}_A \otimes \mathcal{H}_B$, and compute the entanglement entropy , $S_A$, across this division:
\[S_A = -\operatorname{Tr}(\rho_A \log \rho_A).\]
Here, $\rho_A$ describes the state of $A$ when we ignore $B$ (mathematically, we trace out subsystem $B$). For our pure state $\lvert\psi\rangle$, if this entropy is nonzero, we say the two parts are entangled; otherwise, they are not. In the latter case, the state is a tensor product state , $\lvert\psi\rangle = \lvert\psi_A\rangle \otimes \lvert\psi_B\rangle.$
We propose to instead think of entanglement as a property of the state that lets us perform some action on one part, $A$, and undo it with a corresponding action on the other part, $B$. Mathematically, we describe “doing something to parts of a state” using group actions. Here, these are unitary transformations (basically rotation matrices). Calling the two operations $g_A$ and $g_B$, we can translate the idea into math:
\[g_A \lvert\psi\rangle = g_B^{-1}\lvert\psi\rangle,\]
where $g_A$ acts only on $A$, and $g_B$ acts only on $B$. In other words, acting on one side can have the same effect as acting on the other. Roughly, collecting these coordinated transformations gives us the entanglement group
Our prescription extends naturally to multipartite systems, capturing entanglement structure that entropy alone does not fully describe. Using this language, we propose a way to characterize and classify multipartite entanglement , show how monogamy of entanglement arises, identify this entanglement structure in a range of quantum algorithms and protocols, and extend the formulation to mixed states.
* For the precise definition see our paper.
