Add NaFlexViT support to CoCa2 and train a cc12m test model

Author: rwightman

src/open_clip/coca_model.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass
 
 from .transformer import (
+    AttentionalPooler,
     LayerNormFp32,
     LayerNorm,
     QuickGELU,
@@ -83,10 +84,32 @@ def __init__(
         text_cfg = CLIPTextCfg(**text_cfg) if isinstance(text_cfg, dict) else text_cfg
         vision_cfg = CLIPVisionCfg(**vision_cfg) if isinstance(vision_cfg, dict) else vision_cfg
         if vision_cfg.timm_model_name:
+            # timm towers are supported in token mode only, with model-level attentional pooling:
+            # the tower returns {'pooled', 'patch_tokens', 'patch_valid'} and the paper poolers
+            # (generative n_queries + contrastive) live on the model, pooling the trunk tokens with
+            # validity masking. Masking terminates at the pooler: the caption decoder cross-attends
+            # the fixed-count pooled queries, so no context_valid threading is needed downstream
+            # (contrast MaMMUT, which cross-attends raw patch tokens and threads image_embs_valid).
+            if not vision_cfg.output_tokens:
+                raise ValueError(
+                    "CoCa with a timm vision tower requires vision_cfg.output_tokens=true "
+                    "(token mode) so model-level attentional pooling can consume trunk tokens.")
+            if vision_cfg.attentional_pool not in ('cascade', 'parallel'):
+                raise ValueError(
+                    "CoCa with a timm vision tower requires vision_cfg.attentional_pool "
+                    "'cascade' or 'parallel' (paper pooling); the caption decoder consumes the "
+                    "pooled queries, not raw patch tokens.")
+        if getattr(text_cfg, 'text_arch', None) == 'modern' and \
+                getattr(text_cfg, 'attention_mode', None) == 'bidirectional':
+            # The caption decoder consumes the tower's contextualized token embeddings; under
+            # bidirectional attention position i carries future-token information the decoder's
+            # causal mask cannot undo, so the caption loss can copy and generation would mismatch
+            # training. Fail fast rather than train a silently-broken caption objective. (Masked
+            # 'mean' pooling for the contrastive readout works fine over a causal tower.)
             raise ValueError(
-                "CoCa does not support timm vision towers: caption cross-attention requires token projection "
-                "and validity handling that only MaMMUT's timm path currently provides."
-            )
+                "CoCa does not support a bidirectional modern text tower "
+                "(text_cfg.attention_mode='bidirectional'): its token embeddings leak future "
+                "tokens into the caption decoder. Use attention_mode='causal'.")
 
         self.text = _build_text_tower(
             embed_dim=embed_dim,
@@ -108,6 +131,40 @@ def __init__(
             cast_dtype=cast_dtype,
         )
 
+        if vision_cfg.timm_model_name:
+            # model-level paper poolers over trunk tokens (see the token-mode note above)
+            norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+            trunk_dim = self.visual.trunk.num_features
+            self.attn_pool_type = vision_cfg.attentional_pool
+            self.pool_norm = norm_layer(trunk_dim)  # pre-pool norm, mirrors native ln_post-on-width
+            self.attn_pool = AttentionalPooler(
+                embed_dim,
+                trunk_dim,
+                n_head=vision_cfg.attn_pooler_heads,
+                n_queries=vision_cfg.attn_pooler_queries,
+            )
+            self.attn_pool_contrastive = AttentionalPooler(
+                embed_dim,
+                embed_dim if self.attn_pool_type == 'cascade' else trunk_dim,
+                n_head=vision_cfg.attn_pooler_heads,
+                n_queries=1,
+            )
+            with torch.no_grad():  # match VisionTransformer.init_parameters pooler query init
+                nn.init.normal_(self.attn_pool.query, std=embed_dim ** -0.5)
+                nn.init.normal_(self.attn_pool_contrastive.query, std=embed_dim ** -0.5)
+            # Paper pooling replaces the tower's pooled readout entirely; its parameters
+            # (trunk fc_norm + TimmModel head proj) would never receive grad and trip DDP's
+            # unused-parameter check. Remove them rather than carry dead weights. (MaMMUT keeps
+            # them: it uses the tower's pooled output as the contrastive latent.)
+            self.visual.head = nn.Identity()
+            if getattr(self.visual.trunk, 'fc_norm', None) is not None:
+                self.visual.trunk.fc_norm = nn.Identity()
+        else:
+            self.attn_pool_type = ''
+            self.pool_norm = None
+            self.attn_pool = None
+            self.attn_pool_contrastive = None
+
         self.text_decoder = _build_text_decoder_tower(
             vocab_size,
             multimodal_cfg=multimodal_cfg,
@@ -140,7 +197,20 @@ def set_grad_checkpointing(self, enable: bool = True, impl: str = 'inline'):
         self.text_decoder.set_grad_checkpointing(enable, impl=impl)
 
     def _encode_image(self, images, normalize: bool = True):
-        image_latent, tokens_embs = self.visual(images)
+        out = self.visual(images)
+        if isinstance(out, dict):
+            # timm token mode: {'pooled', 'patch_tokens', 'patch_valid'} -- masked model-level
+            # pooling; the tower's own 'pooled' readout is unused (paper pooling replaces it).
+            tokens = self.pool_norm(out['patch_tokens'])
+            patch_valid = out['patch_valid']
+            tokens_embs = self.attn_pool(tokens, key_valid=patch_valid)
+            if self.attn_pool_type == 'cascade':
+                # pooled queries are fixed-count and all valid -> no mask from here on
+                image_latent = self.attn_pool_contrastive(tokens_embs)[:, 0]
+            else:  # parallel: contrastive pooler reads the (masked) trunk tokens directly
+                image_latent = self.attn_pool_contrastive(tokens, key_valid=patch_valid)[:, 0]
+        else:
+            image_latent, tokens_embs = out
         image_latent = F.normalize(image_latent, dim=-1) if normalize else image_latent
         return image_latent, tokens_embs
 

src/open_clip/model_configs/coca2-moderntext-naflex_ViT-B-32.json

@@ -0,0 +1,59 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "timm_model_name": "naflexvit_base_patch16_gap",
+        "timm_model_kwargs": {"patch_size": 32, "pre_norm": true, "pos_embed_grid_size": [8, 8]},
+        "timm_pool": "avg",
+        "timm_proj": "linear",
+        "output_tokens": true,
+        "image_size": 256,
+        "image_seq_len": 64,
+        "attentional_pool": "parallel",
+        "attn_pooler_queries": 128,
+        "attn_pooler_heads": 8
+    },
+    "text_cfg": {
+        "text_arch": "modern",
+        "context_length": 128,
+        "variable_text": true,
+        "vocab_size": 50260,
+        "pad_id": 50258,
+        "bos_id": 50259,
+        "eos_id": 50257,
+        "width": 512,
+        "heads": 8,
+        "layers": 12,
+        "pool_type": "mean",
+        "attention_mode": "causal",
+        "pos_embed": "rope",
+        "mlp_type": "swiglu",
+        "norm_type": "rmsnorm",
+        "norm_eps": 1e-06,
+        "qk_norm": true,
+        "attn_gated": true,
+        "output_tokens": true,
+        "tokenizer_type": "tiktoken",
+        "tiktoken_name": "r50k_base",
+        "tokenizer_kwargs": {
+            "clean": "whitespace"
+        }
+    },
+    "multimodal_cfg": {
+        "text_arch": "modern",
+        "context_length": 128,
+        "vocab_size": 50260,
+        "pad_id": 50258,
+        "bos_id": 50259,
+        "eos_id": 50257,
+        "width": 512,
+        "heads": 8,
+        "layers": 12,
+        "pos_embed": "rope",
+        "mlp_type": "swiglu",
+        "norm_type": "rmsnorm",
+        "norm_eps": 1e-06,
+        "qk_norm": true,
+        "attn_gated": true
+    },
+    "custom_text": true
+}

src/open_clip/transformer.py

@@ -264,11 +264,16 @@ def __init__(
         self.ln_q = norm_layer(d_model)
         self.ln_k = norm_layer(context_dim)
 
-    def forward(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor, key_valid: Optional[torch.Tensor] = None):
+        """key_valid: optional [B, L] bool/int validity (True/1 = real K/V token). Padded keys are
+        excluded from pooling attention -- required for NaFlex/variable-length token inputs where
+        trailing positions are padding, not content."""
         N = x.shape[0]
         x = self.ln_k(x)
         q = self.ln_q(self.query)
-        out = self.attn(q.unsqueeze(0).expand(N, -1, -1), k_x=x, v_x=x)
+        # [B, 1, 1, L] bool key mask, broadcast over heads/queries (incl. all-invalid-row guard)
+        attn_mask = context_attn_mask_from_valid(key_valid)
+        out = self.attn(q.unsqueeze(0).expand(N, -1, -1), k_x=x, v_x=x, attn_mask=attn_mask)
         return out
 
 

tests/test_mammut.py

@@ -460,8 +460,10 @@ def test_mammut_modern_factory_create():
     model = open_clip.create_model('mammut2-moderntext_ViT-B-32')
     assert isinstance(model, MaMMUT)
     assert isinstance(model.text, ModernMultimodalDecoder)
+    # config switched to tiktoken r50k / ctx 128 / qk-norm in f326982
     tokenizer = open_clip.get_tokenizer('mammut2-moderntext_ViT-B-32')
-    assert tokenizer.context_length == 77
+    assert tokenizer.context_length == 128
+    assert model.text.token_embedding.weight.shape[0] == 50260
 
 
 def test_mammut_modern_cross_sublayer_placement():