Optimize DeMemWM memory retrieval and remove diagnostics

algorithms/worldmem/dememwm/algorithm.py

@@ -1,4 +1,3 @@
-
 from __future__ import annotations
 
 import math
@@ -10,13 +9,13 @@ from einops import rearrange
 
 from .cache import StreamingCache
 from .compression import CausalConv3DDynamicCompressor, SpatialConv2DMemoryProjector, latent_patch_tokens, spatial_pool_tokens
-from .diagnostics import summarize_eval_ablation_diagnostics, summarize_noise_bucket_diagnostics, summarize_revisit_diagnostics
 from .injection import InjectionAdapter
 from .memory import CausalMemoryBank, MemoryBankQuery, stack_record_tokens
 from .negatives import apply_revisit_eval_corruption
-from .retrieval import deterministic_revisit_retrieval
-from .schedules import EVAL_CORRUPTION_BRANCHES, compute_stream_gates, denoising_fraction_from_noise_levels, noise_bucket_from_denoising_fraction, noise_bucket_from_noise_levels, noise_bucket_ids_from_noise_levels, normalize_eval_ablation_branch, resolve_curriculum
+from .retrieval import batched_revisit_select_positions, deterministic_revisit_retrieval
+from .schedules import EVAL_CORRUPTION_BRANCHES, compute_stream_gates, denoising_fraction_from_noise_levels, normalize_eval_ablation_branch, resolve_curriculum
 from .types import MemoryRecord, MemorySourceType, MemoryStreamTensors
+from ..df_video import convert_to_plucker
 
 
 class MemoryDiTMixin:
@@ -36,35 +35,9 @@ class MemoryDiTMixin:
     strict_key_substrings = (
         ".memory_token_cross_attn.",
     )
-    _TRAIN_DIAGNOSTIC_LOG_KEYS = frozenset({
-        "revisit_candidate_frame_count",
-        "revisit_pose_preselect_input_count",
-        "revisit_pose_preselect_selected_count",
-        "revisit_exact_fov_candidate_count",
-        "valid_revisit_frame_count",
-        "no_valid_revisit_count",
-        "revisit_selected_frame_count",
-        "revisit_frame_fov_overlap_mean",
-        "revisit_best_selected_frame_fov_overlap_mean",
-        "revisit_best_selected_plucker_overlap_mean",
-        "revisit_best_selected_gap_frames_mean",
-        "revisit_gate_raw",
-        "revisit_gate_eff",
-        "revisit_learned_gate_mean",
-        "revisit_effective_gate_mean",
-        "generated_history_proxy_prob",
-        "noise_bucket_target_count",
-        "noise_bucket_high_target_count",
-        "noise_bucket_mid_target_count",
-        "noise_bucket_low_target_count",
-    })
-    _VALIDATION_DIAGNOSTIC_LOG_KEYS = _TRAIN_DIAGNOSTIC_LOG_KEYS | frozenset({
-        "cache_records",
-        "cache_slots",
-    })
 
     def _memory_cfg(self):
-        return getattr(self.cfg, "dememwm", None)
+        return getattr(getattr(self, "cfg", None), "dememwm", None)
 
     def _cfg_get(self, obj, name, default):
         if obj is None:
@@ -84,8 +57,6 @@ class MemoryDiTMixin:
         except Exception:
             return False
 
-    def _stage_policy_cfg(self):
-        return self._cfg_get(self._memory_cfg(), "stage_policy", None)
 
     def _eval_ablation_cfg(self):
         return self._cfg_get(self._memory_cfg(), "eval_ablation", None)
@@ -99,7 +70,7 @@ class MemoryDiTMixin:
         branch = normalize_eval_ablation_branch(self._cfg_get(cfg, "branch", "A_plus_D_plus_R_normal"))
         return enabled, branch
 
-    def _effective_gate_state(self, denoising_fraction: float | None = None, noise_bucket: str | None = None) -> dict:
+    def _effective_gate_state(self, denoising_fraction: float | None = None) -> dict:
         memory_cfg = self._memory_cfg()
         anchor_cfg = self._cfg_get(memory_cfg, "anchor", None)
         dynamic_cfg = self._cfg_get(memory_cfg, "dynamic", None)
@@ -110,12 +81,9 @@ class MemoryDiTMixin:
         revisit_config_enabled = self._stream_enabled(revisit_cfg)
         curriculum_state = self._curriculum_state()
         eval_ablation_enabled, eval_ablation_branch = self._eval_ablation_state()
-        debug_force = bool(self._cfg_get(memory_cfg, "debug_force_all_streams", False))
-        resolved_noise_bucket = noise_bucket or noise_bucket_from_denoising_fraction(denoising_fraction)
         gates = compute_stream_gates(
             curriculum_state.stage,
             denoising_fraction=denoising_fraction,
-            debug_force_all_streams=debug_force,
             anchor_gate=float(self._cfg_get(injection_cfg, "anchor_gate", 1.0)),
             dynamic_gate=float(self._cfg_get(injection_cfg, "dynamic_gate", 1.0)),
             revisit_gate=float(self._cfg_get(injection_cfg, "revisit_gate", 1.0)),
@@ -139,7 +107,6 @@ class MemoryDiTMixin:
         return {
             "curriculum_state": curriculum_state,
             "gates": gates,
-            "resolved_noise_bucket": resolved_noise_bucket,
             "anchor_config_enabled": anchor_config_enabled,
             "dynamic_config_enabled": dynamic_config_enabled,
             "revisit_config_enabled": revisit_config_enabled,
@@ -152,14 +119,13 @@ class MemoryDiTMixin:
             "force_revisit_on": bool(eval_ablation_enabled and eval_ablation_branch == "R_forced_on"),
         }
 
-    def _validate_config_contract(self) -> dict:
+    def _validate_config_contract(self) -> None:
         if bool(getattr(self, "_dememwm_contract_validated", False)):
-            return getattr(self, "_last_dememwm_config_diagnostics", {})
+            return
         memory_cfg = self._memory_cfg()
         if memory_cfg is None:
             self._dememwm_contract_validated = True
-            self._last_dememwm_config_diagnostics = {}
-            return {}
+            return
 
         stale_sections = [name for name in ("ablation", "memory", "loss", "abstention") if self._cfg_has(memory_cfg, name)]
         if stale_sections:
@@ -193,10 +159,8 @@ class MemoryDiTMixin:
         if not bool(self._cfg_get(revisit_cfg, "deterministic_pose_retrieval", True)):
             raise ValueError("final DeMemWM requires deterministic FOV/Plucker revisit retrieval")
         fov_overlap_threshold = self._cfg_get(revisit_cfg, "fov_overlap_threshold", 0.30)
-        if fov_overlap_threshold is not None:
-            fov_overlap_threshold = float(fov_overlap_threshold)
-            if fov_overlap_threshold < 0.0:
-                raise ValueError("dememwm.revisit.fov_overlap_threshold must be non-negative")
+        if fov_overlap_threshold is not None and float(fov_overlap_threshold) < 0.0:
+            raise ValueError("dememwm.revisit.fov_overlap_threshold must be non-negative")
         high_quality_fov_threshold = float(self._cfg_get(revisit_cfg, "high_quality_fov_threshold", 0.70))
         if high_quality_fov_threshold < 0.0:
             raise ValueError("dememwm.revisit.high_quality_fov_threshold must be non-negative")
@@ -222,9 +186,6 @@ class MemoryDiTMixin:
             value = int(self._cfg_get(revisit_cfg, field_name, default))
             if value <= 0:
                 raise ValueError(f"dememwm.revisit.{field_name} must be positive")
-        stage_policy_cfg = self._stage_policy_cfg()
-        if not bool(self._cfg_get(stage_policy_cfg, "noise_bucket_logging", True)):
-            raise ValueError("final DeMemWM keeps noise_bucket logging enabled")
         proxy_cfg = self._generated_history_proxy_cfg()
         proxy_max_prob = float(self._cfg_get(proxy_cfg, "max_prob", 0.0))
         proxy_dropout_prob = float(self._cfg_get(proxy_cfg, "dropout_prob", 0.0))
@@ -240,18 +201,7 @@ class MemoryDiTMixin:
             raise ValueError("dememwm.generated_history_proxy.ramp_steps must be non-negative")
         eval_ablation_cfg = self._eval_ablation_cfg()
         normalize_eval_ablation_branch(self._cfg_get(eval_ablation_cfg, "branch", "A_plus_D_plus_R_normal"))
-
-        diagnostics = {
-            "dynamic_exclude_latest_local_frames": exclude_latest_local_frames,
-            "revisit_deterministic_fov_plucker_retrieval": True,
-            "revisit_local_context_exclusion_frames": self._local_context_exclusion_frames(),
-            "revisit_fov_overlap_threshold": -1.0 if fov_overlap_threshold is None else fov_overlap_threshold,
-            "revisit_plucker_weight": plucker_weight,
-            "stage_policy_noise_bucket_logging": True,
-        }
         self._dememwm_contract_validated = True
-        self._last_dememwm_config_diagnostics = diagnostics
-        return diagnostics
 
     def _stream_enabled(self, stream_cfg) -> bool:
         return bool(self._cfg_get(stream_cfg, "enabled", True))
@@ -294,25 +244,17 @@ class MemoryDiTMixin:
         source_is_generated: torch.Tensor | None,
         context_frame_count: int | None = None,
         target_start_frame: int | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         cfg = self._generated_history_proxy_cfg()
         prob = self._generated_history_proxy_prob()
         noise_std = float(self._cfg_get(cfg, "noise_std", 0.0))
         dropout_prob = float(self._cfg_get(cfg, "dropout_prob", 0.0))
-        diagnostics = {
-            "generated_history_proxy_enabled": bool(self._cfg_get(cfg, "enabled", False)),
-            "generated_history_proxy_prob": float(prob),
-            "generated_history_proxy_noise_std": float(noise_std),
-            "generated_history_proxy_dropout_prob": float(dropout_prob),
-            "generated_history_proxy_frame_count": 0,
-            "generated_history_proxy_frame_fraction": 0.0,
-        }
         if source_is_generated is None:
             source_is_generated = torch.zeros(source_latents.shape[:2], device=source_latents.device, dtype=torch.bool)
         else:
             source_is_generated = source_is_generated.to(device=source_latents.device, dtype=torch.bool)
         if prob <= 0.0 or source_latents.numel() == 0:
-            return source_latents, source_is_generated, diagnostics
+            return source_latents, source_is_generated
 
         eligible_mask = torch.ones(source_latents.shape[:2], device=source_latents.device, dtype=torch.bool)
         if context_frame_count is not None or target_start_frame is not None:
@@ -322,12 +264,8 @@ class MemoryDiTMixin:
             if target_start_frame is not None:
                 eligible_mask &= frame_positions < max(0, int(target_start_frame))
         proxy_mask = (torch.rand(source_latents.shape[:2], device=source_latents.device) < prob) & eligible_mask
-        proxy_count = int(proxy_mask.detach().long().sum().item())
-        total_count = max(1, int(proxy_mask.numel()))
-        diagnostics["generated_history_proxy_frame_count"] = proxy_count
-        diagnostics["generated_history_proxy_frame_fraction"] = float(proxy_count / total_count)
-        if proxy_count == 0:
-            return source_latents, source_is_generated, diagnostics
+        if not proxy_mask.any():
+            return source_latents, source_is_generated
 
         corrupt_latents = source_latents.clone()
         frame_mask = proxy_mask[:, :, None, None, None].to(dtype=corrupt_latents.dtype)
@@ -342,7 +280,7 @@ class MemoryDiTMixin:
             corrupt_latents = torch.where(dropout_mask, corrupt_latents.new_zeros(()), corrupt_latents)
         source_is_generated = source_is_generated.clone()
         source_is_generated |= proxy_mask
-        return corrupt_latents, source_is_generated, diagnostics
+        return corrupt_latents, source_is_generated
 
     def _checkpoint_cfg(self):
         return self._cfg_get(self._memory_cfg(), "checkpoint", None)
@@ -397,62 +335,21 @@ class MemoryDiTMixin:
 
     def _apply_freeze_policy(self, optimizer=None, step: int | None = None):
         state = self._curriculum_state(step)
-
-        # Keep DDP's trainable graph stable: DiT params stay requires_grad=True
-        # from step 0 and are frozen by optimizer LR=0 until the full stage.
-        # Re-walk only when curriculum diagnostics can change.
         freeze_key = (state.stage, state.dit_train_state, state.freeze_vae)
-        last_key = getattr(self, "_last_freeze_key", None)
-        if last_key != freeze_key:
-            trainable_tensors = {
-                "dememwm_modules": 0,
-                "memory_adapters": 0,
-                "full_dit": 0,
-                "excluded_frozen": 0,
-            }
-            trainable_scalars = {key: 0 for key in trainable_tensors}
-            requires_grad_tensors = {key: 0 for key in trainable_tensors}
-            requires_grad_scalars = {key: 0 for key in trainable_tensors}
+        if getattr(self, "_last_freeze_key", None) != freeze_key:
             for name, param in self.named_parameters():
                 group_name = self._param_group_name(name, state)
-                should_train = self._group_trainable(group_name, state)
                 if group_name == "excluded_frozen" or (name.startswith("vae.") and state.freeze_vae):
-                    should_train = False
-                    should_require_grad = False
+                    param.requires_grad_(False)
                 else:
-                    should_require_grad = True
-                param.requires_grad_(should_require_grad)
-                if should_train:
-                    trainable_tensors[group_name] = trainable_tensors.get(group_name, 0) + 1
-                    trainable_scalars[group_name] = trainable_scalars.get(group_name, 0) + int(param.numel())
-                if should_require_grad:
-                    requires_grad_tensors[group_name] = requires_grad_tensors.get(group_name, 0) + 1
-                    requires_grad_scalars[group_name] = requires_grad_scalars.get(group_name, 0) + int(param.numel())
+                    param.requires_grad_(True)
             self._last_freeze_key = freeze_key
-            self._last_trainable_tensors = trainable_tensors
-            self._last_trainable_scalars = trainable_scalars
-            self._last_requires_grad_tensors = requires_grad_tensors
-            self._last_requires_grad_scalars = requires_grad_scalars
-        else:
-            trainable_tensors = getattr(self, "_last_trainable_tensors", {})
-            trainable_scalars = getattr(self, "_last_trainable_scalars", {})
-            requires_grad_tensors = getattr(self, "_last_requires_grad_tensors", {})
-            requires_grad_scalars = getattr(self, "_last_requires_grad_scalars", {})
 
         if optimizer is not None:
             for param_group in optimizer.param_groups:
                 group_name = param_group.get("name", "")
                 trainable = self._group_trainable(group_name, state)
                 param_group["lr"] = self._group_lr(group_name, state) if trainable else 0.0
-
-        diagnostics = state.diagnostics()
-        for group_name in ("dememwm_modules", "memory_adapters", "full_dit"):
-            diagnostics[f"trainable_tensors_{group_name}"] = trainable_tensors.get(group_name, 0)
-            diagnostics[f"trainable_params_{group_name}"] = trainable_scalars.get(group_name, 0)
-            diagnostics[f"requires_grad_tensors_{group_name}"] = requires_grad_tensors.get(group_name, 0)
-            diagnostics[f"requires_grad_params_{group_name}"] = requires_grad_scalars.get(group_name, 0)
-            diagnostics[f"optimizer_lr_{group_name}"] = self._group_lr(group_name, state) if self._group_trainable(group_name, state) else 0.0
-        self._last_dememwm_freeze_diagnostics = diagnostics
         return state
 
     def configure_optimizers(self):
@@ -1101,7 +998,7 @@ class MemoryDiTMixin:
         plucker_weight: float,
         revisit_retrieval_kwargs: dict | None,
         token_patch_size: int,
-    ) -> tuple[list[CausalMemoryBank], list[CausalMemoryBank], int, dict]:
+    ) -> tuple[list[CausalMemoryBank], list[CausalMemoryBank], int, list[list[tuple[MemoryRecord, ...]]] | None, dict | None]:
         if committed_latents.ndim != 5:
             raise ValueError("committed_latents must have shape (T_src,B,C,H,W)")
         T_src, B, _, H, W = committed_latents.shape
@@ -1126,12 +1023,6 @@ class MemoryDiTMixin:
         revisit_banks: list[CausalMemoryBank] = []
         dummy_tokens = committed_latents.new_zeros((1, hidden_size))
         dummy_mask = torch.ones((1,), device=stream_device, dtype=torch.bool)
-        preselection_candidate_count = 0
-        preselection_valid_candidate_label_count = 0
-        preselection_selected_count = 0
-        projected_anchor_frames = 0
-        projected_revisit_frames = 0
-        projected_revisit_records = 0
         retrieval_kwargs = dict(revisit_retrieval_kwargs or {})
 
         # Pre-convert pose tensors to stream_device once so that the
@@ -1188,6 +1079,175 @@ class MemoryDiTMixin:
         def _pose_subset(positions: torch.Tensor, batch_idx: int):
             return _tensor_subset(pose, positions, batch_idx)
 
