Skip to content

feat: add CAPI clustering head and cross-attention predictor - #1995

Open
lorinczszabolcs wants to merge 1 commit into
lightly-ai:masterfrom
lorinczszabolcs:capi-modules
Open

feat: add CAPI clustering head and cross-attention predictor#1995
lorinczszabolcs wants to merge 1 commit into
lightly-ai:masterfrom
lorinczszabolcs:capi-modules

Conversation

@lorinczszabolcs

@lorinczszabolcs lorinczszabolcs commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Adds the CAPI building blocks for #1985: CAPIProjectionHead in heads.py (L2-normalizes the input, then a linear layer to num_clusters prototypes; a weight_norm option gives the student head unit-norm prototypes, while the teacher clustering head leaves it off and uses a bias), CAPIPredictorTIMM, a cross-attention decoder with rotary position embeddings via timm.layers.RotaryEmbeddingCat, gated behind timm_vit_available() like the other *_timm modules, and random_inverse_block_mask in models/utils.py (inverse-block masking: a contiguous block is kept visible, the rest masked, with optional roll).

Head and predictor put in the library since it matches the existing pattern (heads.py already holds the method heads, and MaskedVisionTransformerDecoder is a library-level predictor), and the RoPE predictor is reusable rather than CAPI-only wiring. The masking util sits with the existing masking utilities, which all live in models/utils.py.

Follow up of #1994.

@lorinczszabolcs

Copy link
Copy Markdown
Contributor Author

@gabrielfruet Small deviation from what we agreed on in the issue: we agreed only loss + transform in the lib, and wiring goes in the examples.

This PR also puts the clustering head and predictor in the library, since heads.py already holds every method's head and MAEDecoderTIMM is a library-level predictor, so they felt like reusable blocks. I can move them into the example instead if you wish.

@lorinczszabolcs

lorinczszabolcs commented Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

On a related note, CAPI puts RoPE inside the encoder and runs it over the masked sequence, so a faithful version needs a masking-aware RoPE encoder, which would be a separate component (I think DINOv3 #1881, would need it as well). For now I've kept the encoder as a standard masked ViT with RoPE only in the predictor.

@gabrielfruet

gabrielfruet commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Tip: When you're stacking PRs you should change the target merge branch to the one you stacked on top. It's easier for me to review :)

  • I took a look and that's not possible, unfortunately :(

@lorinczszabolcs

lorinczszabolcs commented Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

Tip: When you're stacking PRs you should change the target merge branch to the one you stacked on top. It's easier for me to review :)

  • I took a look and that's not possible, unfortunately :(

Yes, unfortunately that's the cross-fork limitation I flagged in #1985. Let me rebase so it's easier to review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants