From 3657caa154cd485d58b9a3730b0fa2651ce6e956 Mon Sep 17 00:00:00 2001 From: Haiyu Wu Date: Sat, 25 Jul 2026 18:28:10 -0400 Subject: [PATCH] [Feature] Improve VISReg loss numerical stability and defaults - VISRegLoss: clamp std to a minimum instead of adding 1e-6, so the scale term and normalisation stay unbiased for well-conditioned features while still guarding against division-by-zero on collapsed dimensions. - VISRegLoss: keep the Gaussian quantile target in fp32 rather than casting it to the input dtype, for a more stable shape term under 16-mixed. - VISReg: change the default convex weight lamb from 0.02 to 0.9. --- stable_pretraining/methods/visreg.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/stable_pretraining/methods/visreg.py b/stable_pretraining/methods/visreg.py index f8cfe2556..caa235b15 100644 --- a/stable_pretraining/methods/visreg.py +++ b/stable_pretraining/methods/visreg.py @@ -85,13 +85,13 @@ def __init__( self._cached_B = -1 self._cached_target = None - def _get_target(self, B: int, device, dtype) -> torch.Tensor: + def _get_target(self, B: int, device) -> torch.Tensor: """Theoretical standard-normal quantiles for ``B`` sorted samples.""" if self._cached_B != B: q = torch.linspace(1, B, B, device=device, dtype=torch.float32) / (B + 1) self._cached_target = torch.erfinv(2 * q - 1).mul_(math.sqrt(2)) self._cached_B = B - return self._cached_target.to(device=device, dtype=dtype) + return self._cached_target.to(device=device) def forward(self, z: torch.Tensor) -> torch.Tensor: """:param z: Embeddings [V, B, D] (views, batch, dim). @@ -104,13 +104,13 @@ def forward(self, z: torch.Tensor) -> torch.Tensor: center_loss = mu.pow(2).mean() z_centered = z - mu - std = z_centered.norm(dim=1).div(math.sqrt(B)) + 1e-6 + std = z_centered.norm(dim=1).div(math.sqrt(B)).clamp(min=1e-6) scale_loss = (std - 1.0).pow(2).mean() z_norm = z_centered / std.detach().unsqueeze(1) W = F.normalize(torch.randn(D, self.K, device=z.device, dtype=z.dtype), dim=0) p_sorted = (z_norm @ W).sort(dim=1).values - target = self._get_target(B, z.device, z.dtype).view(1, B, 1) + target = self._get_target(B, z.device).view(1, B, 1) shape_loss = (p_sorted - target).pow(2).mean() return ( @@ -207,7 +207,7 @@ def __init__( lambda_scale: float = 1.0, lambda_shape: float = 1.0, lambda_center: float = 1.0, - lamb: float = 0.02, + lamb: float = 0.9, pretrained: bool = False, drop_path_rate: float = 0.1, ):