+        fast_source_pose_shape = (
+            pose is not None
+            and pose.ndim >= 3
+            and ((pose.shape[0] == T_src and pose.shape[1] == B) or (pose.shape[0] == B and pose.shape[1] == T_src))
+        )
+        fast_target_pose_shape = (
+            target_pose is not None
+            and target_pose.ndim >= 3
+            and ((target_pose.shape[0] == T_tgt and target_pose.shape[1] == B) or (target_pose.shape[0] == B and target_pose.shape[1] == T_tgt))
+        )
+        if fast_source_pose_shape and fast_target_pose_shape:
+            selection = batched_revisit_select_positions(
+                source_frame_indices,
+                pose,
+                target_frame_indices,
+                target_pose,
+                topk=revisit_max_frames,
+                exclude_local_context_frames=exclude_local_context_frames,
+                fov_overlap_threshold=fov_overlap_threshold,
+                plucker_weight=plucker_weight,
+                fov_half_h=float(retrieval_kwargs.get("fov_half_h", 52.5)),
+                fov_half_v=float(retrieval_kwargs.get("fov_half_v", 37.5)),
+                fov_yaw_samples=int(retrieval_kwargs.get("fov_yaw_samples", 25)),
+                fov_pitch_samples=int(retrieval_kwargs.get("fov_pitch_samples", 20)),
+                fov_depth_samples=int(retrieval_kwargs.get("fov_depth_samples", 20)),
+                fov_radius=float(retrieval_kwargs.get("fov_radius", 30.0)),
+                plucker_grid_h=int(retrieval_kwargs.get("plucker_grid_h", 4)),
+                plucker_grid_w=int(retrieval_kwargs.get("plucker_grid_w", 4)),
+                plucker_focal_length=float(retrieval_kwargs.get("plucker_focal_length", 0.35)),
+                pose_preselect_topk=retrieval_kwargs.get("pose_preselect_topk", 64),
+            )
+            selected_records_by_target: list[list[tuple[MemoryRecord, ...]]] = [
+                [tuple() for _ in range(T_tgt)] for _ in range(B)
+            ]
+
+            for batch_idx in range(B):
+                anchor_bank = CausalMemoryBank()
+                revisit_bank = CausalMemoryBank()
+                src_frames = source_frame_indices[:, batch_idx]
+                if generated is None:
+                    non_generated = torch.ones_like(src_frames, dtype=torch.bool)
+                else:
+                    non_generated = ~generated[:, batch_idx]
+
+                source_positions = torch.nonzero(non_generated, as_tuple=False).flatten()
+                anchor_positions = source_positions[:0].to(device=stream_device, dtype=torch.long)
+                if anchor_indices and source_positions.numel() > 0:
+                    if anchor_diverse:
+                        anchor_source_positions = source_positions[source_positions < self._context_frame_count()]
+                        if anchor_source_positions.numel() > 0:
+                            anchor_pose = _pose_subset(anchor_source_positions, batch_idx)
+                            anchor_positions = self._select_diverse_anchor_positions(
+                                anchor_source_positions, anchor_pose, len(anchor_indices)
+                            ).to(device=stream_device, dtype=torch.long)
+                    else:
+                        selected_anchor_positions = []
+                        for anchor_idx in anchor_indices:
+                            if 0 <= int(anchor_idx) < source_positions.numel():
+                                selected_anchor_positions.append(source_positions[int(anchor_idx)])
+                        if selected_anchor_positions:
+                            anchor_positions = torch.stack(selected_anchor_positions).to(device=stream_device, dtype=torch.long)
+                if anchor_positions.numel() > 0:
+                    anchor_projected = self._project_latent_patch_tokens(
+                        committed_latents.index_select(0, anchor_positions)[:, batch_idx:batch_idx + 1],
+                        self.dememwm_anchor_proj,
+                        token_patch_size,
+                    )[0]
+                    for local_idx, source_pos in enumerate(anchor_positions):
+                        source_pos_i = int(source_pos.item())
+                        anchor_tokens = self._spatial_pool_tokens(anchor_projected[local_idx], anchor_pool_h, anchor_pool_w, src_h, src_w)
+                        n_slots = anchor_tokens.shape[0]
+                        record_mask = torch.ones((n_slots,), device=stream_device, dtype=torch.bool)
+                        anchor_bank.add_prefix_anchors(
+                            anchor_tokens.unsqueeze(0),
+                            record_mask.unsqueeze(0),
+                            src_frames[source_pos_i:source_pos_i + 1],
+                            slots_per_anchor=n_slots,
+                        )
+
+                selected_b = selection.selected_positions[batch_idx]
+                selected_mask_b = selection.selected_mask[batch_idx]
+                selected_fov_b = selection.selected_fov_overlap[batch_idx]
+                selected_plucker_b = selection.selected_plucker_overlap[batch_idx]
+                selected_gap_b = selection.selected_gap_frames[batch_idx]
+                high_quality_fov_threshold = float(retrieval_kwargs.get("high_quality_fov_threshold", 0.70))
+                metadata_by_position: dict[int, dict] = {}
+                for target_idx in range(T_tgt):
+                    for slot_idx in range(selected_b.shape[1]):
+                        if not bool(selected_mask_b[target_idx, slot_idx].detach().item()):
+                            continue
+                        source_pos_i = int(selected_b[target_idx, slot_idx].detach().item())
+                        if source_pos_i < 0:
+                            continue
+                        frame = int(src_frames[source_pos_i].detach().item())
+                        fov_overlap = float(selected_fov_b[target_idx, slot_idx].detach().item())
+                        plucker_overlap = float(selected_plucker_b[target_idx, slot_idx].detach().item())
+                        gap_frames = float(selected_gap_b[target_idx, slot_idx].detach().item())
+                        existing = metadata_by_position.get(source_pos_i)
+                        if existing is not None:
+                            existing_rank = (
+                                float(existing.get("dememwm_selected_frame_fov_overlap", 0.0)),
+                                -float(existing.get("dememwm_selected_gap_frames", 1.0e9)),
+                            )
+                            new_rank = (fov_overlap, -gap_frames)
+                            if existing_rank >= new_rank:
+                                continue
+                        metadata_by_position[source_pos_i] = {
+                            "dememwm_selected_revisit_fov_overlap": fov_overlap,
+                            "dememwm_selected_revisit_plucker_overlap": plucker_overlap,
+                            "dememwm_selected_gap_frames": gap_frames,
+                            "dememwm_selected_frame_index": frame,
+                            "dememwm_selected_frame_fov_overlap": fov_overlap,
+                            "dememwm_selected_frame_fov_threshold": high_quality_fov_threshold,
+                            "dememwm_selected_frame_passes_high_quality": bool(fov_overlap >= high_quality_fov_threshold),
+                        }
+                flat_selected = selected_b[selected_b >= 0].to(device=stream_device, dtype=torch.long)
+                unique_positions = torch.unique(flat_selected, sorted=True) if flat_selected.numel() > 0 else flat_selected
+                records_by_position: dict[int, MemoryRecord] = {}
+                if unique_positions.numel() > 0:
+                    revisit_projected = self._project_latent_patch_tokens(
+                        committed_latents.index_select(0, unique_positions)[:, batch_idx:batch_idx + 1],
+                        self.dememwm_revisit_proj,
+                        token_patch_size,
+                    )[0]
+                    for local_idx, source_pos in enumerate(unique_positions):
+                        source_pos_i = int(source_pos.item())
+                        frame_index = src_frames[source_pos_i]
+                        frame = int(frame_index.detach().item())
+                        is_generated = False if generated is None else bool(generated[source_pos_i, batch_idx].detach().item())
+                        source_type = MemorySourceType.GENERATED if is_generated else MemorySourceType.PREFIX_GT
+                        prefix = "generated" if is_generated else "prefix"
+                        frame_tokens = self._spatial_pool_tokens(revisit_projected[local_idx], revisit_pool_h, revisit_pool_w, src_h, src_w)
+                        frame_mask = torch.ones((frame_tokens.shape[0],), device=stream_device, dtype=torch.bool)
+                        record = MemoryRecord(
+                            tokens=frame_tokens,
+                            mask=frame_mask,
+                            source_start=frame,
+                            source_end=frame + 1,
+                            frame_indices=frame_index.reshape(1).to(device=stream_device),
+                            pose=_pose_subset(source_pos.reshape(1), batch_idx),
+                            source_type=source_type,
+                            is_generated=is_generated,
+                            chunk_id=f"{prefix}_revisit_b{batch_idx}_f{frame}",
+                            metadata=metadata_by_position.get(source_pos_i, {}),
+                        )
+                        revisit_bank.add_record(record)
+                        records_by_position[source_pos_i] = record
+
+                for target_idx in range(T_tgt):
+                    target_records: list[MemoryRecord] = []
+                    for source_pos in selected_b[target_idx]:
+                        source_pos_i = int(source_pos.detach().item())
+                        if source_pos_i < 0:
+                            continue
+                        record = records_by_position.get(source_pos_i)
+                        if record is not None:
+                            target_records.append(record)
+                    selected_records_by_target[batch_idx][target_idx] = tuple(target_records)
+
+                anchor_banks.append(anchor_bank)
+                revisit_banks.append(revisit_bank)
+
+            fast_revisit_stats = {
+                "best_selected_fov_overlap": selection.best_selected_fov_overlap.to(device=stream_device),
+                "best_selected_plucker_overlap": selection.best_selected_plucker_overlap.to(device=stream_device),
+                "best_selected_gap_frames": selection.best_selected_gap_frames.to(device=stream_device),
+            }
+            return anchor_banks, revisit_banks, tokens_per_frame, selected_records_by_target, fast_revisit_stats
+
         def _candidate_record(
             *,
             batch_idx: int,
@@ -1238,7 +1298,6 @@ class MemoryDiTMixin:
                     if selected_anchor_positions:
                         anchor_positions = torch.stack(selected_anchor_positions).to(device=stream_device, dtype=torch.long)
             if anchor_positions.numel() > 0:
-                projected_anchor_frames += int(anchor_positions.numel())
                 anchor_projected = self._project_latent_patch_tokens(
                     committed_latents.index_select(0, anchor_positions)[:, batch_idx:batch_idx + 1],
                     self.dememwm_anchor_proj,
@@ -1258,9 +1317,7 @@ class MemoryDiTMixin:
 
             candidate_records: list[MemoryRecord] = []
             candidate_positions: dict[str, torch.Tensor] = {}
-            src_frames_cpu = src_frames.detach().cpu()
-            target_frames_cpu = target_frame_indices[:, batch_idx].detach().cpu().to(dtype=torch.long)
-            latest_valid_source_frame_exclusive = int(target_frames_cpu.max().item()) - int(exclude_local_context_frames)
+            latest_valid_source_frame_exclusive = target_frame_indices[:, batch_idx].amax() - int(exclude_local_context_frames)
             for prefix, positions, source_type, is_generated in (
                 ("prefix", source_positions, MemorySourceType.PREFIX_GT, False),
                 (
@@ -1270,14 +1327,15 @@ class MemoryDiTMixin:
                     True,
                 ),
             ):
-                if positions.numel() == 0 or latest_valid_source_frame_exclusive <= 0:
+                if positions.numel() == 0:
                     continue
-                positions_cpu = positions.detach().cpu().to(dtype=torch.long)
-                for frame_position_cpu in positions_cpu:
-                    frame = int(src_frames_cpu[int(frame_position_cpu.item())].item())
-                    if frame >= latest_valid_source_frame_exclusive:
-                        continue
-                    frame_position = frame_position_cpu.reshape(1).to(device=stream_device, dtype=torch.long)
+                positions = positions.to(device=stream_device, dtype=torch.long)
+                frame_values = src_frames.index_select(0, positions).to(device=stream_device)
+                valid_positions = positions[frame_values < latest_valid_source_frame_exclusive]
+                valid_frames = src_frames.index_select(0, valid_positions) if valid_positions.numel() else src_frames[:0]
+                for frame_position, frame_tensor in zip(valid_positions.unbind(0), valid_frames.unbind(0)):
+                    frame = int(frame_tensor.item())
+                    frame_position = frame_position.reshape(1)
                     record_id = f"{prefix}_revisit_b{batch_idx}_f{frame}"
                     candidate_positions[record_id] = frame_position
                     candidate_records.append(_candidate_record(
@@ -1301,12 +1359,9 @@ class MemoryDiTMixin:
                     exclude_local_context_frames=exclude_local_context_frames,
                     fov_overlap_threshold=fov_overlap_threshold,
                     plucker_weight=plucker_weight,
-                    target_video_id=_target_video_id(batch_idx, target_idx),
+                    target_video_id=None,
                     **retrieval_kwargs,
                 )
-                preselection_candidate_count += int(result.diagnostics.get("revisit_candidate_frame_count", result.diagnostics.get("revisit_candidate_count", 0)))
-                preselection_valid_candidate_label_count += int(result.diagnostics.get("valid_candidate_label_count", 0))
-                preselection_selected_count += int(result.diagnostics.get("revisit_selected_frame_count", result.diagnostics.get("revisit_selected_count", 0)))
                 for selected_record in result.records:
                     if selected_record.chunk_id is None:
                         continue
@@ -1319,8 +1374,6 @@ class MemoryDiTMixin:
                     continue
                 record_id = str(record.chunk_id)
                 frame_position = candidate_positions[record_id]
-                projected_revisit_records += 1
-                projected_revisit_frames += int(frame_position.numel())
                 revisit_projected = self._project_latent_patch_tokens(
                     committed_latents.index_select(0, frame_position)[:, batch_idx:batch_idx + 1],
                     self.dememwm_revisit_proj,
@@ -1344,17 +1397,7 @@ class MemoryDiTMixin:
             anchor_banks.append(anchor_bank)
             revisit_banks.append(revisit_bank)
 
-        diagnostics = {
-            "preselected_anchor_projected_frame_count": projected_anchor_frames,
-            "preselected_revisit_projected_frame_count": projected_revisit_frames,
-            "preselected_revisit_projected_frame_record_count": projected_revisit_records,
-            "preselected_revisit_candidate_frame_count": preselection_candidate_count,
-            "preselected_revisit_candidate_count": preselection_candidate_count,
-            "preselected_revisit_valid_candidate_label_count": preselection_valid_candidate_label_count,
-            "preselected_revisit_selected_frame_count": preselection_selected_count,
-            "preselected_revisit_selected_count": preselection_selected_count,
-        }
-        return anchor_banks, revisit_banks, tokens_per_frame, diagnostics
+        return anchor_banks, revisit_banks, tokens_per_frame, None, None
 
     def _causal_cached_revisit_records(
         self,
@@ -1486,8 +1529,6 @@ class MemoryDiTMixin:
         target_video_ids=None,
         source_is_generated: torch.Tensor | None = None,
         denoising_fraction: float | None = None,
-        noise_bucket: str | None = None,
-        noise_bucket_ids: torch.Tensor | None = None,
         streaming_cache: StreamingCache | None = None,
     ) -> MemoryStreamTensors:
         if target_frame_indices is None:
@@ -1499,10 +1540,9 @@ class MemoryDiTMixin:
         dynamic_cfg = self._cfg_get(memory_cfg, "dynamic", None)
         revisit_cfg = self._cfg_get(memory_cfg, "revisit", None)
         injection_cfg = self._cfg_get(memory_cfg, "injection", None)
-        contract_diag = self._validate_config_contract()
+        self._validate_config_contract()
         gate_state = self._effective_gate_state(
             denoising_fraction=denoising_fraction,
-            noise_bucket=noise_bucket,
         )
         anchor_config_enabled = gate_state["anchor_config_enabled"]
         dynamic_config_enabled = gate_state["dynamic_config_enabled"]
@@ -1510,7 +1550,6 @@ class MemoryDiTMixin:
         curriculum_state = gate_state["curriculum_state"]
         eval_ablation_enabled = gate_state["eval_ablation_enabled"]
         eval_ablation_branch = gate_state["eval_ablation_branch"]
-        resolved_noise_bucket = gate_state["resolved_noise_bucket"]
         gates = gate_state["gates"]
         anchor_effective_enabled = gate_state["anchor_effective_enabled"]
         dynamic_effective_enabled = gate_state["dynamic_effective_enabled"]
@@ -1565,12 +1604,12 @@ class MemoryDiTMixin:
             "plucker_grid_w": int(self._cfg_get(revisit_cfg, "plucker_grid_w", 4)),
             "plucker_focal_length": float(self._cfg_get(revisit_cfg, "plucker_focal_length", 0.35)),
         }
-        preselection_diag = {}
+        preselected_revisit_records_by_target: list[list[tuple[MemoryRecord, ...]]] | None = None
+        preselected_revisit_stats: dict | None = None
         use_cache_revisit_records = False
         revisit_record_batches: list[tuple[MemoryRecord, ...]] | None = None
 
         cache = streaming_cache if streaming_cache is not None and getattr(streaming_cache, "enabled", False) else None
-        cache_diag = cache.diagnostics("cache") if cache is not None else {"cache_enabled": False, "cache_records": 0, "cache_slots": 0, "cache_evictions": 0, "cache_resets": 0}
         if committed_latents is not None:
             stream_device = committed_latents.device
             stream_dtype = committed_latents.dtype
@@ -1638,7 +1677,7 @@ class MemoryDiTMixin:
             B = committed_latents.shape[1]
             hidden_size = int(self._cfg_get(injection_cfg, "dit_hidden_size", 1024))
             target_pose_source = target_pose if target_pose is not None else pose
-            anchor_banks, revisit_banks, tokens_per_frame, preselection_diag = self._build_preselected_causal_memory_banks(
+            anchor_banks, revisit_banks, tokens_per_frame, preselected_revisit_records_by_target, preselected_revisit_stats = self._build_preselected_causal_memory_banks(
                 committed_latents,
                 source_frame_indices.to(device=stream_device),
                 None if source_is_generated is None else source_is_generated.to(device=stream_device, dtype=torch.bool),
@@ -1710,29 +1749,21 @@ class MemoryDiTMixin:
             dynamic_num_slots = self.dememwm_dynamic_compressor.tokens_per_target(_fallback_h, _fallback_w)
             dynamic_tokens = torch.zeros((B, T_tgt, dynamic_num_slots, hidden_size), device=stream_device, dtype=stream_dtype)
             dynamic_mask = torch.zeros((B, T_tgt, dynamic_num_slots), device=stream_device, dtype=torch.bool)
-            dynamic_diag = {
-                "selected_source_count": torch.zeros((B, T_tgt), dtype=torch.long, device=stream_device),
-                "max_source_frame": torch.full((B, T_tgt), -1, dtype=torch.long, device=stream_device),
-                "generated_source_fraction": torch.zeros((B, T_tgt), dtype=torch.float32, device=stream_device),
-                "dynamic_min_gap_to_target_per_target": torch.full((B, T_tgt), -1, dtype=torch.long, device=stream_device),
-                "dynamic_max_gap_to_target_per_target": torch.full((B, T_tgt), -1, dtype=torch.long, device=stream_device),
-                "dynamic_overlap_with_c_short_count_per_target": torch.zeros((B, T_tgt), dtype=torch.long, device=stream_device),
-                "dynamic_exclude_latest_local_frames": dynamic_recent_exclusion_frames,
-            }
         else:
             # Pre-select dynamic source frame positions using only frame index metadata
             # before touching latents, so we pass a small slice instead of the full
             # 1000-frame tensor to the compressor.
             _dfi = dynamic_frame_indices.to(device=stream_device)
             _max_src = self.dememwm_dynamic_compressor.max_source_frames
-            _needed: list[int] = []
+            _needed_tensors: list[torch.Tensor] = []
             for _b in range(B):
                 for _j in range(T_tgt):
-                    _target = int(target_frame_indices[_j, _b].item())
+                    _target = target_frame_indices[_j, _b]
                     _valid = (_dfi[:, _b] < _target - dynamic_recent_exclusion_frames).nonzero(as_tuple=False).flatten()
-                    _needed.extend(_valid[-_max_src:].tolist())
-            if _needed:
-                _needed_idx = torch.tensor(sorted(set(_needed)), device=stream_device, dtype=torch.long)
+                    if _valid.numel() > 0:
+                        _needed_tensors.append(_valid[-_max_src:])
+            if _needed_tensors:
+                _needed_idx = torch.unique(torch.cat(_needed_tensors, dim=0), sorted=True).to(device=stream_device, dtype=torch.long)
                 _dynamic_latents_small = dynamic_latents.index_select(0, _needed_idx)
                 _dynamic_fi_small = _dfi.index_select(0, _needed_idx)
                 _dynamic_pose_small = dynamic_pose.index_select(0, _needed_idx) if dynamic_pose is not None else None
@@ -1745,7 +1776,7 @@ class MemoryDiTMixin:
                 _dynamic_fi_small = _dfi[:0]
                 _dynamic_pose_small = dynamic_pose[:0] if dynamic_pose is not None else None
                 _dynamic_gen_small = None
-            dynamic_tokens, dynamic_mask, dynamic_diag = self.dememwm_dynamic_compressor(
+            dynamic_tokens, dynamic_mask = self.dememwm_dynamic_compressor(
                 _dynamic_latents_small,
                 _dynamic_fi_small,
                 _dynamic_pose_small,
@@ -1754,18 +1785,6 @@ class MemoryDiTMixin:
                 exclude_latest_local_frames=dynamic_recent_exclusion_frames,
             )
 
-        dynamic_min_gap_tensor = torch.as_tensor(
-            dynamic_diag.get("dynamic_min_gap_to_target_per_target", torch.full((B, T_tgt), -1, device=stream_device)),
-            device=stream_device,
-        )
-        dynamic_max_gap_tensor = torch.as_tensor(
-            dynamic_diag.get("dynamic_max_gap_to_target_per_target", torch.full((B, T_tgt), -1, device=stream_device)),
-            device=stream_device,
-        )
-        dynamic_gap_valid = dynamic_min_gap_tensor >= 0
-        dynamic_min_gap_to_target = int(dynamic_min_gap_tensor[dynamic_gap_valid].min().item()) if dynamic_gap_valid.any() else -1
-        dynamic_max_gap_valid = dynamic_max_gap_tensor >= 0
-        dynamic_max_gap_to_target = int(dynamic_max_gap_tensor[dynamic_max_gap_valid].max().item()) if dynamic_max_gap_valid.any() else -1
         def _target_tensor_or_none(tensor: torch.Tensor | None, batch_idx: int, target_idx: int):
             if tensor is None or tensor.ndim < 2:
                 return None
@@ -1804,15 +1823,11 @@ class MemoryDiTMixin:
         revisit_mask_rows = []
         revisit_max_rows = []
         valid_revisit_mask = torch.zeros((B, T_tgt), device=stream_device, dtype=torch.bool)
-        revisit_candidate_count = torch.zeros((B, T_tgt), device=stream_device, dtype=torch.float32)
-        revisit_selected_count = torch.zeros((B, T_tgt), device=stream_device, dtype=torch.float32)
         revisit_best_selected_fov_overlap = torch.zeros((B, T_tgt), device=stream_device, dtype=torch.float32)
         revisit_best_selected_plucker_overlap = torch.zeros((B, T_tgt), device=stream_device, dtype=torch.float32)
         revisit_selected_gap_frames = torch.full((B, T_tgt), -1.0, device=stream_device, dtype=torch.float32)
         eval_corrupted_revisit_mask = torch.zeros((B, T_tgt), device=stream_device, dtype=torch.bool)
-        revisit_causal_max = torch.full((B, T_tgt), -1, device=stream_device, dtype=torch.long)
         eval_corruption_enabled = bool(eval_ablation_enabled and eval_ablation_branch in EVAL_CORRUPTION_BRANCHES)
-        revisit_result_diagnostics = []
         projected_revisit_record_cache: dict[tuple[int, str, int, int, int, bool], MemoryRecord] = {}
         if revisit_record_batches is None:
             revisit_record_batches = [tuple(bank.records) for bank in revisit_banks]
@@ -1823,49 +1838,65 @@ class MemoryDiTMixin:
             batch_max_rows = []
             for target_idx in range(T_tgt):
                 target_frame = int(target_frame_indices[target_idx, batch_idx].item())
-                if use_cache_revisit_records:
-                    candidate_records = self._causal_cached_revisit_records(
-                        revisit_record_batches[batch_idx],
-                        target_frame,
-                    )
+                if preselected_revisit_records_by_target is not None:
+                    selected_records = list(preselected_revisit_records_by_target[batch_idx][target_idx])
+                    if preselected_revisit_stats is not None:
+                        revisit_best_selected_fov_overlap[batch_idx, target_idx] = torch.as_tensor(
+                            preselected_revisit_stats["best_selected_fov_overlap"][batch_idx, target_idx],
+                            device=stream_device,
+                            dtype=torch.float32,
+                        )
+                        revisit_best_selected_plucker_overlap[batch_idx, target_idx] = torch.as_tensor(
+                            preselected_revisit_stats["best_selected_plucker_overlap"][batch_idx, target_idx],
+                            device=stream_device,
+                            dtype=torch.float32,
+                        )
+                        revisit_selected_gap_frames[batch_idx, target_idx] = torch.as_tensor(
+                            preselected_revisit_stats["best_selected_gap_frames"][batch_idx, target_idx],
+                            device=stream_device,
+                            dtype=torch.float32,
+                        )
                 else:
-                    candidate_records = revisit_bank.query(
-                        MemoryBankQuery(
-                            target_frame=target_frame,
-                            include_generated=True,
+                    if use_cache_revisit_records:
+                        candidate_records = self._causal_cached_revisit_records(
+                            revisit_record_batches[batch_idx],
+                            target_frame,
                         )
+                    else:
+                        candidate_records = revisit_bank.query(
+                            MemoryBankQuery(
+                                target_frame=target_frame,
+                                include_generated=True,
+                            )
+                        )
+                    result = deterministic_revisit_retrieval(
+                        candidate_records,
+                        target_frame=target_frame,
+                        target_pose=_target_tensor_or_none(target_pose_source, batch_idx, target_idx),
+                        target_summary=None,
+                        topk=revisit_max_frames,
+                        exclude_local_context_frames=revisit_context_window_exclusion_frames,
+                        fov_overlap_threshold=fov_overlap_threshold,
+                        plucker_weight=plucker_weight,
+                        target_video_id=None,
+                        **revisit_retrieval_kwargs,
                     )
-                result = deterministic_revisit_retrieval(
-                    candidate_records,
-                    target_frame=target_frame,
-                    target_pose=_target_tensor_or_none(target_pose_source, batch_idx, target_idx),
-                    target_summary=None,
-                    topk=revisit_max_frames,
-                    exclude_local_context_frames=revisit_context_window_exclusion_frames,
-                    fov_overlap_threshold=fov_overlap_threshold,
-                    plucker_weight=plucker_weight,
-                    target_video_id=_target_video_id_or_none(batch_idx, target_idx),
-                    **revisit_retrieval_kwargs,
-                )
-                selected_records = result.records
-                if use_cache_revisit_records and selected_records:
-                    selected_records = self._project_streaming_revisit_records(
-                        cache=cache,
-                        batch_idx=batch_idx,
-                        records=selected_records,
-                        device=stream_device,
-                        dtype=stream_dtype,
-                        token_patch_size=token_patch_size,
-                        revisit_pool_h=revisit_pool_h,
-                        revisit_pool_w=revisit_pool_w,
-                        projection_cache=projected_revisit_record_cache,
-                    )
-                revisit_result_diagnostics.append(result.diagnostics)
-                revisit_candidate_count[batch_idx, target_idx] = float(result.diagnostics.get("revisit_candidate_frame_count", result.diagnostics.get("revisit_candidate_count", 0)))
-                revisit_selected_count[batch_idx, target_idx] = float(result.diagnostics.get("revisit_selected_frame_count", result.diagnostics.get("revisit_selected_count", 0)))
-                revisit_best_selected_fov_overlap[batch_idx, target_idx] = float(result.diagnostics.get("best_selected_fov_overlap", 0.0))
-                revisit_best_selected_plucker_overlap[batch_idx, target_idx] = float(result.diagnostics.get("best_selected_plucker_overlap", 0.0))
-                revisit_selected_gap_frames[batch_idx, target_idx] = float(result.diagnostics.get("best_selected_gap_frames", -1))
+                    selected_records = result.records
+                    if use_cache_revisit_records and selected_records:
+                        selected_records = self._project_streaming_revisit_records(
+                            cache=cache,
+                            batch_idx=batch_idx,
+                            records=selected_records,
+                            device=stream_device,
+                            dtype=stream_dtype,
+                            token_patch_size=token_patch_size,
+                            revisit_pool_h=revisit_pool_h,
+                            revisit_pool_w=revisit_pool_w,
+                            projection_cache=projected_revisit_record_cache,
+                        )
+                    revisit_best_selected_fov_overlap[batch_idx, target_idx] = result.best_selected_fov_overlap.to(device=stream_device, dtype=torch.float32)
+                    revisit_best_selected_plucker_overlap[batch_idx, target_idx] = result.best_selected_plucker_overlap.to(device=stream_device, dtype=torch.float32)
+                    revisit_selected_gap_frames[batch_idx, target_idx] = result.best_selected_gap_frames.to(device=stream_device, dtype=torch.float32)
                 revisit_bank.assert_causal(target_frame, selected_records)
                 if selected_records:
                     valid_revisit_mask[batch_idx, target_idx] = True
@@ -1876,7 +1907,6 @@ class MemoryDiTMixin:
                     stream_device,
                     stream_dtype,
                 )
-                revisit_causal_max[batch_idx, target_idx] = max_source_frame
                 if eval_corruption_enabled:
                     stream_tokens, was_corrupted = apply_revisit_eval_corruption(
                         tokens=stream_tokens,
@@ -1932,8 +1962,6 @@ class MemoryDiTMixin:
         if not revisit_stage_config_enabled:
             revisit_mask = torch.zeros_like(revisit_mask)
             valid_revisit_mask = torch.zeros_like(valid_revisit_mask)
-            revisit_candidate_count = torch.zeros_like(revisit_candidate_count)
-            revisit_selected_count = torch.zeros_like(revisit_selected_count)
             revisit_best_selected_fov_overlap = torch.zeros_like(revisit_best_selected_fov_overlap)
             revisit_best_selected_plucker_overlap = torch.zeros_like(revisit_best_selected_plucker_overlap)
             revisit_selected_gap_frames = torch.full_like(revisit_selected_gap_frames, -1.0)
@@ -1941,85 +1969,6 @@ class MemoryDiTMixin:
             valid_revisit_eff_mask = torch.zeros_like(valid_revisit_eff_mask)
             revisit_gate_raw = torch.zeros_like(revisit_gate_raw)
             revisit_gate = torch.zeros_like(revisit_gate)
-        no_valid_revisit_mask = (~valid_revisit_mask) if revisit_stage_config_enabled else torch.zeros_like(valid_revisit_mask)
-        revisit_diag = summarize_revisit_diagnostics(revisit_result_diagnostics, valid_revisit_mask)
-        causal_violation_count = 0
-        for source_max in (anchor_max, dynamic_diag.get("max_source_frame"), revisit_causal_max):
-            if source_max is None:
-                continue
-            source_max_t = torch.as_tensor(source_max, device=target_frame_indices.device)
-            valid = source_max_t >= 0
-            if valid.any():
-                causal_violation_count += int((source_max_t[valid] >= target_frame_indices.transpose(0, 1)[valid]).sum().item())
-        diagnostics = {
-            **curriculum_state.diagnostics(),
-            **getattr(self, "_last_dememwm_freeze_diagnostics", {}),
-            **contract_diag,
-            **cache_diag,
-            **preselection_diag,
-            **revisit_diag,
-            "dememwm_stage": gates.stage,
-            "dememwm_gate_reason": gates.reason,
-            "anchor_config_enabled": anchor_config_enabled,
-            "dynamic_config_enabled": dynamic_config_enabled,
-            "revisit_config_enabled": revisit_config_enabled,
-            "anchor_effective_enabled": anchor_effective_enabled,
-            "dynamic_effective_enabled": dynamic_effective_enabled,
-            "revisit_effective_enabled": revisit_effective_enabled,
-            "revisit_stage_config_enabled": revisit_stage_config_enabled,
-            "revisit_gate_raw": revisit_gate_raw.detach(),
-            "revisit_gate_eff": revisit_gate.detach() if torch.is_tensor(revisit_gate) else torch.tensor(float(revisit_gate)),
-            "no_valid_revisit_mask": no_valid_revisit_mask,
-            "valid_revisit_eff_mask": valid_revisit_eff_mask,
-            "revisit_candidate_frame_count_per_target": revisit_candidate_count,
-            "revisit_selected_frame_count_per_target": revisit_selected_count,
-            "revisit_best_selected_fov_overlap_per_target": revisit_best_selected_fov_overlap,
-            "revisit_best_selected_plucker_overlap_per_target": revisit_best_selected_plucker_overlap,
-            "revisit_selected_gap_frames_per_target": revisit_selected_gap_frames,
-            "revisit_learned_gate_mean": float(revisit_gate_raw.detach().float().mean().item()) if revisit_gate_raw.numel() else 0.0,
-            "revisit_effective_gate_mean": float(torch.as_tensor(revisit_gate, device=stream_device).float().mean().item()),
-            **summarize_noise_bucket_diagnostics(
-                noise_bucket=resolved_noise_bucket,
-                valid_revisit_mask=valid_revisit_mask,
-                no_valid_revisit_mask=no_valid_revisit_mask,
-                noise_bucket_ids=noise_bucket_ids,
-            ),
-            **summarize_eval_ablation_diagnostics(
-                enabled=eval_ablation_enabled,
-                branch=eval_ablation_branch,
-                valid_revisit_mask=valid_revisit_mask,
-                no_valid_revisit_mask=no_valid_revisit_mask,
-                eval_corrupted_revisit_mask=eval_corrupted_revisit_mask if eval_corruption_enabled else None,
-            ),
-            "token_patch_size": token_patch_size,
-            "tokens_per_frame": tokens_per_frame,
-            "anchor_token_slots": int(anchor_tokens.shape[-2]),
-            "anchor_target_slots": anchor_num_tokens,
-            "anchor_pool_h": anchor_pool_h,
-            "anchor_pool_w": anchor_pool_w,
-            "dynamic_token_slots": int(dynamic_tokens.shape[-2]),
-            "dynamic_target_slots": int(dynamic_tokens.shape[-2]),
-            "dynamic_min_gap_to_target": dynamic_min_gap_to_target,
-            "dynamic_max_gap_to_target": dynamic_max_gap_to_target,
-            "dynamic_exclude_latest_local_frames": dynamic_recent_exclusion_frames,
-            "revisit_token_slots": int(revisit_tokens.shape[-2]),
-            "revisit_target_slots": revisit_target_slots,
-            "revisit_local_context_exclusion_frames": revisit_context_window_exclusion_frames,
-            "revisit_pool_h": revisit_pool_h,
-            "revisit_pool_w": revisit_pool_w,
-            "revisit_max_frames": revisit_max_frames,
-            "anchor_valid_tokens_per_target_max": int(anchor_mask.sum(dim=-1).max().item()) if anchor_mask.numel() else 0,
-            "dynamic_valid_tokens_per_target_max": int(dynamic_mask.sum(dim=-1).max().item()) if dynamic_mask.numel() else 0,
-            "revisit_valid_tokens_per_target_max": int(revisit_mask.sum(dim=-1).max().item()) if revisit_mask.numel() else 0,
-            "causal_violation_count": causal_violation_count,
-            "anchor_max_source_frame": anchor_max,
-            "dynamic_max_source_frame": dynamic_diag.get("max_source_frame"),
-            "revisit_max_source_frame": revisit_max,
-            "dynamic_generated_source_fraction": dynamic_diag.get("generated_source_fraction"),
-        }
-        if eval_corruption_enabled:
-            diagnostics["eval_corrupted_revisit_mask"] = eval_corrupted_revisit_mask
-
         return MemoryStreamTensors(
             anchor_tokens=anchor_tokens,
             anchor_mask=anchor_mask,
@@ -2032,19 +1981,18 @@ class MemoryDiTMixin:
             revisit_gate=revisit_gate,
             revisit_gate_raw=revisit_gate_raw,
             valid_revisit_mask=valid_revisit_mask,
-            no_valid_revisit_mask=no_valid_revisit_mask,
-            diagnostics=diagnostics,
+            revisit_best_selected_fov_overlap=revisit_best_selected_fov_overlap,
+            revisit_best_selected_plucker_overlap=revisit_best_selected_plucker_overlap,
+            revisit_selected_gap_frames=revisit_selected_gap_frames,
         )
 
     def _refresh_stream_gates(
         self,
         streams: MemoryStreamTensors,
         denoising_fraction: float | None = None,
-        noise_bucket: str | None = None,
     ) -> MemoryStreamTensors:
         gate_state = self._effective_gate_state(
             denoising_fraction=denoising_fraction,
-            noise_bucket=noise_bucket,
         )
         gates = gate_state["gates"]
         device = streams.anchor_tokens.device
@@ -2056,20 +2004,17 @@ class MemoryDiTMixin:
         else:
             valid_revisit_mask = valid_revisit_mask.to(device=device, dtype=torch.bool)
 
-        diagnostics = dict(streams.diagnostics)
-
-        def _diagnostic_tensor(name: str, fill_value: float = 0.0) -> torch.Tensor:
-            value = diagnostics.get(name)
+        def _gate_feature(value: torch.Tensor | None, fill_value: float = 0.0) -> torch.Tensor:
             if value is None:
                 return torch.full((B, T_tgt), float(fill_value), device=device, dtype=torch.float32)
-            tensor = torch.as_tensor(value, device=device, dtype=torch.float32)
+            tensor = value.to(device=device, dtype=torch.float32)
             if tensor.ndim == 0:
                 return torch.full((B, T_tgt), float(tensor.item()), device=device, dtype=torch.float32)
             return tensor.expand((B, T_tgt))
 
-        revisit_best_selected_fov_overlap = _diagnostic_tensor("revisit_best_selected_fov_overlap_per_target")
-        revisit_best_selected_plucker_overlap = _diagnostic_tensor("revisit_best_selected_plucker_overlap_per_target")
-        revisit_selected_gap_frames = _diagnostic_tensor("revisit_selected_gap_frames_per_target", -1.0)
+        revisit_best_selected_fov_overlap = _gate_feature(streams.revisit_best_selected_fov_overlap)
+        revisit_best_selected_plucker_overlap = _gate_feature(streams.revisit_best_selected_plucker_overlap)
+        revisit_selected_gap_frames = _gate_feature(streams.revisit_selected_gap_frames, -1.0)
 
         anchor_effective_enabled = gate_state["anchor_effective_enabled"]
         dynamic_effective_enabled = gate_state["dynamic_effective_enabled"]
@@ -2086,53 +2031,16 @@ class MemoryDiTMixin:
                 best_selected_plucker_overlap=revisit_best_selected_plucker_overlap,
                 selected_gap_frames=revisit_selected_gap_frames,
             ).to(device=device, dtype=dtype)
-        valid_revisit_eff_mask = valid_revisit_mask
         if not revisit_stage_config_enabled or gate_state["force_revisit_off"]:
             revisit_gate = torch.zeros_like(revisit_gate_raw)
         elif gate_state["force_revisit_on"]:
-            revisit_gate = valid_revisit_eff_mask.to(device=device, dtype=dtype) * torch.ones_like(revisit_gate_raw)
+            revisit_gate = valid_revisit_mask.to(device=device, dtype=dtype) * torch.ones_like(revisit_gate_raw)
         else:
-            revisit_gate = valid_revisit_eff_mask.to(device=device, dtype=dtype) * revisit_gate_raw * float(gates.revisit_gate)
+            revisit_gate = valid_revisit_mask.to(device=device, dtype=dtype) * revisit_gate_raw * float(gates.revisit_gate)
         if not revisit_stage_config_enabled:
             valid_revisit_mask = torch.zeros_like(valid_revisit_mask)
-            valid_revisit_eff_mask = torch.zeros_like(valid_revisit_eff_mask)
             revisit_gate_raw = torch.zeros_like(revisit_gate_raw)
             revisit_gate = torch.zeros_like(revisit_gate)
-        no_valid_revisit_mask = (~valid_revisit_mask) if revisit_stage_config_enabled else torch.zeros_like(valid_revisit_mask)
-        eval_corrupted_revisit_mask = diagnostics.get("eval_corrupted_revisit_mask")
-        if eval_corrupted_revisit_mask is not None:
-            eval_corrupted_revisit_mask = torch.as_tensor(eval_corrupted_revisit_mask, device=device, dtype=torch.bool)
-        revisit_effective_enabled = bool(revisit_stage_config_enabled and (revisit_gate > 0).any().item())
-        diagnostics.update(gate_state["curriculum_state"].diagnostics())
-        diagnostics.update({
-            "dememwm_stage": gates.stage,
-            "dememwm_gate_reason": gates.reason,
-            "anchor_config_enabled": gate_state["anchor_config_enabled"],
-            "dynamic_config_enabled": gate_state["dynamic_config_enabled"],
-            "revisit_config_enabled": gate_state["revisit_config_enabled"],
-            "anchor_effective_enabled": anchor_effective_enabled,
-            "dynamic_effective_enabled": dynamic_effective_enabled,
-            "revisit_effective_enabled": revisit_effective_enabled,
-            "revisit_stage_config_enabled": revisit_stage_config_enabled,
-            "revisit_gate_raw": revisit_gate_raw.detach(),
-            "revisit_gate_eff": revisit_gate.detach() if torch.is_tensor(revisit_gate) else torch.tensor(float(revisit_gate)),
-            "no_valid_revisit_mask": no_valid_revisit_mask,
-            "valid_revisit_eff_mask": valid_revisit_eff_mask,
-            "revisit_learned_gate_mean": float(revisit_gate_raw.detach().float().mean().item()) if revisit_gate_raw.numel() else 0.0,
-            "revisit_effective_gate_mean": float(revisit_gate.detach().float().mean().item()) if revisit_gate.numel() else 0.0,
-        })
-        diagnostics.update(summarize_noise_bucket_diagnostics(
-            noise_bucket=gate_state["resolved_noise_bucket"],
-            valid_revisit_mask=valid_revisit_mask,
-            no_valid_revisit_mask=no_valid_revisit_mask,
-        ))
-        diagnostics.update(summarize_eval_ablation_diagnostics(
-            enabled=gate_state["eval_ablation_enabled"],
-            branch=gate_state["eval_ablation_branch"],
-            valid_revisit_mask=valid_revisit_mask,
-            no_valid_revisit_mask=no_valid_revisit_mask,
-            eval_corrupted_revisit_mask=eval_corrupted_revisit_mask,
-        ))
         return replace(
             streams,
             anchor_gate=anchor_gate,
@@ -2140,84 +2048,90 @@ class MemoryDiTMixin:
             revisit_gate=revisit_gate,
             revisit_gate_raw=revisit_gate_raw,
             valid_revisit_mask=valid_revisit_mask,
-            no_valid_revisit_mask=no_valid_revisit_mask,
-            diagnostics=diagnostics,
+            revisit_best_selected_fov_overlap=revisit_best_selected_fov_overlap,
+            revisit_best_selected_plucker_overlap=revisit_best_selected_plucker_overlap,
+            revisit_selected_gap_frames=revisit_selected_gap_frames,
         )
 
-    def _streams_to_kwargs(self, streams: MemoryStreamTensors) -> tuple[dict, dict]:
-        memory_kwargs, diagnostics = self.dememwm_injection_adapter(streams, device=streams.anchor_tokens.device, dtype=streams.anchor_tokens.dtype)
-        return memory_kwargs, diagnostics
+    def _training_window_bounds(self, total_frames: int, device: torch.device) -> tuple[int, int]:
+        total_frames = int(total_frames)
+        window = int(getattr(self, "n_tokens", total_frames) or total_frames)
+        if total_frames <= 0:
+            return 0, 0
+        if window <= 0 or total_frames <= window:
+            return 0, total_frames
+        max_start = max(0, total_frames - window)
+        min_start = min(max(0, self._context_frame_count()), max_start)
+        if max_start <= min_start:
+            start = min_start
+        else:
+            start = int(torch.randint(min_start, max_start + 1, (1,), device=device).item())
+        return start, start + window
+
+    def _streams_to_kwargs(self, streams: MemoryStreamTensors) -> dict:
+        return self.dememwm_injection_adapter(
+            streams,
+            device=streams.anchor_tokens.device,
+            dtype=streams.anchor_tokens.dtype,
+        )
 
-    def build_memory_kwargs(self, *args, **kwargs) -> tuple[dict, dict]:
+    def build_memory_kwargs(self, *args, **kwargs) -> dict:
         streams = self.build_memory_streams(*args, **kwargs)
         return self._streams_to_kwargs(streams)
 
-    def _memory_adapter_delta_diagnostics(self) -> dict:
-        dit_model = getattr(getattr(self, "diffusion_model", None), "model", None)
-        diagnostics_fn = getattr(dit_model, "memory_adapter_delta_diagnostics", None)
-        if diagnostics_fn is None:
-            return {}
-        return diagnostics_fn()
-
-    def _log_memory_diagnostics(self, namespace: str, diagnostics: dict) -> None:
-        if namespace == "training/dememwm":
-            allowed_keys = self._TRAIN_DIAGNOSTIC_LOG_KEYS
-        elif namespace.endswith("/dememwm"):
-            allowed_keys = self._VALIDATION_DIAGNOSTIC_LOG_KEYS
-        else:
-            allowed_keys = None
-        for key, value in diagnostics.items():
-            if allowed_keys is not None and key not in allowed_keys:
-                continue
-            if isinstance(value, str) or value is None:
-                continue
-            if torch.is_tensor(value):
-                if value.numel() > 0:
-                    self.log(f"{namespace}/{key}", value.float().mean().item(), prog_bar=False, sync_dist=True)
-            elif isinstance(value, (bool, int, float)):
-                self.log(f"{namespace}/{key}", float(value), prog_bar=False, sync_dist=True)
-
-    def _training_pose_condition(self, xs, pose_conditions, c2w_mat, frame_idx):
-        from ..df_video import convert_to_plucker
-        image_height, image_width = self._image_size(xs)
+    def _training_pose_condition(
+        self,
+        pose_conditions: torch.Tensor,
+        c2w_mat: torch.Tensor,
+        frame_idx: torch.Tensor,
+        *,
+        dtype: torch.dtype,
+        image_width: int,
+        image_height: int,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         if self.use_plucker:
             if self.relative_embedding:
+                memory_condition_length = max(0, int(self.memory_condition_length or 0))
+                if memory_condition_length:
+                    memory_c2w = c2w_mat[-memory_condition_length:]
+                    memory_frame_idx = frame_idx[-memory_condition_length:]
+                else:
+                    memory_c2w = c2w_mat[:0]
+                    memory_frame_idx = frame_idx[:0]
                 input_pose_condition = []
                 frame_idx_list = []
-                ref_c2w = c2w_mat[-self.memory_condition_length:] if self.memory_condition_length else c2w_mat[:0]
-                ref_idx = frame_idx[-self.memory_condition_length:] if self.memory_condition_length else frame_idx[:0]
-                for i in range(c2w_mat.shape[0]):
+                for target_idx in range(c2w_mat.shape[0]):
                     input_pose_condition.append(
                         convert_to_plucker(
-                            torch.cat([c2w_mat[i:i + 1], ref_c2w]).clone(),
+                            torch.cat([c2w_mat[target_idx:target_idx + 1], memory_c2w]).clone(),
                             0,
                             focal_length=self.focal_length,
-                            image_height=image_height, image_width=image_width
-                        ).to(xs.dtype)
+                            image_width=image_width,
+                            image_height=image_height,
+                        ).to(dtype)
+                    )
+                    frame_idx_list.append(
+                        torch.cat([
+                            frame_idx[target_idx:target_idx + 1] - frame_idx[target_idx:target_idx + 1],
+                            memory_frame_idx - frame_idx[target_idx:target_idx + 1],
+                        ]).clone()
                     )
-                    frame_idx_list.append(torch.cat([frame_idx[i:i + 1] - frame_idx[i:i + 1], ref_idx - frame_idx[i:i + 1]]).clone())
                 return torch.cat(input_pose_condition), torch.cat(frame_idx_list)
-            return convert_to_plucker(
-                c2w_mat, 0, focal_length=self.focal_length,
-                image_height=image_height, image_width=image_width
-            ).to(xs.dtype), frame_idx
-        return pose_conditions.to(xs.dtype), None
-
-    def _training_window_bounds(self, total_frames: int, device: torch.device) -> tuple[int, int]:
-        total_frames = max(0, int(total_frames))
-        n_tokens = max(1, min(int(self.n_tokens), total_frames))
-        max_start = max(0, total_frames - n_tokens)
-        if max_start == 0:
-            return 0, n_tokens
-        context_start = self._context_frame_count()
-        min_start = min(context_start, max_start)
-        if min_start == max_start:
-            return min_start, min_start + n_tokens
-        start = int(torch.randint(min_start, max_start + 1, (1,), device=device).item())
-        return start, start + n_tokens
+            return (
+                convert_to_plucker(
+                    c2w_mat,
+                    0,
+                    focal_length=self.focal_length,
+                    image_width=image_width,
+                    image_height=image_height,
+                ).to(dtype),
+                frame_idx,
+            )
+        return pose_conditions.to(dtype), None
 
     def training_step(self, batch, batch_idx):
         xs, conditions, pose_conditions, c2w_mat, frame_idx = self._preprocess_batch(batch)
+        image_height, image_width = self._image_size(xs)
         xs = self._as_latents(xs)
 
         # Randomly select a contiguous n_tokens denoising window inside the long
@@ -2231,7 +2145,12 @@ class MemoryDiTMixin:
         frame_idx_window = frame_idx[start:end]
 
         input_pose_condition, frame_idx_list = self._training_pose_condition(
-            xs_window, pose_conditions[start:end], c2w_mat[start:end], frame_idx_window
+            pose_conditions[start:end],
+            c2w_mat[start:end],
+            frame_idx_window,
+            dtype=xs_window.dtype,
+            image_width=image_width,
+            image_height=image_height,
         )
 
         noise_levels = self._generate_noise_levels(xs_window)
@@ -2239,17 +2158,15 @@ class MemoryDiTMixin:
             noise_levels[-self.memory_condition_length:] = self.diffusion_model.stabilization_level
             conditions_window[-self.memory_condition_length:] *= 0
         source_is_generated = torch.zeros(frame_idx.shape, device=frame_idx.device, dtype=torch.bool)
-        memory_source_latents, source_is_generated, proxy_diagnostics = self._apply_generated_history_proxy(
+        memory_source_latents, source_is_generated = self._apply_generated_history_proxy(
             xs,
             source_is_generated,
             context_frame_count=self._context_frame_count(),
             target_start_frame=start,
         )
         timesteps = int(getattr(self, "timesteps", 0) or 0)
-        training_noise_bucket = noise_bucket_from_noise_levels(noise_levels, timesteps)
-        training_noise_bucket_ids = noise_bucket_ids_from_noise_levels(noise_levels, timesteps)
         training_denoising_fraction = denoising_fraction_from_noise_levels(noise_levels, timesteps)
-        memory_kwargs, diagnostics = self.build_memory_kwargs(
+        memory_kwargs = self.build_memory_kwargs(
             memory_source_latents,
             frame_idx,
             target_frame_indices=frame_idx_window,
@@ -2259,10 +2176,7 @@ class MemoryDiTMixin:
             target_action=conditions_window,
             source_is_generated=source_is_generated,
             denoising_fraction=training_denoising_fraction,
-            noise_bucket=training_noise_bucket,
-            noise_bucket_ids=None if training_noise_bucket_ids is None else training_noise_bucket_ids.transpose(0, 1),
         )
-        diagnostics.update(proxy_diagnostics)
         _, loss = self.diffusion_model(
             xs_window,
             conditions_window,
@@ -2272,19 +2186,19 @@ class MemoryDiTMixin:
             frame_idx=frame_idx_list,
             **memory_kwargs,
         )
-        diagnostics.update(self._memory_adapter_delta_diagnostics())
         if self.memory_condition_length:
             loss = loss[:-self.memory_condition_length]
         loss_denoise = self.reweight_loss(loss, None)
         loss_total = loss_denoise
-        diagnostics["training_window_start"] = int(start)
-        diagnostics["training_window_end"] = int(end)
-        diagnostics["training_window_size"] = int(end - start)
-        diagnostics["loss_denoise"] = float(loss_denoise.detach().item())
-        diagnostics["loss_total"] = float(loss_total.detach().item())
         if batch_idx % 20 == 0:
-            self.log("training/loss", loss_total.detach().cpu())
-            self._log_memory_diagnostics("training/dememwm", diagnostics)
+            revisit_gate = memory_kwargs.get("memory_retrieval_gate")
+            if torch.is_tensor(revisit_gate):
+                revisit_gate_value = revisit_gate.detach().float().mean()
+            else:
+                revisit_gate_value = loss_total.detach().new_tensor(0.0 if revisit_gate is None else float(revisit_gate))
+            self.log("training/loss", loss_total.detach(), prog_bar=True, sync_dist=True)
+            self.log("training/denoise_loss", loss_denoise.detach(), prog_bar=False, sync_dist=True)
+            self.log("training/revisit_gate", revisit_gate_value, prog_bar=False, sync_dist=True)
         return {"loss": loss_total}
 
     def validation_step(self, batch, batch_idx, namespace="validation"):
@@ -2308,7 +2222,6 @@ class MemoryDiTMixin:
         streaming_cache = self._new_streaming_cache(video_id=f"{namespace}:{batch_idx}")
         cached_until = 0
         pbar = tqdm(total=n_frames, initial=curr_frame, desc="Sampling")
-        last_diagnostics = None
         while curr_frame < n_frames:
             if streaming_cache is not None and curr_frame > cached_until:
                 new_generated = torch.zeros(frame_idx[cached_until:curr_frame].shape, dtype=torch.bool, device=frame_idx.device)
@@ -2375,7 +2288,7 @@ class MemoryDiTMixin:
                 from_noise_levels, to_noise_levels = self._prepare_noise_levels(scheduling_matrix, m, curr_frame, batch_size, memory_condition_length)
                 denoise_frac = float(m + 1) / max(float(scheduling_matrix.shape[0] - 1), 1.0)
                 step_streams = self._refresh_stream_gates(memory_streams, denoising_fraction=denoise_frac)
-                memory_kwargs, last_diagnostics = self._streams_to_kwargs(step_streams)
+                memory_kwargs = self._streams_to_kwargs(step_streams)
                 xs_pred[start_frame:] = self.diffusion_model.sample_step(
                     xs_pred[start_frame:].to(input_condition.device),
                     input_condition,
@@ -2405,12 +2318,8 @@ class MemoryDiTMixin:
                     action=conditions[cached_until:curr_frame],
                 )
                 cached_until = curr_frame
-                if last_diagnostics is not None:
-                    last_diagnostics.update(streaming_cache.diagnostics("cache"))
             pbar.update(horizon)
         pbar.close()
-        if last_diagnostics is not None:
-            self._log_memory_diagnostics(f"{namespace}/dememwm", last_diagnostics)
         xs_pred = self.decode(xs_pred[n_context_frames:].to(conditions.device))
         xs_decode = self.decode(xs[n_context_frames:].to(conditions.device))
         self.validation_step_outputs.append((xs_pred.detach().cpu(), xs_decode.detach().cpu()))
@@ -2440,10 +2349,7 @@ class MemoryDiTMixin:
     # Compatibility aliases for old DeMemWM test and experiment call sites.
     dememwm_strict_key_prefixes = strict_key_prefixes
     dememwm_strict_key_substrings = strict_key_substrings
-    _DEMEMWM_TRAIN_DIAGNOSTIC_LOG_KEYS = _TRAIN_DIAGNOSTIC_LOG_KEYS
-    _DEMEMWM_VALIDATION_DIAGNOSTIC_LOG_KEYS = _VALIDATION_DIAGNOSTIC_LOG_KEYS
     _dememwm_cfg = _memory_cfg
-    _dememwm_stage_policy_cfg = _stage_policy_cfg
     _dememwm_eval_ablation_cfg = _eval_ablation_cfg
     _dememwm_generated_history_proxy_cfg = _generated_history_proxy_cfg
     _dememwm_eval_ablation_state = _eval_ablation_state
@@ -2476,8 +2382,6 @@ class MemoryDiTMixin:
     _dememwm_refresh_stream_gates = _refresh_stream_gates
     _dememwm_streams_to_kwargs = _streams_to_kwargs
     build_dememwm_memory_kwargs = build_memory_kwargs
-    _dememwm_memory_adapter_delta_diagnostics = _memory_adapter_delta_diagnostics
-    _log_dememwm_diagnostics = _log_memory_diagnostics
     _dememwm_training_window_bounds = _training_window_bounds
     strict_dememwm_checkpoint_key_check = strict_checkpoint_key_check
 

algorithms/worldmem/dememwm/cache.py

@@ -486,23 +486,6 @@ class StreamingCache:
             ).to(device=device)
         return latents, frame_indices, generated, pose
 
-    def diagnostics(self, prefix: str = "cache") -> dict[str, Any]:
-        return {
-            f"{prefix}_enabled": bool(self.enabled),
-            f"{prefix}_records": int(self.record_count),
-            f"{prefix}_anchor_records": int(self.records_count("anchor")),
-            f"{prefix}_revisit_records": int(self.records_count("revisit")),
-            f"{prefix}_slots": int(self.slot_count),
-            f"{prefix}_raw_frame_slots": int(self.raw_frame_slots),
-            f"{prefix}_raw_segments": int(self.raw_segment_count),
-            f"{prefix}_evictions": int(self.evictions),
-            f"{prefix}_resets": int(self.reset_count),
-            f"{prefix}_capacity_exceeded": int(self.capacity_exceeded_count),
-            f"{prefix}_device": self.device,
-            f"{prefix}_current_video_id": self.current_video_id,
-            f"{prefix}_clear_between_videos": bool(self.clear_between_videos),
-            f"{prefix}_no_evict": bool(self.no_evict),
-        }
 
 
 DeMemWMStreamingCache = StreamingCache

algorithms/worldmem/dememwm/compression.py

@@ -163,7 +163,8 @@ class CausalConv3DDynamicCompressor(nn.Module):
         target_frame_indices: torch.Tensor,
         source_is_generated: Optional[torch.Tensor] = None,
         exclude_latest_local_frames: Optional[int] = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, dict]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        del pose, source_is_generated
         if latents.ndim != 5:
             raise ValueError("latents must have shape (T_src,B,C,H,W)")
         exclude_latest_local_frames = (
@@ -175,86 +176,78 @@ class CausalConv3DDynamicCompressor(nn.Module):
         p = self.patch_size
         if H % p != 0 or W % p != 0:
             raise ValueError(f"latent H,W=({H},{W}) must be divisible by patch_size={p}")
+        if frame_indices.shape != (T_src, B):
+            raise ValueError("frame_indices must have shape (T_src,B)")
         if target_frame_indices.ndim == 1:
             target_frame_indices = target_frame_indices[:, None].expand(-1, B)
-        T_tgt = target_frame_indices.shape[0]
+        if target_frame_indices.ndim != 2 or target_frame_indices.shape[1] != B:
+            raise ValueError("target_frame_indices must have shape (T_tgt,B)")
+
         device = latents.device
-        generated_flags = None if source_is_generated is None else source_is_generated.to(device=device, dtype=torch.bool)
+        frame_indices = frame_indices.to(device=device)
+        target_frame_indices = target_frame_indices.to(device=device)
+        T_tgt = target_frame_indices.shape[0]
         n_spatial = (H // p) * (W // p)
         T_out = self._temporal_output_count()
         num_slots = T_out * n_spatial
         output_time_idx = self._output_time_indices(device)
-        selected_source_count = torch.zeros((B, T_tgt), dtype=torch.long, device=device)
-        max_source_frame = torch.full((B, T_tgt), -1, dtype=torch.long, device=device)
-        generated_source_fraction = torch.zeros((B, T_tgt), dtype=torch.float32, device=device)
-        min_gap = torch.full((B, T_tgt), -1, dtype=torch.long, device=device)
-        max_gap = torch.full((B, T_tgt), -1, dtype=torch.long, device=device)
-        output_rows, mask_rows = [], []
-        for b in range(B):
-            src_frames_b = frame_indices[:, b]
-            tgt_outputs, tgt_masks = [], []
-            for j in range(T_tgt):
-                target = int(target_frame_indices[j, b].item())
-                valid_idx = (
-                    src_frames_b < target - exclude_latest_local_frames
-                ).nonzero(as_tuple=False).flatten()
-                if valid_idx.numel() == 0:
-                    tgt_outputs.append(latents.new_zeros(num_slots, self.dit_hidden_size))
-                    tgt_masks.append(torch.zeros(num_slots, device=device, dtype=torch.bool))
-                    continue
-                selected_frames = src_frames_b.index_select(0, valid_idx)
-                order = torch.argsort(selected_frames)
-                valid_idx = valid_idx.index_select(0, order)[-self.max_source_frames:]
-                selected_frames = src_frames_b.index_select(0, valid_idx)
-                selected_source_count[b, j] = int(selected_frames.numel())
-                max_source_frame[b, j] = selected_frames.max()
-                gaps = target - selected_frames
-                min_gap[b, j] = gaps.min()
-                max_gap[b, j] = gaps.max()
-                if generated_flags is not None:
-                    generated = generated_flags.index_select(0, valid_idx)[:, b]
-                    generated_source_fraction[b, j] = generated.float().mean()
-                chunk = latents[valid_idx, b]
-                real_mask = torch.ones((chunk.shape[0],), device=device, dtype=torch.bool)
-                if chunk.shape[0] < self.max_source_frames:
-                    pad = chunk.new_zeros(self.max_source_frames - chunk.shape[0], C, H, W)
-                    chunk = torch.cat([pad, chunk], dim=0)
-                    real_mask = torch.cat([
-                        torch.zeros((pad.shape[0],), device=device, dtype=torch.bool),
-                        real_mask,
-                    ])
-                inp = chunk.clone()
-                inp[1:] = chunk[1:] - chunk[:-1]
-                x = inp.permute(1, 0, 2, 3).unsqueeze(0)        # (1,C,T,H,W)
-                x = F.pad(x, (0, 0, 0, 0, self.causal_pad, 0))  # left-pad time
-                x = self.conv3d(x)                               # (1,D,T_out,H//p,W//p)
-                x = x.squeeze(0).permute(1, 2, 3, 0)            # (T_out,H//p,W//p,D)
-                x = self.out_norm(x)
-                tokens = x.reshape(num_slots, self.dit_hidden_size)
-                clamped_time_idx = output_time_idx.clamp(min=0, max=self.max_source_frames - 1)
-                temporal_mask = (
-                    (output_time_idx >= 0)
-                    & (output_time_idx < self.max_source_frames)
-                    & real_mask.index_select(0, clamped_time_idx)
-                )
-                mask = temporal_mask[:, None].expand(T_out, n_spatial).reshape(num_slots)
-                tgt_outputs.append(tokens)
-                tgt_masks.append(mask)
-            output_rows.append(torch.stack(tgt_outputs))
-            mask_rows.append(torch.stack(tgt_masks))
-        out_tokens = torch.stack(output_rows)
-        out_mask = torch.stack(mask_rows)
-        diagnostics = {
-            "num_dynamic_slots": num_slots,
-            "dynamic_T_out": T_out,
-            "dynamic_n_spatial": n_spatial,
-            "dynamic_temporal_left_pad": self.causal_pad,
-            "dynamic_output_time_indices": output_time_idx,
-            "selected_source_count": selected_source_count,
-            "max_source_frame": max_source_frame,
-            "generated_source_fraction": generated_source_fraction,
-            "dynamic_min_gap_to_target_per_target": min_gap,
-            "dynamic_max_gap_to_target_per_target": max_gap,
-            "dynamic_exclude_latest_local_frames": exclude_latest_local_frames,
-        }
-        return out_tokens, out_mask, diagnostics
+        if T_src == 0:
+            out_tokens = latents.new_zeros((B, T_tgt, num_slots, self.dit_hidden_size))
+            out_mask = torch.zeros((B, T_tgt, num_slots), device=device, dtype=torch.bool)
+            return out_tokens, out_mask
+
+        source_frames = frame_indices.transpose(0, 1).contiguous()
+        target_frames = target_frame_indices.transpose(0, 1).contiguous()
+        valid = source_frames[:, None, :] < (target_frames[:, :, None] - int(exclude_latest_local_frames))
+        valid_flat = valid.reshape(B * T_tgt, T_src)
+        source_frames_flat = source_frames[:, None, :].expand(B, T_tgt, T_src).reshape(B * T_tgt, T_src)
+
+        topk = min(int(self.max_source_frames), T_src)
+        rank = source_frames_flat.to(dtype=torch.float64).masked_fill(~valid_flat, -float("inf"))
+        top = torch.topk(rank, k=topk, dim=1, largest=True, sorted=True)
+        selected_idx = top.indices.flip(dims=(1,))
+        selected_valid = torch.isfinite(top.values).flip(dims=(1,))
+        if topk < self.max_source_frames:
+            pad_count = self.max_source_frames - topk
+            selected_idx = torch.cat([
+                torch.zeros((B * T_tgt, pad_count), device=device, dtype=torch.long),
+                selected_idx,
+            ], dim=1)
+            selected_valid = torch.cat([
+                torch.zeros((B * T_tgt, pad_count), device=device, dtype=torch.bool),
+                selected_valid,
+            ], dim=1)
+
+        selected_idx_clamped = selected_idx.to(device=device, dtype=torch.long).clamp(min=0, max=max(0, T_src - 1))
+        has_valid = selected_valid.any(dim=1)
+        batch_ids = torch.arange(B, device=device, dtype=torch.long).repeat_interleave(T_tgt)
+        latents_by_batch = latents.permute(1, 0, 2, 3, 4).contiguous()
+        latents_per_query = latents_by_batch.index_select(0, batch_ids)
+        gather_idx = selected_idx_clamped.reshape(B * T_tgt, self.max_source_frames, 1, 1, 1).expand(
+            -1, -1, C, H, W
+        )
+        chunk = torch.gather(latents_per_query, 1, gather_idx)
+        chunk = torch.where(
+            selected_valid[:, :, None, None, None],
+            chunk,
+            torch.zeros((), device=device, dtype=latents.dtype),
+        )
+
+        inp = chunk.clone()
+        inp[:, 1:] = chunk[:, 1:] - chunk[:, :-1]
+        x = inp.permute(0, 2, 1, 3, 4)
+        x = F.pad(x, (0, 0, 0, 0, self.causal_pad, 0))
+        x = self.conv3d(x)
+        x = self.out_norm(x.permute(0, 2, 3, 4, 1))
+        tokens_flat = x.reshape(B * T_tgt, num_slots, self.dit_hidden_size)
+        tokens_flat = torch.where(has_valid[:, None, None], tokens_flat, torch.zeros_like(tokens_flat))
+        out_tokens = tokens_flat.reshape(B, T_tgt, num_slots, self.dit_hidden_size)
+
+        clamped_time_idx = output_time_idx.clamp(min=0, max=self.max_source_frames - 1)
+        temporal_mask = (
+            (output_time_idx >= 0)
+            & (output_time_idx < self.max_source_frames)
+            & selected_valid.index_select(1, clamped_time_idx)
+        )
+        out_mask = temporal_mask[:, :, None].expand(B * T_tgt, T_out, n_spatial).reshape(B, T_tgt, num_slots)
+        return out_tokens, out_mask

algorithms/worldmem/dememwm/diagnostics.py

@@ -1,172 +0,0 @@
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-
-from .schedules import EVAL_ABLATION_BRANCH_TO_ID, NOISE_BUCKETS, NOISE_BUCKET_TO_ID, normalize_eval_ablation_branch, normalize_noise_bucket
-
-
-_REVISIT_LABEL_SOURCE = "deterministic_fov_coverage_plucker"
-
-
-def tensor_valid_fraction(mask: torch.Tensor | None) -> float:
-    if mask is None or mask.numel() == 0:
-        return 0.0
-    return float(mask.detach().bool().float().mean().item())
-
-
-def gate_stats(gate: torch.Tensor | float | int | None) -> dict[str, float]:
-    if gate is None:
-        return {"mean": 0.0, "min": 0.0, "max": 0.0}
-    if not torch.is_tensor(gate):
-        value = float(gate)
-        return {"mean": value, "min": value, "max": value}
-    g = gate.detach().float()
-    return {"mean": float(g.mean().item()), "min": float(g.min().item()), "max": float(g.max().item())}
-
-
-def summarize_stream(name: str, tokens: torch.Tensor | None, mask: torch.Tensor | None, gate: torch.Tensor | float | None) -> dict[str, Any]:
-    return {f"{name}_tokens_shape": None if tokens is None else tuple(tokens.shape), f"{name}_valid_fraction": tensor_valid_fraction(mask), f"{name}_valid_tokens": 0 if mask is None else int(mask.detach().bool().sum().item()), f"{name}_gate": gate_stats(gate)}
-
-
-def assert_no_future_sources(target_frame: int, max_source_frame: int | torch.Tensor) -> None:
-    max_src = int(max_source_frame.detach().max().item()) if torch.is_tensor(max_source_frame) else int(max_source_frame)
-    if max_src >= int(target_frame):
-        raise AssertionError(f"DeMemWM memory source {max_src} is not causal for target {target_frame}")
-
-
-def _collect_values(result_diagnostics: list[dict[str, Any]], key: str) -> list[float]:
-    values: list[float] = []
-    for diag in result_diagnostics:
-        for value in diag.get(key, []) or []:
-            values.append(float(value))
-    return values
-
-
-def _value_stats(values: list[float], prefix: str) -> dict[str, float]:
-    if not values:
-        return {f"{prefix}_mean": 0.0, f"{prefix}_min": 0.0, f"{prefix}_max": 0.0}
-    return {
-        f"{prefix}_mean": float(sum(values) / len(values)),
-        f"{prefix}_min": float(min(values)),
-        f"{prefix}_max": float(max(values)),
-    }
-
-
-def summarize_revisit_diagnostics(result_diagnostics: list[dict[str, Any]], valid_revisit_mask: torch.Tensor | None) -> dict[str, Any]:
-    target_count = len(result_diagnostics)
-    candidate_count = sum(int(diag.get("revisit_candidate_frame_count", diag.get("revisit_candidate_count", diag.get("candidate_count", 0)))) for diag in result_diagnostics)
-    candidate_count_mean = float(candidate_count / target_count) if target_count else 0.0
-    valid_candidate_label_count = sum(int(diag.get("valid_candidate_label_count", diag.get("valid_candidate_count", 0))) for diag in result_diagnostics)
-    pose_preselect_input_count = sum(int(diag.get("revisit_pose_preselect_input_count", 0)) for diag in result_diagnostics)
-    pose_preselect_selected_count = sum(int(diag.get("revisit_pose_preselect_selected_count", 0)) for diag in result_diagnostics)
-    exact_fov_candidate_count = sum(int(diag.get("revisit_exact_fov_candidate_count", 0)) for diag in result_diagnostics)
-    valid_count = sum(int(diag.get("valid_revisit_frame_count", diag.get("valid_revisit_count", diag.get("valid_candidate_count", 0)))) for diag in result_diagnostics)
-    valid_count_mean = float(valid_count / target_count) if target_count else 0.0
-    selected_count = sum(int(diag.get("revisit_selected_frame_count", diag.get("revisit_selected_count", diag.get("selected_count", 0)))) for diag in result_diagnostics)
-    no_valid_count = sum(int(diag.get("no_valid_revisit_count", 0)) for diag in result_diagnostics)
-    abstained_count = sum(int(diag.get("revisit_abstained_count", int(bool(diag.get("abstained", False))))) for diag in result_diagnostics)
-    selected_gaps = [int(diag["revisit_min_gap_to_target"]) for diag in result_diagnostics if int(diag.get("revisit_min_gap_to_target", -1)) >= 0]
-    diagnostics: dict[str, Any] = {
-        "revisit_candidate_frame_count": candidate_count_mean,
-        "revisit_candidate_count": candidate_count_mean,
-        "valid_candidate_label_count": int(valid_candidate_label_count),
-        "revisit_pose_preselect_input_count": float(pose_preselect_input_count / target_count) if target_count else 0.0,
-        "revisit_pose_preselect_selected_count": float(pose_preselect_selected_count / target_count) if target_count else 0.0,
-        "revisit_exact_fov_candidate_count": float(exact_fov_candidate_count / target_count) if target_count else 0.0,
-        "valid_revisit_frame_count": valid_count_mean,
-        "valid_revisit_count": valid_count_mean,
-        "no_valid_revisit_count": int(no_valid_count),
-        "valid_revisit_mask_fraction": tensor_valid_fraction(valid_revisit_mask),
-        "revisit_selected_frame_count": int(selected_count),
-        "revisit_selected_count": int(selected_count),
-        "revisit_abstained_count": int(abstained_count),
-        "revisit_min_gap_to_target": int(min(selected_gaps)) if selected_gaps else -1,
-        "revisit_label_source": _REVISIT_LABEL_SOURCE,
-    }
-    frame_fov_values = _collect_values(result_diagnostics, "frame_fov_overlap_values")
-    if not frame_fov_values:
-        frame_fov_values = _collect_values(result_diagnostics, "fov_overlap_values")
-    diagnostics.update(_value_stats(frame_fov_values, "revisit_frame_fov_overlap"))
-    diagnostics.update(_value_stats(frame_fov_values, "revisit_fov_overlap"))
-    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "plucker_overlap_values"), "revisit_plucker_overlap"))
-    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_fov_overlap_values"), "revisit_best_selected_fov_overlap"))
-    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_plucker_overlap_values"), "revisit_best_selected_plucker_overlap"))
-    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_gap_frame_values"), "revisit_best_selected_gap_frames"))
-    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "best_selected_frame_fov_overlap_values"), "revisit_best_selected_frame_fov_overlap"))
-    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "selected_frame_fov_overlap_values"), "revisit_selected_frame_fov_overlap"))
-    diagnostics.update(_value_stats(_collect_values(result_diagnostics, "selected_incremental_fov_overlap_values"), "revisit_incremental_fov_overlap"))
-    return diagnostics
-
-
-def summarize_noise_bucket_diagnostics(
-    *,
-    noise_bucket: str | None,
-    valid_revisit_mask: torch.Tensor | None,
-    no_valid_revisit_mask: torch.Tensor | None,
-    noise_bucket_ids: torch.Tensor | None = None,
-) -> dict[str, Any]:
-    bucket = normalize_noise_bucket(noise_bucket)
-    diagnostics: dict[str, Any] = {
-        "noise_bucket": bucket,
-        "noise_bucket_id": int(NOISE_BUCKET_TO_ID[bucket]),
-    }
-    for candidate in NOISE_BUCKETS:
-        diagnostics[f"noise_bucket_is_{candidate}"] = int(bucket == candidate)
-
-    valid = torch.zeros(0, dtype=torch.bool) if valid_revisit_mask is None else valid_revisit_mask.detach().bool().reshape(-1).cpu()
-    no_valid = torch.zeros_like(valid) if no_valid_revisit_mask is None else no_valid_revisit_mask.detach().bool().reshape(-1).cpu()
-    target_count = int(valid.numel())
-    diagnostics["noise_bucket_target_count"] = target_count
-    if noise_bucket_ids is None:
-        target_bucket_ids = torch.full((target_count,), int(NOISE_BUCKET_TO_ID[bucket]), dtype=torch.long)
-    else:
-        target_bucket_ids = noise_bucket_ids.detach().long().reshape(-1).cpu()
-        if int(target_bucket_ids.numel()) != target_count:
-            raise ValueError(
-                f"noise_bucket_ids has {int(target_bucket_ids.numel())} targets, expected {target_count}"
-            )
-
-    for bucket_name in NOISE_BUCKETS:
-        bucket_mask = target_bucket_ids == int(NOISE_BUCKET_TO_ID[bucket_name])
-        diagnostics[f"noise_bucket_{bucket_name}_target_count"] = int(bucket_mask.long().sum().item())
-
-    mask_specs = (
-        ("valid_revisit", valid),
-        ("no_valid_revisit", no_valid),
-    )
-    for mask_name, mask in mask_specs:
-        for bucket_name in NOISE_BUCKETS:
-            bucket_mask = target_bucket_ids == int(NOISE_BUCKET_TO_ID[bucket_name])
-            count = int((mask & bucket_mask).long().sum().item()) if mask.numel() else 0
-            diagnostics[f"{mask_name}_noise_bucket_{bucket_name}_count"] = count
-    return diagnostics
-
-
-def summarize_eval_ablation_diagnostics(
-    *,
-    enabled: bool,
-    branch: str | None,
-    valid_revisit_mask: torch.Tensor | None,
-    no_valid_revisit_mask: torch.Tensor | None,
-    eval_corrupted_revisit_mask: torch.Tensor | None,
-) -> dict[str, Any]:
-    branch = normalize_eval_ablation_branch(branch)
-    valid = torch.zeros(0, dtype=torch.bool) if valid_revisit_mask is None else valid_revisit_mask.detach().bool().reshape(-1).cpu()
-    no_valid = torch.zeros_like(valid) if no_valid_revisit_mask is None else no_valid_revisit_mask.detach().bool().reshape(-1).cpu()
-    corrupted = torch.zeros_like(valid) if eval_corrupted_revisit_mask is None else eval_corrupted_revisit_mask.detach().bool().reshape(-1).cpu()
-    true_revisit = valid & (~corrupted)
-    diagnostics: dict[str, Any] = {
-        "eval_ablation_enabled": bool(enabled),
-        "eval_ablation_branch": branch,
-        "eval_ablation_branch_id": int(EVAL_ABLATION_BRANCH_TO_ID[branch]),
-        "eval_bucket_true_revisit_count": int(true_revisit.long().sum().item()),
-        "eval_bucket_no_valid_revisit_count": int(no_valid.long().sum().item()),
-        "eval_bucket_corrupted_memory_count": int(corrupted.long().sum().item()),
-    }
-    total = max(int(valid.numel()), 1)
-    diagnostics["eval_bucket_true_revisit_fraction"] = float(diagnostics["eval_bucket_true_revisit_count"] / total)
-    diagnostics["eval_bucket_no_valid_revisit_fraction"] = float(diagnostics["eval_bucket_no_valid_revisit_count"] / total)
-    diagnostics["eval_bucket_corrupted_memory_fraction"] = float(diagnostics["eval_bucket_corrupted_memory_count"] / total)
-    return diagnostics

