LeJEPA for time series / video

[chemeris]

Very inspiring work on SIGReg/LeJEPA!

I noticed that your demo GIFs are videos, but the code and the paper only mention images. We're training a time-series encoder using JEPA, which is more similar to video than to images due to its time-domain nature, and we're now trying to apply LeJEPA instead of the classical JEPA. I would appreciate your thoughts on how to apply LeJEPA when the time dimension is present.

1. In your video demos, do you train the encoder on images and then apply it to the video frame-by-frame? I.e. is it a true video encoder, or just an image encoder applied to video?
2. My thinking is that in the case of video or time series, we would need a predictor like in V-JEPA, with SIGReg replacing EMA+StopGrad. Maybe-maybe we can make the predictor much simpler (an MLP instead of a transformer), but I doubt we can do without it. And then the training objective will be prediction, not invariance. Does this match your intuition? And will the SIGReg work with the Prediction objective instead of invariance?
3. Our preliminary experiments show that when training a time series JEPA with SIGReg, time-dimension embeddings collapse, so we applied SIGReg twice - across all embeddings of each batch sample individually, and across aggregated embeddings between batch samples. This seems to prevent the collapse, but we haven't finished the downstream task validations yet (it's a bit more developed than in CV, unfortunately).

PS If anyone else is interested in time-series self-supervised JEPA / representation learning - I'd be very interested to chat.

[RandallBalestriero]

Hello!

Thank you for your question.

1. Yes the gif is a video, but we process (and produce the masking) on a frame-by-frame basis. No training or inference is done on video.
2. yes I think for video application (or time-series) you want to leverage that extra data/info during pretraining. In that case having a predictor (next frame or tubelet prediction like VJEPA) would be a good starting point. But yes, basically reuse the exact VJEPA setup and replace stop-gradient/teacher-student/EMA with SIGReg at the output of the encoder
3. that's a good point. The current formulation may not prevent collapse across time depending on how you form the input to the sigreg loss. Doing two losses is a very good starting point, another one could be to do apply sigreg per video (each frame is a sample). This is part of the research question we are hoping to figure out theoretically to reduce the amount of trial-and-error. But so far the best bet is to try a few variations--we are super excited to see which performance you will get and what works and what doesn't!

I do think of a few people! Can you shoot me an email at rbalestr@brown.edu?

[chemeris]

Thank you, Randall.

A theoretical grounding would be awesome!

My current understanding of the time-domain collapse is the following:

1. A constant is a reasonable local minimum for the MAE/MSE loss we're using across the time domain (in the hidden space, of course).
2. The model essentially learns an "ID vector" for each time series to which it converges.
3. Since it's not a full collapse, but rather an "ID vector" collapse where each sample in a batch receives its own "ID", SIGReg drops to low values happily.

I.e. it looks like this (rough sketch):

[[1, 1, 1, 1, 1, 1, 1]  // time series sample 1
[2, 2, 2, 2, 2, 2, 2]   // time series sample 2
[3, 3, 3, 3, 3, 3, 3]   // time series sample 3
...]

This needs more research, but this is my current understanding of what's happening.

Now, an interesting question is what's better - two uni-dimensional SIGReg's or some form of multi-dimensional SIGReg (if it can be formulated).

[cvoscode]

Hey,

I think its a very cool and simple approach, thank you!

I am currently also experimenting with time series forecasting and LeJEPA. I am using overlapping windows and some augmentations to create views. The set up strictly follows LeJEPA (mse and SigReg Loss and a Probe loss).

I am finding that the inv and the sigreg loss go low fast in the beginning and then stay constant. The forecasting probe shows similar characteristics.

However I am wondering if I do anything wrong, since there is no progress after a while, I played with learning rate and different inits. Maybe my augmentation is to heavy? Or do you have other ideas?

[RandallBalestriero]

Interesting! There is a possibility that it comes from sub-optimal optimizer and lr scheduler. Or the capacity may not be strong enough? Overall though it is OK to have both of them nonzero, perhaps it is not possible to predict better (there is uncertainty?). Playing with different lambda should show a case where sigreg can go to 0 (but maybe the inv term goes up more) and one case where the inv term goes to 0 but sigreg goes up. IMO that's a quick way to debug. If you can't reach those extreme cases just by varying lambda, then you know there is something else which is wrong in the pipeline!

[cvoscode]

@RandallBalestriero Thanks for your quick response. I played around with lambda and a lambda schedule to improve on both but the run you see is the best I can get it to work, also used different learning rate schedulers.

