LeJEPA

LeJEPA 0 is a self-supervised learning method that learns image representations by enforcing invariance between multiple augmented views of the same image while regularizing the projected embeddings with SIGReg (Sketched Isotropic Gaussian Regularization). The objective combines a simple mean-squared invariance term with a SIGReg penalty that drives the projected features toward an isotropic Gaussian, removing the need for negative samples, momentum encoders, or stop-gradients.

Key Components

• Multi-view projections: LeJEPA uses a shared backbone and projection head to produce embeddings for several global views and a set of local (smaller) views.

• Invariance loss: Each local view’s projection is pulled toward the centroid of the global views’ projections using a mean-squared distance.

• SIGReg: A sliced, Epps-Pulley based regularizer that compares the empirical characteristic function of the projected embeddings to that of an isotropic Gaussian. It is the only regularizer in the LeJEPA objective.

• Projection head: A multi-layer perceptron with BatchNorm and ReLU (lightly.models.modules.LeJEPAProjectionHead) maps backbone features into the projection space where the loss is computed.

Good to Know

• No negatives or momentum encoder: Unlike contrastive (SimCLR, MoCo) or self-distillation (DINO, BYOL) methods, LeJEPA only needs positive views and a shared online encoder.

• SIGReg replaces VICReg-style variance/covariance terms: The Gaussian matching objective is more principled and avoids the heuristic hyperparameters of variance and covariance regularizers.

• Compatible with DINO-style multi-crop transforms: The example uses lightly.transforms.DINOTransform with global and local crops, matching the LeJEPA paper’s training setup.

Reference:

0

LeJEPA, 2025

1

LeJEPA reference implementation

Note

LeJEPA requires TIMM to be installed

pip install "lightly[timm]"

PyTorchLightningLightning Distributed

This example can be run from the command line with:

python lightly/examples/pytorch/lejepa.py

# This example requires the following dependencies to be installed:
# pip install lightly[timm]

# Note: The model and training settings do not follow the reference settings
# from the paper. The settings are chosen such that the example can easily be
# run on a small dataset with a single GPU.

import torch
import torchvision
from timm.models.vision_transformer import vit_small_patch16_224
from torch import Tensor
from torch.nn import Module
from torch.optim import AdamW

from lightly.loss import LeJEPALoss
from lightly.models.modules import LeJEPAProjectionHead
from lightly.transforms.dino_transform import DINOTransform


def _get_backbone_output_dim(backbone: Module) -> int:
    """Get the output dimension of a backbone by passing a dummy input through it."""
    with torch.inference_mode():
        dummy_input = torch.zeros(1, 3, 224, 224)
        output = backbone(dummy_input)
        output_dim = output.shape[1]
    return output_dim


class LeJEPA(Module):
    def __init__(
        self,
    ) -> None:
        super().__init__()

        self.backbone = vit_small_patch16_224(
            pretrained=False,
            pos_embed="learn",
            num_classes=0,
            dynamic_img_size=True,
            drop_path_rate=0.1,
        )

        backbone_out_dims = _get_backbone_output_dim(self.backbone)
        self.projection_head = LeJEPAProjectionHead(input_dim=backbone_out_dims)

    def forward(self, x: Tensor) -> Tensor:
        emb = self.backbone(x)
        proj = self.projection_head(emb)
        return proj


model = LeJEPA()

transform = DINOTransform(
    global_crop_scale=(0.3, 1),
    local_crop_scale=(0.05, 0.3),
    gaussian_blur=(0.5, 0.5, 0.5),
    n_local_views=6,
)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

dataset = torchvision.datasets.VOCDetection(
    "datasets/pascal_voc",
    download=True,
    transform=transform,
)
# Or create a dataset from a folder containing images or videos.
# dataset = LightlyDataset("path/to/folder")

dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    drop_last=True,
    num_workers=0,
)

# Create the loss functions.
lejepa_criterion = LeJEPALoss()

# Move loss to correct device because it also contains parameters.
lejepa_criterion = lejepa_criterion.to(device)

optimizer = AdamW(model.parameters(), lr=5e-4, weight_decay=5e-2)

epochs = 50
num_batches = len(dataloader)

print("Starting Training")
for epoch in range(epochs):
    total_loss = 0
    for batch_idx, batch in enumerate(dataloader):
        views = batch[0]
        views = [view.to(device) for view in views]
        global_views = views[:2]
        local_views = views[2:]

        global_proj = torch.stack([model(view) for view in global_views])
        local_proj = torch.stack([model(view) for view in local_views])

        loss = lejepa_criterion(local_proj=local_proj, global_proj=global_proj)
        total_loss += loss.detach()
        print(
            f"Epoch [{epoch + 1}/{epochs}] "
            f"Batch [{batch_idx + 1}/{num_batches}] "
            f"Loss [{loss.item():.4f}]"
        )

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    avg_loss = total_loss / num_batches
    print(f"epoch: {epoch:>02}, loss: {avg_loss:.5f}")