Velox: Spatial Memory Makes World Models Efficient

AI summary

Overview

Research area: Computer Vision — efficient long-horizon video generation, specifically video world models and long-context attention.

Technical level: Advanced. The paper assumes familiarity with diffusion transformers, KV caching, FlashAttention-style tiled execution, camera geometry (frusta, unprojection, oriented bounding boxes), and space-filling curves.

Scope: A training-free inference-time attention router that uses camera frustum slab correspondences from streaming geometry to bound the growing KV-attention cost in chunked video world models, evaluated on runtime scaling over 48 chunks and the 158-case WBench navigation split.

What This Paper Is About

Chunked video world models generate long rollouts by letting each new chunk attend to all previously generated frames through a growing KV cache, which makes later chunks progressively slower even when the current camera only sees a small part of the past scene. Velox asks whether the model can instead assign each region of the current chunk its own small, geometry-relevant slice of history, rather than retrieving whole past frames or attending to everything. The goal is to keep a pretrained generator unchanged while removing the late-rollout attention bottleneck without materially degrading quality.

Key Contributions

1. A training-free frustum-slab attention router. Velox represents historical KV blocks as camera frustum slabs with depth intervals, and uses frustum slab overlap to route each current query block to a compact set of historical KV blocks. No fine-tuning or new conditioning branch is required.

2. A soft geometric fusion rule with an unbiased fallback. Rather than imposing a hard geometric top-K, Velox uses oriented-bounding-box (OBB) frustum slab overlap to softly modulate each layer's own model-derived block importance. Regions with insufficient geometric coverage (e.g. newly revealed sky) keep the model-derived route without geometric bias.

3. A systems implementation. Velox is built with global-aligned Gilbert KV storage (replacing an earlier chunk-local order plus a separate cross-chunk seam table), a native Triton indexed block-sparse attention kernel, and asynchronous frustum slab construction running on the same H200 as denoising with a two-chunk freshness lag.

4. An empirical audit of the geometry/importance relationship. Offline affinity dumps quantify how often frustum slab overlap agrees with the model's own preferred blocks, motivating soft fusion over hard replacement.

Main Findings

• Attention scaling is nearly flat. Over 48-chunk timing traces, dense FlashAttention grows with retained context while Velox stays nearly flat. At chunk 48, per-call attention latency drops from approximately 34.2 ms to 7.8 ms, a 4.36× speedup.

• End-to-end denoising improves. Under the same single-H200 asynchronous setting, chunk denoising time falls from approximately 7.19 s to 2.96 s, a 2.43× speedup, with geometry construction sharing the same GPU.

• Quality stays close to dense attention. On all 158 cases of the WBench navigation split, Velox remains within 0.15 API-free average points of dense LingBot-Fast-Cam and improves the 19-metric overall average from 75.52 to 75.60.

• Geometry beats matched sparse routing. Against a model-only block-sparse baseline with the same routing budget, soft geometric fusion raises the API-free average from 73.11 to 73.51, the overall average from 75.01 to 75.60, and spatial consistency from 71.8 to 74.3.

• Geometry and model importance agree well above chance. At τ = 0.1, layer 14 reaches geometric recall R_geo = 0.7302 and fixed-budget overlap R_50 = 0.8909, versus a random expectation of 50/143 = 0.350. Agreement varies by layer, which the authors cite as the reason to keep layer-specific distributions rather than a layer-invariant geometric set.

• Fusion strength was calibrated, not hand-tuned blindly. Over 48 pooled-importance grids, λ = 2 and g_ref = 0.2 raise fused top-80 positive-overlap recall from 0.602 to 0.852 while changing 9.0% of selected blocks. Positive overlap values have median 0.211, which motivates rescaling rather than multiplying by raw IoU.

• Qualitative drift is reduced. In the illustrated rollouts (bedroom, figurine, bamboo, desert), dense attention progressively departs from the input frame while Velox retains the bedspread, desk, armchair, figurine colors, bamboo path, and sand ripples. The authors note these cases are selected for illustration and that the benchmark scores are the primary evidence.

Methodology in Plain English

The starting point is a standard autoregressive video world model that generates one chunk at a time and keeps every previous chunk in a KV cache. Velox changes only one thing at inference: which historical KV blocks each region of the current chunk is allowed to look at.

Step 1 — Summarize the past as slabs, not samples. After a chunk is decoded, a reconstruction backend produces depth, confidence, masks, and camera poses. Velox unprojects reliable depth samples into 3D, projects them into the chunk's latent cameras, buckets them by block, and stores each block as one or more "frustum slabs": a source-camera token rectangle plus a robust depth interval (low and high percentiles) plus validity statistics. Validity gates require enough samples, enough covered tokens, bounded relative depth spread, and a minimum slab thickness. Storing slabs rather than exact 3D point identities makes the representation robust to artifacts in generated video.

Step 2 — Get depth for the current chunk. The current chunk is noisy during denoising, so Velox cannot reconstruct its depth directly. Instead it caches the first and last pixel-frame depths of recent reconstructed chunks and forward-warps them into the current chunk's cameras. For each target token, the newest cached frame with at least a minimum number of warped samples wins, and the median of those samples becomes the query depth. If no cached frame reaches that floor, the token has no depth and can fall back to model-derived routing.

Step 3 — Score overlap in world space. Each query slab and historical slab is approximated as a camera-facing oriented bounding box. The geometric score is an intersection-over-union between the query box and the best-matching historical box, computed with a deterministic OBB grid after a separating-axis test. A separate coverage score counts what fraction of tokens in a query block fall inside a valid historical slab both spatially and in depth.

