MSN (Masked Siamese Networks) for ViT

[sayakpaul]

What does this PR do?

Adds the MSN checkpoints for ViT. MSN shines in the few-shot regimes which would benefit real-world use cases. Later we could add a pre-training script so that people can actually perform pre-training with MSN with their own datasets.

Closes #18758

Who can review?

@sgugger @NielsRogge @amyeroberts

TODO

• Add documentation
• Add rest of the files for repo consistency
• Host MSN weights on the Facebook org on HF Hub (@NielsRogge ?)
• Change the checkpoint paths wherever needed

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sayakpaul]

@NielsRogge, after studying the pretraining script of MSN thoroughly I am still unsure of how to put together a ViTMSNForPretraining similar to ViTMAEForPreTraining. There are multiple moving pieces that I think are best off residing inside a standalone pretraining script:

• A target encoder updated with EMA.
• Learnable prototypes that are needed to compute the final MSN loss.
• Target sharpening amongst other things.

Both the EMA and sharpening components operate with their own schedules.

Given this, I think it's best to resort to a separate pre-training script and use this model for feature extraction and fine-tuning.

There's an ongoing discussion around releasing the weights of the linear classification layers and fine-tuned models. So when that's available, we could directly support those via ViTMSNForImageClassification. Regardless, I am happy to add a ViTMSNForImageClassification for easy access.

What do you think?

[sgugger]

Thanks for your PR! It would be great to have the ViTMSNForImageClassification even if there are no released weights for image classification, so users can already fine-tune the main checkpoint if they want.

For pretraining, if multiple new pieces are needed, maybe it could go in a research project at first, where you can add more modules?

[sayakpaul]

For pretraining, if multiple new pieces are needed, maybe it could go in a research project at first, where you can add more modules?

Sounds good to me.

Thanks for your PR! It would be great to have the ViTMSNForImageClassification even if there are no released weights for image classification, so users can already fine-tune the main checkpoint if they want.

Sure, I will continue the work from here on then. Thank you!

[sayakpaul]

This is here so as to ensure the classification head params are always initialized using the same init. Since we don't have classification head params from the authors yet, this is needed to ensure the tests are passing.

[sayakpaul]

@sgugger @NielsRogge @amyeroberts ready for review.

[sayakpaul]

@sgugger @NielsRogge @amyeroberts a friendly nudge on the PR.

[sgugger]

A couple more comments, then the weights should be transferred to the right org before we merge this PR.

[sayakpaul]

@sgugger addressed your comments. After the weights are transferred to the right org, I will open a PR there adding README.

[amyeroberts]

Very nice PR ❤️ All LGTM.

Thanks for your work adding this model 💪

[amyeroberts]

Nice comment :)

[sayakpaul]

Thanks! I am contacting the model author Mido to check if they can release the classification params before we merge. Will keep y'all posted.

[ydshieh]

Hi @sayakpaul . First, thank you for this PR 🤗 .

The doctest for this model is currently failing, as

transformers/src/transformers/models/vit_msn/modeling_vit_msn.py

Line 646 in 7e84723

>>> print(model.config.id2label[predicted_label])

this outputs the predicted label, but there is no expected value provided.

The config has LABEL_0 ... LABEL_999 in id2label, but I feel it should be the actual labels for the COCO dataset.

Could you take a look for this config, as well as the missing expected outputs for the doctest? Thank you!

Here is the failing doctest job:

https://github.com/huggingface/transformers/actions/runs/3109562462/jobs/5039877349

[sayakpaul]

The config has LABEL_0 ... LABEL_999 in id2label, but I feel it should be the actual labels for the COCO dataset.

The model was trained on ImageNet-1k.

I will add the expected outputs. Thanks for flagging it.