feat: adding stochastic positional embeddings

lightly/models/modules/ijepa.py

@@ -79,6 +79,11 @@ def __init__(
             torch.from_numpy(predictor_pos_embed).float().unsqueeze(0)
         )
 
+        self.use_stop = kwargs.get(
+            "use_stop", False
+        )  # pass use stop embeddings as additional args, default to False
+        self.noise_std = kwargs.get("noise_std", 0.25)  # default 0.25
+
     @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 +139,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 +150,21 @@ 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 as the projector
+        pos_embs = utils.add_stochastic_positional_noise(
+            pos_embs,
+            self.predictor_embed,
+            noise_dim,
+            noise_std=self.noise_std,
+            enabled=self.use_stop,
+        )
+
         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

@@ -61,6 +61,8 @@ 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),
+        use_stop: bool = False,
+        noise_std: float = 0.25,
     ):
         """Initializes the IJEPAPredictorTIMM with the specified dimensions."""
         super().__init__()
@@ -97,6 +99,9 @@ def __init__(
             ]
         )
 
+        self.use_stop = use_stop
+        self.noise_std = noise_std
+
     def forward(
         self,
         x: Tensor,
@@ -123,6 +128,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.predictor_pos_embed.repeat(B, 1, 1)
@@ -136,6 +142,17 @@ def forward(
         pred_tokens = self.mask_token.repeat(pos_embs.size(0), pos_embs.size(1), 1)
 
         pred_tokens += pos_embs
+
+        # we add the stochastic positional embedding here:
+        # use self.predictor_embed as the projector
+        pred_tokens = utils.add_stochastic_positional_noise(
+            pred_tokens,
+            self.predictor_embed,
+            noise_dim,
+            noise_std=self.noise_std,
+            enabled=self.use_stop,
+        )
+
         x = x.repeat(len_masks, 1, 1)
         x = torch.cat([x, pred_tokens], dim=1)
 

lightly/models/utils.py

@@ -1316,3 +1316,46 @@ 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: Module,
+    noise_dim: int,
+    noise_std: float = 0.25,
+    enabled: bool = False,
+) -> 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:
+            Matrix A used to project gaussian noise to the pos_embedding
+            dimension.
+        noise_dim:
+            Dimension of the sampled gaussian noise before projection.
+        noise_std:
+            Standard deviation of the gaussian noise.
+        enabled:
+            If False, returns ``pos_embeddings`` unchanged.
+
+    Returns:
+        Positional embeddings with optional gaussian noise added.
+    """
+    if not enabled or 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 + projection(noise)

tests/models/modules/test_ijepa_timm.py

@@ -14,8 +14,10 @@
 from lightly.models.modules import IJEPAPredictorTIMM
 
 
-class TestIJEPAPredictorTIMM(unittest.TestCase):
-    def test_init(self) -> None:
+class TestIJEPAPredictorTIMM:
+    @pytest.mark.parametrize("use_stop", [True, False])
+    @pytest.mark.parametrize("noise_std", [0.0, 0.1])
+    def test_init(self, use_stop: bool, noise_std: float) -> None:
         IJEPAPredictorTIMM(
             num_patches=196,
             depth=2,
@@ -26,10 +28,17 @@ def test_init(self) -> None:
             mlp_ratio=4.0,
             proj_drop_rate=0.0,
             attn_drop_rate=0.0,
+            use_stop=use_stop,
+            noise_std=noise_std,
         )
 
     def _test_forward(
-        self, device: torch.device, batch_size: int = 4, seed: int = 0
+        self,
+        device: torch.device,
+        use_stop: bool,
+        noise_std: float,
+        batch_size: int = 4,
+        seed: int = 0,
     ) -> None:
         torch.manual_seed(seed)
         num_patches = 196  # 14x14 patches
@@ -48,6 +57,8 @@ def _test_forward(
             mlp_ratio=4.0,
             proj_drop_rate=0.0,
             attn_drop_rate=0.0,
+            use_stop=use_stop,
+            noise_std=noise_std,
         ).to(device)
 
         x = torch.randn(batch_size, num_patches, mlp_dim, device=device)
@@ -56,16 +67,16 @@ def _test_forward(
 
         predictions = predictor(x, masks_x, masks)
 
-        # output shape must be correct
-        expected_shape = [batch_size, num_patches, mlp_dim]
-        self.assertListEqual(list(predictions.shape), expected_shape)
+        assert list(predictions.shape) == [batch_size, num_patches, mlp_dim]
+        assert torch.all(torch.isfinite(predictions))
 
-        # output must have reasonable numbers
-        self.assertTrue(torch.all(torch.isfinite(predictions)))
+    @pytest.mark.parametrize("use_stop", [True, False])
+    @pytest.mark.parametrize("noise_std", [0.0, 0.1])
+    def test_forward(self, use_stop: bool, noise_std: float) -> None:
+        self._test_forward(torch.device("cpu"), use_stop, noise_std)
 
-    def test_forward(self) -> None:
-        self._test_forward(torch.device("cpu"))
-
-    @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available.")
-    def test_forward_cuda(self) -> None:
-        self._test_forward(torch.device("cuda"))
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available.")
+    @pytest.mark.parametrize("use_stop", [True, False])
+    @pytest.mark.parametrize("noise_std", [0.0, 0.1])
+    def test_forward_cuda(self, use_stop: bool, noise_std: float) -> None:
+        self._test_forward(torch.device("cuda"), use_stop, noise_std)

tests/utils/test_stochastic_positional_embedding.py

@@ -0,0 +1,31 @@
+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=projection,
+        noise_dim=8,
+        enabled=False,
+    )
+
+    assert torch.equal(out, pos_embeddings)
+
+
+def test_add_stochastic_positional_noise_enabled_shape() -> 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=projection,
+        noise_dim=8,
+        enabled=True,
+    )
+
+    assert out.shape == pos_embeddings.shape