Question regarding JEPA loss from LeJEPA and original JEPA

[NoTody]

Hi, @RandallBalestriero

Thanks for the great work. I'm wondering if there are any specific reasons on why LeJEPA used DINO style transformation for representation consistency. (e.g. I-JEPA, V-JEPA have masked images/videos as target, but LeJEPA doesn't seem to have them according to https://github.com/rbalestr-lab/lejepa/blob/main/MINIMAL.md?) Did you experiment on both and found DINO style (global/local views) augmentation give better performance?

[RandallBalestriero]

Great question! So we mostly wanted something general that works for any architecture, so we picked the DINO-style transforms that is identical (almost) across ViTs, resnets, efficientnet, vgg, ..... But there is absolutely no other reason, and in fact we are looking for some interested people to push a IJEPA version. Or rather a version of this https://arxiv.org/abs/2308.00566 which is supposed to be an improved IJEPA, where we keep everything the same in terms of augmentation and prediction task, but simply replace teacher-student/EMA by SIGReg :)

[gia-guar]

Hi there, I am trying to integrate the SIGReg Loss with a classic I-JEPA algorithm. I am working on another dataset, but so far the losses seem to go down. Here is my code:

    def forward(self, imgs, masks_enc, masks_pred):

        def _pool_tokens(x: torch.Tensor) -> torch.Tensor:
            return x.mean(dim=-2)  # (..., N, D) -> (..., D)

        with torch.autocast(device_type="cuda", dtype=self.amp_dtype, enabled=self.use_amp):

            # ---------- compute encoder targets with stop-grad
            with torch.no_grad():
                h_full = self.encoder(imgs)                         # (B, N_patches, D)
                h_full = F.layer_norm(h_full, (h_full.size(-1),))
                b = h_full.size(0)
                h_pred = apply_masks(h_full, masks_pred)   # no grad path 
                h_pred = repeat_interleave_batch(h_pred, b, repeat=len(masks_enc))

            # ---------- context -> predict targets
            z = self.encoder(imgs, masks_enc)                   # (B*Nm_enc, Nctx, D)
            pred = self.predictor(z, masks_enc, masks_pred)
            
            # F.smooth_l1_loss(z, h)
            predictive_loss = self.loss_fn(pred, h_pred)

            # ---------- SIGReg on pooled context
            z_pooled = _pool_tokens(z).reshape(-1, z.size(-1))  # (N_views_ctx, D)

            proj = z_pooled.clone()

            sigreg_loss = self.sigreg(proj)

            loss = (1.0 - self.lambd) * predictive_loss + self.lambd * sigreg_loss
            self.last_loss = (predictive_loss.item(), sigreg_loss.item())

        return loss

This is a TSNE of the global representation of each 2D Image (averaged embeddings).

For further details, refer to the I-JEPA public repo.

[lukas-kuhn]

Interesting, I also came up with nearly the exact same code as you @gia-guar but I found that the loss is much more volatile then it would be for augmentation view-based LeJEPA. I am training on Imagenette and a Ridge probe on val is also fluctuating wildly.

Did you notice the same?

I tried different lambdas from 1e-3 to 5e-2 as well as different learning rates from 1e-3 to 1e-5, with and without schedule and warmup.

Here are the first 13 epochs of a current run on Imagenette with a deit_tiny:

Edit: Additionally a quick check with standard LeJEPA shows that Ridge AUC on val is much higher (+10%) even after only 5 epochs of training.

[gia-guar]

Yes I also observe the same behavior. I noticed that including SIGReg loss increased the gradient by several orders of magnitude. With I-JEPA I was having gradient norms between 1e-4 and 1e-2, currently I am looking at gradients between 1e-2 and 20. Not yet sure if it's a problem but can explain why we see volatility. I am playing around with learning rate and gradient clipping to see if that leads to smoother losses. I excluded already floating point precision.

Another aspect worth mentioning is that I would not necessarily expect the same smoothness as the vanilla I-JEPA since a lot of its variability comes from the EMA schedule, which we don't use here.

[chemeris]

I'm working on a similar setup for time series (see issue #27 ). While some of the properties are different, the basic setup is the same with ViT context encoder/predictor/target encoder, EMA, and Stop-Grad. So I think the issues are similar.

1. I was initially thinking that we can avoid Stop-Grad, but it doesn't look like the case. My understanding is that we then essentially "dumb down" the teacher (target encoder) to match "guesses" of the predictor. This setup pretty much immediately collapses. The only way I can think of to fix this is to make the setup symmetric - generate many views (by using different masks), predict, and "pull" them to each other, while SIGReg'ing. This requires quite a large number of views to work, and the validation linear probe AUC is still only 0.72 vs 0.92 for the standard JEPA, even though it starts faster:

2. If we remove the EMA and do not introduce SIGReg, the model still trains, but the validation linear probe AUC flattens out at 0.65 (vs 0.72 for the symmetric approach and 0.92 for the normal JEPA). This is the baseline from which SIGReg should work (light brown line):

3. Now, if I add SIGReg on the context encoder output, i.e. Context SIGReg + Stop-Grad, no EMA, the best I was able to get is better than without it (option 2) but slightly worse than the symmetric approach (1), and still much worse than the traditional JEPA:

