#!/usr/bin/env python3 """Analyze timestep and cross-chunk feature reuse in Self-Forcing. The script runs the released four-step causal checkpoint without changing its outputs. Forward hooks capture sampled DiT hidden states, residual updates, the final velocity, and a low-dimensional dense projection used for optical-flow alignment. It then produces: * within-chunk and cross-chunk feature-pair metrics; * leave-one-prompt-out channel-wise probes for conditional chunk information; * shuffled, wrong-step, distant-chunk, zero, and noise controls; * motion-stratified raw/global/dense-flow alignment measurements. Each prompt is saved independently, so interrupted generation can be resumed. """ from __future__ import annotations import argparse import csv import json import math import os import random import sys import time from collections import defaultdict from pathlib import Path from typing import Any, Iterable def _preparse_gpu() -> str: parser = argparse.ArgumentParser(add_help=False) parser.add_argument("--gpu", default="2") args, _ = parser.parse_known_args() os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) return str(args.gpu) PHYSICAL_GPU = _preparse_gpu() import cv2 import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch import torch.nn.functional as F from omegaconf import OmegaConf REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from pipeline import CausalInferencePipeline from utils.misc import set_seed EXPECTED_LATENT_HEIGHT = 60 EXPECTED_LATENT_WIDTH = 104 FRAME_TOKEN_HEIGHT = 30 FRAME_TOKEN_WIDTH = 52 FRAME_SEQ_LENGTH = FRAME_TOKEN_HEIGHT * FRAME_TOKEN_WIDTH def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Self-Forcing cross-chunk feature-cache analysis" ) parser.add_argument("--gpu", default=PHYSICAL_GPU) parser.add_argument( "--config_path", type=Path, default=Path("configs/self_forcing_dmd.yaml") ) parser.add_argument( "--checkpoint_path", type=Path, default=Path("checkpoints/self_forcing_dmd.pt"), ) parser.add_argument( "--prompt_path", type=Path, default=Path("prompts/MovieGenVideoBench_extended.txt"), ) parser.add_argument("--output_dir", type=Path, required=True) parser.add_argument("--num_prompts", type=int, default=3) parser.add_argument("--num_frames", type=int, default=21) parser.add_argument("--seed", type=int, default=20260728) parser.add_argument( "--same_seed", action="store_true", help="Reset every prompt to --seed for paired cross-model evaluation.", ) parser.add_argument( "--layers", type=int, nargs="+", default=[0, 9, 19, 29] ) parser.add_argument("--max_tokens", type=int, default=256) parser.add_argument("--projection_dim", type=int, default=64) parser.add_argument("--ridge", type=float, default=1e-4) parser.add_argument("--use_ema", action="store_true", default=True) parser.add_argument("--no_ema", action="store_false", dest="use_ema") parser.add_argument("--overwrite", action="store_true") parser.add_argument( "--cosine_only", action="store_true", help="Generate pair metrics/heatmap only; skip probes and motion analysis.", ) parser.add_argument( "--analysis_only", action="store_true", help="Skip model loading and analyze existing per-prompt snapshots.", ) parser.add_argument( "--save_preview", action="store_true", help="Save a compact MP4 preview when torchvision video IO is available.", ) args = parser.parse_args() if args.num_frames % 3 != 0: parser.error("--num_frames must be divisible by the configured 3-frame chunk") if args.num_prompts < 3 and not args.cosine_only: parser.error("--num_prompts must be at least 3 for held-out/shuffle controls") return args def resolve_path(path: Path) -> Path: return path if path.is_absolute() else REPO_ROOT / path def write_csv(path: Path, rows: list[dict[str, Any]]) -> None: if not rows: return path.parent.mkdir(parents=True, exist_ok=True) fields: list[str] = [] for row in rows: for key in row: if key not in fields: fields.append(key) with path.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore") writer.writeheader() writer.writerows(rows) def read_prompts(path: Path, count: int) -> list[str]: with path.open("r", encoding="utf-8") as handle: prompts = [line.strip() for line in handle if line.strip()] if len(prompts) < count: raise ValueError(f"Requested {count} prompts, found {len(prompts)} in {path}") return prompts[:count] def regular_grid_indices( frames: int, height: int, width: int, max_tokens: int, device: torch.device ) -> tuple[torch.Tensor, torch.Tensor]: total = frames * height * width if max_tokens >= total: coords = torch.cartesian_prod( torch.arange(frames, device=device), torch.arange(height, device=device), torch.arange(width, device=device), ) else: per_frame = max(1, max_tokens // frames) h_count = max( 1, min(height, int(round(math.sqrt(per_frame * height / width)))) ) w_count = max(1, min(width, per_frame // h_count)) while frames * h_count * w_count > max_tokens and w_count > 1: w_count -= 1 while frames * h_count * w_count > max_tokens and h_count > 1: h_count -= 1 hs = ( torch.linspace(0, height - 1, h_count, device=device) .round() .long() .unique() ) ws = ( torch.linspace(0, width - 1, w_count, device=device) .round() .long() .unique() ) coords = torch.cartesian_prod( torch.arange(frames, device=device), hs, ws ) flat = ( coords[:, 0] * height * width + coords[:, 1] * width + coords[:, 2] ) return flat.long(), coords.long() class FeatureRecorder: def __init__( self, model: torch.nn.Module, layers: list[int], denoising_timesteps: Iterable[float], num_frame_per_block: int, max_tokens: int, projection_dim: int, ) -> None: self.model = model self.layers = sorted(set(int(value) for value in layers)) self.timesteps = [float(value) for value in denoising_timesteps] self.num_frame_per_block = int(num_frame_per_block) self.max_tokens = int(max_tokens) self.projection_dim = int(projection_dim) self.projection_layer = self.layers[-1] self.records: dict[str, dict[str, torch.Tensor]] = defaultdict(dict) self.projected: dict[str, torch.Tensor] = {} self.projected_by_layer: dict[str, torch.Tensor] = {} self.sample_coords: dict[str, torch.Tensor] = {} self.current: dict[str, Any] = {"active": False} self.handles: list[Any] = [] self._projection_cache: dict[tuple[int, str], torch.Tensor] = {} self._register() def _register(self) -> None: self.handles.append( self.model.register_forward_pre_hook(self._model_pre_hook, with_kwargs=True) ) self.handles.append(self.model.register_forward_hook(self._model_output_hook)) for layer in self.layers: if layer < 0 or layer >= len(self.model.blocks): raise ValueError( f"Layer {layer} outside model block range 0..{len(self.model.blocks)-1}" ) self.handles.append( self.model.blocks[layer].register_forward_hook( self._make_block_hook(layer) ) ) def close(self) -> None: for handle in self.handles: handle.remove() self.handles.clear() def reset(self) -> None: self.records = defaultdict(dict) self.projected = {} self.projected_by_layer = {} self.sample_coords = {} self.current = {"active": False} def _model_pre_hook( self, _module: torch.nn.Module, _args: tuple[Any, ...], kwargs: dict[str, Any] ) -> None: timestep = kwargs.get("t") current_start = int(kwargs.get("current_start", 0) or 0) if not isinstance(timestep, torch.Tensor) or timestep.numel() == 0: self.current = {"active": False} return value = float(timestep.detach().float().reshape(-1)[0].item()) distances = [abs(value - expected) for expected in self.timesteps] step = int(np.argmin(distances)) if distances[step] > 0.5: self.current = {"active": False, "timestep": value} return frames = int(timestep.shape[-1]) if timestep.ndim > 1 else 1 if frames != self.num_frame_per_block: self.current = {"active": False, "timestep": value} return start_frame = current_start // FRAME_SEQ_LENGTH chunk = start_frame // self.num_frame_per_block self.current = { "active": True, "chunk": int(chunk), "step": step, "timestep": value, "frames": frames, } def _key(self) -> str: return f"{self.current['chunk']}:{self.current['step']}" def _hidden_indices( self, tokens: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: frames = int(self.current["frames"]) if tokens.shape[1] != frames * FRAME_TOKEN_HEIGHT * FRAME_TOKEN_WIDTH: raise ValueError( f"Unexpected hidden token count {tokens.shape[1]} for {frames} frames" ) return regular_grid_indices( frames, FRAME_TOKEN_HEIGHT, FRAME_TOKEN_WIDTH, self.max_tokens, tokens.device, ) def _projection(self, dim: int, device: torch.device) -> torch.Tensor: cache_key = (dim, str(device)) if cache_key not in self._projection_cache: generator = torch.Generator(device="cpu").manual_seed(20260728 + dim) signs = torch.randint( 0, 2, (dim, self.projection_dim), generator=generator, dtype=torch.int8, ) projection = ( signs.float().mul_(2).sub_(1).div_(math.sqrt(self.projection_dim)) ) self._projection_cache[cache_key] = projection.to(device) return self._projection_cache[cache_key] def _make_block_hook(self, layer: int): def hook( _module: torch.nn.Module, inputs: tuple[torch.Tensor, ...], output: torch.Tensor, ) -> None: if not self.current.get("active", False): return if not inputs or not isinstance(output, torch.Tensor): return key = self._key() hidden_input = inputs[0] indices, coords = self._hidden_indices(output) hidden = output[0].index_select(0, indices) delta = (output - hidden_input)[0].index_select(0, indices) self.records[f"block_{layer}_hidden"][key] = ( hidden.detach().to(dtype=torch.float16, device="cpu") ) self.records[f"block_{layer}_delta"][key] = ( delta.detach().to(dtype=torch.float16, device="cpu") ) self.sample_coords["hidden"] = coords.detach().cpu() if layer in self.layers: projection = self._projection(output.shape[-1], output.device) dense = torch.matmul(output[0].float(), projection) frames = int(self.current["frames"]) dense = dense.reshape( frames, FRAME_TOKEN_HEIGHT, FRAME_TOKEN_WIDTH, self.projection_dim, ) dense_cpu = dense.detach().to(dtype=torch.float16, device="cpu") self.projected_by_layer[f"{layer}:{key}"] = dense_cpu # Preserve the legacy key layout for existing final-layer analyses. if layer == self.projection_layer: self.projected[key] = dense_cpu return hook def _model_output_hook( self, _module: torch.nn.Module, _inputs: tuple[Any, ...], output: torch.Tensor, ) -> None: if not self.current.get("active", False): return if not isinstance(output, torch.Tensor) or output.ndim != 5: return batch, channels, frames, height, width = output.shape if batch != 1: raise ValueError(f"Analysis expects batch size 1, got {batch}") tokens = output.permute(0, 2, 3, 4, 1).reshape( batch, frames * height * width, channels ) indices, coords = regular_grid_indices( frames, height, width, self.max_tokens, output.device ) sampled = tokens[0].index_select(0, indices) self.records["dit_output"][self._key()] = sampled.detach().to( dtype=torch.float16, device="cpu" ) self.sample_coords["dit_output"] = coords.detach().cpu() def state_dict(self) -> dict[str, Any]: return { "timesteps": self.timesteps, "layers": self.layers, "projection_layer": self.projection_layer, "projection_dim": self.projection_dim, "records": {stage: dict(values) for stage, values in self.records.items()}, "projected": dict(self.projected), "projected_by_layer": dict(self.projected_by_layer), "sample_coords": dict(self.sample_coords), } def build_pipeline(args: argparse.Namespace) -> tuple[CausalInferencePipeline, Any]: config = OmegaConf.load(resolve_path(args.config_path)) default_config = OmegaConf.load(REPO_ROOT / "configs/default_config.yaml") config = OmegaConf.merge(default_config, config) device = torch.device("cuda") pipeline = CausalInferencePipeline(config, device=device) checkpoint = torch.load( resolve_path(args.checkpoint_path), map_location="cpu", weights_only=False ) state_key = "generator_ema" if args.use_ema else "generator" pipeline.generator.load_state_dict(checkpoint[state_key]) del checkpoint pipeline = pipeline.to(dtype=torch.bfloat16) pipeline.text_encoder.to(device=device) pipeline.generator.to(device=device) pipeline.vae.to(device=device) pipeline.eval() return pipeline, config def downsample_anchors(video: torch.Tensor, latent_frames: int) -> np.ndarray: value = video[0].detach().float().cpu() frame_count = value.shape[0] if frame_count == latent_frames: indices = np.arange(latent_frames) else: indices = np.linspace(0, frame_count - 1, latent_frames).round().astype(int) anchors = value[indices].permute(0, 2, 3, 1).clamp(0, 1).numpy() result = [] for frame in anchors: frame_u8 = np.uint8(np.round(frame * 255.0)) result.append(cv2.resize(frame_u8, (416, 240), interpolation=cv2.INTER_AREA)) return np.stack(result) def maybe_save_preview(path: Path, anchors: np.ndarray) -> None: try: from torchvision.io import write_video repeated = np.repeat(anchors, 4, axis=0) tensor = torch.from_numpy(repeated) write_video(str(path), tensor, fps=16) except Exception as error: print(f"[preview] skipped: {error}", flush=True) @torch.inference_mode() def generate_snapshots(args: argparse.Namespace) -> list[Path]: output_dir = args.output_dir runs_dir = output_dir / "runs" runs_dir.mkdir(parents=True, exist_ok=True) prompts = read_prompts(resolve_path(args.prompt_path), args.num_prompts) expected_paths = [runs_dir / f"prompt_{index:02d}.pt" for index in range(len(prompts))] missing = [ path for path in expected_paths if args.overwrite or not path.exists() ] if not missing: print("[generation] all prompt snapshots already exist", flush=True) return expected_paths pipeline, config = build_pipeline(args) denoising_timesteps = [ float(value) for value in pipeline.denoising_step_list.detach().cpu().tolist() ] recorder = FeatureRecorder( model=pipeline.generator.model, layers=args.layers, denoising_timesteps=denoising_timesteps, num_frame_per_block=pipeline.num_frame_per_block, max_tokens=args.max_tokens, projection_dim=args.projection_dim, ) metadata = { "physical_gpu": args.gpu, "config_path": str(resolve_path(args.config_path)), "checkpoint_path": str(resolve_path(args.checkpoint_path)), "use_ema": args.use_ema, "num_prompts": args.num_prompts, "num_frames": args.num_frames, "num_frame_per_block": pipeline.num_frame_per_block, "denoising_timesteps": denoising_timesteps, "layers": args.layers, "max_tokens": args.max_tokens, "projection_dim": args.projection_dim, "seed": args.seed, "dtype": "bfloat16", } (output_dir / "experiment_config.json").write_text( json.dumps(metadata, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) try: for index, (prompt, path) in enumerate(zip(prompts, expected_paths)): if path.exists() and not args.overwrite: print(f"[generation] skip existing {path.name}", flush=True) continue recorder.reset() run_seed = args.seed if args.same_seed else args.seed + index set_seed(run_seed) noise = torch.randn( 1, args.num_frames, 16, EXPECTED_LATENT_HEIGHT, EXPECTED_LATENT_WIDTH, device="cuda", dtype=torch.bfloat16, ) torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() start = time.perf_counter() print( f"[generation] prompt {index + 1}/{len(prompts)} seed={run_seed}", flush=True, ) video, latents = pipeline.inference( noise=noise, text_prompts=[prompt], return_latents=True, initial_latent=None, low_memory=False, ) torch.cuda.synchronize() elapsed = time.perf_counter() - start peak_gib = torch.cuda.max_memory_allocated() / (1024**3) anchors = downsample_anchors(video, args.num_frames) state = { "run_index": index, "prompt": prompt, "seed": run_seed, "elapsed_s": elapsed, "peak_gpu_gib": peak_gib, "num_frames": args.num_frames, "num_frame_per_block": pipeline.num_frame_per_block, "latents": latents[0].detach().to( dtype=torch.float16, device="cpu" ), **recorder.state_dict(), } torch.save(state, path) np.savez_compressed(path.with_suffix(".anchors.npz"), frames=anchors) if args.save_preview: maybe_save_preview(path.with_suffix(".mp4"), anchors) print( f"[generation] saved {path.name}: {elapsed:.1f}s, peak={peak_gib:.1f} GiB", flush=True, ) del video, latents, noise, state pipeline.vae.model.clear_cache() torch.cuda.empty_cache() finally: recorder.close() return expected_paths def load_runs(paths: list[Path]) -> list[dict[str, Any]]: runs = [] for path in paths: if not path.exists(): raise FileNotFoundError(path) run = torch.load(path, map_location="cpu", weights_only=False) anchor_path = path.with_suffix(".anchors.npz") if not anchor_path.exists(): raise FileNotFoundError(anchor_path) run["anchors"] = np.load(anchor_path, allow_pickle=False)["frames"] run["path"] = str(path) runs.append(run) return runs def feature(run: dict[str, Any], stage: str, chunk: int, step: int) -> torch.Tensor: return run["records"][stage][f"{chunk}:{step}"].float() def projected_feature( run: dict[str, Any], chunk: int, step: int ) -> torch.Tensor: return run["projected"][f"{chunk}:{step}"].float() def available_chunks(run: dict[str, Any], stage: str) -> list[int]: return sorted( {int(key.split(":")[0]) for key in run["records"][stage].keys()} ) def available_steps(run: dict[str, Any], stage: str) -> list[int]: return sorted( {int(key.split(":")[1]) for key in run["records"][stage].keys()} ) def pair_metrics( reference: torch.Tensor, target: torch.Tensor, compute_cka: bool = True, ) -> dict[str, float]: if reference.shape != target.shape: raise ValueError(f"Pair shape mismatch: {reference.shape} vs {target.shape}") eps = 1e-8 x = reference.float() y = target.float() xf = x.reshape(-1) yf = y.reshape(-1) diff = yf - xf cosine = F.cosine_similarity(xf[None], yf[None], dim=1, eps=eps)[0] xc_flat = xf - xf.mean() yc_flat = yf - yf.mean() centered_cosine = torch.dot(xc_flat, yc_flat) / ( torch.linalg.vector_norm(xc_flat) * torch.linalg.vector_norm(yc_flat) + eps ) token_cosine = F.cosine_similarity(x, y, dim=1, eps=eps) rel_l2 = diff.square().mean().sqrt() / (xf.square().mean().sqrt() + eps) nmse = diff.square().mean() / (yc_flat.square().mean() + eps) if compute_cka: xc = x - x.mean(dim=0, keepdim=True) yc = y - y.mean(dim=0, keepdim=True) gram_x = xc @ xc.T gram_y = yc @ yc.T cka = (gram_x * gram_y).sum() / ( (gram_x.square().sum() * gram_y.square().sum()).sqrt() + eps ) else: cka = torch.tensor(float("nan")) quantiles = torch.quantile( token_cosine, torch.tensor([0.1, 0.5, 0.9], dtype=token_cosine.dtype), ) return { "cosine": float(cosine), "centered_cosine": float(centered_cosine), "linear_cka": float(cka), "rel_l2": float(rel_l2), "nmse": float(nmse), "token_cosine_mean": float(token_cosine.mean()), "token_cosine_p10": float(quantiles[0]), "token_cosine_p50": float(quantiles[1]), "token_cosine_p90": float(quantiles[2]), "reference_rms": float(xf.square().mean().sqrt()), "target_rms": float(yf.square().mean().sqrt()), } def collect_pair_rows( runs: list[dict[str, Any]], cosine_only: bool = False ) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] stages = sorted(runs[0]["records"]) if cosine_only: stages = [ stage for stage in stages if stage.startswith("block_") and stage.endswith("_hidden") ] for run_index, run in enumerate(runs): shuffled_run = runs[(run_index + 1) % len(runs)] for stage in stages: chunks = available_chunks(run, stage) steps = available_steps(run, stage) def append( comparison: str, ref_run: dict[str, Any], ref_chunk: int, ref_step: int, target_chunk: int, target_step: int, ) -> None: metrics = pair_metrics( feature(ref_run, stage, ref_chunk, ref_step), feature(run, stage, target_chunk, target_step), compute_cka=not cosine_only, ) rows.append( { "run": run_index, "comparison": comparison, "stage": stage, "reference_chunk": ref_chunk, "target_chunk": target_chunk, "reference_step": ref_step, "target_step": target_step, **metrics, } ) for chunk in chunks: for step in steps[1:]: append( "within_adjacent", run, chunk, step - 1, chunk, step, ) if chunk < 1: continue for step in steps: append( "cross_same", run, chunk - 1, step, chunk, step, ) if cosine_only: continue append( "cross_video_shuffle", shuffled_run, chunk - 1, step, chunk, step, ) if step > 0: append( "cross_wrong_step", run, chunk - 1, step - 1, chunk, step, ) if chunk > 1: append( "cross_distant", run, chunk - 2, step, chunk, step, ) return rows def predictor_columns( name: str, within: torch.Tensor, cross: torch.Tensor, distant: torch.Tensor, wrong: torch.Tensor, batch: torch.Tensor, noise_seed: int, ) -> list[torch.Tensor]: ones = torch.ones_like(within) shifted = cross.roll(shifts=max(1, cross.shape[0] // 2), dims=0) generator = torch.Generator(device="cpu").manual_seed(noise_seed) noise = torch.randn( cross.shape, generator=generator, dtype=cross.dtype ) noise = noise * cross.std(dim=0, keepdim=True).clamp_min(1e-6) noise = noise + cross.mean(dim=0, keepdim=True) mapping = { "within_affine": [within, ones], "within_quadratic": [within, within.square(), ones], "cross_affine": [cross, ones], "fusion_same": [within, cross, ones], "fusion_distant": [within, distant, ones], "fusion_token_shift": [within, shifted, ones], "fusion_wrong_step": [within, wrong, ones], "fusion_batch_shuffle": [within, batch, ones], "fusion_zero": [within, torch.zeros_like(cross), ones], "fusion_noise": [within, noise, ones], } return mapping[name] PROBE_NAMES = [ "within_affine", "within_quadratic", "cross_affine", "fusion_same", "fusion_distant", "fusion_token_shift", "fusion_wrong_step", "fusion_batch_shuffle", "fusion_zero", "fusion_noise", ] def gather_probe_data( runs: list[dict[str, Any]], run_indices: list[int], stage: str, step: int, ) -> dict[str, torch.Tensor]: buckets: dict[str, list[torch.Tensor]] = defaultdict(list) for run_index in run_indices: run = runs[run_index] other = runs[(run_index + 1) % len(runs)] chunks = available_chunks(run, stage) for chunk in chunks: # Use c >= 2 for every probe so correct, distant, and all other # controls are evaluated on exactly the same target tokens. if chunk < 2: continue buckets["target"].append(feature(run, stage, chunk, step)) buckets["within"].append(feature(run, stage, chunk, step - 1)) buckets["cross"].append(feature(run, stage, chunk - 1, step)) buckets["distant"].append(feature(run, stage, chunk - 2, step)) buckets["wrong"].append(feature(run, stage, chunk - 1, step - 1)) buckets["batch"].append(feature(other, stage, chunk - 1, step)) if not buckets: raise ValueError(f"No probe data for stage={stage}, step={step}") return {key: torch.cat(values, dim=0).float() for key, values in buckets.items()} def fit_channelwise_probe( columns: list[torch.Tensor], target: torch.Tensor, ridge: float, ) -> torch.Tensor: design = torch.stack(columns, dim=-1).double() y = target.double() gram = torch.einsum("ndp,ndq->dpq", design, design) rhs = torch.einsum("ndp,nd->dp", design, y) feature_count = gram.shape[-1] diagonal_scale = ( gram.diagonal(dim1=-2, dim2=-1).mean(dim=-1).clamp_min(1e-8) ) regularizer = ( torch.eye(feature_count, dtype=gram.dtype)[None] * (ridge * diagonal_scale)[:, None, None] ) regularizer[:, -1, -1] = 0.0 try: weights = torch.linalg.solve(gram + regularizer, rhs.unsqueeze(-1)).squeeze(-1) except torch.linalg.LinAlgError: weights = ( torch.linalg.pinv(gram + regularizer) @ rhs.unsqueeze(-1) ).squeeze(-1) return weights.float() def apply_channelwise_probe( columns: list[torch.Tensor], weights: torch.Tensor ) -> torch.Tensor: design = torch.stack(columns, dim=-1).float() return torch.einsum("ndp,dp->nd", design, weights) def prediction_metrics(prediction: torch.Tensor, target: torch.Tensor) -> dict[str, float]: eps = 1e-8 pred = prediction.float() y = target.float() error = pred - y mse = error.square().mean() variance = (y - y.mean()).square().mean() nrmse = mse.sqrt() / (variance.sqrt() + eps) r2 = 1.0 - mse / (variance + eps) cosine = F.cosine_similarity( pred.reshape(1, -1), y.reshape(1, -1), dim=1, eps=eps )[0] return { "mse": float(mse), "nrmse": float(nrmse), "r2": float(r2), "cosine": float(cosine), } def run_conditional_probes( runs: list[dict[str, Any]], ridge: float ) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] stages = sorted(runs[0]["records"]) steps = available_steps(runs[0], stages[0]) for stage in stages: for step in steps[1:]: for held_out in range(len(runs)): train_indices = [index for index in range(len(runs)) if index != held_out] train = gather_probe_data(runs, train_indices, stage, step) test = gather_probe_data(runs, [held_out], stage, step) for probe_name in PROBE_NAMES: train_columns = predictor_columns( probe_name, train["within"], train["cross"], train["distant"], train["wrong"], train["batch"], noise_seed=1000 + held_out * 100 + step, ) test_columns = predictor_columns( probe_name, test["within"], test["cross"], test["distant"], test["wrong"], test["batch"], noise_seed=2000 + held_out * 100 + step, ) weights = fit_channelwise_probe( train_columns, train["target"], ridge=ridge ) prediction = apply_channelwise_probe(test_columns, weights) rows.append( { "held_out_run": held_out, "stage": stage, "step": step, "probe": probe_name, "train_tokens": int(train["target"].shape[0]), "test_tokens": int(test["target"].shape[0]), **prediction_metrics(prediction, test["target"]), } ) baseline_lookup = { (row["held_out_run"], row["stage"], row["step"], row["probe"]): row for row in rows if row["probe"] in {"within_affine", "within_quadratic"} } for row in rows: key = (row["held_out_run"], row["stage"], row["step"]) for baseline_name in ("within_affine", "within_quadratic"): baseline = baseline_lookup[(*key, baseline_name)] row[f"mse_reduction_vs_{baseline_name}"] = ( baseline["mse"] - row["mse"] ) / max(baseline["mse"], 1e-12) row[f"r2_gain_vs_{baseline_name}"] = row["r2"] - baseline["r2"] return rows def gray(frame: np.ndarray) -> np.ndarray: return cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) def farneback(source: np.ndarray, target: np.ndarray) -> np.ndarray: return cv2.calcOpticalFlowFarneback( gray(source), gray(target), None, pyr_scale=0.5, levels=4, winsize=21, iterations=5, poly_n=7, poly_sigma=1.5, flags=0, ) def resize_flow(flow: np.ndarray, height: int, width: int) -> torch.Tensor: source_height, source_width = flow.shape[:2] resized = cv2.resize(flow, (width, height), interpolation=cv2.INTER_AREA) resized[..., 0] *= width / source_width resized[..., 1] *= height / source_height return torch.from_numpy(resized).permute(2, 0, 1).float() def warp_feature( source: torch.Tensor, target_to_source_flow: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: height, width, _ = source.shape yy, xx = torch.meshgrid( torch.arange(height, dtype=torch.float32), torch.arange(width, dtype=torch.float32), indexing="ij", ) sample_x = xx + target_to_source_flow[0] sample_y = yy + target_to_source_flow[1] grid = torch.stack( [ 2.0 * sample_x / max(width - 1, 1) - 1.0, 2.0 * sample_y / max(height - 1, 1) - 1.0, ], dim=-1, )[None] value = source.permute(2, 0, 1)[None].float() warped = F.grid_sample( value, grid, mode="bilinear", padding_mode="zeros", align_corners=True )[0].permute(1, 2, 0) mask = ( (sample_x >= 0) & (sample_x <= width - 1) & (sample_y >= 0) & (sample_y <= height - 1) ) return warped, mask def masked_cosine( left: torch.Tensor, right: torch.Tensor, mask: torch.Tensor | None = None ) -> float: value = F.cosine_similarity(left.float(), right.float(), dim=-1, eps=1e-8) if mask is not None: if not bool(mask.any()): return float("nan") value = value[mask] return float(value.mean()) def collect_motion_rows(runs: list[dict[str, Any]]) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] for run_index, run in enumerate(runs): anchors = run["anchors"] chunk_size = int(run["num_frame_per_block"]) chunk_count = int(run["num_frames"]) // chunk_size steps = sorted( {int(key.split(":")[1]) for key in run["projected"].keys()} ) for chunk in range(1, chunk_count): previous_anchor_index = chunk * chunk_size - 1 previous_frame = anchors[previous_anchor_index] boundary_flow = farneback( previous_frame, anchors[chunk * chunk_size] ) median_flow = np.median( boundary_flow.reshape(-1, 2), axis=0 ) camera_motion = float(np.linalg.norm(median_flow)) residual = boundary_flow - median_flow[None, None] object_motion = float( np.linalg.norm(residual, axis=-1).mean() ) total_motion = float( np.linalg.norm(boundary_flow, axis=-1).mean() ) for step in steps: source_map = projected_feature(run, chunk - 1, step)[-1] target_map = projected_feature(run, chunk, step) same_slot_source = projected_feature(run, chunk - 1, step) per_slot: list[dict[str, float]] = [] for slot in range(chunk_size): target = target_map[slot] same_slot = same_slot_source[slot] raw_boundary = source_map target_frame = anchors[chunk * chunk_size + slot] backward = farneback(target_frame, previous_frame) feature_flow = resize_flow( backward, FRAME_TOKEN_HEIGHT, FRAME_TOKEN_WIDTH ) global_flow = torch.zeros_like(feature_flow) global_flow[0].fill_(float(np.median(feature_flow[0].numpy()))) global_flow[1].fill_(float(np.median(feature_flow[1].numpy()))) global_aligned, global_mask = warp_feature( source_map, global_flow ) flow_aligned, flow_mask = warp_feature(source_map, feature_flow) per_slot.append( { "same_slot_cosine": masked_cosine(target, same_slot), "boundary_raw_cosine": masked_cosine( target, raw_boundary ), "global_aligned_cosine": masked_cosine( target, global_aligned, global_mask ), "flow_aligned_cosine": masked_cosine( target, flow_aligned, flow_mask ), "valid_flow_ratio": float(flow_mask.float().mean()), } ) rows.append( { "run": run_index, "chunk": chunk, "step": step, "total_motion": total_motion, "camera_motion": camera_motion, "object_motion": object_motion, **{ key: float(np.nanmean([item[key] for item in per_slot])) for key in per_slot[0] }, } ) motion_values = np.asarray([row["total_motion"] for row in rows]) if len(motion_values) >= 3: low, high = np.quantile(motion_values, [1 / 3, 2 / 3]) for row in rows: value = row["total_motion"] row["motion_bin"] = "low" if value <= low else "high" if value > high else "medium" return rows def group_mean( rows: list[dict[str, Any]], keys: list[str], metrics: list[str] ) -> list[dict[str, Any]]: groups: dict[tuple[Any, ...], list[dict[str, Any]]] = defaultdict(list) for row in rows: groups[tuple(row[key] for key in keys)].append(row) result = [] for group, values in sorted(groups.items(), key=lambda item: tuple(map(str, item[0]))): output = {key: value for key, value in zip(keys, group)} output["count"] = len(values) for metric in metrics: finite = [ float(row[metric]) for row in values if metric in row and math.isfinite(float(row[metric])) ] output[metric] = float(np.mean(finite)) if finite else float("nan") result.append(output) return result def paired_probe_reduction( rows: list[dict[str, Any]], stage: str, reference_probe: str, candidate_probe: str = "fusion_same", ) -> dict[str, Any]: lookup = { (int(row["held_out_run"]), int(row["step"]), row["probe"]): row for row in rows if row["stage"] == stage and row["probe"] in {reference_probe, candidate_probe} } pairs = sorted( { (held_out, step) for held_out, step, probe in lookup if probe == candidate_probe and (held_out, step, reference_probe) in lookup } ) reductions = [] for held_out, step in pairs: reference = float(lookup[(held_out, step, reference_probe)]["mse"]) candidate = float(lookup[(held_out, step, candidate_probe)]["mse"]) reductions.append((reference - candidate) / max(reference, 1e-12)) return { "reference": reference_probe, "paired_count": len(reductions), "mean_mse_reduction": ( float(np.mean(reductions)) if reductions else float("nan") ), "median_mse_reduction": ( float(np.median(reductions)) if reductions else float("nan") ), "wins": int(sum(value > 0 for value in reductions)), } def plot_similarity(pair_rows: list[dict[str, Any]], output_dir: Path) -> None: stages = [ stage for stage in sorted({row["stage"] for row in pair_rows}) if stage.endswith("_hidden") ] comparisons = [ "within_adjacent", "cross_same", "cross_wrong_step", "cross_distant", "cross_video_shuffle", ] steps = sorted({int(row["target_step"]) for row in pair_rows}) fig, axes = plt.subplots( len(stages), 2, figsize=(12, max(3.2, 2.8 * len(stages))), squeeze=False ) for stage_index, stage in enumerate(stages): for metric_index, metric in enumerate(["cosine", "nmse"]): matrix = np.full((len(comparisons), len(steps)), np.nan) for row_index, comparison in enumerate(comparisons): for col_index, step in enumerate(steps): values = [ float(row[metric]) for row in pair_rows if row["stage"] == stage and row["comparison"] == comparison and int(row["target_step"]) == step ] if values: matrix[row_index, col_index] = np.mean(values) ax = axes[stage_index, metric_index] image = ax.imshow( matrix, aspect="auto", cmap="viridis_r" if metric == "nmse" else "viridis", ) ax.set_title(f"{stage}: {metric}") ax.set_xticks(range(len(steps)), labels=steps) ax.set_yticks(range(len(comparisons)), labels=comparisons) ax.set_xlabel("target denoising step") fig.colorbar(image, ax=ax, fraction=0.03) fig.tight_layout() fig.savefig(output_dir / "feature_redundancy_heatmap.png", dpi=180) plt.close(fig) def plot_probe_gain(probe_rows: list[dict[str, Any]], output_dir: Path) -> None: aggregate = group_mean( probe_rows, ["stage", "step", "probe"], ["mse", "nrmse", "r2", "mse_reduction_vs_within_quadratic"], ) stages = [ stage for stage in sorted({row["stage"] for row in aggregate}) if stage.endswith("_hidden") ] probes = [ "fusion_same", "fusion_distant", "fusion_token_shift", "fusion_wrong_step", "fusion_batch_shuffle", ] steps = sorted({int(row["step"]) for row in aggregate}) fig, axes = plt.subplots( len(stages), 1, figsize=(9, max(3.4, 3.0 * len(stages))), squeeze=False ) for stage_index, stage in enumerate(stages): ax = axes[stage_index, 0] for probe in probes: values = [] for step in steps: match = [ row for row in aggregate if row["stage"] == stage and int(row["step"]) == step and row["probe"] == probe ] values.append( match[0]["mse_reduction_vs_within_quadratic"] if match else np.nan ) ax.plot(steps, values, marker="o", label=probe) ax.axhline(0, color="black", linewidth=0.8) ax.set_title(stage) ax.set_ylabel("MSE reduction vs within quadratic") ax.set_xlabel("target denoising step") ax.legend(fontsize=8) ax.grid(alpha=0.25) fig.tight_layout() fig.savefig(output_dir / "conditional_chunk_gain.png", dpi=180) plt.close(fig) def plot_motion(motion_rows: list[dict[str, Any]], output_dir: Path) -> None: fig, axes = plt.subplots(1, 2, figsize=(12, 4.5)) axes[0].scatter( [row["total_motion"] for row in motion_rows], [row["boundary_raw_cosine"] for row in motion_rows], s=18, alpha=0.65, label="boundary raw", ) axes[0].scatter( [row["total_motion"] for row in motion_rows], [row["flow_aligned_cosine"] for row in motion_rows], s=18, alpha=0.65, label="dense-flow aligned", ) axes[0].set_xlabel("optical-flow magnitude") axes[0].set_ylabel("feature cosine") axes[0].legend() axes[0].grid(alpha=0.25) aggregate = group_mean( motion_rows, ["motion_bin"], [ "same_slot_cosine", "boundary_raw_cosine", "global_aligned_cosine", "flow_aligned_cosine", ], ) bins = ["low", "medium", "high"] metrics = [ "same_slot_cosine", "boundary_raw_cosine", "global_aligned_cosine", "flow_aligned_cosine", ] width = 0.18 x = np.arange(len(bins)) for metric_index, metric in enumerate(metrics): values = [] for name in bins: match = [row for row in aggregate if row["motion_bin"] == name] values.append(match[0][metric] if match else np.nan) axes[1].bar( x + (metric_index - 1.5) * width, values, width=width, label=metric.replace("_cosine", ""), ) axes[1].set_xticks(x, bins) axes[1].set_ylabel("mean feature cosine") axes[1].set_xlabel("motion bin") axes[1].legend(fontsize=8) axes[1].grid(axis="y", alpha=0.25) fig.tight_layout() fig.savefig(output_dir / "motion_alignment_analysis.png", dpi=180) plt.close(fig) def markdown_table( rows: list[dict[str, Any]], columns: list[str], digits: int = 4 ) -> str: lines = [ "| " + " | ".join(columns) + " |", "|" + "|".join(["---"] * len(columns)) + "|", ] for row in rows: values = [] for column in columns: value = row.get(column, "") if isinstance(value, float): values.append(f"{value:.{digits}f}") else: values.append(str(value)) lines.append("| " + " | ".join(values) + " |") return "\n".join(lines) def build_report( runs: list[dict[str, Any]], pair_rows: list[dict[str, Any]], probe_rows: list[dict[str, Any]], motion_rows: list[dict[str, Any]], output_dir: Path, ) -> None: final_stage = f"block_{runs[0]['projection_layer']}_hidden" pair_summary = group_mean( [ row for row in pair_rows if row["stage"] == final_stage ], ["comparison"], ["cosine", "linear_cka", "rel_l2", "nmse", "token_cosine_p10"], ) probe_summary = group_mean( [ row for row in probe_rows if row["stage"] == final_stage and row["probe"] in set(PROBE_NAMES) ], ["probe"], [ "nrmse", "r2", "mse_reduction_vs_within_affine", "mse_reduction_vs_within_quadratic", ], ) motion_summary = group_mean( motion_rows, ["motion_bin"], [ "total_motion", "same_slot_cosine", "boundary_raw_cosine", "global_aligned_cosine", "flow_aligned_cosine", ], ) control_order = [ "within_affine", "within_quadratic", "fusion_distant", "fusion_wrong_step", "fusion_token_shift", "fusion_batch_shuffle", "fusion_zero", "fusion_noise", ] control_summary = [ paired_probe_reduction( probe_rows, stage=final_stage, reference_probe=control, ) for control in control_order ] control_lookup = {row["reference"]: row for row in control_summary} affine = control_lookup["within_affine"] shuffled = control_lookup["fusion_batch_shuffle"] conditional_statement = ( f"在主层 `{final_stage}` 上，加入正确的前一 chunk 同 timestep 特征，" f"相对 within-only affine probe 平均降低 held-out MSE " f"{100*affine['mean_mse_reduction']:.2f}%，并在 " f"{affine['wins']}/{affine['paired_count']} 个 " f"prompt–timestep 配对中取得改善。相对参数量一致的跨视频 " f"shuffle 对照，MSE 平均降低 " f"{100*shuffled['mean_mse_reduction']:.2f}%。" ) pair_lookup = {row["comparison"]: row for row in pair_summary} same_pair = pair_lookup.get("cross_same") random_pair = pair_lookup.get("cross_video_shuffle") pair_statement = "" if same_pair is not None and random_pair is not None: pair_statement = ( f"正确相邻 chunk 的平均 cosine/CKA 为 " f"{same_pair['cosine']:.4f}/{same_pair['linear_cka']:.4f}，" f"跨视频 shuffle 为 " f"{random_pair['cosine']:.4f}/{random_pair['linear_cka']:.4f}。" ) total_motion = np.asarray( [float(row["total_motion"]) for row in motion_rows], dtype=np.float64 ) raw_cosine = np.asarray( [float(row["boundary_raw_cosine"]) for row in motion_rows], dtype=np.float64, ) flow_cosine = np.asarray( [float(row["flow_aligned_cosine"]) for row in motion_rows], dtype=np.float64, ) motion_correlation = ( float(np.corrcoef(total_motion, raw_cosine)[0, 1]) if len(motion_rows) > 1 else float("nan") ) flow_gain = float(np.mean(flow_cosine - raw_cosine)) flow_wins = int(np.sum(flow_cosine > raw_cosine)) motion_statement = ( f"运动强度与未对齐跨 chunk cosine 的 Pearson 相关系数为 " f"{motion_correlation:.3f}；dense-flow 对齐平均恢复 " f"{flow_gain:.4f} cosine，并在 {flow_wins}/{len(motion_rows)} " f"个 chunk–timestep 样本上改善。" ) report = f"""# Self-Forcing Feature Cache 分析结果 ## 实验配置 - Prompts：{len(runs)} - 每条视频 latent frames：{runs[0]['num_frames']} - 每个 AR chunk latent frames：{runs[0]['num_frame_per_block']} - Denoising timesteps：{runs[0]['timesteps']} - Hook layers：{runs[0]['layers']} - 主分析 stage：`{final_stage}` ## 核心观察 {conditional_statement} {pair_statement} {motion_statement} 以下结果是 3 条 prompt 的 pilot 分析。条件 probe 按 prompt 留一测试， 统计单位是 held-out prompt 与 timestep；不能将 token 数量解释为独立视频样本数， 也不能据此声称数据集级统计显著性。 ## 主层特征配对 {markdown_table(pair_summary, ['comparison', 'count', 'cosine', 'linear_cka', 'rel_l2', 'nmse', 'token_cosine_p10'])} ## 条件 Probe Probe 仅使用 chunk index `c ≥ 2` 的目标 chunk，使 `correct`、`c-2 distant` 及其他控制组在完全相同的 token 上比较。 {markdown_table(probe_summary, ['probe', 'count', 'nrmse', 'r2', 'mse_reduction_vs_within_affine', 'mse_reduction_vs_within_quadratic'])} ### `fusion_same` 的成对控制实验 正值表示正确前一 chunk 同 timestep 输入的 MSE 更低。 {markdown_table(control_summary, ['reference', 'paired_count', 'mean_mse_reduction', 'median_mse_reduction', 'wins'])} ## 运动与对齐 {markdown_table(motion_summary, ['motion_bin', 'count', 'total_motion', 'same_slot_cosine', 'boundary_raw_cosine', 'global_aligned_cosine', 'flow_aligned_cosine'])} 这里的 `global_aligned` 是由光流中位数估计的全局平移对齐， `flow_aligned` 是 dense optical-flow oracle；本轮未实现 homography。 ## 本轮结论边界 - 已完成：四步 DMD 模型的 hidden/residual-delta 特征采集、相似度与 nMSE/CKA、linear/ridge 条件 probe、负对照、运动分桶和光流对齐。 - 尚未完成：接入用户现有的小型非线性预测网络、真实 cache 替换干预、 最终视频质量与端到端加速评估、规模化多视频置信区间。 - 因而当前结果支持“前一 chunk 提供额外且具有空间对应性的条件信息”， 但还不能单独证明最终生成质量或真实加速收益。 ## 输出文件 - `feature_pair_metrics.csv` - `feature_pair_summary.csv` - `conditional_probe_folds.csv` - `conditional_probe_summary.csv` - `motion_alignment_metrics.csv` - `motion_alignment_summary.csv` - `feature_redundancy_heatmap.png` - `conditional_chunk_gain.png` - `motion_alignment_analysis.png` """ (output_dir / "REPORT.md").write_text(report, encoding="utf-8") def analyze(args: argparse.Namespace, paths: list[Path]) -> None: print("[analysis] loading snapshots", flush=True) runs = load_runs(paths) pair_rows = collect_pair_rows(runs, cosine_only=args.cosine_only) if args.cosine_only: pair_summary = group_mean( pair_rows, ["comparison", "stage", "target_step"], [ "cosine", "centered_cosine", "linear_cka", "rel_l2", "nmse", "token_cosine_mean", "token_cosine_p10", "token_cosine_p50", "token_cosine_p90", ], ) write_csv(args.output_dir / "feature_pair_metrics.csv", pair_rows) write_csv(args.output_dir / "feature_pair_summary.csv", pair_summary) plot_similarity(pair_rows, args.output_dir) summary = { "runs": len(runs), "pair_rows": len(pair_rows), "cosine_only": True, } (args.output_dir / "summary.json").write_text( json.dumps(summary, indent=2) + "\n", encoding="utf-8" ) print(f"[analysis] cosine-only complete: {args.output_dir}", flush=True) return probe_rows = run_conditional_probes(runs, ridge=args.ridge) motion_rows = collect_motion_rows(runs) pair_summary = group_mean( pair_rows, ["comparison", "stage", "target_step"], [ "cosine", "centered_cosine", "linear_cka", "rel_l2", "nmse", "token_cosine_mean", "token_cosine_p10", "token_cosine_p50", "token_cosine_p90", ], ) probe_summary = group_mean( probe_rows, ["stage", "step", "probe"], [ "mse", "nrmse", "r2", "cosine", "mse_reduction_vs_within_affine", "mse_reduction_vs_within_quadratic", "r2_gain_vs_within_affine", "r2_gain_vs_within_quadratic", ], ) motion_summary = group_mean( motion_rows, ["motion_bin", "step"], [ "total_motion", "camera_motion", "object_motion", "same_slot_cosine", "boundary_raw_cosine", "global_aligned_cosine", "flow_aligned_cosine", "valid_flow_ratio", ], ) write_csv(args.output_dir / "feature_pair_metrics.csv", pair_rows) write_csv(args.output_dir / "feature_pair_summary.csv", pair_summary) write_csv(args.output_dir / "conditional_probe_folds.csv", probe_rows) write_csv(args.output_dir / "conditional_probe_summary.csv", probe_summary) write_csv(args.output_dir / "motion_alignment_metrics.csv", motion_rows) write_csv(args.output_dir / "motion_alignment_summary.csv", motion_summary) plot_similarity(pair_rows, args.output_dir) plot_probe_gain(probe_rows, args.output_dir) plot_motion(motion_rows, args.output_dir) build_report(runs, pair_rows, probe_rows, motion_rows, args.output_dir) summary = { "runs": len(runs), "pair_rows": len(pair_rows), "probe_rows": len(probe_rows), "motion_rows": len(motion_rows), "report": str(args.output_dir / "REPORT.md"), } (args.output_dir / "summary.json").write_text( json.dumps(summary, indent=2) + "\n", encoding="utf-8" ) print(f"[analysis] complete: {args.output_dir / 'REPORT.md'}", flush=True) def main() -> None: args = parse_args() args.config_path = resolve_path(args.config_path) args.checkpoint_path = resolve_path(args.checkpoint_path) args.prompt_path = resolve_path(args.prompt_path) args.output_dir = resolve_path(args.output_dir) args.output_dir.mkdir(parents=True, exist_ok=True) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.set_grad_enabled(False) run_paths = [ args.output_dir / "runs" / f"prompt_{index:02d}.pt" for index in range(args.num_prompts) ] if not args.analysis_only: run_paths = generate_snapshots(args) analyze(args, run_paths) if __name__ == "__main__": main()