diff --git a/lightly/models/modules/ijepa.py b/lightly/models/modules/ijepa.py index 7889dadd1..6332df79a 100644 --- a/lightly/models/modules/ijepa.py +++ b/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 if self.training else 0.0, + ) + 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) diff --git a/lightly/models/modules/ijepa_timm.py b/lightly/models/modules/ijepa_timm.py index 80385a2ed..962de87df 100644 --- a/lightly/models/modules/ijepa_timm.py +++ b/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 if self.training else 0.0, + ) + 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) diff --git a/lightly/models/utils.py b/lightly/models/utils.py index 19df3e001..31ae9dbd2 100644 --- a/lightly/models/utils.py +++ b/lightly/models/utils.py @@ -1560,3 +1560,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) diff --git a/tests/models/modules/test_ijepa_timm.py b/tests/models/modules/test_ijepa_timm.py index eafa09066..d728abb21 100644 --- a/tests/models/modules/test_ijepa_timm.py +++ b/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,6 +98,7 @@ def test_init(self) -> None: mlp_ratio=4.0, proj_drop_rate=0.0, attn_drop_rate=0.0, + noise_std=noise_std, ) @pytest.mark.parametrize("device", ["cpu", "cuda"]) @@ -111,6 +113,7 @@ def test_forward(self, device: str) -> None: predictor_embed_dim = 128 depth = 3 num_heads = 2 + noise_std = 0.1 predictor = IJEPAPredictorTIMM( num_patches=num_patches, @@ -122,6 +125,7 @@ def test_forward(self, device: str) -> None: 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) diff --git a/tests/utils/test_stochastic_positional_embedding.py b/tests/utils/test_stochastic_positional_embedding.py new file mode 100644 index 000000000..0fd8b4252 --- /dev/null +++ b/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)