Exp Suggestions : LeJEPA outperforms DINOv3 on global tasks but underperforms on segmentation (local tasks)

[gajeshladhar]

Hi @RandallBalestriero ,

I trained LeJEPA (ViT-L/16) initialized from DINOv3 ViT-L on the Core-Five training dataset (https://huggingface.co/datasets/gajeshladhar/core-five) for 48 hours on A100.

LeJEPA outperforms DINOv3 by ~2% on global classification benchmark tasks, but underperforms by ~1.5–2% on segmentation tasks, even after trying patch-wise/local-loss modifications in SIGReg.

Looking for clarification on whether this behavior is expected or if there are recommended settings to improve local spatial representation.

Experiment Details

1. Training Setup

• Backbone: ViT-L/16
• Initialization: DINOv3 ViT-L weights
• Training Dataset: Core-Five (HighRes EO Images)
• Hardware: A100 GPU
• Training Duration: 48 hours
• Objective: As mentioned in the paper & minimal setup code

2. Global Classification (Linear Probe) Performance

LeJEPA outperformed DINOv3 by ~2% on:

• AID Scene Classification
• UC Merced
• RESISC-45
• RSSCN7

(Metric: Top-1 accuracy)

3. Segmentation / Dense Prediction Performance

LeJEPA underperformed DINOv3 by ~1.5–2% on:

• INRIA Aerial Image Labeling
• Massachusetts Roads Dataset

(Metric: IoU / F1)

4. Attempts to Improve Local Feature Accuracy

Tried the following without improvement:

• Patch-wise local loss similar to iBOT

Segmentation accuracy remained lower than DINOv3.

Questions for Authors

1. Is LeJEPA inherently optimized for global semantics over local spatial accuracy?
2. Any recommended hyperparameters or architectural tweaks for dense prediction tasks?
3. Were internal segmentation benchmarks run, and can details be shared?

[RandallBalestriero]

Hello!
Thanks for your interest and for trying out LeJEPA! That is a great question that we are also thinking about. My intuition is that improving dense prediction task may be done by

1. considering all the tokens (not just CLS) as samples to be fed into SIGReg, i.e., we want all the tokens to be different iid samples from N(0,I).
2. Another possible setup is to concat (along the dim) the different tokens and then use those as input to SIGReg.

I think your idea to have a iBOT like loss is also valid, but because we don't have teacher/student/EMA to prevent collapse on the non-CLS tokens, there is nothing preventing the tokens to collapse (except for CLS which is fed to SIGReg). So perhaps combining your iBOT idea with 1. or 2. could be a good starting point (or even 1. or 2. alone), would love to see what you can get from it, keep us posted!

[gajeshladhar]

Hi @RandallBalestriero,
Thank you for the detailed reply & suggestions, really appreciate you taking the time!

I tried both suggested approaches:

1. Feeding all tokens (not just CLS) into SIGReg

2. Concatenating all tokens and applying SIGReg globally

Unfortunately, neither approach improved dense prediction performance. In my experiments:

Segmentation accuracy remained ~1.5–2% below DINOv3

Global classification dropped slightly (~0.3–0.5%) compared to baseline LeJEPA trained on only CLS

Below is the pseudo-code I used for the global + local losses:

proj, proj_patch = model(batch)

# proj shape : (N,V,D) 
# proj_patch shape : (N,V,H,W,D) (H*W = number of tokens)

# --- Global LeJEPA loss ---
sigreg_loss_global = sigreg(proj)
inv_loss_global = (proj.transpose(0,1).mean(0) - proj.transpose(0,1)).square().mean()
lejepa_loss_global = sigreg_loss_global * lambd + inv_loss_global * (1 - lambd)

# --- Local LeJEPA loss over tokens ---
sigreg_loss_local = sigreg(proj_patch)
inv_loss_local = (proj_patch.transpose(0,1).mean(0) - proj_patch.transpose(0,1)).square().mean()
lejepa_loss_local = sigreg_loss_local * lambd + inv_loss_local * (1 - lambd)

# --- Total ---
lejepa_loss = lejepa_loss_local + lejepa_loss_global
loss = lejepa_loss

Next steps:
I am currently experimenting with Gram Anchoring (using pretrained DINOv3 as a reference to stabilize local spatial structure). If that still does not close the gap on dense tasks, I will try integrating an iBOT-style token-level objective as you suggested.

I’ll keep you posted with results!
— Gajesh

[niklas030]

Any Updates? @gajeshladhar

[gajeshladhar]

Hi,
quick update.

I trained LeJEPA on top of a DINOv3 baseline, using SIGReg (LeJEPA loss). The main observation so far:

      1. Global representations improve (+1.0%)
      2. Local / dense performance remains the bottleneck

Local tokens don’t preserve spatial structure compared to DINOv3

I tried:

   1. Applying SIGReg on all tokens (not just CLS)
   2. Global + local SIGReg (token-level + CLS)
   3. Variants of Gram Anchoring (using pretrained DINOv3 as reference)

So far, these stabilize training but don’t fully close the gap on dense tasks vs DINOv3.

Code, training configs, and inference codebase are here: https://huggingface.co/gajeshladhar/core-jepa

— Gajesh

[niklas030]

Ah cool, as far as I remember, they used the DINO model as an early teacher. It’s probably impossible to get the weights at this epoch. What I also saw: in order to get more dense features, they used different deeper layers as an anchor and used SAM logits as an embedding to predict. https://arxiv.org/pdf/2504.13181 or here is a implementation for Dense-Jepa Features but its with the MIM pretask.

[RandallBalestriero]

Thank you for the update! Yes, I beleive that to have stronger dense feature, a prediction task (e.g. ibot style) must be employed. SIGReg on its own just requires to have uncollapsed representation which you can achieve by any entangled feature. The role of the prediction loss is to disentangle those features and thus must be employed alongside SIGReg (also I of course abuse the terms around disentanglement).
Did you try by any chance some inter-token prediction loss as used in DINO + SIGreg to prevent collapse on the tokens?

[gajeshladhar]

Thanks @RandallBalestriero for the suggestion.

Yes, I also agree that SIGReg alone only enforces non-collapse and does not necessarily encourage disentanglement at the token level. In my experiments so far, I mainly tried token-level Gram anchoring on top of DINOv3 (frozen anchor) features, with the intuition that constraining token statistics might act as a weak proxy for token prediction, but i still think it likely isn’t sufficient to drive disentanglement for dense features.

I haven’t yet tried an explicit inter-token prediction loss (iBOT-style) yet, I’ll definitely explore that next.

I’ll report back once I get time to experiment & I have results.

[gajeshladhar]

Hi @RandallBalestriero ,

I tried an iBOT-style patch loss on top of LeJEPA, but downstream performance on local tasks (INRIA building segmentations and many others using linear probe) still doesn’t surpass DINOv3 (dinov3 is +2%). Global features improve slightly (lejepa is +2.3%), but local representation remain weaker.

Experiment idea: each student patch token should be explainable by the token population of another view of the same image (soft inter-token prediction). For stability, I used frozen DINOv3 patch tokens as the teacher, assuming iBOT benefits from a strong, fixed teacher.

iBOT Loss :

def _ibot_loss(student_tokens, teacher_tokens, tau=0.1, tau_t=0.04):
    student_tokens = F.normalize(student_tokens, dim=-1)
    teacher_tokens = F.normalize(teacher_tokens, dim=-1)

    sim_s_t = student_tokens @ teacher_tokens.transpose(-1, -2)
    P_s = F.log_softmax(sim_s_t / tau, dim=-1)

    sim_t_t = teacher_tokens @ teacher_tokens.transpose(-1, -2)
    Nt = sim_t_t.size(-1)
    sim_t_t = sim_t_t.masked_fill(torch.eye(Nt, device=sim_t_t.device).bool(), -1e9)
    P_t = F.softmax(sim_t_t / tau_t, dim=-1)

    W = F.softmax(sim_s_t / tau, dim=-1)
    P_t_soft = W @ P_t
    return -(P_t_soft * P_s).sum(dim=-1).mean()

Applied across views (student = LeJEPA, teacher = frozen DINOv3):

def get_ibot_jepa_dino_loss(proj_patch_jepa, proj_patch_dino):
    loss, count = 0.0, 0
    N, V, D, H, W = proj_patch_jepa.shape
    for s in range(V):
        for t in range(V):
            if s == t: continue
            z_s = proj_patch_jepa[:, s].permute(0,2,3,1).reshape(N, H*W, D)
            z_t = proj_patch_dino[:, t].permute(0,2,3,1).reshape(N, H*W, D)
            loss += _ibot_loss(z_s, z_t)
            count += 1
    return loss / count

Colab Notebook (iBOT on LeJEPA experiments):
https://colab.research.google.com/drive/1d2Wx-8ZRyhRKtv2a3FDzxEndNoqE-7_8

I’d love your thoughts on:

whether a frozen DINOv3 teacher is the right choice here,
if this soft token matching misses a key ingredient (e.g. Sinkhorn or hard assignment) or
if iBOT-style objectives are fundamentally misaligned with JEPA-style training for dense EO segmentation ?

Thanks a lot, any feedback would be really helpful.