LeJEPA

[MohamedFarag21]

Hi all,

I have found this today, and I believe it will be a very strong addition to the library!

Looks amazing from my perspective

leJEPA

[yutong-xiang-97]

Hi @MohamedFarag21,

Thank you so much for sharing the paper! We will see if we can put this on our agenda :)

[guarin]

Feel also free to open a PR if you would like to contribute. Loss functions and transforms are usually a great start. Once we have those we can put everything together and run benchmarks. Given that they benchmarked on rather small models and datasets this paper should be simpler to reproduce.

[gabrielfruet]

I will work on that if possible.

Planning to also add support for CNN-JEPA + LeJEPA.

Since i already implemented SparK, CNN-JEPA will be easier.

[IgorSusmelj]

Remaining tasks to finish LeJepa. I suggest to do the changes in 3 pull requests if possible to make reviews easier. It should not be much work :)

1) Core LeJEPA Building Blocks

• Add lightly/models/modules/lejepa.py with reusable LeJEPA components:
  • LeJEPAEncoder (wrapper around backbone + projector forward)
  • lejepa_invariance_loss(proj: Tensor) -> Tensor implementing (proj.mean(0) - proj).square().mean()
  • LeJEPALoss (combines SIGReg + invariance with lambda)
• Add a dedicated projector head in lightly/models/modules/heads.py:
  • LeJEPAProjectionHead (MLP style consistent with other methods, e.g. BN/ReLU blocks)
• Export new classes/functions:
  • Update lightly/models/modules/init.py to expose LeJEPA module + projector
• Add/extend tests together with this PR:
  • projector forward/backward shape test
  • combined LeJEPA loss wiring test (SIGReg + inv + lambda)

2) LeJEPA Transforms & Data Integration

• Add lightly/transforms/lejepa_transform.py:
  • LeJEPATransform returning 2 global + 6 local views
  • Parameters aligned to LeJEPA defaults (global/local crop sizes and scales configurable)
  • Reuse existing transform primitives (prefer DINOViewTransform / MultiViewTransform)
• Export transform:
  • Update lightly/transforms/init.py to expose LeJEPATransform
• Ensure integration behavior matches other transform building blocks:
  • Compatible output format for existing dataset/dataloader flows
  • No new collator required for LeJEPA (multi-view transform path is sufficient)
• Tests policy for this PR:
  • Add tests only if new logic is introduced
  • If transform is thin reuse of existing DINO-style components, no extra test required

3) Examples & Documentation

• Add PyTorch example:
  • examples/pytorch/lejepa.py with full training loop
  • Includes: backbone, LeJEPA projector, SIGReg, invariance loss, lambda mix, optimizer/scheduler config
• Add PyTorch Lightning example:
  • examples/pytorch_lightning/lejepa.py in the same style as other methods
• Add docs page:
  • docs/source/examples/lejepa.rst
  • Link both example scripts and short usage notes
• Register docs entry:
  • Add lejepa.rst to docs/source/examples/models.rst toctree
• Other update:
  • Mention LeJEPA in readme

[gabrielfruet]

@IgorSusmelj Since i'm already developing the OBB support for LightlyTrain i'll stick with it for now. After that i can take a look into some of these new methods to the LightlySSL repo. WDYT?

If you prefer, I could also switch priorities and focus on this first.

[RecreationalMath]

Hi @IgorSusmelj, I'm interested in picking this up and wanted to clarify one point before starting.

1) Core LeJEPA Building Blocks

• • Add lightly/models/modules/lejepa.py with reusable LeJEPA components:

    • LeJEPAEncoder (wrapper around backbone + projector forward)

The task list includes a LeJEPAEncoder (wrapper around backbone + projector forward). Currently, no other simple method (VICReg, Barlow Twins, SimCLR, etc.) has a library-level encoder module, they all compose backbone + projection head inline in example scripts. The complex methods (I-JEPA, MAE, DINO) do have dedicated modules because they have non-trivial architecture logic (masking, momentum, student/teacher).

Since LeJEPA's architecture is similarly simple (just backbone -> projector), is there a specific reason you'd like this as a library export?

[RecreationalMath]

Here is the link to the open-source code for the paper mentioned in the issue description (forgot to include this in my previous comment).

[RecreationalMath]

I suggest to do the changes in 3 pull requests if possible to make reviews easier.

PR 1 under review: #1916
PR 2 & 3 to follow after PR 1 is reviewed and accepted.