How to change MSN loss to PMSN loss? (from paper "The Hidden Uniform Cluster Prior in Self-Supervised Learning")

[faris-k]

I couldn't find an official code release for this paper arxiv.org/abs/2210.07277, in which an extension is proposed to MSN to allow arbitrary feature priors.

It looks like the main difference is a change of a single term in the loss function. Is that correct? How would I implement the changes mentioned in the PMSN paper?

Previous MSN loss:

New loss: Prior Matching for Siamese Networks, PMSN:

[faris-k]

In case anyone stumbles upon this, the lightly package has a nice implementation of PMSN. I believe they're using a regularization weight of $\lambda = 1$.

https://github.com/lightly-ai/lightly/blob/ddfed3c4dc03a8d2722df24bfa537d24ac80bde6/examples/pytorch_lightning/pmsn.py

https://github.com/lightly-ai/lightly/blob/ddfed3c4dc03a8d2722df24bfa537d24ac80bde6/lightly/loss/pmsn_loss.py