FAR teaches world models which memories to recall by peeking at the future during training

Learning What to Recall: Adaptive Multi-Cue Episodic Memory for World Models

Beomsu Kim, Chieh-Hsin Lai, Bac Nguyen, Amir Bar, Jong Chul Ye, Yuki Mitsufuji

cs.LG, cs.AI, cs.CV

2026-09-28

FAR trains a future-blind memory retriever with the observed future, adaptively weighing time, pose, vision and audio cues, beating hand-designed recall on three benchmarks.

What problem this solves

World models predict future frames from the current observation and an action, but much of what determines those frames has already left the field of view. Navigate back along a loop and the model needs the streets it saw on the way out; open a fridge and it needs to remember whether a tomato went inside. The standard fix is an external episodic memory that keeps every past observation and feeds a few retrieved frames to the generator. The hard part is the retrieval itself. Existing systems hard-code it: Navigation World Models takes the most recent frames, WorldMem retrieves by field-of-view overlap with the goal pose, LongLive-RAG by visual-embedding similarity. Each rule fails somewhere predictable. In an elbow corridor the pose-overlapping memory may sit behind a wall; similar corridor segments are indistinguishable by appearance; after an interaction changes the world, the geometrically best-matched frame depicts a stale state and the model hallucinates. FAR's position is that both questions, which memories help this prediction and which cues should find them, should be learned.

Method

FAR (Future-Aware Recall; KAIST with Sony) imports a recipe from retrieval-augmented language models. EMDR2 trained a discrete document retriever from downstream reader likelihood; FAR applies the same latent-variable treatment to an agent's own history, exploiting an asymmetry unique to world models: at inference the future is unknown, but during training it has already happened and is visible. Four moving pieces:

• Predictive utility as the teacher. Condition the world model on each candidate memory alone and score how well it predicts the realized future. Exact likelihood is too expensive for a video diffusion model, so FAR uses the negative diffusion prediction loss over four noise levels, with all candidates sharing the same perturbations to keep comparisons fair. Proposition 1 supplies the theory: the conditional mutual information between the recalled context and the future is lower-bounded by the world model's conditional log-likelihood.
• Posterior distillation. Add the retriever's own score to the utility and softmax to get a future-aware posterior; the inference-time retriever, which never sees the future, is trained to match this fixed target by KL with a stop-gradient. The paper argues against supervising with utility alone: that implicitly swaps the query-dependent prior for a uniform one and discards relevance already inferred from the current query.
• Adaptive multi-cue fusion. Each cue (time, pose, vision, audio) gets its own scorer. Scores are standardized over the current memory bank, then combined with weights produced by a learned gate from per-cue score statistics, so the same environment can trust pose on one query and audio on the next.
• Engineering. Vision and audio encoders are contrastively pretrained then frozen, so historical keys are computed once and cached; only a light query-side adapter trains jointly. Recall is chunk-structured: memory is partitioned into temporal chunks, each chunk contributes at most one frame, and the top-K representatives are recalled, so K slots never collapse onto adjacent redundant frames.

A deliberate boundary: cues are used only to select memories, never as extra conditions to the generator. SoundSpaces audio picks frames; generation still consumes frames plus pose, time, and action.

Results

LoopNav (Minecraft loop navigation; outbound leg is the memory pool, return leg is a 128-frame rollout):

Same cues, learned retrieval, clearly better. The ablation completes the chain: temporal baseline 13.23 PSNR, WorldMem 15.37, FAR with a frozen pretrained ViT encoder 16.25, plus a query adapter 16.48, plus metadata fused at a fixed 0.5 weight 16.91, and with learned fusion weights 17.27 PSNR and 0.082 DreamSim. That last step is the point: adaptive weighting buys about 0.35 PSNR over fixed equal fusion.

SoundSpaces (scanned indoor scenes with binaural spatial audio): with dense periodic scans, metadata alone is strong and audio adds little; with sparse endpoint-only scans and long returns, the meta+audio advantage grows with trajectory length. In the qualitative corridor example, temporal, pose, and geometric retrieval all pick views obstructed by walls, while the audio cue retrieves a memory with clear sight of the goal region.

AI2-THOR (interactive kitchens where objects move and states go stale): on state accuracy, the fraction of changed states rendered correctly rather than stale, multi-cue FAR dominates both baselines, most clearly on container reveals, where the current observation says nothing about hidden contents. A counterfactual pair: same closed-fridge view, same opening action, different histories; FAR renders the tomato only when it was placed there, while WorldMem hallucinates it at its stale table location. On the off-scene dynamics task, predicting a patrolling agent visible only intermittently through a doorway: temporal 41.9%, WorldMem 32.4%, FAR with metadata only 67.1%, multi-cue 94.6%.

Why it matters

For anyone building long-video generators or world models, this is a clean demonstration that memory retrieval can be a learned component, with the training objective inherited from EMDR2 rather than invented. Two transferable ideas: the train-time versus inference information gap (the answer is visible during training) can be converted into supervision well beyond world models; and the cue-fusion interface lets a new modality join by adding one scorer, leaving the generator untouched. Code is at sony/far on GitHub.

Scale deserves honesty: the experiments are hundreds to tens of thousands of controlled simulation episodes at up to 360x640 resolution, with a 4-frame training context (12 at eval). This is not yet external memory for a frontier video model.

Limitations

The authors list four: only recall is learned, while memory writing, compression, and forgetting are untouched; the method needs informative cues and pays extra DiT forward passes for utility evaluation; the diffusion loss underweights small semantic changes (patched with change masks in AI2-THOR); everything runs in simulation. Two more concerns from reading closely. AI2-THOR-dyn injects task-specific machinery, including second-agent position supervision and silhouette-weighted losses, directly into the retriever reward, and the paper never disentangles how much of the 94.6% is framework versus task supervision. LoopNav and AI2-THOR-dyn use episode-level splits where train and test share scenes, so scene-disjoint generalization is established only on SoundSpaces and AI2-THOR. MemLearner (ECCV 2026), a contemporary learned retriever, appears only in the conceptual comparison table, not as an experimental baseline.

Terms

• World model: a generative model that predicts future observations from the current observation and actions; used as a learned environment simulator.
• Episodic memory: an external store that keeps each past observation individually addressable, unlike persistent-state memory that compresses history into one evolving representation.
• Retrieval cue: a signal used to locate memories; here, timestamps, camera pose, visual embeddings, and audio embeddings.
• EMDR2: a retrieval-augmented-LM method that trains a discrete retriever from downstream reader likelihood; FAR's training objective descends from it.
• WorldMem: a NeurIPS 2025 world-model memory system that retrieves frames by field-of-view overlap; the strongest hand-designed baseline here.
• DreamSim / LPIPS / PSNR: image-similarity metrics ranging from pixel-level (PSNR) to perceptual (DreamSim).

Source

• arXiv 2609.34677v1
• PDF
• Hugging Face
• GitHub · https://github.com/sony/far

Related papers

• Two empowerment encoders separate control from noise, with a 1,000× sensitivity gap
• Heavy-tailed data and gradient noise drive LLM loss spikes; a convex toy model predicts them
• Test-time ridge regression: +12 RULER points, 55.6% fewer visual tokens
• Muon-trained Moonlight gains 7.9 HumanEval points over its AdamW twin at 1.2T
• Parallax: a learned covariance branch cuts pretraining perplexity 17%, but only under Muon
• Muon beats Adam on ImageNet-C and typo text; ViT effective rank 6.88 vs 1.86

All paper explainers