feat: add CAPI clustering head and cross-attention predictor - #1995
feat: add CAPI clustering head and cross-attention predictor#1995lorinczszabolcs wants to merge 1 commit into
Conversation
aca60e5 to
d1eb7cd
Compare
|
@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 |
d1eb7cd to
4aa2a8e
Compare
|
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. |
|
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 :)
|
Yes, unfortunately that's the cross-fork limitation I flagged in #1985. Let me rebase so it's easier to review |
c9fd611 to
83c5c65
Compare
Adds the CAPI building blocks for #1985:
CAPIProjectionHeadinheads.py(L2-normalizes the input, then a linear layer tonum_clustersprototypes; aweight_normoption 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 viatimm.layers.RotaryEmbeddingCat, gated behindtimm_vit_available()like the other*_timmmodules, andrandom_inverse_block_maskinmodels/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.pyalready holds the method heads, andMaskedVisionTransformerDecoderis 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 inmodels/utils.py.Follow up of #1994.