algorithms/worldmem/dememwm/injection.py

@@ -1,4 +1,3 @@
-
 from __future__ import annotations
 
 from dataclasses import dataclass
@@ -6,7 +5,6 @@ from typing import Any
 
 import torch
 
-from .diagnostics import summarize_stream
 from .types import MemoryStreamTensors
 
 
@@ -31,7 +29,12 @@ class InjectionAdapter:
             return gate.to(device=device, dtype=dtype)
         return torch.tensor(float(gate), device=device, dtype=dtype)
 
-    def __call__(self, streams: MemoryStreamTensors, device=None, dtype=None) -> tuple[dict[str, Any], dict[str, Any]]:
+    def __call__(
+        self,
+        streams: MemoryStreamTensors,
+        device=None,
+        dtype=None,
+    ) -> dict[str, Any]:
         ref = streams.anchor_tokens
         device = device or ref.device
         dtype = dtype or ref.dtype
@@ -62,22 +65,7 @@ class InjectionAdapter:
             if not revisit_mask.any():
                 kwargs["memory_retrieval_tokens"] = None
                 kwargs["memory_retrieval_mask"] = None
-        diagnostics = dict(streams.diagnostics)
-        diagnostics.update(summarize_stream("anchor", anchor_tokens, anchor_mask, kwargs["memory_anchor_gate"]))
-        diagnostics.update(summarize_stream("dynamic", dynamic_tokens, dynamic_mask, kwargs["memory_dynamic_gate"]))
-        diagnostics.update(summarize_stream("revisit", revisit_tokens, revisit_mask, kwargs["memory_retrieval_gate"]))
-        if streams.revisit_gate_raw is not None:
-            raw_gate = streams.revisit_gate_raw.to(device=device, dtype=dtype)
-            diagnostics["revisit_gate_raw"] = raw_gate
-            diagnostics["revisit_gate_raw_mean"] = float(raw_gate.detach().float().mean().item()) if raw_gate.numel() else 0.0
-            diagnostics["revisit_gate_raw_min"] = float(raw_gate.detach().float().min().item()) if raw_gate.numel() else 0.0
-            diagnostics["revisit_gate_raw_max"] = float(raw_gate.detach().float().max().item()) if raw_gate.numel() else 0.0
-        if streams.no_valid_revisit_mask is not None:
-            diagnostics["no_valid_revisit_mask"] = streams.no_valid_revisit_mask.to(device=device, dtype=torch.bool)
-        max_sources = [v for k, v in streams.diagnostics.items() if k.endswith("max_source_frame")]
-        if max_sources:
-            diagnostics["max_source_frame"] = max(int(torch.as_tensor(v).max().item()) for v in max_sources)
-        return kwargs, diagnostics
+        return kwargs
 
 
 DeMemWMInjectionAdapter = InjectionAdapter

algorithms/worldmem/dememwm/retrieval.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import math
-from dataclasses import replace
+from dataclasses import dataclass, replace
 from typing import Any, Optional
 
 import torch
@@ -9,6 +9,7 @@ import torch
 from .labels import (
     LABEL_SOURCE,
     RevisitCandidateLabel,
+    _angle_diff_degrees,
     _inside_fov_3d_hv,
     _plucker_descriptor,
     _target_fov_points,
@@ -16,24 +17,19 @@ from .labels import (
 from .types import MemoryRecord, RevisitRetrievalResult
 
 
-def _overlap_values(labels, name: str) -> list[float]:
-    values: list[float] = []
-    for label in labels:
-        value = getattr(label, name)
-        if value is not None:
-            values.append(float(value))
-    return values
+@dataclass
+class BatchedRevisitSelectionResult:
+    selected_positions: torch.Tensor
+    selected_mask: torch.Tensor
+    selected_scores: torch.Tensor
+    selected_fov_overlap: torch.Tensor
+    selected_plucker_overlap: torch.Tensor
+    selected_gap_frames: torch.Tensor
+    best_selected_fov_overlap: torch.Tensor
+    best_selected_plucker_overlap: torch.Tensor
+    best_selected_gap_frames: torch.Tensor
 
 
-def _overlap_stats(values: list[float], prefix: str) -> dict[str, float]:
-    if not values:
-        return {f"{prefix}_mean": 0.0, f"{prefix}_min": 0.0, f"{prefix}_max": 0.0}
-    return {
-        f"{prefix}_mean": float(sum(values) / len(values)),
-        f"{prefix}_min": float(min(values)),
-        f"{prefix}_max": float(max(values)),
-    }
-
 
 def _pose_rows(pose) -> torch.Tensor | None:
     if pose is None:
@@ -58,6 +54,348 @@ def _pose_forward(poses: torch.Tensor) -> torch.Tensor:
     )
 
 
+
+def _time_batch_pose_rows(
+    pose: torch.Tensor,
+    *,
+    time: int,
+    batch: int,
+    name: str,
+) -> torch.Tensor:
+    if pose is None:
+        raise ValueError(f"{name} is required for batched DeMemWM revisit retrieval")
+    pose_tensor = pose if torch.is_tensor(pose) else torch.as_tensor(pose, dtype=torch.float32)
+    if pose_tensor.ndim < 3 or pose_tensor.shape[-1] < 5:
+        raise ValueError(f"{name} must have shape (T,B,D) or (B,T,D) with D >= 5")
+    if pose_tensor.shape[0] == time and pose_tensor.shape[1] == batch:
+        pose_tb = pose_tensor
+    elif pose_tensor.shape[0] == batch and pose_tensor.shape[1] == time:
+        pose_tb = pose_tensor.transpose(0, 1)
+    else:
+        raise ValueError(f"{name} must match time/batch dimensions ({time},{batch})")
+    return pose_tb[..., :5].detach().to(dtype=torch.float32)
+
+
+def _time_batch_mask(
+    mask: torch.Tensor | None,
+    *,
+    time: int,
+    batch: int,
+    device: torch.device,
+) -> torch.Tensor:
+    if mask is None:
+        return torch.ones((time, batch), device=device, dtype=torch.bool)
+    mask_tensor = mask if torch.is_tensor(mask) else torch.as_tensor(mask)
+    if mask_tensor.ndim != 2:
+        raise ValueError("source_candidate_mask must have shape (T,B) or (B,T)")
+    if mask_tensor.shape == (time, batch):
+        mask_tb = mask_tensor
+    elif mask_tensor.shape == (batch, time):
+        mask_tb = mask_tensor.transpose(0, 1)
+    else:
+        raise ValueError(f"source_candidate_mask must match time/batch dimensions ({time},{batch})")
+    return mask_tb.to(device=device, dtype=torch.bool)
+
+
+def _target_fov_points_batched(
+    target_poses: torch.Tensor,
+    *,
+    fov_half_h: float,
+    fov_half_v: float,
+    yaw_samples: int,
+    pitch_samples: int,
+    depth_samples: int,
+    radius: float,
+) -> torch.Tensor:
+    yaw_samples = max(1, int(yaw_samples))
+    pitch_samples = max(1, int(pitch_samples))
+    depth_samples = max(1, int(depth_samples))
+    device = target_poses.device
+    dtype = target_poses.dtype
+    if yaw_samples == 1:
+        yaw_offsets = torch.zeros((1,), device=device, dtype=dtype)
+    else:
+        yaw_offsets = torch.linspace(-float(fov_half_h), float(fov_half_h), yaw_samples + 2, device=device, dtype=dtype)[1:-1]
+    if pitch_samples == 1:
+        pitch_offsets = torch.zeros((1,), device=device, dtype=dtype)
+    else:
+        pitch_offsets = torch.linspace(-float(fov_half_v), float(fov_half_v), pitch_samples + 2, device=device, dtype=dtype)[1:-1]
+    if depth_samples == 1:
+        depths = torch.full((1,), float(radius), device=device, dtype=dtype)
+    else:
+        depths = torch.linspace(float(radius) / float(depth_samples), float(radius), depth_samples, device=device, dtype=dtype)
+    depth_grid, pitch_grid, yaw_grid = torch.meshgrid(depths, pitch_offsets, yaw_offsets, indexing="ij")
+    pitch_offsets_flat = pitch_grid.reshape(1, -1)
+    yaw_offsets_flat = yaw_grid.reshape(1, -1)
+    depth = depth_grid.reshape(1, -1)
+    pitch = torch.deg2rad(target_poses[:, 3:4] + pitch_offsets_flat)
+    yaw = torch.deg2rad(target_poses[:, 4:5] + yaw_offsets_flat)
+    cos_pitch = torch.cos(pitch)
+    vectors = torch.stack(
+        [
+            depth * cos_pitch * torch.sin(yaw),
+            depth * torch.sin(pitch),
+            depth * cos_pitch * torch.cos(yaw),
+        ],
+        dim=-1,
+    )
+    return target_poses[:, None, :3] + vectors
+
+
+def _inside_fov_3d_hv_batched(
+    points: torch.Tensor,
+    poses: torch.Tensor,
+    *,
+    fov_half_h: float,
+    fov_half_v: float,
+) -> torch.Tensor:
+    vectors = points[:, None, :, :] - poses[:, :, None, :3]
+    x = vectors[..., 0]
+    y = vectors[..., 1]
+    z = vectors[..., 2]
+    azimuth = torch.atan2(x, z) * (180.0 / math.pi)
+    elevation = torch.atan2(y, torch.sqrt(x.square() + z.square()).clamp_min(1e-8)) * (180.0 / math.pi)
+    diff_azimuth = _angle_diff_degrees(azimuth - poses[:, :, None, 4])
+    diff_elevation = _angle_diff_degrees(elevation - poses[:, :, None, 3])
+    return (diff_azimuth < float(fov_half_h)) & (diff_elevation < float(fov_half_v))
+
+
+def _batched_tie_mask(mask: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
+    neg_inf = torch.full_like(values, -float("inf"))
+    best = torch.where(mask, values, neg_inf).max(dim=1).values
+    return mask & torch.isclose(values, best[:, None], rtol=0.0, atol=1e-12)
+
+
+def batched_revisit_select_positions(
+    source_frame_indices: torch.Tensor,
+    source_pose: torch.Tensor,
+    target_frame_indices: torch.Tensor,
+    target_pose: torch.Tensor,
+    *,
+    source_candidate_mask: torch.Tensor | None = None,
+    topk: int = 2,
+    exclude_local_context_frames: int = 0,
+    fov_overlap_threshold: Optional[float] = 0.30,
+    plucker_weight: float = 0.1,
+    fov_half_h: float = 105.0 / 2.0,
+    fov_half_v: float = 75.0 / 2.0,
+    fov_yaw_samples: int = 25,
+    fov_pitch_samples: int = 20,
+    fov_depth_samples: int = 20,
+    fov_radius: float = 30.0,
+    plucker_grid_h: int = 4,
+    plucker_grid_w: int = 4,
+    plucker_focal_length: float = 0.35,
+    pose_preselect_topk: Optional[int] = 64,
+    query_chunk_size: int = 16,
+) -> BatchedRevisitSelectionResult:
+    if source_frame_indices.ndim != 2:
+        raise ValueError("source_frame_indices must have shape (T_src,B)")
+    if target_frame_indices.ndim == 1:
+        target_frame_indices = target_frame_indices[:, None]
+    if target_frame_indices.ndim != 2:
+        raise ValueError("target_frame_indices must have shape (T_tgt,B)")
+
+    T_src, B = source_frame_indices.shape
+    T_tgt, B_tgt = target_frame_indices.shape
+    if B_tgt != B:
+        raise ValueError("source_frame_indices and target_frame_indices must share batch dimension")
+
+    source_pose_tensor = source_pose if torch.is_tensor(source_pose) else torch.as_tensor(source_pose, dtype=torch.float32)
+    target_pose_tensor = target_pose if torch.is_tensor(target_pose) else torch.as_tensor(target_pose, dtype=torch.float32)
+    device = source_pose_tensor.device
+    if target_pose_tensor.is_cuda:
+        device = target_pose_tensor.device
+    elif source_pose_tensor.is_cuda:
+        device = source_pose_tensor.device
+    source_frames_tb = source_frame_indices.to(device=device)
+    target_frames_tb = target_frame_indices.to(device=device)
+    source_pose_tb = _time_batch_pose_rows(source_pose_tensor.to(device=device), time=T_src, batch=B, name="source_pose")
+    target_pose_tb = _time_batch_pose_rows(target_pose_tensor.to(device=device), time=T_tgt, batch=B, name="target_pose")
+    candidate_mask_tb = _time_batch_mask(source_candidate_mask, time=T_src, batch=B, device=device)
+
+    topk = max(0, int(topk))
+    selected_positions = torch.full((B, T_tgt, topk), -1, device=device, dtype=torch.long)
+    selected_mask = torch.zeros((B, T_tgt, topk), device=device, dtype=torch.bool)
+    selected_scores = torch.zeros((B, T_tgt, topk), device=device, dtype=torch.float32)
+    selected_fov_overlap = torch.zeros((B, T_tgt, topk), device=device, dtype=torch.float32)
+    selected_plucker_overlap = torch.zeros((B, T_tgt, topk), device=device, dtype=torch.float32)
+    selected_gap_frames = torch.full((B, T_tgt, topk), -1.0, device=device, dtype=torch.float32)
+    best_fov = torch.zeros((B, T_tgt), device=device, dtype=torch.float32)
+    best_plucker = torch.zeros((B, T_tgt), device=device, dtype=torch.float32)
+    best_gap = torch.full((B, T_tgt), -1.0, device=device, dtype=torch.float32)
+    if topk == 0 or T_src == 0 or T_tgt == 0 or B == 0:
+        return BatchedRevisitSelectionResult(
+            selected_positions=selected_positions,
+            selected_mask=selected_mask,
+            selected_scores=selected_scores,
+            selected_fov_overlap=selected_fov_overlap,
+            selected_plucker_overlap=selected_plucker_overlap,
+            selected_gap_frames=selected_gap_frames,
+            best_selected_fov_overlap=best_fov,
+            best_selected_plucker_overlap=best_plucker,
+            best_selected_gap_frames=best_gap,
+        )
+
+    source_pose_flat = source_pose_tb.reshape(-1, source_pose_tb.shape[-1])
+    source_forward_tb = _pose_forward(source_pose_flat).reshape(T_src, B, 3)
+    source_desc_tb = _plucker_descriptor(
+        source_pose_flat,
+        grid_h=plucker_grid_h,
+        grid_w=plucker_grid_w,
+        focal_length=plucker_focal_length,
+    ).reshape(T_src, B, -1)
+    target_frames_flat = target_frames_tb.transpose(0, 1).contiguous().reshape(-1)
+    target_pose_flat = target_pose_tb.transpose(0, 1).contiguous().reshape(-1, target_pose_tb.shape[-1])
+    target_forward_flat = _pose_forward(target_pose_flat)
+    target_desc_flat = _plucker_descriptor(
+        target_pose_flat,
+        grid_h=plucker_grid_h,
+        grid_w=plucker_grid_w,
+        focal_length=plucker_focal_length,
+    )
+    batch_ids = torch.arange(B, device=device, dtype=torch.long).repeat_interleave(T_tgt)
+    Q = int(target_frames_flat.numel())
+    chunk_size = max(1, int(query_chunk_size))
+    source_positions = torch.arange(T_src, device=device, dtype=torch.long)
+    pose_topk = None if pose_preselect_topk is None else int(pose_preselect_topk)
+
+    selected_positions_flat_view = selected_positions.reshape(-1, topk)
+    selected_mask_flat_view = selected_mask.reshape(-1, topk)
+    selected_scores_flat_view = selected_scores.reshape(-1, topk)
+    selected_fov_flat_view = selected_fov_overlap.reshape(-1, topk)
+    selected_plucker_flat_view = selected_plucker_overlap.reshape(-1, topk)
+    selected_gap_flat_view = selected_gap_frames.reshape(-1, topk)
+    best_fov_flat_view = best_fov.reshape(-1)
+    best_plucker_flat_view = best_plucker.reshape(-1)
+    best_gap_flat_view = best_gap.reshape(-1)
+
+    for start in range(0, Q, chunk_size):
+        end = min(Q, start + chunk_size)
+        b_idx = batch_ids[start:end]
+        target_frames = target_frames_flat[start:end]
+        target_poses = target_pose_flat[start:end]
+        q = int(end - start)
+
+        source_frames = source_frames_tb.index_select(1, b_idx).transpose(0, 1).contiguous()
+        source_candidates = candidate_mask_tb.index_select(1, b_idx).transpose(0, 1).contiguous()
+        score_valid = source_candidates & (source_frames < (target_frames[:, None] - int(exclude_local_context_frames)))
+
+        if pose_topk is None or pose_topk <= 0 or T_src <= pose_topk:
+            preselect_idx = source_positions.reshape(1, -1).expand(q, -1)
+            preselected_valid = score_valid
+        else:
+            source_poses_q = source_pose_tb.index_select(1, b_idx).permute(1, 0, 2).contiguous()
+            source_forward_q = source_forward_tb.index_select(1, b_idx).permute(1, 0, 2).contiguous()
+            translation_norm = torch.linalg.vector_norm(source_poses_q[:, :, :3] - target_poses[:, None, :3], dim=-1) / max(float(fov_radius), 1e-6)
+            dot = (
+                source_forward_q * target_forward_flat[start:end].reshape(q, 1, 3)
+            ).sum(dim=-1).clamp(-1.0, 1.0)
+            pose_distance = translation_norm + (torch.acos(dot) / math.pi)
+            rank = (
+                pose_distance.to(dtype=torch.float64)
+                - source_frames.to(dtype=torch.float64) * 1e-12
+                + source_frames.to(dtype=torch.float64) * 1e-15
+            )
+            rank = rank.masked_fill(~score_valid, float("inf"))
+            k_pre = min(max(1, pose_topk), T_src)
+            top = torch.topk(rank, k=k_pre, largest=False, sorted=True)
+            preselect_idx = top.indices
+            preselected_valid = torch.isfinite(top.values) & torch.gather(score_valid, 1, preselect_idx)
+
+        K = int(preselect_idx.shape[1])
+        if K == 0:
+            continue
+
+        gather_pose_idx = preselect_idx.unsqueeze(-1).expand(-1, -1, source_pose_tb.shape[-1])
+        source_poses_q = source_pose_tb.index_select(1, b_idx).permute(1, 0, 2).contiguous()
+        selected_poses = torch.gather(source_poses_q, 1, gather_pose_idx)
+        selected_frames = torch.gather(source_frames, 1, preselect_idx)
+
+        points = _target_fov_points_batched(
+            target_poses,
+            fov_half_h=fov_half_h,
+            fov_half_v=fov_half_v,
+            yaw_samples=fov_yaw_samples,
+            pitch_samples=fov_pitch_samples,
+            depth_samples=fov_depth_samples,
+            radius=fov_radius,
+        )
+        inside = _inside_fov_3d_hv_batched(points, selected_poses, fov_half_h=fov_half_h, fov_half_v=fov_half_v)
+        fov_values = inside.float().mean(dim=2)
+
+        source_desc_q = source_desc_tb.index_select(1, b_idx).permute(1, 0, 2).contiguous()
+        gather_desc_idx = preselect_idx.unsqueeze(-1).expand(-1, -1, source_desc_tb.shape[-1])
+        selected_desc = torch.gather(source_desc_q, 1, gather_desc_idx)
+        diff = selected_desc - target_desc_flat[start:end, None, :]
+        plucker_distance = torch.linalg.vector_norm(diff, dim=-1) / math.sqrt(float(diff.shape[-1]))
+        plucker_values = 1.0 / (1.0 + plucker_distance.clamp_min(0.0))
+
+        valid_mask = preselected_valid
+        if fov_overlap_threshold is not None:
+            valid_mask = valid_mask & (fov_values >= float(fov_overlap_threshold))
+
+        remaining = valid_mask.clone()
+        covered = torch.zeros((q, inside.shape[2]), device=device, dtype=torch.bool)
+        chosen_positions = torch.full((q, topk), -1, device=device, dtype=torch.long)
+        chosen_scores = torch.zeros((q, topk), device=device, dtype=torch.float32)
+        chosen_fov = torch.zeros((q, topk), device=device, dtype=torch.float32)
+        chosen_plucker = torch.zeros((q, topk), device=device, dtype=torch.float32)
+        chosen_gap = torch.full((q, topk), -1.0, device=device, dtype=torch.float32)
+        row_idx = torch.arange(q, device=device, dtype=torch.long)
+        gap_values = target_frames[:, None] - selected_frames
+        for slot in range(topk):
+            active = remaining.any(dim=1)
+            gains = (inside & ~covered[:, None, :]).float().mean(dim=2)
+            tied = _batched_tie_mask(remaining, gains)
+            tied = _batched_tie_mask(tied, fov_values)
+            tied = _batched_tie_mask(tied, plucker_values * float(plucker_weight))
+            tied = _batched_tie_mask(tied, -gap_values.to(dtype=torch.float32))
+            tied = _batched_tie_mask(tied, -selected_frames.to(dtype=torch.float32))
+            best_idx = tied.to(dtype=torch.long).argmax(dim=1)
+            chosen_positions[active, slot] = preselect_idx[row_idx[active], best_idx[active]]
+            chosen_scores[active, slot] = gains[row_idx[active], best_idx[active]]
+            chosen_fov[active, slot] = fov_values[row_idx[active], best_idx[active]]
+            chosen_plucker[active, slot] = plucker_values[row_idx[active], best_idx[active]]
+            chosen_gap[active, slot] = gap_values[row_idx[active], best_idx[active]].to(dtype=torch.float32)
+            covered[active] = covered[active] | inside[row_idx[active], best_idx[active]]
+            remaining[row_idx[active], best_idx[active]] = False
+
+        chosen_mask = chosen_positions >= 0
+        chosen_rank = (
+            chosen_fov.to(dtype=torch.float64)
+            + chosen_plucker.to(dtype=torch.float64) * 1e-9
+            - chosen_gap.to(dtype=torch.float64) * 1e-12
+        ).masked_fill(~chosen_mask, -float("inf"))
+        has_choice = chosen_mask.any(dim=1)
+        best_slot = chosen_rank.argmax(dim=1)
+        best_fov_flat = torch.where(has_choice, chosen_fov[row_idx, best_slot], torch.zeros((q,), device=device, dtype=torch.float32))
+        best_plucker_flat = torch.where(has_choice, chosen_plucker[row_idx, best_slot], torch.zeros((q,), device=device, dtype=torch.float32))
+        best_gap_flat = torch.where(has_choice, chosen_gap[row_idx, best_slot], torch.full((q,), -1.0, device=device, dtype=torch.float32))
+
+        selected_positions_flat_view[start:end] = chosen_positions
+        selected_mask_flat_view[start:end] = chosen_mask
+        selected_scores_flat_view[start:end] = chosen_scores
+        selected_fov_flat_view[start:end] = chosen_fov
+        selected_plucker_flat_view[start:end] = chosen_plucker
+        selected_gap_flat_view[start:end] = chosen_gap
+        best_fov_flat_view[start:end] = best_fov_flat
+        best_plucker_flat_view[start:end] = best_plucker_flat
+        best_gap_flat_view[start:end] = best_gap_flat
+
+    return BatchedRevisitSelectionResult(
+        selected_positions=selected_positions,
+        selected_mask=selected_mask,
+        selected_scores=selected_scores,
+        selected_fov_overlap=selected_fov_overlap,
+        selected_plucker_overlap=selected_plucker_overlap,
+        selected_gap_frames=selected_gap_frames,
+        best_selected_fov_overlap=best_fov,
+        best_selected_plucker_overlap=best_plucker,
+        best_selected_gap_frames=best_gap,
+    )
+
+
 def _single_frame_pose(record: MemoryRecord) -> torch.Tensor | None:
     if int(record.frame_indices.numel()) != 1:
         return None
@@ -83,19 +421,9 @@ def _vectorized_frame_candidate_labels(
     plucker_grid_w: int,
     plucker_focal_length: float,
     pose_preselect_topk: Optional[int],
-) -> tuple[list[RevisitCandidateLabel], dict[str, float | int]]:
-    diagnostics: dict[str, float | int] = {
-        "revisit_pose_preselect_input_count": len(records),
-        "revisit_pose_preselect_scored_count": len(records),
-        "revisit_pose_preselect_unscored_count": 0,
-        "revisit_pose_preselect_selected_count": len(records),
-        "revisit_pose_preselect_min_distance": 0.0,
-        "revisit_pose_preselect_max_distance": 0.0,
-        "revisit_exact_fov_candidate_count": len(records),
-        "revisit_vectorized_frame_scorer_used": 1,
-    }
+) -> list[RevisitCandidateLabel]:
     if not records:
-        return [], diagnostics
+        return []
 
     target_poses = _pose_rows(target_pose)
     if target_poses is None:
@@ -137,9 +465,6 @@ def _vectorized_frame_candidate_labels(
         ]
         ranked.sort()
         selected_indices = [idx for *_, idx in ranked[:topk]]
-        diagnostics["revisit_pose_preselect_selected_count"] = len(selected_indices)
-        diagnostics["revisit_pose_preselect_min_distance"] = float(min(distance_values))
-        diagnostics["revisit_pose_preselect_max_distance"] = float(max(distance_values))
 
     selected_tensor = torch.tensor(selected_indices, device=device, dtype=torch.long)
     selected_records = [records[idx] for idx in selected_indices]
@@ -175,7 +500,6 @@ def _vectorized_frame_candidate_labels(
     if fov_overlap_threshold is not None:
         valid_mask = fov_values >= float(fov_overlap_threshold)
 
-    diagnostics["revisit_exact_fov_candidate_count"] = len(selected_records)
     fov_list = [float(value) for value in fov_values.detach().cpu().tolist()]
     plucker_list = [float(value) for value in plucker_values.detach().cpu().tolist()]
     valid_list = [bool(value) for value in valid_mask.detach().cpu().tolist()]
@@ -200,7 +524,7 @@ def _vectorized_frame_candidate_labels(
                 best_frame_fov_overlap=fov_overlap,
             )
         )
-    return labels, diagnostics
+    return labels
 
 
 def _coverage_gain(label: RevisitCandidateLabel, covered_mask: torch.Tensor | None) -> float:
@@ -297,23 +621,6 @@ def _best_selected_label(labels: list[RevisitCandidateLabel]) -> RevisitCandidat
     )
 
 
