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 as the projector
+        pos_embs = utils.add_stochastic_positional_noise(
+            pos_embs,
+            self.predictor_embed,
+            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

@@ -46,6 +46,9 @@ class IJEPAPredictorTIMM(nn.Module):  # type: ignore[misc]
             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__(
@@ -61,6 +64,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__()
@@ -97,6 +101,8 @@ def __init__(
             ]
         )
 
+        self.noise_std = noise_std
+
     def forward(
         self,
         x: Tensor,
@@ -123,6 +129,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)
@@ -133,9 +140,20 @@ def forward(
         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,
+        )
+
         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/utils.py

@@ -1316,3 +1316,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: Module,
+    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: Matrix A used to project Gaussian noise to the positional
+            embedding dimension.
+        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 + projection(noise)

tests/models/modules/test_ijepa_timm.py

@@ -1,9 +1,6 @@
-import unittest
-
 import pytest
 import torch
 
-from lightly.models import utils
 from lightly.utils import dependency
 
 if not dependency.timm_vit_available():
@@ -14,8 +11,9 @@
 from lightly.models.modules import IJEPAPredictorTIMM
 
 
-class TestIJEPAPredictorTIMM(unittest.TestCase):
-    def test_init(self) -> None:
+class TestIJEPAPredictorTIMM:
+    @pytest.mark.parametrize("noise_std", [0.0, 0.1])
+    def test_init(self, noise_std: float) -> None:
         IJEPAPredictorTIMM(
             num_patches=196,
             depth=2,
@@ -26,10 +24,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
@@ -48,6 +51,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)
@@ -56,16 +60,14 @@ 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)
-
-        # output must have reasonable numbers
-        self.assertTrue(torch.all(torch.isfinite(predictions)))
+        assert list(predictions.shape) == [batch_size, num_patches, mlp_dim]
+        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)
 
-    @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("noise_std", [0.0, 0.1])
+    def test_forward_cuda(self, noise_std: float) -> None:
+        self._test_forward(torch.device("cuda"), noise_std)

tests/utils/test_stochastic_positional_embedding.py

@@ -0,0 +1,29 @@
+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, 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=projection,
+        noise_dim=8,
+        noise_std=0.25,
+    )
+
+    assert out.shape == pos_embeddings.shape
+    assert not torch.equal(out, pos_embeddings)