rwightman/coca2-moderntext-naflex_ViT-B-32.cc12m-qk-lr2e3-b4096

Model card for coca2-moderntext-naflex_ViT-B-32.cc12m-qk-lr2e3-b4096

CoCa2 with a NaFlexVit B/32 vision tower and parallel paper pooling: a 128-query generative pooler feeds the caption decoder while an independent 1-query pooler (reading the same validity-masked trunk tokens) produces the contrastive latent. Modern causal text encoder with masked mean pooling, modern multimodal decoder, tiktoken r50k at ctx 128. Matches the bs-4096 MaMMUT2 on caption metrics but gives up the caption-objective I2T retrieval boost (see comparison).

Research checkpoint from a B/32-scale architecture/objective comparison on CC12M (CLIP vs MaMMUT2 vs modern text decoder); not a production model.

Model Details

• —Dataset: pixparse/cc12m-wds
• —Epochs: 32
• —Precision: bf16 amp
• —Weight decay: 0.25
• —Objective: InfoNCE + 2x caption CE + 1e-5 z-loss
• —Vision tower: timm NaFlexVit base, patch 32, NaFlex data pipeline; parallel attentional pooling (128-query generative + 1-query contrastive)
• —Text encoder: modern (RoPE/RMSNorm/SwiGLU/gated attn), qk-norm, causal, masked mean pooling, tiktoken r50k vocab 50260, ctx 128, variable-length text
• —Multimodal decoder: modern, cross-attn every block over the 128 pooled queries
• —Optimizer: adamw, lr 2e-3
• —Batch size: 4096

Model Usage

import torch
import open_clip

model, preprocess = open_clip.create_model_from_pretrained('hf-hub:rwightman/coca2-moderntext-naflex_ViT-B-32.cc12m-qk-lr2e3-b4096')
tokenizer = open_clip.get_tokenizer('hf-hub:rwightman/coca2-moderntext-naflex_ViT-B-32.cc12m-qk-lr2e3-b4096')

image = preprocess(pil_image).unsqueeze(0)
text = tokenizer(['a dog', 'a cat', 'a bird'])
with torch.no_grad():
    image_features = model.encode_image(image)
    text_features = model.encode_text(text)

Caption generation:

with torch.no_grad():
    tokens = model.generate(image, generation_type='beam_search', seq_len=30)
print(tokenizer.decode(tokens[0].cpu().tolist()))

Model Comparison

All ten runs: CC12M, 32 epochs, 2x GPU, bf16 amp, wd 0.25, evaluated at epoch 32 with one harness (ImageNet zero-shot w/ OpenAI 80-prompt ensemble; COCO Karpathy-test retrieval, I2T = any-of-5-captions; caption metrics vs the 5 COCO refs, beam search). Key reads:

• —The CLIP rows are prior runs, not purpose-trained controls: they were the closest available baselines but come from NAdaMuon optimizer experiments (nadamuon @ bs 4096 vs the MaMMUT runs' adamw @ bs ~1024), so cross-family deltas bundle optimizer + batch + loss.
• —Caption objective -> I2T retrieval: every MaMMUT2 run beats both CLIP runs by 1-4 pt I2T R@1 despite 1/4 the contrastive batch and a plainer optimizer.
• —Zero-shot ordering tracks batch/loss/optimizer, not architecture -- confirmed by the bs-4096 moderntext run: at matched batch (still adamw) it lands within 0.4 IN1K of clip-naflex and within 0.5 IN-R of clip-siglip. Batch was most of the CLIP lead.
• —The caption-objective I2T advantage is batch-independent (37.7-38.6 I2T R@1 at both bs 1024 and 4096), while zero-shot, T2I, and caption metrics all scale with batch.
• —...and it needs the shared caption/contrastive pathway, not just the caption loss: CoCa2 with parallel pooling (an independent 1-query contrastive pooler, decoupled from the caption path at the pooler) matches the bs-4096 MaMMUT2 on caption metrics (CIDEr 12.41 vs 12.23) yet drops I2T R@1 to 33.6 -- below both CLIP baselines. Its zero-shot/robustness deltas vs that run are NaFlex-tower-sized (cf. clip-naflex), so the retrieval regression is the pooling topology.
• —NaFlex vision and 1e-5 caption z-loss are both ~free (within the <=0.6 pt repeat-run noise).
• —Modern text tower (bundled with r50k tokenizer + ctx 128) at its best LR: IN1K parity with classic, +V2/+IN-R/+T2I, and clearly better caption metrics. qk-norm buys 2x LR headroom rather than direct quality; 3e-3 is past the stable edge even with it.
• —Caption scores are style-limited: CC12M alt-text register vs COCO references. Absolute CIDEr is low for all runs; relative differences are the signal.