-def _best_selected_frame_label(labels: list[RevisitCandidateLabel]) -> RevisitCandidateLabel | None:
-    frame_labels = [label for label in labels if label.best_frame_fov_overlap is not None]
-    if not frame_labels:
-        return None
-    return max(
-        frame_labels,
-        key=lambda label: (
-            float(label.best_frame_fov_overlap),
-            0.0 if label.fov_overlap is None else float(label.fov_overlap),
-            0.0 if label.plucker_overlap is None else float(label.plucker_overlap),
-            -int(label.gap_to_target),
-            -int(label.record.source_start),
-            str(label.record.chunk_id or ""),
-        ),
-    )
-
-
 def _record_with_selected_frame_metadata(
     label: RevisitCandidateLabel,
     *,
@@ -334,6 +641,11 @@ def _record_with_selected_frame_metadata(
     return replace(label.record, metadata=metadata)
 
 
+
+def _record_frame_id(record: MemoryRecord) -> int:
+    return int(record.source_end) - 1
+
+
 def deterministic_revisit_retrieval(
     records: list[MemoryRecord],
     target_frame: int,
@@ -367,7 +679,7 @@ def deterministic_revisit_retrieval(
         for record in causal_records
         if int(record.source_end) <= target_frame - exclude_local_context_frames
     ]
-    labels, pose_preselect_diagnostics = _vectorized_frame_candidate_labels(
+    labels = _vectorized_frame_candidate_labels(
         score_records,
         target_frame=target_frame,
         target_pose=target_pose,
@@ -383,21 +695,16 @@ def deterministic_revisit_retrieval(
         plucker_focal_length=plucker_focal_length,
         pose_preselect_topk=pose_preselect_topk,
     )
-    exact_fov_candidate_count = int(pose_preselect_diagnostics["revisit_exact_fov_candidate_count"])
     valid_labels = [label for label in labels if label.valid]
-    selected_labels, selected_scores, selected_gains = _select_greedy_coverage(
+    selected_labels, selected_scores, _ = _select_greedy_coverage(
         valid_labels,
         topk=topk,
         plucker_weight=float(plucker_weight),
     )
     best_selected = _best_selected_label(selected_labels)
-    best_selected_frame = _best_selected_frame_label(selected_labels)
     best_selected_fov = 0.0 if best_selected is None or best_selected.fov_overlap is None else float(best_selected.fov_overlap)
     best_selected_plucker = 0.0 if best_selected is None or best_selected.plucker_overlap is None else float(best_selected.plucker_overlap)
     best_selected_gap = -1 if best_selected is None else int(best_selected.gap_to_target)
-    best_selected_frame_fov = 0.0 if best_selected_frame is None else float(best_selected_frame.best_frame_fov_overlap)
-    best_selected_frame_index = -1 if best_selected_frame is None or best_selected_frame.best_frame_index is None else int(best_selected_frame.best_frame_index)
-    high_quality_selected = int(best_selected_frame is not None and best_selected_frame_fov >= float(high_quality_fov_threshold))
     selected_records = [
         _record_with_selected_frame_metadata(label, high_quality_fov_threshold=float(high_quality_fov_threshold))
         for label in selected_labels
@@ -405,71 +712,11 @@ def deterministic_revisit_retrieval(
     score_device = selected_records[0].tokens.device if selected_records else torch.device("cpu")
     scores = torch.tensor(selected_scores, dtype=torch.float32, device=score_device)
 
-    fov_values = _overlap_values(valid_labels, "fov_overlap")
-    plucker_values = _overlap_values(valid_labels, "plucker_overlap")
-    selected_gaps = [label.gap_to_target for label in selected_labels]
-    selected_frame_fov_values = [
-        float(label.best_frame_fov_overlap)
-        for label in selected_labels
-        if label.best_frame_fov_overlap is not None
-    ]
-    diagnostics = {
-        "target_frame": int(target_frame),
-        "candidate_count": len(causal_records),
-        "candidate_frame_count": len(causal_records),
-        "valid_candidate_count": len(valid_labels),
-        "revisit_exact_fov_candidate_count": exact_fov_candidate_count,
-        "valid_candidate_frame_count": len(valid_labels),
-        "valid_candidate_label_count": len(valid_labels),
-        "selected_count": len(selected_records),
-        "selected_frame_count": len(selected_records),
-        "revisit_candidate_frame_count": len(causal_records),
-        "revisit_candidate_count": len(causal_records),
-        "valid_revisit_frame_count": len(valid_labels),
-        "valid_revisit_count": len(valid_labels),
-        "no_valid_revisit_count": int(len(valid_labels) == 0),
-        "valid_revisit_mask": int(len(valid_labels) > 0),
-        "revisit_abstained_count": int(len(selected_records) == 0),
-        "abstained": bool(len(selected_records) == 0),
-        "revisit_selected_frame_count": len(selected_records),
-        "revisit_selected_count": len(selected_records),
-        "revisit_min_gap_to_target": int(min(selected_gaps)) if selected_gaps else -1,
-        "best_selected_fov_overlap": best_selected_fov,
-        "best_selected_plucker_overlap": best_selected_plucker,
-        "best_selected_gap_frames": best_selected_gap,
-        "best_selected_frame_index": best_selected_frame_index,
-        "best_selected_frame_fov_overlap": best_selected_frame_fov,
-        "best_selected_frame_passes_high_quality": high_quality_selected,
-        "high_quality_selected_revisit": high_quality_selected,
-        "high_quality_fov_threshold": float(high_quality_fov_threshold),
-        "revisit_label_source": LABEL_SOURCE,
-        "selected_frame_ids": [int(record.max_source_frame) for record in selected_records],
-        "selected_frame_record_ids": [record.chunk_id for record in selected_records],
-        "selected_ranges": [(record.source_start, record.source_end) for record in selected_records],
-        "frame_fov_overlap_values": fov_values,
-        "fov_overlap_values": fov_values,
-        "plucker_overlap_values": plucker_values,
-        "best_selected_fov_overlap_values": [] if best_selected is None else [best_selected_fov],
-        "best_selected_plucker_overlap_values": [] if best_selected is None else [best_selected_plucker],
-        "best_selected_gap_frame_values": [] if best_selected is None else [best_selected_gap],
-        "best_selected_frame_fov_overlap_values": [] if best_selected_frame is None else [best_selected_frame_fov],
-        "selected_frame_fov_overlap_values": selected_frame_fov_values,
-        "selected_incremental_fov_overlap_values": selected_gains,
-        "selected_revisit_scores": selected_scores,
-        **pose_preselect_diagnostics,
-    }
-    diagnostics.update(_overlap_stats(fov_values, "revisit_frame_fov_overlap"))
-    diagnostics.update(_overlap_stats(fov_values, "revisit_fov_overlap"))
-    diagnostics.update(_overlap_stats(plucker_values, "revisit_plucker_overlap"))
-    diagnostics.update(_overlap_stats(diagnostics["best_selected_fov_overlap_values"], "revisit_best_selected_fov_overlap"))
-    diagnostics.update(_overlap_stats(diagnostics["best_selected_plucker_overlap_values"], "revisit_best_selected_plucker_overlap"))
-    diagnostics.update(_overlap_stats(diagnostics["best_selected_gap_frame_values"], "revisit_best_selected_gap_frames"))
-    diagnostics.update(_overlap_stats(diagnostics["best_selected_frame_fov_overlap_values"], "revisit_best_selected_frame_fov_overlap"))
-    diagnostics.update(_overlap_stats(selected_frame_fov_values, "revisit_selected_frame_fov_overlap"))
-    diagnostics.update(_overlap_stats(selected_gains, "revisit_incremental_fov_overlap"))
     return RevisitRetrievalResult(
         records=selected_records,
         scores=scores,
         selected_frame_ids=[int(record.max_source_frame) for record in selected_records],
-        diagnostics=diagnostics,
+        best_selected_fov_overlap=torch.as_tensor(best_selected_fov, dtype=torch.float32, device=score_device),
+        best_selected_plucker_overlap=torch.as_tensor(best_selected_plucker, dtype=torch.float32, device=score_device),
+        best_selected_gap_frames=torch.as_tensor(float(best_selected_gap), dtype=torch.float32, device=score_device),
     )

algorithms/worldmem/dememwm/schedules.py

@@ -8,8 +8,6 @@ import torch
 
 from .types import StreamGateState
 
-NOISE_BUCKETS = ("high", "mid", "low")
-NOISE_BUCKET_TO_ID = {name: idx for idx, name in enumerate(NOISE_BUCKETS)}
 EVAL_ABLATION_BRANCHES = (
     "memory_off",
     "A_only",
@@ -41,46 +39,6 @@ def _clamp01(value: float) -> float:
     return max(0.0, min(1.0, float(value)))
 
 
-def noise_bucket_from_denoising_fraction(denoising_fraction: float | None) -> str:
-    if denoising_fraction is None:
-        return "mid"
-    frac = _clamp01(float(denoising_fraction))
-    if frac < (1.0 / 3.0):
-        return "high"
-    if frac < (2.0 / 3.0):
-        return "mid"
-    return "low"
-
-
-def noise_bucket_from_noise_levels(noise_levels: torch.Tensor | None, timesteps: int | None) -> str:
-    if noise_levels is None or timesteps is None or int(timesteps) <= 1:
-        return "mid"
-    noise_fraction = _clamp01(float(noise_levels.detach().float().mean().item()) / float(int(timesteps) - 1))
-    if noise_fraction >= (2.0 / 3.0):
-        return "high"
-    if noise_fraction >= (1.0 / 3.0):
-        return "mid"
-    return "low"
-
-
-def noise_bucket_ids_from_noise_levels(noise_levels: torch.Tensor | None, timesteps: int | None) -> torch.Tensor | None:
-    if noise_levels is None or timesteps is None or int(timesteps) <= 1:
-        return None
-    noise_fraction = noise_levels.detach().float() / float(int(timesteps) - 1)
-    bucket_ids = torch.full_like(noise_levels, NOISE_BUCKET_TO_ID["mid"], dtype=torch.long)
-    bucket_ids = torch.where(
-        noise_fraction >= (2.0 / 3.0),
-        torch.full_like(bucket_ids, NOISE_BUCKET_TO_ID["high"]),
-        bucket_ids,
-    )
-    bucket_ids = torch.where(
-        noise_fraction < (1.0 / 3.0),
-        torch.full_like(bucket_ids, NOISE_BUCKET_TO_ID["low"]),
-        bucket_ids,
-    )
-    return bucket_ids
-
-
 def denoising_fraction_from_noise_levels(noise_levels: torch.Tensor | None, timesteps: int | None) -> float | None:
     if noise_levels is None or timesteps is None or int(timesteps) <= 1:
         return None
@@ -97,12 +55,6 @@ def normalize_eval_ablation_branch(branch: str | None) -> str:
     return branch
 
 
-def normalize_noise_bucket(noise_bucket: str | None) -> str:
-    if noise_bucket in NOISE_BUCKET_TO_ID:
-        return str(noise_bucket)
-    return "mid"
-
-
 _STAGE_ENABLES = {
     'stage_1': (True, True, True),
     'stage_2': (True, True, True),
@@ -131,22 +83,6 @@ class CurriculumState:
     def dit_full_trainable(self) -> bool:
         return self.dit_train_state == "full"
 
-    def diagnostics(self) -> dict[str, Any]:
-        return {
-            "dememwm_global_step": self.global_step,
-            "dememwm_curriculum_enabled": self.enabled,
-            "dememwm_stage": self.stage,
-            "curriculum_anchor_enabled": self.anchor_enabled,
-            "curriculum_dynamic_enabled": self.dynamic_enabled,
-            "curriculum_revisit_enabled": self.revisit_enabled,
-            "dit_train_state": self.dit_train_state,
-            "dit_full_trainable": self.dit_full_trainable,
-            "freeze_vae": self.freeze_vae,
-            "lr_dememwm_modules": self.dememwm_lr,
-            "lr_memory_adapters": self.memory_adapter_lr,
-            "lr_full_dit": self.full_dit_lr,
-        }
-
 
 def _cfg_get(obj: Any, name: str, default: Any) -> Any:
     return getattr(obj, name, default) if obj is not None else default
@@ -209,9 +145,7 @@ DeMemWMCurriculumState = CurriculumState
 resolve_dememwm_curriculum = resolve_curriculum
 
 
-def compute_stream_gates(stage: str, denoising_fraction: float | None = None, debug_force_all_streams: bool = False, anchor_gate: float = 1.0, dynamic_gate: float = 1.0, revisit_gate: float = 1.0) -> StreamGateState:
-    if debug_force_all_streams:
-        return StreamGateState(True, True, True, float(anchor_gate), float(dynamic_gate), float(revisit_gate), stage, "debug_force_all_streams")
+def compute_stream_gates(stage: str, denoising_fraction: float | None = None, anchor_gate: float = 1.0, dynamic_gate: float = 1.0, revisit_gate: float = 1.0) -> StreamGateState:
     if stage not in _STAGE_ENABLES:
         raise ValueError(f"unknown DeMemWM stage: {stage}")
     a_on, d_on, r_on = _STAGE_ENABLES[stage]

algorithms/worldmem/dememwm/types.py

@@ -1,4 +1,3 @@
-
 from __future__ import annotations
 
 from dataclasses import dataclass, field
@@ -74,8 +73,9 @@ class MemoryStreamTensors:
     revisit_gate: torch.Tensor | float
     revisit_gate_raw: torch.Tensor | None = None
     valid_revisit_mask: torch.Tensor | None = None
-    no_valid_revisit_mask: torch.Tensor | None = None
-    diagnostics: dict[str, Any] = field(default_factory=dict)
+    revisit_best_selected_fov_overlap: torch.Tensor | None = None
+    revisit_best_selected_plucker_overlap: torch.Tensor | None = None
+    revisit_selected_gap_frames: torch.Tensor | None = None
 
 
 @dataclass(frozen=True)
@@ -95,4 +95,6 @@ class RevisitRetrievalResult:
     records: list[MemoryRecord]
     scores: torch.Tensor
     selected_frame_ids: list[int]
-    diagnostics: dict[str, Any]
+    best_selected_fov_overlap: torch.Tensor
+    best_selected_plucker_overlap: torch.Tensor
+    best_selected_gap_frames: torch.Tensor

algorithms/worldmem/models/dit.py

@@ -166,12 +166,6 @@ class MemoryTokenCrossAttention(nn.Module):
         self.memory_type_embed = nn.Embedding(num_memory_types, hidden_size)
         self.memory_type_scale = nn.Parameter(torch.ones(num_memory_types, hidden_size))
         self.memory_type_gate = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, num_memory_types, bias=True))
-        self.last_gate_mean = None
-        self.last_delta_ratio = None
-        self.last_valid_fraction = None
-        self.last_type_gate_mean = None
-        for type_name in MEMORY_TYPE_NAMES[:num_memory_types]:
-            setattr(self, f"last_type_gate_{type_name}_mean", None)
         nn.init.normal_(self.memory_type_embed.weight, std=0.02)
         self.reset_identity_init()
 
@@ -236,19 +230,11 @@ class MemoryTokenCrossAttention(nn.Module):
             type_scale = type_scale.unsqueeze(0)
         return memory_tokens * type_scale + type_embed
 
-    def _store_type_gate_diagnostics(self, stage_gate):
-        with torch.no_grad():
-            detached = stage_gate.detach().float()
-            self.last_type_gate_mean = detached.mean()
-            for type_idx, type_name in enumerate(MEMORY_TYPE_NAMES[: self.num_memory_types]):
-                setattr(self, f"last_type_gate_{type_name}_mean", detached[..., type_idx].mean())
-
     def _type_stage_gate(self, c, memory_tokens, memory_type_ids):
         if memory_type_ids is None:
             return None
         memory_type_ids = memory_type_ids.to(device=memory_tokens.device, dtype=torch.long)
         stage_gate = torch.sigmoid(self.memory_type_gate(c)).to(memory_tokens.dtype)
-        self._store_type_gate_diagnostics(stage_gate)
         if memory_tokens.dim() == 4:
             batch_size, num_frames, num_tokens = memory_tokens.shape[:3]
             if memory_type_ids.dim() == 1:
@@ -329,33 +315,6 @@ class MemoryTokenCrossAttention(nn.Module):
                 gate_tensor = gate_tensor.unsqueeze(-1)
         return gate_tensor
 
-    def _store_diagnostics(self, output, base, gate_msa, gate_mlp, valid_rows):
-        with torch.no_grad():
-            batch_size, num_frames = base.shape[:2]
-            gate_values = torch.cat(
-                [gate_msa.detach().float().abs(), gate_mlp.detach().float().abs()],
-                dim=-1,
-            )
-            gate_mask = self._gate_valid_mask(
-                valid_rows,
-                batch_size,
-                num_frames,
-                dtype=gate_values.dtype,
-                device=gate_values.device,
-            )
-            if gate_mask is not None:
-                gate_values = gate_values * gate_mask
-                self.last_valid_fraction = valid_rows.detach().float().mean()
-                valid_count = (gate_mask.sum() * gate_values.shape[-1]).clamp_min(1.0)
-                self.last_gate_mean = gate_values.sum() / valid_count
-            else:
-                self.last_valid_fraction = base.detach().new_tensor(1.0, dtype=torch.float32)
-                self.last_gate_mean = gate_values.mean()
-
-            delta_norm = (output.detach().float() - base.detach().float()).norm()
-            base_norm = base.detach().float().norm()
-            self.last_delta_ratio = delta_norm / (base_norm + 1e-6)
-
     def forward(
         self,
         x,
@@ -437,7 +396,6 @@ class MemoryTokenCrossAttention(nn.Module):
         if residual_gate_tensor is not None:
             mlp_delta = mlp_delta * residual_gate_tensor
         output = output + mlp_delta
-        self._store_diagnostics(output, residual_base, m_gate_msa, m_gate_mlp, valid_rows)
         if return_delta:
             return attn_delta + mlp_delta
         return output
@@ -767,38 +725,6 @@ class DiT(nn.Module):
                 if memory_adapter is not None:
                     memory_adapter.reset_identity_init()
 
-    def memory_adapter_delta_diagnostics(self):
-        diagnostics = {}
-        ratios = []
-        type_gate_values = {type_name: [] for type_name in MEMORY_TYPE_NAMES}
-        shared_type_gate_values = []
-        for block in self.blocks:
-            adapter = getattr(block, "memory_token_cross_attn", None)
-            if adapter is None:
-                continue
-            ratio = getattr(adapter, "last_delta_ratio", None)
-            if ratio is not None:
-                ratios.append(torch.as_tensor(ratio).detach().float())
-            type_gate = getattr(adapter, "last_type_gate_mean", None)
-            if type_gate is not None:
-                shared_type_gate_values.append(torch.as_tensor(type_gate).detach().float())
-            for type_name in MEMORY_TYPE_NAMES:
-                value = getattr(adapter, f"last_type_gate_{type_name}_mean", None)
-                if value is not None:
-                    type_gate_values[type_name].append(torch.as_tensor(value).detach().float())
-        if ratios:
-            values = torch.stack(ratios)
-            diagnostics["memory_adapter_delta_ratio_max"] = float(values.max().item())
-            diagnostics["memory_adapter_delta_ratio_mean"] = float(values.mean().item())
-        if shared_type_gate_values:
-            values = torch.stack(shared_type_gate_values)
-            diagnostics["memory_adapter_type_gate_mean"] = float(values.mean().item())
-        for type_name, values_list in type_gate_values.items():
-            if values_list:
-                values = torch.stack(values_list)
-                diagnostics[f"memory_adapter_type_gate_{type_name}_mean"] = float(values.mean().item())
-        return diagnostics
-
     def unpatchify(self, x):
         """
         x: (N, H, W, patch_size**2 * C)

configurations/algorithm/dememwm_memory_dit.yaml

@@ -15,7 +15,6 @@ log_video: false
 dememwm:
   enabled: true
   training_stage: stage_1  # fallback only when curriculum.enabled=false
-  debug_force_all_streams: false
   curriculum:
     enabled: true
     full_stage_start_step: 60000
@@ -66,8 +65,6 @@ dememwm:
     plucker_focal_length: 0.35
     compress:
       downsample_ratio: 4
-  stage_policy:
-    noise_bucket_logging: true
   eval_ablation:
     enabled: false
     branch: A_plus_D_plus_R_normal

scripts/dememwm_full_eval.slurm

@@ -118,7 +118,6 @@ EVAL_ARGS=(
   "++algorithm.context_frames=${CONTEXT_FRAMES}"
   "++algorithm.log_video=${LOG_VIDEO}"
   "++algorithm.diffusion.sampling_timesteps=${SAMPLING_TIMESTEPS}"
-  "++algorithm.dememwm.debug_force_all_streams=false"
   "++algorithm.dememwm.training_stage=stage_2"
   "++algorithm.dememwm.anchor.enabled=true"
   "++algorithm.dememwm.anchor.anchor_indices=[0,1,2,3]"
@@ -139,7 +138,6 @@ EVAL_ARGS=(
   "++algorithm.dememwm.revisit.plucker_weight=0.10"
   "++algorithm.dememwm.revisit.max_frames=${REVISIT_MAX_FRAMES}"
   "++algorithm.dememwm.revisit.compress.downsample_ratio=${REVISIT_DOWNSAMPLE_RATIO}"
-  "++algorithm.dememwm.stage_policy.noise_bucket_logging=true"
   "++algorithm.dememwm.eval_ablation.enabled=true"
   "++algorithm.dememwm.eval_ablation.branch=${ABLATION_BRANCH}"
   "++algorithm.dememwm.cache.enabled=true"

scripts/dememwm_full_train.slurm

@@ -52,7 +52,6 @@ srun python -m main \
     ++algorithm.context_frames=100 \
     ++algorithm.log_video=true \
     ++algorithm.diffusion.sampling_timesteps=20 \
-    ++algorithm.dememwm.debug_force_all_streams=false \
     ++algorithm.dememwm.generated_history_proxy.enabled=true \
     ++algorithm.dememwm.generated_history_proxy.start_step=40000 \
     ++algorithm.dememwm.generated_history_proxy.ramp_steps=40000 \
@@ -77,7 +76,6 @@ srun python -m main \
     ++algorithm.dememwm.revisit.plucker_weight=0.10 \
     ++algorithm.dememwm.revisit.max_frames=2 \
     ++algorithm.dememwm.revisit.compress.downsample_ratio=3 \
-    ++algorithm.dememwm.stage_policy.noise_bucket_logging=true \
     ++algorithm.dememwm.cache.enabled=true \
     ++algorithm.dememwm.cache.device=cpu \
     ++algorithm.dememwm.cache.keep_raw_latents=all \

tests/test_dememwm_compression.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn.functional as F
 from dememwm_import_helper import install_dememwm_namespace
 
 install_dememwm_namespace()
@@ -18,16 +19,77 @@ def small_compressor(**kwargs):
     )
 
 
+def legacy_dynamic_forward(comp, latents, frame_indices, target_frame_indices, source_is_generated=None, exclude_latest_local_frames=None):
+    del source_is_generated
+    exclude_latest_local_frames = (
+        comp.exclude_latest_local_frames
+        if exclude_latest_local_frames is None
+        else int(exclude_latest_local_frames)
+    )
+    T_src, B, C, H, W = latents.shape
+    if target_frame_indices.ndim == 1:
+        target_frame_indices = target_frame_indices[:, None].expand(-1, B)
+    T_tgt = target_frame_indices.shape[0]
+    device = latents.device
+    n_spatial = (H // comp.patch_size) * (W // comp.patch_size)
+    T_out = comp._temporal_output_count()
+    num_slots = T_out * n_spatial
+    output_time_idx = comp._output_time_indices(device)
+    output_rows, mask_rows = [], []
+    for b in range(B):
+        src_frames_b = frame_indices[:, b]
+        tgt_outputs, tgt_masks = [], []
+        for j in range(T_tgt):
+            target = int(target_frame_indices[j, b].item())
+            valid_idx = (src_frames_b < target - exclude_latest_local_frames).nonzero(as_tuple=False).flatten()
+            if valid_idx.numel() == 0:
+                tgt_outputs.append(latents.new_zeros(num_slots, comp.dit_hidden_size))
+                tgt_masks.append(torch.zeros(num_slots, device=device, dtype=torch.bool))
+                continue
+            selected_frames = src_frames_b.index_select(0, valid_idx)
+            order = torch.argsort(selected_frames)
+            valid_idx = valid_idx.index_select(0, order)[-comp.max_source_frames:]
+            chunk = latents[valid_idx, b]
+            real_mask = torch.ones((chunk.shape[0],), device=device, dtype=torch.bool)
+            if chunk.shape[0] < comp.max_source_frames:
+                pad = chunk.new_zeros(comp.max_source_frames - chunk.shape[0], C, H, W)
+                chunk = torch.cat([pad, chunk], dim=0)
+                real_mask = torch.cat([
+                    torch.zeros((pad.shape[0],), device=device, dtype=torch.bool),
+                    real_mask,
+                ])
+            inp = chunk.clone()
+            inp[1:] = chunk[1:] - chunk[:-1]
+            x = inp.permute(1, 0, 2, 3).unsqueeze(0)
+            x = F.pad(x, (0, 0, 0, 0, comp.causal_pad, 0))
+            x = comp.conv3d(x)
+            x = x.squeeze(0).permute(1, 2, 3, 0)
+            x = comp.out_norm(x)
+            tokens = x.reshape(num_slots, comp.dit_hidden_size)
+            clamped_time_idx = output_time_idx.clamp(min=0, max=comp.max_source_frames - 1)
+            temporal_mask = (
+                (output_time_idx >= 0)
+                & (output_time_idx < comp.max_source_frames)
+                & real_mask.index_select(0, clamped_time_idx)
+            )
+            mask = temporal_mask[:, None].expand(T_out, n_spatial).reshape(num_slots)
+            tgt_outputs.append(tokens)
+            tgt_masks.append(mask)
+        output_rows.append(torch.stack(tgt_outputs))
+        mask_rows.append(torch.stack(tgt_masks))
+    return torch.stack(output_rows), torch.stack(mask_rows)
+
+
 def test_dynamic_compressor_shapes_and_budget():
     comp = small_compressor(exclude_latest_local_frames=0)
     latents = torch.randn(4, 2, 3, 2, 2)
     frame_indices = torch.arange(4)[:, None].repeat(1, 2)
     target = torch.tensor([[1, 2], [4, 4]])
-    tokens, mask, diag = comp(latents, frame_indices, None, target)
+    tokens, mask = comp(latents, frame_indices, None, target)
     assert tokens.shape == (2, 2, 2, 8)
     assert mask.shape == (2, 2, 2)
     assert mask[0, 0].any()
-    assert diag["selected_source_count"].max().item() <= 4
+    assert mask.sum(dim=-1).max().item() <= tokens.shape[2]
 
 
 def test_dynamic_compressor_abstains_without_old_enough_sources():
@@ -35,23 +97,23 @@ def test_dynamic_compressor_abstains_without_old_enough_sources():
     latents = torch.randn(2, 1, 3, 2, 2)
     frame_indices = torch.tensor([[5], [6]])
     target = torch.tensor([[8]])
-    tokens, mask, diag = comp(latents, frame_indices, None, target)
+    tokens, mask = comp(latents, frame_indices, None, target)
     assert tokens.shape == (1, 1, 2, 8)
+    assert not tokens.any()
     assert not mask.any()
-    assert diag["max_source_frame"].item() == -1
-    assert diag["dynamic_min_gap_to_target_per_target"].item() == -1
 
 
-def test_dynamic_compressor_reports_generated_fraction_and_no_future():
+def test_dynamic_compressor_ignores_generated_flags_and_excludes_future_sources():
     comp = small_compressor(exclude_latest_local_frames=0)
     latents = torch.randn(3, 1, 3, 2, 2)
     frame_indices = torch.tensor([[0], [2], [5]])
     generated = torch.tensor([[False], [True], [True]])
     target = torch.tensor([[3]])
-    _, mask, diag = comp(latents, frame_indices, None, target, generated)
+    tokens, mask = comp(latents, frame_indices, None, target, generated)
+    expected_tokens, expected_mask = legacy_dynamic_forward(comp, latents, frame_indices, target, generated)
     assert mask.any()
-    assert diag["max_source_frame"].item() == 2
-    assert 0.0 < diag["generated_source_fraction"].item() < 1.0
+    assert torch.allclose(tokens, expected_tokens, atol=1e-6, rtol=1e-6)
+    assert torch.equal(mask, expected_mask)
 
 
 def test_dynamic_compressor_excludes_c_short_overlap_and_keeps_shape():
@@ -59,13 +121,12 @@ def test_dynamic_compressor_excludes_c_short_overlap_and_keeps_shape():
     latents = torch.randn(5, 1, 3, 2, 2)
     frame_indices = torch.tensor([[0], [1], [2], [3], [4]])
     target = torch.tensor([[5]])
-    tokens, mask, diag = comp(latents, frame_indices, None, target)
+    tokens, mask = comp(latents, frame_indices, None, target)
+    expected_tokens, expected_mask = legacy_dynamic_forward(comp, latents, frame_indices, target)
     assert tokens.shape == (1, 1, 2, 8)
     assert mask.any()
-    assert diag["max_source_frame"].item() == 2
-    assert diag["dynamic_min_gap_to_target_per_target"].item() == 3
-    assert diag["dynamic_max_gap_to_target_per_target"].item() == 5
-    assert diag["dynamic_exclude_latest_local_frames"] == 2
+    assert torch.allclose(tokens, expected_tokens, atol=1e-6, rtol=1e-6)
+    assert torch.equal(mask, expected_mask)
 
 
 def test_cache_materialize_raw_latents_excludes_c_short_overlap():
@@ -91,7 +152,7 @@ def test_dynamic_compressor_preserves_grad_to_trainable_parts():
     latents = torch.randn(4, 1, 3, 2, 2)
     frame_indices = torch.arange(4)[:, None]
     target = torch.tensor([[4]])
-    tokens, mask, _ = comp(latents, frame_indices, None, target)
+    tokens, mask = comp(latents, frame_indices, None, target)
     assert mask.any()
     tokens.square().sum().backward()
     grads = [
@@ -107,6 +168,54 @@ def test_dynamic_compressor_selects_only_recent_valid_sources():
     latents = torch.randn(20, 1, 3, 2, 2)
     frame_indices = torch.arange(20)[:, None]
     target = torch.tensor([[10]])
-    _, mask, diag = comp(latents, frame_indices, None, target)
+    tokens, mask = comp(latents, frame_indices, None, target)
+    expected_tokens, expected_mask = legacy_dynamic_forward(comp, latents, frame_indices, target)
     assert mask.any()
-    assert diag["selected_source_count"].item() == 4
+    assert torch.allclose(tokens, expected_tokens, atol=1e-6, rtol=1e-6)
+    assert torch.equal(mask, expected_mask)
+
+
+def test_dynamic_compressor_batched_matches_legacy_loop():
+    torch.manual_seed(11)
+    comp = small_compressor(exclude_latest_local_frames=2)
+    latents = torch.randn(6, 3, 3, 2, 2)
+    frame_indices = torch.tensor([
+        [0, 4, 1],
+        [3, 1, 5],
+        [7, 8, 9],
+        [2, 6, 11],
+        [12, 3, 13],
+        [15, 10, 4],
+    ])
+    generated = torch.tensor([
+        [False, True, False],
+        [True, False, False],
+        [False, True, True],
+        [True, True, False],
+        [False, False, True],
+        [True, False, True],
+    ])
+    targets = torch.tensor([
+        [5, 5, 6],
+        [9, 9, 12],
+        [20, 11, 3],
+    ])
+
+    expected_tokens, expected_mask = legacy_dynamic_forward(
+        comp, latents, frame_indices, targets, generated
+    )
+    tokens, mask = comp(latents, frame_indices, None, targets, generated)
+
+    assert torch.allclose(tokens, expected_tokens, atol=1e-6, rtol=1e-6)
+    assert torch.equal(mask, expected_mask)
+
+
+def test_dynamic_compressor_handles_empty_source_tensor():
+    comp = small_compressor(exclude_latest_local_frames=2)
+    latents = torch.randn(0, 2, 3, 2, 2)
+    frame_indices = torch.empty(0, 2, dtype=torch.long)
+    target = torch.tensor([[5, 6], [7, 8]])
+    tokens, mask = comp(latents, frame_indices, None, target)
+    assert tokens.shape == (2, 2, 2, 8)
+    assert not tokens.any()
+    assert not mask.any()

tests/test_dememwm_config_static.py

@@ -7,7 +7,8 @@ def test_config_is_distinct_standalone_memory_dit_path():
     assert "base_video_dit" in text
     assert "memory_token_cross_attention: true" in text
     assert "dememwm:" in text
-    assert "debug_force_all_streams" in text
+    assert "diagnostics:" not in text
+    assert "noise_bucket_logging" not in text
     assert "ssm_memory" not in text
     assert "ssm_memory_ckpt_path" not in text
 
@@ -42,7 +43,6 @@ def test_current_config_contract_is_explicit_and_has_no_stale_sections():
         "plucker_grid_h: 4",
         "plucker_grid_w: 4",
         "plucker_focal_length: 0.35",
-        "noise_bucket_logging: true",
         "eval_ablation:",
         "branch: A_plus_D_plus_R_normal",
         "generated_history_proxy:",
@@ -64,6 +64,8 @@ def test_current_config_contract_is_explicit_and_has_no_stale_sections():
         "min_gap_frames",
         "max_chunks",
         "chunk_frames",
+        "diagnostics:",
+        "noise_bucket_logging",
     ):
         assert forbidden not in text
 
@@ -77,7 +79,6 @@ def test_full_scripts_use_consumed_contract_overrides():
         "algorithm.dememwm.revisit.fov_pitch_samples=20",
         "algorithm.dememwm.revisit.fov_depth_samples=20",
         "algorithm.dememwm.revisit.plucker_weight=0.10",
-        "algorithm.dememwm.stage_policy.noise_bucket_logging=true",
         "algorithm.dememwm.cache.keep_compressed_records=true",
     ]
     stale = [
@@ -95,6 +96,8 @@ def test_full_scripts_use_consumed_contract_overrides():
         "algorithm.dememwm.revisit.min_score",
         "algorithm.dememwm.revisit.generated_penalty",
         "algorithm.dememwm.rollout.",
+        "algorithm.dememwm.diagnostics",
+        "algorithm.dememwm.stage_policy.noise_bucket_logging",
     ]
     expected_by_script = {
         "scripts/dememwm_full_train.slurm": [
@@ -120,10 +123,9 @@ def test_algorithm_consumes_final_contract_guards_and_revisit_geometry_args():
         "_validate_config_contract",
         "deterministic_pose_retrieval",
         "exclude_latest_local_frames",
-        "noise_bucket_logging",
         "anchor_effective_enabled",
         "dynamic_effective_enabled",
-        "revisit_effective_enabled",
+        "revisit_stage_config_enabled",
         "stale DeMemWM config fields",
         "revisit_retrieval_kwargs",
         "fov_half_h",
@@ -135,56 +137,58 @@ def test_algorithm_consumes_final_contract_guards_and_revisit_geometry_args():
         assert token in text
     assert '_cfg_get(revisit_cfg, "topk"' not in text
     assert "lambda_abstain" not in text
+    assert "noise_bucket" not in text
+    assert "diagnostics" not in text
 
 
 def test_revisit_retrieval_is_deterministic_fov_plucker_contract():
     retrieval = Path("algorithms/worldmem/dememwm/retrieval.py").read_text()
     labels = Path("algorithms/worldmem/dememwm/labels.py").read_text()
     algorithm = Path("algorithms/worldmem/dememwm/algorithm.py").read_text()
-    diagnostics = Path("algorithms/worldmem/dememwm/diagnostics.py").read_text()
     for token in [
         "exclude_local_context_frames",
         "fov_overlap_threshold",
         "plucker_weight",
         "high_quality_fov_threshold",
-        "best_selected_frame_fov_overlap",
+        "dememwm_selected_frame_fov_overlap",
+        "best_selected_fov_overlap",
+        "best_selected_plucker_overlap",
+        "best_selected_gap_frames",
         "deterministic_fov_coverage_plucker",
         "valid_revisit_mask",
-        "revisit_candidate_frame_count",
-        "valid_candidate_label_count",
-        "valid_revisit_frame_count",
-        "no_valid_revisit_count",
-        "revisit_selected_frame_count",
-        "revisit_frame_fov_overlap",
-        "revisit_abstained_count",
+        "batched_revisit_select_positions",
     ]:
-        assert token in retrieval + labels + algorithm + diagnostics
+        assert token in retrieval + labels + algorithm
     assert "same_video" not in retrieval + labels
     assert "wrong_video" not in retrieval + labels
     for stale in ["time_weight", "pose_weight", "latent_weight", "generated_penalty", "min_score"]:
         assert f'self._cfg_get(revisit_cfg, "{stale}"' not in algorithm
 
 
-def test_dynamic_compressor_excludes_c_short_contract():
+def test_dynamic_compressor_excludes_c_short_contract_without_diagnostics():
     compression = Path("algorithms/worldmem/dememwm/compression.py").read_text()
     cache = Path("algorithms/worldmem/dememwm/cache.py").read_text()
     algorithm = Path("algorithms/worldmem/dememwm/algorithm.py").read_text()
     for token in [
         "exclude_latest_local_frames",
-        "src_frames_b < target - exclude_latest_local_frames",
+        "source_frames[:, None, :] < (target_frames[:, :, None] - int(exclude_latest_local_frames))",
+        "_local_context_exclusion_frames",
+    ]:
+        assert token in compression + cache + algorithm
+    for removed in [
         "dynamic_min_gap_to_target_per_target",
         "dynamic_max_gap_to_target_per_target",
         "dynamic_exclude_latest_local_frames",
-        "_local_context_exclusion_frames",
+        "selected_source_count",
+        "generated_source_fraction",
     ]:
-        assert token in compression + cache + algorithm
+        assert removed not in compression + algorithm
     assert "src_frames_b < target, as_tuple=False" not in compression
     assert "src < int(target), as_tuple=False" not in cache
 
 
-def test_eval_ablation_and_noise_bucket_logging_contracts():
+def test_eval_ablation_contracts_without_diagnostic_summaries():
     schedules = Path("algorithms/worldmem/dememwm/schedules.py").read_text()
-    diagnostics = Path("algorithms/worldmem/dememwm/diagnostics.py").read_text()
     algorithm = Path("algorithms/worldmem/dememwm/algorithm.py").read_text()
     for branch in [
         "memory_off",
@@ -202,15 +206,17 @@ def test_eval_ablation_and_noise_bucket_logging_contracts():
         "local_context_overlap_fake_revisit",
     ]:
         assert branch in schedules
-    for token in [
-        "noise_bucket_from_denoising_fraction",
-        "noise_bucket_from_noise_levels",
+    assert "apply_revisit_eval_corruption" in algorithm
+    for removed in [
         "summarize_noise_bucket_diagnostics",
-        "noise_bucket_id",
         "summarize_eval_ablation_diagnostics",
         "eval_bucket_true_revisit_count",
         "eval_bucket_no_valid_revisit_count",
         "eval_bucket_corrupted_memory_count",
-        "apply_revisit_eval_corruption",
+        "noise_bucket_id",
     ]:
-        assert token in schedules + diagnostics + algorithm
+        assert removed not in schedules + algorithm
+
+
+def test_dememwm_diagnostics_module_removed():
+    assert not Path("algorithms/worldmem/dememwm/diagnostics.py").exists()

tests/test_dememwm_dit_extension_static.py

@@ -73,17 +73,15 @@ def test_shared_memory_attention_zero_revisit_gate_matches_anchor_only():
     assert torch.allclose(out_anchor, out_packed, atol=1e-5)
 
 
-def test_memory_cross_attention_type_ids_record_stage_gates():
+def test_memory_cross_attention_type_ids_apply_stage_gates():
     attn = MemoryTokenCrossAttention(hidden_size=8, num_heads=2)
     x = torch.randn(1, 2, 1, 1, 8)
     c = torch.randn(1, 2, 8)
     mem = torch.randn(1, 2, 3, 8)
     mask = torch.ones(1, 2, 3, dtype=torch.bool)
     type_ids = torch.tensor([0, 1, 2])
-    _ = attn(x, c, mem, mask, memory_type_ids=type_ids, memory_token_gate=torch.ones(1, 2, 3))
-    assert torch.is_tensor(attn.last_type_gate_anchor_mean)
-    assert torch.is_tensor(attn.last_type_gate_dynamic_mean)
-    assert torch.is_tensor(attn.last_type_gate_revisit_mean)
+    out = attn(x, c, mem, mask, memory_type_ids=type_ids, memory_token_gate=torch.ones(1, 2, 3))
+    assert out.shape == x.shape
 
 
 def test_dit_accepts_dynamic_rank4_tokens_and_all_false_masks_without_nan():
@@ -120,7 +118,7 @@ def test_diffusion_methods_accept_option_c_kwargs_by_signature():
 
 
 
-def test_fresh_memory_cross_attention_identity_init_delta_ratio_is_zero():
+def test_fresh_memory_cross_attention_identity_init_delta_is_zero():
     attn = MemoryTokenCrossAttention(hidden_size=8, num_heads=2)
     x = torch.randn(1, 2, 1, 1, 8)
     c = torch.randn(1, 2, 8)
@@ -128,10 +126,9 @@ def test_fresh_memory_cross_attention_identity_init_delta_ratio_is_zero():
     mask = torch.ones(1, 2, 3, dtype=torch.bool)
     delta = attn(x, c, mem, mask, return_delta=True, residual_gate=torch.ones(1, 2, 1))
     assert torch.allclose(delta, torch.zeros_like(delta), atol=1e-6)
-    assert float(attn.last_delta_ratio.item()) <= 1e-7
 
 
-def test_fresh_dit_memory_on_matches_memory_off_and_reports_delta_ratio():
+def test_fresh_dit_memory_on_matches_memory_off():
     model = DiT(input_h=4, input_w=4, patch_size=2, in_channels=2, hidden_size=32, depth=1, num_heads=4, action_cond_dim=0, max_frames=2, reference_length=1, memory_token_cross_attention=True)
     x = torch.randn(1, 2, 2, 4, 4)
     t = torch.zeros(1, 2, dtype=torch.long)
@@ -152,6 +149,4 @@ def test_fresh_dit_memory_on_matches_memory_off_and_reports_delta_ratio():
         memory_dynamic_gate=torch.ones(1, 2, 1),
         memory_retrieval_gate=torch.ones(1, 2, 1),
     )
-    diagnostics = model.memory_adapter_delta_diagnostics()
     assert torch.allclose(out_on, out_off, atol=1e-6)
-    assert diagnostics["memory_adapter_delta_ratio_max"] <= 1e-7

tests/test_dememwm_eval_ablation.py

@@ -7,7 +7,6 @@ from dememwm_import_helper import install_dememwm_namespace
 install_dememwm_namespace()
 from algorithms.worldmem.dememwm.algorithm import MemoryDiTMixin
 from algorithms.worldmem.dememwm.compression import CausalConv3DDynamicCompressor
-from algorithms.worldmem.dememwm.diagnostics import summarize_eval_ablation_diagnostics
 from algorithms.worldmem.dememwm.schedules import (
     EVAL_ABLATION_BRANCHES,
     EVAL_ABLATION_BRANCH_TO_ID,
@@ -48,25 +47,6 @@ def test_wave9_branch_registry_is_exact_and_validated():
         normalize_eval_ablation_branch("ratio_sweep")
 
 
-def test_eval_ablation_diagnostics_bucket_counts():
-    diag = summarize_eval_ablation_diagnostics(
-        enabled=True,
-        branch="wrong_pose",
-        valid_revisit_mask=torch.tensor([[True, True, True, False]]),
-        no_valid_revisit_mask=torch.tensor([[False, False, False, True]]),
-        eval_corrupted_revisit_mask=torch.tensor([[False, True, True, False]]),
-    )
-    assert diag["eval_ablation_enabled"] is True
-    assert diag["eval_ablation_branch"] == "wrong_pose"
-    assert diag["eval_ablation_branch_id"] == EVAL_ABLATION_BRANCH_TO_ID["wrong_pose"]
-    assert diag["eval_bucket_true_revisit_count"] == 1
-    assert diag["eval_bucket_no_valid_revisit_count"] == 1
-    assert diag["eval_bucket_corrupted_memory_count"] == 2
-    assert diag["eval_bucket_true_revisit_fraction"] == pytest.approx(0.25)
-    assert diag["eval_bucket_no_valid_revisit_fraction"] == pytest.approx(0.25)
-    assert diag["eval_bucket_corrupted_memory_fraction"] == pytest.approx(0.5)
-
-
 class ConstantGate(torch.nn.Module):
     def __init__(self, value: float):
         super().__init__()
@@ -83,7 +63,6 @@ class DummyDeMemWM(MemoryDiTMixin):
             dememwm=types.SimpleNamespace(
                 enabled=True,
                 training_stage="stage_2",
-                debug_force_all_streams=False,
                 token_patch_size=2,
                 curriculum=types.SimpleNamespace(enabled=False),
                 anchor=types.SimpleNamespace(
@@ -108,7 +87,6 @@ class DummyDeMemWM(MemoryDiTMixin):
                     max_frames=2,
                     compress=types.SimpleNamespace(pool_h=1, pool_w=1),
                 ),
-                stage_policy=types.SimpleNamespace(noise_bucket_logging=True),
                 eval_ablation=types.SimpleNamespace(enabled=True, branch=branch),
                 generated_history_proxy=types.SimpleNamespace(enabled=False),
                 injection=types.SimpleNamespace(dit_hidden_size=8, anchor_gate=1.0, dynamic_gate=1.0, revisit_gate=1.0),
@@ -190,12 +168,11 @@ def test_eval_ablation_forced_revisit_controls_are_isolated_to_eval_branch():
     assert torch.allclose(normal.revisit_gate, torch.full_like(normal.revisit_gate, 0.25))
     assert torch.count_nonzero(forced_off.revisit_gate).item() == 0
     assert torch.equal(forced_on.revisit_gate, forced_on.valid_revisit_mask.to(dtype=forced_on.revisit_gate.dtype))
-    assert forced_on.diagnostics["eval_ablation_branch"] == "R_forced_on"
 
 
 def test_eval_ablation_corruption_branch_marks_corrupted_revisit_without_zeroing_gate():
     wrong_pose = _streams("wrong_pose")
     assert wrong_pose.valid_revisit_mask.all()
+    assert wrong_pose.revisit_mask.any()
     assert torch.allclose(wrong_pose.revisit_gate, torch.full_like(wrong_pose.revisit_gate, 0.25))
-    assert wrong_pose.diagnostics["eval_bucket_corrupted_memory_count"] == int(wrong_pose.valid_revisit_mask.numel())
-    assert wrong_pose.diagnostics["eval_bucket_true_revisit_count"] == 0
+    assert wrong_pose.revisit_best_selected_fov_overlap.shape == wrong_pose.valid_revisit_mask.shape

tests/test_dememwm_freeze_policy.py

@@ -62,8 +62,6 @@ def test_dit_freeze_keeps_requires_grad_stable_and_zeroes_optimizer_lr():
     assert frozen_state.dit_train_state == "frozen"
     assert full_dit_params
     assert all(param.requires_grad for param in full_dit_params)
-    assert model._last_dememwm_freeze_diagnostics["trainable_tensors_full_dit"] == 0
-    assert model._last_dememwm_freeze_diagnostics["requires_grad_tensors_full_dit"] == len(full_dit_params)
     assert all(not param.requires_grad for param in model.vae.parameters())
 
     for param in full_dit_params:
@@ -85,7 +83,6 @@ def test_dit_freeze_keeps_requires_grad_stable_and_zeroes_optimizer_lr():
 
     assert full_state.dit_train_state == "full"
     assert all(param.requires_grad for param in full_dit_params)
-    assert model._last_dememwm_freeze_diagnostics["trainable_tensors_full_dit"] == len(full_dit_params)
     assert lr_by_name["full_dit"] == 1.0e-5
     assert all(not param.requires_grad for param in model.vae.parameters())
 

tests/test_dememwm_generated_history_proxy.py

@@ -68,7 +68,7 @@ def test_generated_history_proxy_corrupts_only_returned_memory_source_and_marks_
     source_is_generated = torch.zeros(4, 1, dtype=torch.bool)
 
     torch.manual_seed(123)
-    corrupted, generated, diagnostics = model._apply_generated_history_proxy(
+    corrupted, generated = model._apply_generated_history_proxy(
         source_latents,
         source_is_generated,
     )
@@ -77,8 +77,7 @@ def test_generated_history_proxy_corrupts_only_returned_memory_source_and_marks_
     assert not torch.equal(corrupted, source_latents)
     assert generated.all()
     assert not source_is_generated.any()
-    assert diagnostics["generated_history_proxy_frame_count"] == 4
-    assert diagnostics["generated_history_proxy_frame_fraction"] == 1.0
+    assert generated.sum().item() == 4
 
 
 
@@ -88,7 +87,7 @@ def test_generated_history_proxy_respects_context_prefix_and_target_window_bound
     source_is_generated = torch.zeros(8, 1, dtype=torch.bool)
 
     torch.manual_seed(123)
-    corrupted, generated, diagnostics = model._apply_generated_history_proxy(
+    corrupted, generated = model._apply_generated_history_proxy(
         source_latents,
         source_is_generated,
         context_frame_count=3,
@@ -104,8 +103,7 @@ def test_generated_history_proxy_respects_context_prefix_and_target_window_bound
     assert torch.equal(corrupted[:3], source_latents[:3])
     assert not torch.equal(corrupted[3:6], source_latents[3:6])
     assert torch.equal(corrupted[6:], source_latents[6:])
-    assert diagnostics["generated_history_proxy_frame_count"] == 3
-    assert diagnostics["generated_history_proxy_frame_fraction"] == 3 / 8
+    assert generated.sum().item() == 3
 
 
 def test_generated_proxy_frames_skip_prefix_anchors_but_remain_revisit_sources():

tests/test_dememwm_injection_static.py

@@ -18,23 +18,19 @@ def _streams(dtype=torch.float32):
         anchor_gate=1.0,
         dynamic_gate=torch.ones(2, 3, 1) * 0.5,
         revisit_gate=0.0,
-        diagnostics={"selected_revisit_frame_record_ids": ["c1"], "dynamic_max_source_frame": torch.tensor(2)},
     )
 
 
-def test_injection_kwarg_names_masks_dtype_and_diagnostics():
-    kwargs, diag = InjectionAdapter()(_streams(), dtype=torch.float64)
+def test_injection_kwarg_names_masks_and_dtype():
+    kwargs = InjectionAdapter()(_streams(), dtype=torch.float64)
     assert set(kwargs) == {"memory_tokens", "memory_token_mask", "memory_dynamic_tokens", "memory_dynamic_mask", "memory_retrieval_tokens", "memory_retrieval_mask", "memory_anchor_gate", "memory_dynamic_gate", "memory_retrieval_gate"}
     assert kwargs["memory_tokens"].dtype == torch.float64
     assert kwargs["memory_dynamic_mask"].dtype == torch.bool
-    assert diag["anchor_valid_tokens"] == 6
-    assert diag["dynamic_valid_fraction"] > 0.0
-    assert diag["selected_revisit_frame_record_ids"] == ["c1"]
-    assert diag["max_source_frame"] == 2
+    assert kwargs["memory_retrieval_tokens"].dtype == torch.float64
 
 
 def test_injection_omit_disabled_streams():
-    kwargs, _ = InjectionAdapter(omit_disabled=True)(_streams())
+    kwargs = InjectionAdapter(omit_disabled=True)(_streams())
     assert kwargs["memory_retrieval_tokens"] is None
     assert kwargs["memory_retrieval_mask"] is None
     assert kwargs["memory_dynamic_tokens"] is not None

tests/test_dememwm_noise_bucket.py

@@ -1,102 +1,19 @@
-import torch
-from dememwm_import_helper import install_dememwm_namespace
-
-install_dememwm_namespace()
-from algorithms.worldmem.dememwm.algorithm import MemoryDiTMixin
-from algorithms.worldmem.dememwm.diagnostics import summarize_noise_bucket_diagnostics, summarize_revisit_diagnostics
-
-
-def test_revisit_diagnostics_report_mean_counts_per_target():
-    diag = summarize_revisit_diagnostics(
-        [
-            {"valid_revisit_frame_count": 6, "revisit_candidate_frame_count": 8, "revisit_selected_frame_count": 2},
-            {"valid_revisit_frame_count": 3, "revisit_candidate_frame_count": 5, "revisit_selected_frame_count": 1},
-            {"valid_revisit_frame_count": 0, "revisit_candidate_frame_count": 2, "revisit_selected_frame_count": 0, "no_valid_revisit_count": 1},
-        ],
-        valid_revisit_mask=torch.tensor([[True, True, False]]),
-    )
-    assert diag["revisit_candidate_frame_count"] == 5.0
-    assert diag["revisit_candidate_count"] == 5.0
-    assert diag["valid_revisit_frame_count"] == 3.0
-    assert diag["valid_revisit_count"] == 3.0
-    assert diag["revisit_selected_frame_count"] == 3
-    assert diag["no_valid_revisit_count"] == 1
-
-
-def test_noise_bucket_diagnostics_include_valid_and_no_valid_counts():
-    diag = summarize_noise_bucket_diagnostics(
-        noise_bucket="high",
-        valid_revisit_mask=torch.tensor([[True, True, False]]),
-        no_valid_revisit_mask=torch.tensor([[False, False, True]]),
-    )
-    assert diag["noise_bucket"] == "high"
-    assert diag["noise_bucket_id"] == 0
-    assert diag["noise_bucket_is_high"] == 1
-    assert diag["noise_bucket_is_mid"] == 0
-    assert diag["noise_bucket_high_target_count"] == 3
-    assert diag["noise_bucket_mid_target_count"] == 0
-    assert diag["valid_revisit_noise_bucket_high_count"] == 2
-    assert diag["no_valid_revisit_noise_bucket_high_count"] == 1
-
-
-def test_noise_bucket_diagnostics_count_per_target_bucket_ids():
-    diag = summarize_noise_bucket_diagnostics(
-        noise_bucket="mid",
-        noise_bucket_ids=torch.tensor([[0, 1, 2]]),
-        valid_revisit_mask=torch.tensor([[True, True, False]]),
-        no_valid_revisit_mask=torch.tensor([[False, False, True]]),
-    )
-    assert diag["noise_bucket"] == "mid"
-    assert diag["noise_bucket_id"] == 1
-    assert diag["noise_bucket_is_mid"] == 1
-    assert diag["noise_bucket_high_target_count"] == 1
-    assert diag["noise_bucket_mid_target_count"] == 1
-    assert diag["noise_bucket_low_target_count"] == 1
-    assert diag["valid_revisit_noise_bucket_high_count"] == 1
-    assert diag["valid_revisit_noise_bucket_mid_count"] == 1
-    assert diag["valid_revisit_noise_bucket_low_count"] == 0
-    assert diag["no_valid_revisit_noise_bucket_low_count"] == 1
-
-
-def test_noise_bucket_log_allowlist_keeps_target_counts_only():
-    keys = MemoryDiTMixin._TRAIN_DIAGNOSTIC_LOG_KEYS
-    for key in (
-        "anchor_valid_fraction",
-        "dynamic_valid_fraction",
-        "revisit_valid_fraction",
-        "valid_revisit_mask_fraction",
-        "revisit_candidate_count",
-        "valid_revisit_count",
-        "revisit_selected_count",
-        "revisit_fov_overlap_mean",
-        "revisit_incremental_fov_overlap_mean",
-        "revisit_plucker_overlap_mean",
-        "causal_violation_count",
-        "noise_bucket_id",
-        "noise_bucket_is_high",
-        "noise_bucket_is_mid",
-        "noise_bucket_is_low",
-        "revisit_raw_gate_mean",
-        "valid_revisit_noise_bucket_high_count",
-        "valid_revisit_noise_bucket_mid_count",
-        "valid_revisit_noise_bucket_low_count",
-        "no_valid_revisit_noise_bucket_high_count",
-        "no_valid_revisit_noise_bucket_mid_count",
-        "no_valid_revisit_noise_bucket_low_count",
-    ):
-        assert key not in keys
-    for key in [
-        "noise_bucket_target_count",
-        "noise_bucket_high_target_count",
-        "noise_bucket_mid_target_count",
-        "noise_bucket_low_target_count",
-        "revisit_candidate_frame_count",
-        "valid_revisit_frame_count",
-        "revisit_selected_frame_count",
-        "revisit_frame_fov_overlap_mean",
-        "revisit_best_selected_frame_fov_overlap_mean",
-        "revisit_best_selected_plucker_overlap_mean",
-        "revisit_best_selected_gap_frames_mean",
-        "revisit_learned_gate_mean",
+from pathlib import Path
+
+
+def test_training_logging_keeps_only_core_scalars():
+    algorithm = Path("algorithms/worldmem/dememwm/algorithm.py").read_text()
+    expected_logs = [
+        "training/loss",
+        "training/denoise_loss",
+        "training/revisit_gate",
+    ]
+    for key in expected_logs:
+        assert key in algorithm
+    for removed_key in [
+        "training/dynamic_max_source_frame",
+        "training/revisit_valid_count",
+        "training/noise_bucket_id",
+        "training/eval_bucket_true_revisit_count",
     ]:
-        assert key in keys
+        assert removed_key not in algorithm

tests/test_dememwm_preselection.py

@@ -107,7 +107,7 @@ def test_diverse_anchor_selection_uses_context_frames_not_literal_limit():
     frame_indices = torch.arange(8)[:, None]
     poses = torch.zeros((8, 1, 5), dtype=torch.float32)
     target_pose = torch.zeros((1, 1, 5), dtype=torch.float32)
-    anchor_banks, _, _, diag = harness._build_preselected_causal_memory_banks(
+    anchor_banks, _, _, _, _ = harness._build_preselected_causal_memory_banks(
         committed_latents=latents,
         source_frame_indices=frame_indices,
         source_is_generated=None,
@@ -133,7 +133,7 @@ def test_diverse_anchor_selection_uses_context_frames_not_literal_limit():
     )
 
     assert [int(record.frame_indices.item()) for record in anchor_banks[0].records] == [0, 1]
-    assert diag["preselected_anchor_projected_frame_count"] == 2
+    assert harness.project_call_lengths == [2]
 
 
 def test_streaming_diverse_anchor_selection_uses_context_frames():
@@ -168,7 +168,7 @@ def test_preselected_memory_banks_project_only_selected_frames():
     target_frame_indices = torch.tensor([[10], [11]])
     poses = torch.zeros((20, 1, 5), dtype=torch.float32)
     target_pose = torch.zeros((2, 1, 5), dtype=torch.float32)
-    anchor_banks, revisit_banks, tokens_per_frame, diag = harness._build_preselected_causal_memory_banks(
+    anchor_banks, revisit_banks, tokens_per_frame, selected_by_target, stats = harness._build_preselected_causal_memory_banks(
         committed_latents=latents,
         source_frame_indices=frame_indices,
         source_is_generated=None,
@@ -195,10 +195,10 @@ def test_preselected_memory_banks_project_only_selected_frames():
     assert tokens_per_frame == 1
     assert len(anchor_banks[0].records) == 4
     assert len(revisit_banks[0].records) == 3
-    assert diag["preselected_anchor_projected_frame_count"] == 4
-    assert diag["preselected_revisit_projected_frame_count"] == 3
-    assert diag["preselected_revisit_projected_frame_record_count"] == 3
-    assert harness.project_call_lengths == [4, 1, 1, 1]
+    assert selected_by_target is not None
+    assert stats is not None
+    assert [len(records) for records in selected_by_target[0]] == [2, 2]
+    assert harness.project_call_lengths == [4, 3]
     assert 20 not in harness.project_call_lengths
 
 
@@ -221,7 +221,7 @@ def test_preselected_revisit_projects_best_fov_frame_not_latest():
     )
     poses = pose_rows[:, None, :]
 
-    _, revisit_banks, _, _ = harness._build_preselected_causal_memory_banks(
+    _, revisit_banks, _, selected_by_target, _ = harness._build_preselected_causal_memory_banks(
         committed_latents=latents,
         source_frame_indices=frame_indices,
         source_is_generated=None,
@@ -246,6 +246,7 @@ def test_preselected_revisit_projects_best_fov_frame_not_latest():
         token_patch_size=2,
     )
 
+    assert selected_by_target is not None
     assert len(revisit_banks[0].records) == 1
     assert revisit_banks[0].records[0].metadata["dememwm_selected_frame_index"] == 1
     assert harness.project_call_values == [[1.0]]

tests/test_dememwm_retrieval.py

@@ -91,21 +91,19 @@ def test_revisit_candidates_require_causal_c_short_gap():
         exclude_local_context_frames=4,
     )
     assert [r.max_source_frame for r in result.records] == [1]
-    assert result.diagnostics["revisit_candidate_frame_count"] == 2
-    assert result.diagnostics["revisit_candidate_count"] == 2
-    assert result.diagnostics["valid_revisit_frame_count"] == 1
-    assert result.diagnostics["valid_revisit_count"] == 1
-    assert result.diagnostics["valid_candidate_label_count"] == 1
-    assert result.diagnostics["revisit_min_gap_to_target"] == 5
-    assert result.diagnostics["revisit_vectorized_frame_scorer_used"] == 1
+    assert result.selected_frame_ids == [1]
+    assert result.scores.numel() == 1
+    assert result.best_selected_gap_frames.item() == pytest.approx(5.0)
 
 
 def test_revisit_abstains_when_no_valid_candidate():
     result = deterministic_revisit_retrieval([rec(2, 2), rec(3, 3)], target_frame=6, topk=2, exclude_local_context_frames=4)
     assert result.records == []
-    assert result.diagnostics["abstained"] is True
-    assert result.diagnostics["valid_revisit_mask"] == 0
-    assert result.diagnostics["no_valid_revisit_count"] == 1
+    assert result.selected_frame_ids == []
+    assert result.scores.numel() == 0
+    assert result.best_selected_fov_overlap.item() == pytest.approx(0.0)
+    assert result.best_selected_plucker_overlap.item() == pytest.approx(0.0)
+    assert result.best_selected_gap_frames.item() == pytest.approx(-1.0)
 
 
 def test_revisit_retrieval_rejects_non_vectorized_inputs():
@@ -150,14 +148,11 @@ def test_fov_threshold_filters_candidates_without_action():
         exclude_local_context_frames=4,
         topk=4,
     )
-    assert result.diagnostics["selected_frame_record_ids"] == ["c0"]
-    assert result.diagnostics["valid_revisit_frame_count"] == 1
-    assert result.diagnostics["valid_revisit_count"] == 1
-    assert result.diagnostics["best_selected_fov_overlap"] == 1.0
-    assert result.diagnostics["revisit_best_selected_fov_overlap_max"] == 1.0
-    assert result.diagnostics["best_selected_gap_frames"] == 10
-    assert result.diagnostics["revisit_fov_overlap_max"] == 1.0
-    assert result.diagnostics["revisit_plucker_overlap_max"] > 0.0
+    assert [record.chunk_id for record in result.records] == ["c0"]
+    assert result.selected_frame_ids == [0]
+    assert result.best_selected_fov_overlap.item() == pytest.approx(1.0)
+    assert result.best_selected_plucker_overlap.item() > 0.0
+    assert result.best_selected_gap_frames.item() == pytest.approx(10.0)
 
 
 def test_pose_preselect_uses_local_position_and_view_direction_before_fov():
@@ -177,13 +172,8 @@ def test_pose_preselect_uses_local_position_and_view_direction_before_fov():
         pose_preselect_topk=1,
     )
 
-    assert result.diagnostics["selected_frame_record_ids"] == ["near_same_direction"]
-    assert result.diagnostics["revisit_pose_preselect_input_count"] == 3
-    assert result.diagnostics["revisit_pose_preselect_scored_count"] == 3
-    assert result.diagnostics["revisit_pose_preselect_selected_count"] == 1
-    assert result.diagnostics["revisit_exact_fov_candidate_count"] == 1
-    assert result.diagnostics["revisit_vectorized_frame_scorer_used"] == 1
-    assert abs(result.diagnostics["revisit_pose_preselect_min_distance"] - (1.0 / 30.0)) < 1e-6
+    assert [record.chunk_id for record in result.records] == ["near_same_direction"]
+    assert result.selected_frame_ids == [2]
 
 
 def test_selected_frame_carries_frame_metadata_for_projection():
@@ -197,15 +187,14 @@ def test_selected_frame_carries_frame_metadata_for_projection():
         topk=1,
     )
 
-    assert result.diagnostics["selected_frame_record_ids"] == ["frame_1"]
+    assert [record.chunk_id for record in result.records] == ["frame_1"]
     assert result.selected_frame_ids == [1]
     assert result.records[0].metadata["dememwm_selected_frame_index"] == 1
+    assert result.records[0].metadata["dememwm_selected_frame_fov_overlap"] == pytest.approx(1.0)
     assert result.records[0].metadata["dememwm_selected_frame_passes_high_quality"] is True
-    assert result.diagnostics["best_selected_frame_index"] == 1
-    assert result.diagnostics["best_selected_frame_fov_overlap"] == 1.0
 
 
-def test_high_quality_threshold_is_selected_target_diagnostic_only():
+def test_high_quality_threshold_marks_selected_frame_metadata():
     result = deterministic_revisit_retrieval(
         [rec(0, 0, pose=[0.0, 0.0, 0.0, 0.0, 0.0])],
         target_frame=10,
@@ -215,9 +204,10 @@ def test_high_quality_threshold_is_selected_target_diagnostic_only():
         exclude_local_context_frames=4,
         topk=1,
     )
-    assert result.diagnostics["selected_frame_record_ids"] == ["c0"]
-    assert result.diagnostics["valid_revisit_count"] == 1
-    assert 0.30 <= result.diagnostics["best_selected_fov_overlap"] < 0.70
+    assert [record.chunk_id for record in result.records] == ["c0"]
+    assert len(result.records) == 1
+    assert 0.30 <= result.best_selected_fov_overlap.item() < 0.70
+    assert result.records[0].metadata["dememwm_selected_frame_passes_high_quality"] is False
 
 
 def test_video_metadata_does_not_filter_revisit_candidates():
@@ -233,8 +223,8 @@ def test_video_metadata_does_not_filter_revisit_candidates():
         exclude_local_context_frames=4,
         topk=4,
     )
-    assert result.diagnostics["selected_frame_record_ids"] == ["c1", "c0"]
-    assert result.diagnostics["valid_revisit_count"] == 2
+    assert [record.chunk_id for record in result.records] == ["c1", "c0"]
+    assert result.selected_frame_ids == [1, 0]
 
 
 def test_tie_breaking_is_overlap_then_age_then_source_then_record_id():
@@ -244,4 +234,4 @@ def test_tie_breaking_is_overlap_then_age_then_source_then_record_id():
         rec(2, 2, pose=[0.0, 0.0, 0.0, 0.0, 0.0], chunk_id="c"),
     ]
     result = deterministic_revisit_retrieval(records, target_frame=10, target_pose=torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0]), exclude_local_context_frames=4, topk=3)
-    assert result.diagnostics["selected_frame_record_ids"] == ["c", "a", "b"]
+    assert [record.chunk_id for record in result.records] == ["c", "a", "b"]

tests/test_dememwm_schedules.py

@@ -4,9 +4,7 @@ from types import SimpleNamespace
 
 from algorithms.worldmem.dememwm.schedules import (
     compute_stream_gates,
-    noise_bucket_from_denoising_fraction,
-    noise_bucket_from_noise_levels,
-    noise_bucket_ids_from_noise_levels,
+    denoising_fraction_from_noise_levels,
     resolve_curriculum,
 )
 
@@ -43,37 +41,16 @@ def test_two_stage_curriculum_switches_at_full_stage_start():
     assert stage_1.anchor_enabled and stage_1.dynamic_enabled and stage_1.revisit_enabled
     assert stage_1.dit_train_state == "frozen"
     assert not hasattr(stage_1, "dit_late_blocks_trainable")
-    assert all("late" not in key for key in stage_1.diagnostics())
     assert stage_2.stage == "stage_2"
     assert stage_2.anchor_enabled and stage_2.dynamic_enabled and stage_2.revisit_enabled
     assert stage_2.dit_train_state == "full"
 
 
-def test_debug_force_all_streams_overrides_stage():
-    gates = compute_stream_gates("stage_1", debug_force_all_streams=True)
-    assert gates.anchor_enabled and gates.dynamic_enabled and gates.revisit_enabled
-    assert gates.reason == "debug_force_all_streams"
-
-
 def test_unknown_stage_fails():
     with pytest.raises(ValueError):
         compute_stream_gates("unknown")
 
 
-def test_noise_bucket_from_denoising_fraction():
-    assert noise_bucket_from_denoising_fraction(0.0) == "high"
-    assert noise_bucket_from_denoising_fraction(0.5) == "mid"
-    assert noise_bucket_from_denoising_fraction(1.0) == "low"
-
-
-def test_noise_bucket_from_training_noise_levels():
-    import torch
-    assert noise_bucket_from_noise_levels(torch.tensor([9, 8]), 10) == "high"
-    assert noise_bucket_from_noise_levels(torch.tensor([5, 4]), 10) == "mid"
-    assert noise_bucket_from_noise_levels(torch.tensor([1, 0]), 10) == "low"
-
-
-def test_noise_bucket_ids_from_training_noise_levels():
+def test_denoising_fraction_from_training_noise_levels():
     import torch
-    bucket_ids = noise_bucket_ids_from_noise_levels(torch.tensor([[9, 4, 0]]), 10)
-    assert bucket_ids.tolist() == [[0, 1, 2]]
+    assert denoising_fraction_from_noise_levels(torch.tensor([9, 0]), 10) == pytest.approx(0.5)

train_dememwm_full_berzelius.sh

@@ -51,7 +51,6 @@ srun python -m main \
     ++algorithm.context_frames=100 \
     ++algorithm.log_video=true \
     ++algorithm.diffusion.sampling_timesteps=20 \
-    ++algorithm.dememwm.debug_force_all_streams=false \
     ++algorithm.dememwm.generated_history_proxy.enabled=true \
     ++algorithm.dememwm.generated_history_proxy.start_step=40000 \
     ++algorithm.dememwm.generated_history_proxy.ramp_steps=40000 \
@@ -76,7 +75,6 @@ srun python -m main \
     ++algorithm.dememwm.revisit.plucker_weight=0.10 \
     ++algorithm.dememwm.revisit.max_frames=2 \
     ++algorithm.dememwm.revisit.compress.downsample_ratio=3 \
-    ++algorithm.dememwm.stage_policy.noise_bucket_logging=true \
     ++algorithm.dememwm.cache.enabled=true \
     ++algorithm.dememwm.cache.device=cpu \
     ++algorithm.dememwm.cache.keep_raw_latents=all \