Faster-WAM in LeRobot: from 300 cache reads to 10

A world action model, or WAM, is usually a robot policy built on a video model. Fast-WAM (March 2026) is fast compared to the WAMs that generate future frames before they act. It learns to predict video in training, but when the robot acts it runs the 30-layer Wan2.2 video model once on the current frame and caches every layer's keys and values. Then a 30-layer action expert denoises the next chunk of 32 actions in 10 steps, and at every step action layer $j$ reads only video layer $j$'s cache. Faster-WAM (August 2026) asks if that expert needs to be deep. Its answer is no: one action layer that reads a mix of all 30 caches stays competitive.

$$\begin{aligned}\text{Fast-WAM:}\quad & 30 \times 10 = 300 \text{ cache reads}\\ \text{Faster-WAM:}\quad & 1 \times 10 = 10 \text{ cache reads}\end{aligned}$$

Three changes

The paper's ablation (Figure 3) cuts Fast-WAM to one action layer that reads only the last video layer, then adds three changes one at a time. LIBERO is a simulated tabletop benchmark, and LIBERO-Plus perturbs the same scenes.

The paper copies Fast-WAM's rates from earlier papers, and only the last row is a three-run mean. On RoboTwin 2.0 it reports 91.8% for Fast-WAM and 90.5% for Faster-WAM.

Lite KV-Fusion

Faster-WAM mixes all 30 caches into one, once per chunk, with a learned weight per layer and head:

$$\beta_{h\ell}=\frac{e^{\alpha_{h\ell}}}{\sum_{\ell'=1}^{30} e^{\alpha_{h\ell'}}}$$

$$\widehat K_h=\sum_{\ell=1}^{30}\beta_{h\ell}\,K_{h\ell}, \qquad \widehat V_h=\sum_{\ell=1}^{30}\beta_{h\ell}\,V_{h\ell}$$

Here $\ell$ is the video layer, $h$ is one of 24 heads, the $\alpha_{h\ell}$ are 720 learned numbers, and $K_{h\ell}$ and $V_{h\ell}$ are the cached keys and values. The $\alpha$ start at zero, so every $\beta$ starts at $1/30$, and in the paper's Figure 4 (LIBERO) they stay within about 10% of it, read by eye. The authors still call the mix non-uniform, with different heads leaning on different depths.

So how does the action expert use that cache? At every denoising step, the 32 action tokens, one per action in the chunk, attend to the fused cache and to each other in one softmax per head:

$$O_h=\operatorname{softmax}\!\left(\frac{Q_h\,[\widehat K_h;\,K^{a}_h]^{\top}}{\sqrt{D}}\right)[\widehat V_h;\,V^{a}_h]$$

$Q_h$, $K^{a}_h$ and $V^{a}_h$ are the queries, keys and values the action layer computes from the 32 action tokens, $[\,\cdot\,;\cdot\,]$ stacks the video rows on top of the action rows, and $D=128$ is the size of one head. Since it's a single softmax, each action query splits its attention between the scene and the other actions. Only the action rows change from one step to the next, while $\widehat K_h$ and $\widehat V_h$ stay fixed for the whole chunk. Each time an action layer evaluates this, that's one of the cache reads counted at the top. Fast-WAM evaluates it in all 30 action layers at every step, with layer $j$ reading video layer $j$'s own $K_{hj}$ and $V_{hj}$ instead of the fused pair. The paper also puts a mask inside the softmax that hides future frames in training, but at inference the cache only holds the current frame, so there's nothing to hide.

World-model-only conditioning

Text and robot state go into the video model only, whose layers already attend to them. The action layer drops its own cross-attention to them and gets them through the fused cache.

Retracted 1D RoPE

Rotary position encoding (RoPE) tells attention where tokens are. Each head rotates 64 pairs of query and key numbers by angles that grow with position, and the score sees the angle difference. Video keys rotate by frame time, row and column (22, 21 and 21 pairs), while action queries rotate by action index, so in 42 pairs the score subtracts an action index from a row or a column, which has no geometric meaning. Faster-WAM undoes the video rotation on the fused keys, $\bar k(p)=R^{v}(p)^{\top}\,\widehat k(p)$, and the attention above uses these keys, $\bar K_h$, in place of $\widehat K_h$. In pair $r$ the angle goes from

$$\theta_r^{\text{before}}=\omega_r\,p_{b(r)}-\nu_r\,i \qquad\text{to}\qquad \theta_r^{\text{after}}=-\nu_r\,i$$

Here $i$ is the action index, $p=(t,y,x)$ is the video token's frame time, row and column, $b(r)$ picks one of the three, and $\omega_r$ and $\nu_r=10000^{-2r/128}$ are the video and action rates. The frame sits at time 0. For action 5 and a patch in row 3, counting from 0, the first row pair ($r=22$, $\omega=1$, $\nu\approx0.042$) goes from $3-0.21=2.79$ to $-0.21$, now the same for every patch. One undo after the fusion is exact, because all 30 layers rotate a token by the same angle (the paper's Proposition 2.2).

In LeRobot

I opened issue #5048 and PR #5049, which adds four flags to FastWAMConfig, all off by default:

action_expert_layers=1
lite_kv_fusion=True
retracted_1d_rope=True
world_model_only_conditioning=True

You can start from a Fast-WAM checkpoint but have to fine-tune it, since the one-layer expert and the fusion weights train jointly with the video model. A maintainer asked for changes on October 22, I pushed them the same day, and the PR was merged on October 23. AI agents generated the code, tests and draft text, independent agents reviewed it, and the PR says so. No code comes from the authors' repository, whose original code is under a noncommercial license.

Speed

My runs use random video and action weights with the real image encoder (VAE) and text encoder, so they measure speed only: about 2.5x in float32 and 5.3x in bfloat16, end to end.

bfloat16 gains more for two measured reasons. The encoders and the video prefill, the one video pass that fills the cache, run about twice as fast as in float32, while Fast-WAM's eager denoiser gets slower, from 493 to 797 ms, with 29,020 GPU kernel launches per chunk against Faster-WAM's 770. So the 5.3x is mostly an eager-execution ratio on this host, not a ratio of arithmetic. My Fast-WAM loop is also about 5x slower than the paper's (797 against about 150 ms, both eager bfloat16), which is most of why my 5.3x is bigger than the paper's 3.8x.

My take: the action expert was the expensive part, and it doesn't need to be deep. In bfloat16 the denoising loop fell from 797.3 to 23.7 ms, so the text encoder (about 42 ms) and the video prefill (about 88 ms) now each take longer than the whole loop. The next speedup has to come from there, or from compiling the pipeline.

Limits

• Speed only. No trained Faster-WAM policy, no LIBERO run and no robot.
• I haven't timed both models compiled, or the merged code.
• I haven't tried a fixed $1/30$ average in place of the learned weights.
• arXiv 2608.04404 is a different paper with the same Faster-WAM name. A separate post compares WAMs with vision-language-action models.

← All writing