My thinking is that in the absence of EMA, the target gets much noisier, which leads to worse training performance.

I'll play with thy hyperparameters a bit more, but I think we have a more fundamental issue here.

I also did a few tests adding SIGReg to the traditional JEPA as an additional regularizaion on top of Stop-Grad and EMA just to see how it would affect the training, and it makes it start faster, but then falls back to pretty much the same old JEPA curve:

[lukas-kuhn]

I do not use EMA or a teacher/student setup in general, because it seems hacky to me and I was happy that they figured out a way with SIGReg how to not rely on that (and I think neither is @gia-guar).

"My thinking is that in the absence of EMA, the target gets much noisier, which leads to worse training performance."

I think this is what we continue to notice in our pred_loss as well. But I am unsure and will investigate further tonight.

[chemeris]

@lukas-kuhn I call this part of @gia-guar's code "teacher", i.e. an asymmetric part of the setup that sees the full input data:

            # ---------- compute encoder targets with stop-grad
            with torch.no_grad():
                h_full = self.encoder(imgs)                         # (B, N_patches, D)
                h_full = F.layer_norm(h_full, (h_full.size(-1),))
                b = h_full.size(0)
                h_pred = apply_masks(h_full, masks_pred)   # no grad path 
                h_pred = repeat_interleave_batch(h_pred, b, repeat=len(masks_enc))

Is your setup different?

[lukas-kuhn]

But this is the same encoder that we run the context through, just for target creation. So this is not a seperate model that is lagging behind so I didn't call it a teacher but I understand what you mean.

[saarangp]

Hi! I'm also trying to understand how to adapt LeJEPA + Sigreg for timeseries/medical data, so I wanted to comment to make sure I can stay in the loop with this discussion!

One thing I also did (and am unsure about) like in @gia-guar 's code was this bit here:

proj = z_pooled.clone()

sigreg_loss = self.sigreg(proj)

However, in the Minimal.MD
https://github.com/rbalestr-lab/lejepa/blob/eddffdd857a19726342ef6526e01021fb3c32cae/MINIMAL.md?plain=1#L93

sigreg is computed on proj(z) where proj is an MLP (with projection dim = 16). I wonder if this could make a major difference? I will try on my end to see!

[lukas-kuhn]

@saarangp I tested this and it does not seem to make too much of a difference empirically if you use a projector to a lower dimension as SIGReg works fine in higher dimensions as well. In the normal LeJEPA training you can actually skip the projector and train the encoder directly (both MSE and SIGReg).

EDIT: Even though SIGReg works fine on the encoder directly, using a projector MLP does increase subsequent Ridge AUC on Imagenette significantly.

[ForestsKing]

@chemeris I'm encountering a similar situation. I'm applying LeJEPA to my sequence data, but the MCC plateaus at 0.5, while a supervised baseline easily hits 0.95. I suspect this is due to the lack of augmentation options for sequence data (removing augmentations in the minimal working example dropped accuracy from 0.9 to 0.6 for me).

I saw that your standard I-JEPA setup performed much better than the LeJEPA approach. Since I-JEPA is well-suited for data that is hard to augment, I want to try it out. However, as a beginner for JEPA, I'm struggling with the implementation.

I would be extremely grateful if you could share the core code for your I-JEPA setup. Having a working reference would be invaluable to my learning process. If you are willing to share, my email is cswang@bupt.edu.cn. Thank you very much for your time and insights!

[FalsoMoralista]

@ForestsKing The code for I-JEPA is open-source. Perhaps you could adapt it from here.

[ForestsKing]

@FalsoMoralista I appreciate the help. I have actually tried that implementation, but transfer learning on my target dataset yielded no positive results. While the loss trajectory followed the expected pattern (an initial rapid drop, followed by a rise, and then a slow decrease), the accuracy remained stagnant. LeJEPA provided at least some baseline improvement. However, due to the difficulty of applying data augmentation to my specific data, the achievable upper bound is somewhat limited.

[sehunfromdaegu]

@ForestsKing In my experience, the success of pretraining I-JEPA depends on the hyper parameter choice in my experience. I also tried I-JEPA for time series (https://github.com/sehunfromdaegu/ECG_JEPA), but suffered a bit from collapse behavior. With sequencing data, is it also difficult to try cropping?

[ForestsKing]

@sehunfromdaegu That is incredibly valuable work, it is impressive that you have managed such a substantial workload single-handedly!

"After extensive experimentation, I have finally managed to stabilize the I-JEPA training loss. However, even though the model now converges properly, the final results remain unsatisfactory. On our dataset, LeJEPA achieved an MCC of approximately 0.54, which falls short of my target. Consequently, I switched to I-JEPA, given its reputation for being less dependent on data augmentation. However, the performance was even lower, yielding an MCC of only around 0.45. For this preliminary experiment, I trained a 50M parameter model on a single 48GB GPU for roughly 4 hours. While extending the training duration yields some improvement, the gains are marginal. Contrary to my expectations, a model of this capacity should not saturate so quickly and should theoretically be capable of absorbing more data.

I would greatly appreciate any guidance or insights you could offer on this matter.