"""Independent normalized Nyström MMD objectives over multiple video layers. The V-JEPA2 caller concatenates the sampled dense matrices from every selected layer along the token axis. This module splits that transport tensor back into its layer-local matrices, evaluates one frozen-reference MMD per layer, applies detached scalar self-normalization independently, and sums the results. Keeping the objectives separate is important: concatenating layer features inside one kernel would change both the bandwidth and the teacher reference distribution instead of giving each layer equal normalized influence. """ from __future__ import annotations import math from collections.abc import Mapping, Sequence from typing import Any import torch from torch import nn from losses.video_nystrom_mmd import VideoNystromMMD class MultiLayerVideoNystromMMD(nn.Module): """Sum independent, optionally self-normalized per-layer video MMDs.""" def __init__( self, objectives: Mapping[int, VideoNystromMMD], *, self_normalize: bool = True, norm_eps: float = 1.0e-4, layer_weights: Mapping[int, float] | None = None, ) -> None: super().__init__() if not isinstance(objectives, Mapping) or not objectives: raise ValueError("objectives must be a non-empty layer mapping") ordered_layers = tuple(int(layer) for layer in objectives) if len(set(ordered_layers)) != len(ordered_layers): raise ValueError("objective layer indices must be unique") if any(layer < 0 for layer in ordered_layers): raise ValueError("objective layer indices must be non-negative") if tuple(sorted(ordered_layers)) != ordered_layers: raise ValueError("objective layer indices must be sorted") if ( not math.isfinite(float(norm_eps)) or float(norm_eps) <= 0.0 ): raise ValueError("norm_eps must be finite and positive") weights = ( {layer: 1.0 for layer in ordered_layers} if layer_weights is None else {int(layer): float(value) for layer, value in layer_weights.items()} ) if set(weights) != set(ordered_layers): raise ValueError( "layer_weights must contain exactly the objective layers") if any( not math.isfinite(value) or value < 0.0 for value in weights.values() ): raise ValueError( "layer weights must be finite and non-negative") modules = {} sample_shapes = [] for layer in ordered_layers: objective = objectives[layer] if not isinstance(objective, VideoNystromMMD): raise TypeError( "every layer objective must be VideoNystromMMD") if any( parameter.requires_grad for parameter in objective.parameters() ): raise ValueError( "Nyström reference objectives must remain frozen") modules[str(layer)] = objective sample_shapes.append(objective.sample_shape) feature_dims = {shape[1] for shape in sample_shapes} if len(feature_dims) != 1: raise ValueError( "all layer objectives must share one feature dimension") self.objectives = nn.ModuleDict(modules) self.layer_indices = ordered_layers self.layer_token_counts = tuple( int(shape[0]) for shape in sample_shapes) self.feature_dim = int(sample_shapes[0][1]) self.self_normalize = bool(self_normalize) self.norm_eps = float(norm_eps) self.layer_weights = tuple( float(weights[layer]) for layer in ordered_layers) ema_settings = { ( objective.attraction_ema_enabled, objective.attraction_ema_beta, ) for objective in self.objectives.values() } if len(ema_settings) != 1: raise ValueError( "all layer objectives must use the same attraction EMA " "enabled/beta setting") ( self.attraction_ema_enabled, self.attraction_ema_beta, ) = next(iter(ema_settings)) @property def sample_shape(self) -> tuple[int, int]: """Shape of the concatenated GradCache transport matrix.""" return sum(self.layer_token_counts), self.feature_dim @property def landmarks(self) -> torch.Tensor: """Compatibility device/landmark-count view used by GradCache logs.""" return self.objectives[str(self.layer_indices[0])].landmarks @property def reference_hashes(self) -> tuple[str, ...]: return tuple( str(self.objectives[str(layer)].reference_hash) for layer in self.layer_indices ) @property def reference_metadata(self) -> tuple[Mapping[str, Any], ...]: return tuple( self.objectives[str(layer)].reference_metadata for layer in self.layer_indices ) @property def attraction_ema_pending(self) -> bool: pending = tuple( objective.attraction_ema_pending for objective in self.objectives.values() ) if pending and any(value != pending[0] for value in pending[1:]): raise RuntimeError( "per-layer attraction EMA pending state is inconsistent") return bool(pending[0]) if pending else False def split_features( self, generated: torch.Tensor, ) -> tuple[torch.Tensor, ...]: if not isinstance(generated, torch.Tensor): raise TypeError("generated must be a torch.Tensor") if ( generated.ndim != 3 or min(generated.shape) <= 0 or not generated.is_floating_point() ): raise ValueError( "generated must be a non-empty floating [B,N,D] tensor") if not bool(torch.isfinite(generated.detach()).all().item()): raise FloatingPointError( "generated contains non-finite values") expected_tokens, expected_dim = self.sample_shape actual_shape = ( int(generated.shape[1]), int(generated.shape[2])) if actual_shape != (expected_tokens, expected_dim): raise ValueError( "concatenated multi-layer video matrix shape does not match " f"the frozen references: {actual_shape} != " f"{(expected_tokens, expected_dim)}") return tuple(torch.split( generated, self.layer_token_counts, dim=1)) def forward( self, generated: torch.Tensor, *, stage: bool = True, ) -> tuple[torch.Tensor, dict[str, Any]]: layer_features = self.split_features(generated) total = generated.new_zeros((), dtype=torch.float32) raw_sum = total.clone() repulsion_sum = total.clone() attraction_sum = total.clone() current_attraction_sum = total.clone() reference_constant_sum = total.clone() layer_metadata: dict[str, dict[str, torch.Tensor]] = {} global_samples = None ema_updates = None attraction_current_coefficient = None for layer, weight, features in zip( self.layer_indices, self.layer_weights, layer_features, ): objective = self.objectives[str(layer)] raw_mmd, metadata = ( objective(features, stage=stage) if self.attraction_ema_enabled else objective(features) ) if self.self_normalize: # Nyström attraction is approximate, so its finite estimate # can be slightly negative even though exact squared MMD is # non-negative. Normalize by detached magnitude to avoid both # a singularity near -eps and an accidental gradient reversal. denominator = raw_mmd.detach().abs() + self.norm_eps if ( not bool(torch.isfinite(denominator).item()) or float(denominator) == 0.0 ): raise FloatingPointError( f"layer {layer} has an invalid MMD normalization " "denominator") normalized_mmd = raw_mmd / denominator else: normalized_mmd = raw_mmd weighted_mmd = float(weight) * normalized_mmd total = total + weighted_mmd raw_sum = raw_sum + raw_mmd.detach() repulsion_sum = repulsion_sum + metadata["repulsion"] attraction_sum = attraction_sum + metadata["attraction"] if self.attraction_ema_enabled: current_attraction_sum = ( current_attraction_sum + metadata[ "current_attraction"]) reference_constant_sum = ( reference_constant_sum + metadata["reference_constant"]) observed_global_samples = int(metadata["global_samples"]) if global_samples is None: global_samples = observed_global_samples elif observed_global_samples != global_samples: raise RuntimeError( "per-layer MMD objectives disagree on global batch size") if self.attraction_ema_enabled: observed_updates = int( metadata["attraction_ema_updates"]) if ema_updates is None: ema_updates = observed_updates elif observed_updates != ema_updates: raise RuntimeError( "per-layer attraction EMA update counts disagree") observed_coefficient = float( metadata["attraction_current_coefficient"]) if attraction_current_coefficient is None: attraction_current_coefficient = observed_coefficient elif observed_coefficient != attraction_current_coefficient: raise RuntimeError( "per-layer attraction EMA coefficients disagree") layer_metadata[str(layer)] = { **metadata, "raw_mmd": raw_mmd.detach(), "normalized_mmd": normalized_mmd.detach(), "weighted_mmd": weighted_mmd.detach(), "weight": raw_mmd.new_tensor(float(weight)), } if not bool(torch.isfinite(total.detach()).item()): raise FloatingPointError( "combined multi-layer MMD loss is non-finite") combined_metadata = { "layers": layer_metadata, "layer_indices": torch.tensor( self.layer_indices, device=total.device, dtype=torch.long, ), "num_layers": total.new_tensor( len(self.layer_indices), dtype=torch.long), "self_normalized": total.new_tensor( int(self.self_normalize), dtype=torch.long), "normalization_epsilon": total.new_tensor(self.norm_eps), "raw_mmd_sum": raw_sum.detach(), "normalized_mmd_sum": total.detach(), # Aggregated raw-kernel terms preserve the existing generic log # contract. Per-layer diagnostics above are authoritative. "repulsion": repulsion_sum.detach(), "attraction": attraction_sum.detach(), "reference_constant": reference_constant_sum.detach(), "global_samples": total.new_tensor( int(global_samples), dtype=torch.long), } if self.attraction_ema_enabled: combined_metadata.update({ "current_attraction": current_attraction_sum.detach(), "attraction_ema_enabled": total.new_ones( (), dtype=torch.long), "attraction_ema_updates": total.new_tensor( int(ema_updates), dtype=torch.long), "attraction_current_coefficient": total.new_tensor( float(attraction_current_coefficient)), }) return total, combined_metadata def prepare_commit(self) -> None: """Validate every layer before publishing any EMA scalar.""" if not self.attraction_ema_enabled: return for objective in self.objectives.values(): objective.prepare_commit() def commit(self) -> None: """Publish one attraction EMA update for every layer.""" if not self.attraction_ema_enabled: return if not self.attraction_ema_pending: return self.prepare_commit() for objective in self.objectives.values(): objective.commit() def discard(self) -> None: for objective in self.objectives.values(): objective.discard() def attraction_ema_state_dict(self) -> Mapping[str, Any] | None: if not self.attraction_ema_enabled: return None if self.attraction_ema_pending: raise RuntimeError( "refusing to checkpoint pending per-layer attraction EMA") return { "version": 1, "layers": list(self.layer_indices), "states": { str(layer): self.objectives[ str(layer)].attraction_ema_state_dict() for layer in self.layer_indices }, } def validate_attraction_ema_state_dict( self, state: Mapping[str, Any], ) -> Mapping[str, Any]: if not self.attraction_ema_enabled: raise ValueError( "cannot load attraction EMA state while the path is disabled") if ( not isinstance(state, Mapping) or set(state) != {"version", "layers", "states"} or state.get("version") != 1 or tuple(state.get("layers", ())) != self.layer_indices or not isinstance(state.get("states"), Mapping) or set(state["states"]) != {str(layer) for layer in self.layer_indices} ): raise ValueError( "multi-layer attraction EMA checkpoint schema mismatch") # Validate every layer before the first live scalar is mutated. validated = {} for layer in self.layer_indices: objective = self.objectives[str(layer)] validated[str(layer)] = ( objective.validate_attraction_ema_state_dict( state["states"][str(layer)]) ) update_counts = { int(layer_state["updates"].item()) for layer_state in validated.values() } if len(update_counts) != 1: raise ValueError( "multi-layer attraction EMA update counts disagree") return { "version": 1, "layers": list(self.layer_indices), "states": validated, } def load_attraction_ema_state_dict( self, state: Mapping[str, Any], ) -> None: validated = self.validate_attraction_ema_state_dict(state) for layer in self.layer_indices: self.objectives[str(layer)].load_attraction_ema_state_dict( validated["states"][str(layer)])