feat: adding stochastic positional embeddings

lightly/models/modules/ijepa.py

@@ -39,6 +39,9 @@ class IJEPAPredictor(vision_transformer.Encoder):
             Percentage of elements set to zero after the MLP in the transformer.
         attention_dropout:
             Percentage of elements set to zero after the attention head.
+        noise_std:
+            Standard deviation of the Gaussian noise added to positional embeddings.
+            Default ``0.0`` to disable stochastic positional embeddings.
     """
 
     def __init__(
@@ -53,6 +56,7 @@ def __init__(
         dropout: float,
         attention_dropout: float,
         norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+        noise_std: float = 0.0,
         **kwargs,
     ):
         """Initializes the IJEPAPredictor with the specified dimensions."""
@@ -79,6 +83,8 @@ def __init__(
             torch.from_numpy(predictor_pos_embed).float().unsqueeze(0)
         )
 
+        self.noise_std = noise_std
+
     @classmethod
     def from_vit_encoder(cls, vit_encoder, num_patches):
         """Creates an I-JEPA predictor backbone (multi-head attention and layernorm) from a torchvision ViT encoder.
@@ -134,6 +140,7 @@ def forward(self, x, masks_x, masks):
         if not isinstance(masks, list):
             masks = [masks]
 
+        noise_dim = x.shape[-1]
         B = len(x) // len(masks_x)
         x = self.predictor_embed(x)
         x_pos_embed = self.predictor_pos_embed.repeat(B, 1, 1)
@@ -144,9 +151,20 @@ def forward(self, x, masks_x, masks):
         pos_embs = self.predictor_pos_embed.repeat(B, 1, 1)
         pos_embs = utils.apply_masks(pos_embs, masks)
         pos_embs = utils.repeat_interleave_batch(pos_embs, B, repeat=len(masks_x))
+
+        # we add the stochastic positional embedding here:
+        # use self.predictor_embed.weight as the projection matrix
+        pos_embs = utils.add_stochastic_positional_noise(
+            pos_embs,
+            self.predictor_embed.weight,
+            noise_dim,
+            noise_std=self.noise_std,
+        )
+
         pred_tokens = self.mask_token.repeat(pos_embs.size(0), pos_embs.size(1), 1)
 
         pred_tokens += pos_embs
+
         x = x.repeat(len(masks), 1, 1)
         x = torch.cat([x, pred_tokens], dim=1)
 

lightly/models/modules/ijepa_timm.py

@@ -57,6 +57,9 @@ class IJEPAPredictorTIMM(nn.Module):
             Percentage of elements set to zero after the attention head.
         norm_layer:
             Normalization layer.