Step 4 — Fuse softly instead of overriding the model. For each layer, Velox computes a model-derived importance distribution over visible KV blocks from mean query and mean key vectors. For well-covered blocks, overlap values above a noise threshold τ are rescaled and used to multiply that distribution by (1 + λ·ĝ), then renormalize. A block with full overlap gets at most a 1+λ boost; a block with no vetted overlap keeps its original importance. This preserves each pretrained layer's own notion of useful context.

Step 5 — Materialize bounded routes with safeguards. Well-covered query blocks get the previous chunk's KV blocks reserved first, then the top-scoring remaining blocks up to K_cov = 80. Poorly-covered blocks keep the model distribution and take the smallest prefix whose cumulative mass reaches a threshold, capped at K_unc = 150. Both routes are then unioned with a 26-neighbor local window in a global-aligned Gilbert block ordering, plus optional reference-frame blocks, added outside the route budget.

Step 6 — Execute with a real sparse kernel. The measured path is a native Triton indexed block-sparse attention implementation, not a dense call with a Python mask. Because the block size matches FlashAttention-style tiling, skipping a block removes both the QKᵀ tile product and the KV tile load. The Gilbert ordering keeps neighboring image regions contiguous in memory, so local neighborhoods and block-contiguous KV storage share one coordinate frame.

The calibrated configuration is τ = 0.1, g_ref = 0.2, λ = 2, η_cov = 0.9, K_cov = 80, K_unc = 150. The whole stack is built on LingBot-Fast-Cam with a Wan-style video diffusion transformer.

Why This Matters

Impact on research. The paper reframes long-context video generation from "which frames do I retrieve?" to "which past regions should this region attend to?" That is a finer granularity than frame- or view-level memory, and it lands at the same token-block unit that tiled attention kernels already use — meaning the routing decision is directly executable rather than requiring irregular row-wise indexing. The training-free framing also means the approach can be layered onto existing pretrained generators without architectural adaptation, which is a constraint that limits reuse of many warped-rendering injection methods.

Real-world applications:

• Interactive simulation and game-like environments where a user drives a camera through a generated world and expects responsive frame rates at long horizons.
• Embodied AI and robot training, where simulated rollouts must stay visually consistent as an agent revisits the same locations.
• Virtual production and previsualization, where a director explores a scene along a camera path and needs the set to remain stable across the shot.
• Long-horizon content creation and digital twins of physical spaces, where persistent memory of a location is the point rather than a byproduct.

Industry relevance. The bottleneck Velox targets — late-chunk latency growth — is the difference between a world model that demos well for a few seconds and one that can serve an interactive session. A 4.36× attention and 2.43× denoising speedup at chunk 48 on a single H200, with quality held within 0.15 API-free points of dense attention, is the kind of trade-off that matters for inference cost per session. The fact that geometry construction runs asynchronously on the same GPU, rather than requiring a second device, lowers the deployment bar further.

Future Directions

• Closing the quality gap rather than bounding it. Velox is described as remaining "within" 0.15 points of dense attention on the API-free average, not exceeding it. Whether better Q-depth (currently recency-priority multiwarp with a two-chunk freshness lag) narrows that further is an open question.

• Per-layer rather than fixed geometric fusion. The agreement audit shows substantial layer variation (R_geo from 0.6652 at layer 25 to 0.7302 at layer 14), yet the reported configuration applies a fixed fusion setting at every transformer layer. Layer-adaptive λ or τ is a natural next step.

• Reducing the geometry freshness lag. Chunk c normally consumes geometry only through chunk c−2. Reducing this lag, or handling disocclusion and stale warped depths more explicitly, could improve routing in fast-moving or heavily revealing trajectories.

• Generalization beyond the evaluated setting. Results are reported on one inference stack (LingBot-Fast-Cam with a Wan-style transformer) and one benchmark (the 158-case WBench navigation split). Behavior on other world models, other camera trajectories, and non-navigation scene types is not reported.

Target Audience

Researchers and engineers working on video diffusion transformers, long-context attention, and interactive world models — particularly those who care about inference latency at long rollout horizons. It will also interest practitioners in 3D vision and novel-view synthesis, since the core mechanism imports streaming reconstruction geometry (frustum slabs, unprojection, OBB overlap) into the attention-routing problem. Readers without a background in diffusion transformers, KV caching, and camera geometry will find the method section demanding; the abstract and qualitative comparison are accessible to a broader audience.

Authors’ abstract

Video world models turn actions and camera motion into long visual rollouts, making video a natural interface for interactive simulation. Such interaction requires responsive generation, yet coherent rollouts must retain visual memory of what has already been generated. Chunked video diffusion models usually keep this context in growing KV caches, so later chunks become progressively slower. To address this bottleneck, we introduce Velox, a training-free attention router for region-level context selection using frustum slab correspondences from streaming geometry. Velox represents historical blocks as camera frustum slabs with depth intervals and obtains current-block depth through recency-priority multiwarp. For well-covered regions, frustum slab overlap softly modulates each layer's model-derived block importance. Uncovered regions retain the model-derived route without geometric bias. The fixed token-block granularity matches FlashAttention-style tiled execution and reduces the attention term that grows with rollout length. Frustum slab construction runs asynchronously with denoising on the same H200. At chunk 48, Velox accelerates attention by $4.36\times$ and denoising by $2.43\times$. On the full 158-case WBench navigation split, it remains within 0.15 API-free average points of dense LingBot-Fast-Cam. Please see our project page at https://velox-world.com/velox/ for more details.