From e5ab291af5649cedb2fef348273a827a985acd83 Mon Sep 17 00:00:00 2001 From: Szabolcs Botond Lorincz Molnar Date: Fri, 21 Aug 2026 20:28:37 +0200 Subject: [PATCH 1/2] fix: forward only attention kwargs in MaskedCausalBlock MaskedCausalBlock forwarded its full block-level kwargs to MaskedCausalAttention, including arguments the timm Attention constructor does not accept (e.g. mlp_ratio). This made MaskedCausalVisionTransformer, and therefore the AIM model, fail to construct on every supported timm version (0.9.9-1.0.28) with "Attention.__init__() got an unexpected keyword argument 'mlp_ratio'". Forward only the arguments that Attention.__init__ defines, selected via its signature so the fix stays correct as timm evolves. Add a regression test covering the block and the vision transformer. --- .../masked_causal_vision_transformer.py | 10 +- .../test_masked_causal_vision_transformer.py | 98 +++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) create mode 100644 tests/models/modules/test_masked_causal_vision_transformer.py diff --git a/lightly/models/modules/masked_causal_vision_transformer.py b/lightly/models/modules/masked_causal_vision_transformer.py index 71bba1ea9..4e92ae8b4 100644 --- a/lightly/models/modules/masked_causal_vision_transformer.py +++ b/lightly/models/modules/masked_causal_vision_transformer.py @@ -1,3 +1,4 @@ +import inspect from typing import Optional import torch @@ -120,9 +121,16 @@ class for input parameters. *args, **kwargs, ) + # Keep only the kwargs that timm's Attention accepts. Block-level kwargs + # such as mlp_ratio do not apply to the attention layer. + attention_kwargs = { + key: value + for key, value in kwargs.items() + if key in inspect.signature(Attention.__init__).parameters + } self.attn = MaskedCausalAttention( *args, - **kwargs, + **attention_kwargs, ) def forward( diff --git a/tests/models/modules/test_masked_causal_vision_transformer.py b/tests/models/modules/test_masked_causal_vision_transformer.py new file mode 100644 index 000000000..3eaff1473 --- /dev/null +++ b/tests/models/modules/test_masked_causal_vision_transformer.py @@ -0,0 +1,98 @@ +import pytest +import torch + +from lightly.utils import dependency + +if not dependency.timm_vit_available(): + # We do not use pytest.importorskip on module level because it makes mypy unhappy. + pytest.skip("TIMM vision transformer is not available", allow_module_level=True) + +from lightly.models.modules.masked_causal_vision_transformer import ( + MaskedCausalAttention, + MaskedCausalBlock, + MaskedCausalVisionTransformer, +) + +# MaskedCausalAttention supports only fused attention, which needs +# torch.nn.functional.scaled_dot_product_attention (PyTorch >=2.0). +skip_without_fused_attention = pytest.mark.skipif( + not hasattr(torch.nn.functional, "scaled_dot_product_attention"), + reason="MaskedCausalAttention supports only fused attention (PyTorch >=2.0).", +) + + +class TestMaskedCausalBlock: + def test_init__forwards_only_attention_kwargs(self) -> None: + # Block-level kwargs such as mlp_ratio must not reach the attention layer. + block = MaskedCausalBlock(dim=24, num_heads=3, mlp_ratio=4.0, qkv_bias=True) + assert isinstance(block.attn, MaskedCausalAttention) + # The mlp_ratio is still applied to the block's mlp, not dropped. + fc1 = block.mlp.fc1 + assert isinstance(fc1, torch.nn.Linear) + assert fc1.out_features == 24 * 4 + + +class TestMaskedCausalVisionTransformer: + def test_init(self) -> None: + model = MaskedCausalVisionTransformer( + img_size=32, + patch_size=16, + embed_dim=24, + depth=2, + num_heads=3, + mlp_ratio=4.0, + ) + assert all( + isinstance(block.attn, MaskedCausalAttention) for block in model.blocks + ) + + @skip_without_fused_attention + def test_init__aim_config(self) -> None: + # The AIM examples build the backbone without a class token and with + # average pooling. timm asserts global_pool != "token" when class_token is + # False, so both settings are required for the backbone to build. + model = MaskedCausalVisionTransformer( + img_size=32, + patch_size=16, + embed_dim=24, + depth=2, + num_heads=3, + class_token=False, + no_embed_class=True, + global_pool="avg", + ) + assert all( + isinstance(block.attn, MaskedCausalAttention) for block in model.blocks + ) + images = torch.rand(2, 3, 32, 32) + sequence_length = (32 // 16) ** 2 + model.num_prefix_tokens + mask = torch.zeros(2, sequence_length, dtype=torch.bool) + mask[:, model.num_prefix_tokens :] = True + features = model.forward_features(images, mask=mask) + assert features.shape == (2, sequence_length, 24) + + @skip_without_fused_attention + def test_forward(self) -> None: + model = MaskedCausalVisionTransformer( + img_size=32, patch_size=16, embed_dim=24, depth=2, num_heads=3 + ) + images = torch.rand(2, 3, 32, 32) + features = model.forward_features(images) + sequence_length = (32 // 16) ** 2 + model.num_prefix_tokens + assert features.shape == (2, sequence_length, 24) + + @skip_without_fused_attention + def test_forward__with_mask(self) -> None: + model = MaskedCausalVisionTransformer( + img_size=32, patch_size=16, embed_dim=24, depth=2, num_heads=3 + ) + images = torch.rand(2, 3, 32, 32) + sequence_length = (32 // 16) ** 2 + model.num_prefix_tokens + mask = torch.zeros(2, sequence_length, dtype=torch.bool) + mask[:, model.num_prefix_tokens :] = True + + features = model.forward_features(images, mask=mask) + features_no_mask = model.forward_features(images) + assert features.shape == (2, sequence_length, 24) + # The mask switches the patch tokens to causal attention and changes the output. + assert not torch.allclose(features, features_no_mask) From 9ee3d4e8f427815b1a2bc9907ecdd3a134930c29 Mon Sep 17 00:00:00 2001 From: Szabolcs Botond Lorincz Molnar Date: Fri, 21 Aug 2026 20:28:38 +0200 Subject: [PATCH 2/2] fix: set global_pool for the AIM masked causal vision transformer timm's VisionTransformer asserts `class_token or global_pool != 'token'`. The AIM examples and benchmark build MaskedCausalVisionTransformer with class_token=False but did not set global_pool, so construction failed on current timm. AIM's self-supervised path uses forward_features and is not affected by global_pool, so "avg" satisfies the assertion without changing behaviour. Regenerate the AIM example notebooks accordingly. --- benchmarks/imagenet/vitb16/aim.py | 1 + examples/notebooks/pytorch/aim.ipynb | 1 + examples/notebooks/pytorch_lightning/aim.ipynb | 1 + examples/notebooks/pytorch_lightning_distributed/aim.ipynb | 1 + examples/pytorch/aim.py | 1 + examples/pytorch_lightning/aim.py | 1 + examples/pytorch_lightning_distributed/aim.py | 1 + 7 files changed, 7 insertions(+) diff --git a/benchmarks/imagenet/vitb16/aim.py b/benchmarks/imagenet/vitb16/aim.py index 0fc2dc17b..a566e156e 100644 --- a/benchmarks/imagenet/vitb16/aim.py +++ b/benchmarks/imagenet/vitb16/aim.py @@ -31,6 +31,7 @@ def __init__(self, batch_size_per_device: int, num_classes: int) -> None: qk_norm=False, class_token=False, no_embed_class=True, + global_pool="avg", ) utils.initialize_2d_sine_cosine_positional_embedding( pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token diff --git a/examples/notebooks/pytorch/aim.ipynb b/examples/notebooks/pytorch/aim.ipynb index f91b28fee..5dd280ebd 100644 --- a/examples/notebooks/pytorch/aim.ipynb +++ b/examples/notebooks/pytorch/aim.ipynb @@ -110,6 +110,7 @@ " qk_norm=False,\n", " class_token=False,\n", " no_embed_class=True,\n", + " global_pool=\"avg\",\n", ")\n", "model = AIM(vit)" ] diff --git a/examples/notebooks/pytorch_lightning/aim.ipynb b/examples/notebooks/pytorch_lightning/aim.ipynb index 9e4414e67..09738d903 100644 --- a/examples/notebooks/pytorch_lightning/aim.ipynb +++ b/examples/notebooks/pytorch_lightning/aim.ipynb @@ -73,6 +73,7 @@ " qk_norm=False,\n", " class_token=False,\n", " no_embed_class=True,\n", + " global_pool=\"avg\",\n", " )\n", " utils.initialize_2d_sine_cosine_positional_embedding(\n", " pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token\n", diff --git a/examples/notebooks/pytorch_lightning_distributed/aim.ipynb b/examples/notebooks/pytorch_lightning_distributed/aim.ipynb index 7f560368a..a4ec74b15 100644 --- a/examples/notebooks/pytorch_lightning_distributed/aim.ipynb +++ b/examples/notebooks/pytorch_lightning_distributed/aim.ipynb @@ -73,6 +73,7 @@ " qk_norm=False,\n", " class_token=False,\n", " no_embed_class=True,\n", + " global_pool=\"avg\",\n", " )\n", " utils.initialize_2d_sine_cosine_positional_embedding(\n", " pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token\n", diff --git a/examples/pytorch/aim.py b/examples/pytorch/aim.py index e9623c452..d320cd9d4 100644 --- a/examples/pytorch/aim.py +++ b/examples/pytorch/aim.py @@ -57,6 +57,7 @@ def forward(self, images): qk_norm=False, class_token=False, no_embed_class=True, + global_pool="avg", ) model = AIM(vit) diff --git a/examples/pytorch_lightning/aim.py b/examples/pytorch_lightning/aim.py index e21c04812..755f05483 100644 --- a/examples/pytorch_lightning/aim.py +++ b/examples/pytorch_lightning/aim.py @@ -27,6 +27,7 @@ def __init__(self) -> None: qk_norm=False, class_token=False, no_embed_class=True, + global_pool="avg", ) utils.initialize_2d_sine_cosine_positional_embedding( pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token diff --git a/examples/pytorch_lightning_distributed/aim.py b/examples/pytorch_lightning_distributed/aim.py index 87f2671a4..d5d01c4e5 100644 --- a/examples/pytorch_lightning_distributed/aim.py +++ b/examples/pytorch_lightning_distributed/aim.py @@ -27,6 +27,7 @@ def __init__(self) -> None: qk_norm=False, class_token=False, no_embed_class=True, + global_pool="avg", ) utils.initialize_2d_sine_cosine_positional_embedding( pos_embedding=vit.pos_embed, has_class_token=vit.has_class_token