LeJEPA loss missing contributions from global views?
#2 002 ouverte le 25 juil. 2026
Métriques du dépôt
- Stars
- (3 775 étoiles)
- Métriques de merge PR
- (Métriques PR en attente)
Description
I appreciate the addition of LeJEPA to LightlySSL, but I was confused by one aspect of the implementation as I worked through adding it to one of my projects.
In the original paper, I read the loss function as described in pseudocode on p. 11, or with standard math notation on p. 12, as a sum of contributions from all of the image views, local as well as global.
In the LightlySSL implementation of the LeJEPA loss, however, we seem to only include the contributions from the local views.
I suspect this is a bug since it seems odd to only constrain the embeddings associated with the local views during training. The original authors also describe e.g. training ResNets using only global views, and that doesn't make sense as the loss is implemented here.
Assuming my understanding is correct, I believe the fix is simple: https://github.com/lightly-ai/lightly/blob/master/lightly/loss/lejepa_loss.py#L327-L330 should pass all views, local and global, where it currently passes only local_proj, to the underlying SIGReg and invariance losses.