fix: target_encoder.py — respect config.jepa_loss_fn (smooth_l1/mse/cosine) instead of hardcoded MSE

mr_jepa/models/target_encoder.py

@@ -75,7 +75,6 @@ class TargetEncoder(nn.Module):
         """
         # Compute momentum
         if self.config.ema_schedule == "cosine":
-            # Cosine annealing from base to end momentum
             progress = step / max(total_steps, 1)
             momentum = self.config.ema_momentum_end - \
                 (self.config.ema_momentum_end - self.config.ema_momentum_base) * \
@@ -119,7 +118,6 @@ class TargetEncoder(nn.Module):
                 'target_trajectory': [B, K+1, N_s, D] - target states
                 'target_evidence': [B, N_e, D] - target evidence tokens
         """
-        # Target evidence memory
         evidence_output = self.target_evidence_memory(
             visual_tokens=visual_tokens,
             text_tokens=text_tokens,
@@ -129,7 +127,6 @@ class TargetEncoder(nn.Module):
         
         target_evidence = evidence_output['evidence_tokens']
         
-        # Target rollout
         rollout_output = self.target_rollout(
             evidence_tokens=target_evidence,
         )
@@ -149,14 +146,11 @@ class SIGRegLoss(nn.Module):
     
     Uses random projections + Epps-Pulley test statistic.
     SIGReg(Z) = (1/M) Σ_m T(Z @ u_m)
-    
-    where T is the Epps-Pulley univariate normality test.
     """
     
     def __init__(self, hidden_dim: int, num_projections: int = 1024):
         super().__init__()
         self.num_projections = num_projections
-        # Random projection directions (fixed, not learned)
         self.register_buffer(
             'projections',
             F.normalize(torch.randn(hidden_dim, num_projections), dim=0)
@@ -165,26 +159,15 @@ class SIGRegLoss(nn.Module):
     def _epps_pulley_statistic(self, h: torch.Tensor) -> torch.Tensor:
         """
         Compute Epps-Pulley test statistic for univariate normality.
-        
-        T(h) measures how far the distribution of h is from N(0,1).
-        Lower values = more Gaussian.
-        
-        Simplified version: uses moment-based approximation.
+        Simplified: variance + kurtosis penalty.
         """
-        # Standardize
         h_mean = h.mean()
         h_std = h.std() + 1e-6
         h_norm = (h - h_mean) / h_std
         
-        n = h_norm.size(0)
-        
-        # Compute pairwise differences for the EP statistic
-        # EP test: based on characteristic function
-        # Simplified: variance + kurtosis penalty
         variance = h_norm.var()
-        kurtosis = ((h_norm ** 4).mean() - 3).abs()  # Excess kurtosis
+        kurtosis = ((h_norm ** 4).mean() - 3).abs()
         
-        # Penalize deviation from unit variance and zero excess kurtosis
         return (variance - 1.0) ** 2 + 0.5 * kurtosis
     
     def forward(self, z: torch.Tensor) -> torch.Tensor:
@@ -193,7 +176,6 @@ class SIGRegLoss(nn.Module):
         
         Args:
             z: Latent embeddings [B, N, D] or [B*N, D]
-            
         Returns:
             Scalar SIGReg loss
         """
@@ -203,12 +185,10 @@ class SIGRegLoss(nn.Module):
         else:
             z_flat = z
         
-        # Project onto random directions
         projections = z_flat @ self.projections  # [B*N, M]
         
-        # Compute EP statistic for each projection
         losses = []
-        for m in range(min(self.num_projections, 64)):  # Sample subset for efficiency
+        for m in range(min(self.num_projections, 64)):
             losses.append(self._epps_pulley_statistic(projections[:, m]))
         
         return torch.stack(losses).mean()
@@ -220,7 +200,7 @@ class VICRegLoss(nn.Module):
     
     Three terms:
     - Variance: keep feature std above a threshold
-    - Invariance: prediction should match target (already handled by L2)
+    - Invariance: (handled by prediction loss)
     - Covariance: decorrelate features
     """
     
@@ -230,10 +210,6 @@ class VICRegLoss(nn.Module):
         self.cov_weight = cov_weight
     
     def forward(self, z: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            z: [B*N, D] latent embeddings
-        """
         if z.dim() == 3:
             z = z.reshape(-1, z.size(-1))
         
@@ -246,7 +222,6 @@ class VICRegLoss(nn.Module):
         N = z_centered.size(0)
         cov = (z_centered.T @ z_centered) / (N - 1)
         D = cov.size(0)
-        # Off-diagonal elements
         off_diag = cov.flatten()[:-1].view(D - 1, D + 1)[:, 1:].flatten()
         cov_loss = (off_diag ** 2).mean()
         
@@ -257,7 +232,10 @@ class JEPALoss(nn.Module):
     """
     Complete JEPA objective for MR-JEPA.
     
-    L_JEPA = (1/K) Σ_{k=1}^{K} ||z_pred_k - sg(z*_k)||²
+    Supports three loss functions (controlled by config.jepa_loss_fn):
+    - smooth_l1: SmoothL1Loss (hybrid default, robust to outliers)
+    - mse: MSE / L2 (original I-JEPA)
+    - cosine: Cosine similarity loss (purist default)
     
     Plus anti-collapse regularization:
     L_total = L_JEPA + λ * SIGReg(Z) + L_task + α * L_gen
@@ -279,17 +257,29 @@ class JEPALoss(nn.Module):
         target_trajectory: torch.Tensor,       # [B, K+1, N_s, D]
     ) -> torch.Tensor:
         """
-        Compute L2 prediction loss between online and target trajectories.
+        Compute prediction loss between online and target trajectories.
         
         Only compute loss for steps k=1..K (not z₀, which is deterministic).
+        Loss function selected by config.jepa_loss_fn.
         """
         # Skip z₀ (step 0) — only supervise predicted states
         pred = predicted_trajectory[:, 1:]   # [B, K, N_s, D]
-        target = target_trajectory[:, 1:]     # [B, K, N_s, D]
-        
-        # L2 loss per step, averaged
-        loss = F.mse_loss(pred, target.detach())
-        return loss
+        target = target_trajectory[:, 1:].detach()  # [B, K, N_s, D]
+        
+        loss_fn = self.config.jepa_loss_fn
+        
+        if loss_fn == "smooth_l1":
+            return F.smooth_l1_loss(pred, target)
+        elif loss_fn == "mse":
+            return F.mse_loss(pred, target)
+        elif loss_fn == "cosine":
+            # Cosine similarity loss: 1 - cos(pred, target), averaged
+            pred_flat = pred.reshape(-1, pred.size(-1))
+            target_flat = target.reshape(-1, target.size(-1))
+            cos_sim = F.cosine_similarity(pred_flat, target_flat, dim=-1)
+            return (1 - cos_sim).mean()
+        else:
+            raise ValueError(f"Unknown JEPA loss function: {loss_fn}")
     
     def compute_regularization(
         self,
@@ -299,7 +289,6 @@ class JEPALoss(nn.Module):
         reg_loss = torch.tensor(0.0, device=trajectory.device)
         
         if self.config.use_sigreg:
-            # Apply SIGReg to each step's representations
             B, Kp1, N_s, D = trajectory.shape
             for k in range(Kp1):
                 reg_loss = reg_loss + self.sigreg(trajectory[:, k])
@@ -308,9 +297,11 @@ class JEPALoss(nn.Module):
         
         if self.config.use_vicreg:
             B, Kp1, N_s, D = trajectory.shape
+            vicreg_loss = torch.tensor(0.0, device=trajectory.device)
             for k in range(Kp1):
-                reg_loss = reg_loss + self.vicreg(trajectory[:, k])
-            reg_loss = reg_loss / Kp1
+                vicreg_loss = vicreg_loss + self.vicreg(trajectory[:, k])
+            vicreg_loss = vicreg_loss / Kp1
+            reg_loss = reg_loss + vicreg_loss
         
         return reg_loss