I think the data holds some uncertainty (PEMS) but when forecasting I can get only get double the MAE of the current best models for this dataset. (Which is already great, but maybe I can push it a bit further)

I guess my augmentation might be to heavy or the encoder might be not strong enough? But I also tested several combinations of augmentation and encoders. I am interested in what do you think could help, make the encoder even bigger? Or get a stronger projector, since we would have a more complicated transformation of the embeddings before comparing them? Or work on augmentations?

[RandallBalestriero]

Hello @cvoscode !

There are definitely some things that could further help, e.g.,

• different encoder/projector/predictor architectures
• training time and number of views
• more informed data augmentations

But perhaps the first question is, can you assess how much more performance you expect to gain? e.g. would a supervised baseline be just a few points above your current best model or significantly better?

Another very popular trick that gives you a few points for free is to prove not only the last layer of your model but the last K layers (K between 1 to 4 typically). At least on images this consistently gives you 2 points on average!

[zaptrem]

3. across aggregated embeddings between batch samples

@chemeris

Can you elaborate on how this works? Are you taking the mean across time? Or something else? Thanks

[sattwik-sahu]

Hi, I am trying to adapt the JEPA framework to create energy-based value-learners for robots. I am adapting the V-JEPA architecture for this. I wonder how I should apply SIGReg for this. One idea I am thinking of is to train the encoder with V-JEPA as usual, with the frames in sequence and apply the SIGReg loss by sampling a different batch of random frames from different episodes. I'd love to know if there are better alternatives to this approach.

[sehunfromdaegu]

Very inspiring work on SIGReg/LeJEPA!

I noticed that your demo GIFs are videos, but the code and the paper only mention images. We're training a time-series encoder using JEPA, which is more similar to video than to images due to its time-domain nature, and we're now trying to apply LeJEPA instead of the classical JEPA. I would appreciate your thoughts on how to apply LeJEPA when the time dimension is present.

1. In your video demos, do you train the encoder on images and then apply it to the video frame-by-frame? I.e. is it a true video encoder, or just an image encoder applied to video?
2. My thinking is that in the case of video or time series, we would need a predictor like in V-JEPA, with SIGReg replacing EMA+StopGrad. Maybe-maybe we can make the predictor much simpler (an MLP instead of a transformer), but I doubt we can do without it. And then the training objective will be prediction, not invariance. Does this match your intuition? And will the SIGReg work with the Prediction objective instead of invariance?
3. Our preliminary experiments show that when training a time series JEPA with SIGReg, time-dimension embeddings collapse, so we applied SIGReg twice - across all embeddings of each batch sample individually, and across aggregated embeddings between batch samples. This seems to prevent the collapse, but we haven't finished the downstream task validations yet (it's a bit more developed than in CV, unfortunately).

PS If anyone else is interested in time-series self-supervised JEPA / representation learning - I'd be very interested to chat.

In my experience, I-JEPA like architecture gave me really good performance in multi-channel time series (ECG).

[cvoscode]

@sattwik-sahu What I have tested looks like this:
I use multiple overlapping windows t-1, t0 and t+1. I augement the window t0 and use multiple views of t0 for the inv_loss and sigreg. Then I use the result of the projector of all windows and train a predictor to predict t0 from t-1 and t+1 from t0, in order to achieve temporal coherence using mse. However, there might be a smarter solution, but it seems to work. However I just tested it with PEMS08. I am also not sure if we get some data leakage on subsequent forecasting task this way?

You can try to take a look in here: https://github.com/cvoscode/ts-lejepa. Currently, its a bit messy as I experimented a lot, however I would be happy to work on a better solution, maybe following the predictor idea or a smarter one. I you want to connect you can reach out to christoph.a.vos@gmail.com.

[zaptrem]

@sattwik-sahu What I have tested looks like this: I use multiple overlapping windows t-1, t0 and t+1. I augement the window t0 and use multiple views of t0 for the inv_loss and sigreg. Then I use the result of the projector of all windows and train a predictor to predict t0 from t-1 and t+1 from t0, in order to achieve temporal coherence using mse. However, there might be a smarter solution, but it seems to work. However I just tested it with PEMS08. I am also not sure if we get some data leakage on subsequent forecasting task this way?

You can try to take a look in here: https://github.com/cvoscode/ts-lejepa. Currently, its a bit messy as I experimented a lot, however I would be happy to work on a better solution, maybe following the predictor idea or a smarter one. I you want to connect you can reach out to christoph.a.vos@gmail.com.

Why not use I-JEPA/MAE style augmentation? Seems like it would be much better for time series.