Questions about context_loss warmup schedule and weight_distance_loss in V-JEPA 2.1 fine-tuning

[hegene]

Hi, thanks for releasing the V-JEPA 2.1 codebase.

I am fine-tuning V-JEPA 2.1 ViT-L on a private dataset, and I have a question about the context_loss schedule.

From the official code, it looks like the context loss schedule is currently based on fixed iteration numbers:

• 0 - 15000 iterations: only pred_loss
• 15000 - 30000 iterations: linearly add context_loss, with its weight increasing from 0 to 0.5
• after 30000 iterations: keep the context_loss weight fixed

My understanding is that this setting corresponds to the longer training setup described in the V-JEPA 2.1 paper, e.g. around 135k iterations. However, in my fine-tuning setup, the total training length is only around 30k iterations. Therefore, I scaled the context-loss warmup schedule proportionally and used:

• 0 - 3300 iterations: only pred_loss
• 3300 - 6700 iterations: linearly increase context_loss weight
• after 6700 iterations: keep the context-loss weight fixed

However, I observed that the loss curve becomes unstable after around 6700 iterations, right after the context-loss weight reaches its maximum value. The attached figure shows the loss curve from my fine-tuning run.

So I would like to ask for some advice:

1. For shorter fine-tuning runs, is it recommended to scale the context_loss warmup schedule proportionally to the total number of iterations, or should the original absolute iteration schedule be kept?
2. Is there any recommended warmup range or maximum weight for context_loss during fine-tuning, especially when the total number of training iterations is much smaller than the original setting?
3. Should context_loss be enabled together with weight_distance_loss? In other words, is context_loss expected to be used only when weight_distance_loss=True, or are these two options independent?
4. Are there any known stability considerations when enabling context_loss, such as learning rate, EMA momentum, loss normalization, or recommended delay before turning it on?

Any guidance or experience about using context_loss in private-dataset fine-tuning would be very helpful.

Thanks!