+        noise_std:
+            Standard deviation of the Gaussian noise added to positional embeddings.
+            Default ``0.0`` to disable stochastic positional embeddings.
     """
 
     def __init__(
@@ -72,6 +75,7 @@ def __init__(
         proj_drop_rate: float = 0.0,
         attn_drop_rate: float = 0.0,
         norm_layer: Callable[..., nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+        noise_std: float = 0.0,
     ):
         """Initializes the IJEPAPredictorTIMM with the specified dimensions."""
         super().__init__()
@@ -104,6 +108,8 @@ def __init__(
         # remapping their keys on load (see the note in the class docstring).
         self._register_load_state_dict_pre_hook(self._migrate_legacy_state_dict)
 
+        self.noise_std = noise_std
+
     def forward(
         self,
         x: Tensor,
@@ -130,6 +136,7 @@ def forward(
         len_masks_x = len(masks_x) if isinstance(masks_x, list) else 1
         len_masks = len(masks) if isinstance(masks, list) else 1
 
+        noise_dim = x.shape[-1]
         B = len(x) // len_masks_x
         x = self.predictor_embed(x)
         x_pos_embed = self.decoder.pos_embed.repeat(B, 1, 1)
@@ -140,11 +147,21 @@ def forward(
         pos_embs = self.decoder.pos_embed.repeat(B, 1, 1)
         pos_embs = utils.apply_masks(pos_embs, masks)
         pos_embs = utils.repeat_interleave_batch(pos_embs, B, repeat=len_masks_x)
+        # we add the stochastic positional embedding here:
+        # use self.predictor_embed.weight as the projection matrix
+        pos_embs = utils.add_stochastic_positional_noise(
+            pos_embs,
+            self.predictor_embed.weight,
+            noise_dim,
+            noise_std=self.noise_std,
+        )
+
         pred_tokens = self.decoder.mask_token.repeat(
             pos_embs.size(0), pos_embs.size(1), 1
         )
 
         pred_tokens += pos_embs
+
         x = x.repeat(len_masks, 1, 1)
         x = torch.cat([x, pred_tokens], dim=1)
 

lightly/models/utils.py

@@ -1449,3 +1449,40 @@ def apply_masks(x: Tensor, masks: Tensor | list[Tensor]) -> Tensor:
         mask_keep = m.unsqueeze(-1).repeat(1, 1, x.size(-1))
         all_x += [torch.gather(x, dim=1, index=mask_keep)]
     return torch.cat(all_x, dim=0)
+
+
+def add_stochastic_positional_noise(
+    pos_embeddings: Tensor,
+    projection_weight: Tensor,
+    noise_dim: int,
+    noise_std: float = 0.0,
+) -> Tensor:
+    """Adds stochastic noise to positional embeddings.
+
+    - [0]: https://arxiv.org/pdf/2308.00566
+    - [1]: https://github.com/amirbar/StoP/blob/main/src/deit.py
+
+    Args:
+        pos_embeddings: Positional embeddings of shape
+            ``(batch_size, num_tokens, predictor_embed_dim)``.
+        projection_weight: Matrix A used to project Gaussian noise to the positional
+            embedding dimension. Must have shape ``(predictor_embed_dim, noise_dim)``.
+        noise_dim: Dimension of the sampled Gaussian noise before projection.
+        noise_std: Standard deviation of the Gaussian noise. If ``0.0``,
+            returns ``pos_embeddings`` unchanged.
+
+    Returns:
+        Positional embeddings with optional Gaussian noise added.
+    """
+    if noise_std == 0.0:
+        return pos_embeddings
+
+    noise = torch.normal(
+        mean=0.0,
+        std=noise_std,
+        size=(pos_embeddings.shape[0], pos_embeddings.shape[1], noise_dim),
+        device=pos_embeddings.device,
+        dtype=pos_embeddings.dtype,
+    )
+
+    return pos_embeddings + nn.functional.linear(noise, projection_weight, bias=None)

tests/models/modules/test_ijepa_timm.py

@@ -86,7 +86,8 @@ def forward(self, x: Tensor, masks_x: Tensor, masks: Tensor) -> Tensor:
 
 
 class TestIJEPAPredictorTIMM:
-    def test_init(self) -> None:
+    @pytest.mark.parametrize("noise_std", [0.0, 0.1])
+    def test_init(self, noise_std: float) -> None:
         IJEPAPredictorTIMM(
             num_patches=196,
             depth=2,
@@ -97,10 +98,15 @@ def test_init(self) -> None:
             mlp_ratio=4.0,
             proj_drop_rate=0.0,
             attn_drop_rate=0.0,
+            noise_std=noise_std,
         )
 
     def _test_forward(
-        self, device: torch.device, batch_size: int = 4, seed: int = 0
+        self,
+        device: torch.device,
+        noise_std: float,
+        batch_size: int = 4,
+        seed: int = 0,
     ) -> None:
         torch.manual_seed(seed)
         num_patches = 196  # 14x14 patches
@@ -119,6 +125,7 @@ def _test_forward(
             mlp_ratio=4.0,
             proj_drop_rate=0.0,
             attn_drop_rate=0.0,
+            noise_std=noise_std,
         ).to(device)
 
         x = torch.randn(batch_size, num_patches, mlp_dim, device=device)
@@ -134,12 +141,14 @@ def _test_forward(
         # output must have reasonable numbers
         assert torch.all(torch.isfinite(predictions))
 
-    def test_forward(self) -> None:
-        self._test_forward(torch.device("cpu"))
+    @pytest.mark.parametrize("noise_std", [0.0, 0.1])
+    def test_forward(self, noise_std: float) -> None:
+        self._test_forward(torch.device("cpu"), noise_std)
 
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available.")
-    def test_forward_cuda(self) -> None:
-        self._test_forward(torch.device("cuda"))
+    @pytest.mark.parametrize("noise_std", [0.0, 0.1])
+    def test_forward_cuda(self, noise_std: float) -> None:
+        self._test_forward(torch.device("cuda"), noise_std)
 
     def test_migrates_legacy_checkpoint(self) -> None:
         # A checkpoint saved with the pre-refactor parameter names must still load,

tests/utils/test_stochastic_positional_embedding.py

@@ -0,0 +1,32 @@
+import torch
+
+from lightly.models import utils
+
+
+def test_add_stochastic_positional_noise_disabled() -> None:
+    projection = torch.nn.Linear(8, 4)
+    pos_embeddings = torch.randn(2, 3, 4)
+
+    out = utils.add_stochastic_positional_noise(
+        pos_embeddings=pos_embeddings,
+        projection_weight=projection.weight,
+        noise_dim=8,
+        noise_std=0.0,
+    )
+
+    assert torch.equal(out, pos_embeddings)
+
+
+def test_add_stochastic_positional_noise_enabled() -> None:
+    projection = torch.nn.Linear(8, 4)
+    pos_embeddings = torch.randn(2, 3, 4)
+
+    out = utils.add_stochastic_positional_noise(
+        pos_embeddings=pos_embeddings,
+        projection_weight=projection.weight,
+        noise_dim=8,
+        noise_std=0.25,
+    )
+
+    assert out.shape == pos_embeddings.shape
+    assert not torch.equal(out, pos_embeddings)