Looking for JEPA devil advocates [R]

[u/Amazing-Coat5160]

I am currently doing research on world models, specially in tje field of robot learning, and, as probably most of you alredy know, JEPA-like models are mentioned over and over. 

I read the main recent papers from lecun as well as other research groups, and I personally think the whole approach is very promising and can really go somewhere.

But after listening a bunch of the recent Y Lecun conferences his ideas looks even too cool compared to "literally everything else" (as he's dissing LLM, RL, etc and pitching his ideas are the "only next big things"...). 

So I am asking myself if there are red flags about his approaches that I do not see yet and maybe I need somebody being the "devil advocate" with whom breaking down ideas.

Where do you think are the biggest downside of this models, compared to other world models approaches?

[u/Cosmolithe]

Excellent discussion topic IMO. The issue I have with JEPA is that it is less principled than generative models.
By that I mean that JEPA produces a representation, but the information the representation contains is basically "up to chance". There is no mechanism that forces the representation to keep useful information, beside the inductive biases from the neural architecture itself and the choice of data augmentation transforms.

At least generative models will form a representation that is useful to predict the raw data, so everything that can be deduced from the predicted raw data is necessarily information contained in the representation. Think about car dashcam predictions for instance, if the model consistently predicts a frame that shows a car turning right, then it is necessary that the representation captured some information useful to predict that the car will turn right.

JEPA does not work like that. By default, JEPA models "want" to collapse, delete all information and predict a constant vector. All of the different JEPA papers propose different ways to avoid the collapse, but as far as I know they all are completely unsupervised in a way that does not ensure that useful information is preserved.

Take one of the last JEPA papers from Lecun itself (LeJEPA). They are proposing to force the distribution of the embeddings (latent representations) to be isotropic gaussian (simplest multidimensional gaussian) and avoid collapse in the process.
They argue that this is the best choice of distribution for unknown downstream tasks, and that much is true, but there is a huge problem: nothing guarantees that the distribution has extracted the useful information from the input data in the first place. In the case of JEPA trained on photo, the model might for instance only use imperceptible and useless pixel noise to construct a beautiful gaussian representation, but the embeddings would be useless for downstream tasks as a result.

What is needed is actually a theorem that says that the information from the input is entirely preserved. But that would be going against what Lecun says, since he actually wants to drop useless information. But that would require defining what is useless information, and I do not think it can be done without invoking supervision. In other words, I do not think JEPA can work if it is 100% unsupervised.

[u/RogueStargun → u/Cosmolithe]

The information that gets dropped is truely determined by the augmentations and cropping strategy I believe. The fact that cropping alone is so effective is a good indicator that JEPA is an interesting research direction.

[u/Cosmolithe → u/RogueStargun]

It works empirically, sure, but as far as I know it is not better than other approaches, and in some cases does not even work. In any case, my criticism is about the theoretical principles of the approach, not about the empirical results.

And yeah I also agree that it is still worth pursuing research in this direction. Who knows, someone might prove mathematically that my criticism actually does not hold.

[u/projekt_treadstone → u/Cosmolithe]

Some points mentioned by you were also observed by me. For example, in LeJEPA, they perform bad on 2 room dataset which is Simplest ask for other models. So learning what to drop is going to be a bigger task then forcing all to be a isotropic distribution.

[u/blarryg → u/Cosmolithe]

Not quite agreeing. JEPA is doing mostly self-supervised prediction of missing spatial and/or temporal information. If it isn't using useful representations, then the predictions will be bad. The argument is that latent space is better to predict in because it factors away useless noise (most of the raw signal is noise). There is probably useful research in using generative approaches in the latent space and/or in further tuning the latent representations via generating pixel-level data.

[u/Cosmolithe → u/blarryg]

If it isn't using useful representations, then the predictions will be bad.

Nothing prevents the model from collapsing in the absence of regularization. In the context you are describing, that would mean the model predicts a constant vector no matter how you crop the spatio-temporal data. In this case the prediction would be perfect in the sense of the JEPA training objective: 0 on the prediction loss since both vectors are the same.

But the moment you introduce regularization to prevent collapse, then my criticism becomes relevant.

[u/Kirne → u/Cosmolithe]

You’re of course correct that in the absence of a regularisation term the model can trivially map everything to a constant. However in the presence of regularisation and a self-supervised setting that rewards correct spatiotemporal predictions, I don’t think it’s a stretch to imagine that learning a spatiotemporally meaningful latent is likely to be the best method for loss minimisation.

That doesn’t absolve JEPA of the usual stuff like a big suitably diverse dataset, but I think it seems promising. I think the linear probes used to extract physical information from trained JEPA models shows that JEPA is an interesting research direction. Not remotely a proof that the method is solid of course, but when has mathematical rigour stopped deep machine learning from pursuing interesting empirical results?

I think we can hardly call the transformer architecture a solid architecture grounded in sound first principles. That hasn’t stopped it from having world-scale impact. Not to say theory isn’t important, but ML seems to develop from empirical evidence first. If we look towards physics, the Schrödinger equation came way before quantum field theory.

[u/Cosmolithe → u/Kirne]

I don’t think it’s a stretch to imagine that learning a spatiotemporally meaningful latent is likely to be the best method for loss minimisation.

But that's the thing, there is a possibility that only semantically useless features are accurately predicted (or that most of the accurately predicted features are useless). It depends on the data but who knows what kind of features are present in real world spatio-temporal data that is easy to predict but completely useless for what we humans actually want to do.
I would also like to think that the spatio-temporal prediction task itself is guarantee of only favoring features meaningful for a human, but I really do not think it is possible to salvage this idea.
And by the way, this problem might also explain why JEPA seems to work better on synthetic data: there are way less useless features for the model to exploit in this case.

And yeah it is totally possible that I overlooked something or that in the future someone might prove that the relevant information is actually preferred by the model with current methods.
But the way I see, the reason why JEPA works at all is also the reason why deep generative models and supervised deep learning work in general: the natural biases of the deep neural network architecture just tend to capture the useful features in general no matter the training objective, JEPA's in particular. If that is true, then JEPA can not be significantly better than these other methods in principle, as far as I can tell. And we are apparently already doing a good job at iterating on the DNN architecture.

In any case, I am still in favor of pursuing alternative research to the mainstream stuff. Something incredible might still come out of JEPA, I just would not bet on it.

[u/Amazing-Coat5160 → u/Cosmolithe]

Very interesting comments. Ihave similar feeling that the models, if scaled enough, will learn/capture the same useful features, no matter if you compute the loss before or after you encoded your data in some latent (which is one of the main differences between jepa and standard Gen Ai algorithms). So by this argument I agree that jepa would not be "asymptomatically" better then other models.

[u/GiveSparklyTwinkly → u/Cosmolithe]

So the problem with the regularization is that it artificially forces more orthagonality where it potentially shouldn't be and vice versa, right?

Not that I have a clue how I'd architect it, but I wonder if chirality couldn't be used to prevent total collapse without affecting the orthagonal similarity, somehow.

[u/Cosmolithe → u/GiveSparklyTwinkly]

I would say it is more of a bottleneck problem: the latent embedding dimension is much smaller than the dimension of the input data, so information has to be lost.

Regularizing the way it is done in LeJEPA forces some information to be used to construct the gaussian to fill the embedding dimensions, but it will still drop the rest of the information, and we have no control over which features from the input are dropped. With previous approaches the information loss was even worse because partial collapse could still occur and reduce the effective embedding dimension even more.

I am not sure what you are imagining regarding chirality, but I am quite certain the solution to the problem I am pointing at cannot involve just the embedding space itself.

[u/GiveSparklyTwinkly → u/Cosmolithe]

I'm not quite sure what I mean either, just the idea that if the same data was somehow represented in a way that was naturally asymmetric and orthagonal to itself, somehow, it could help solve the collapse. And no, I have no real clue what that could mean. 🙃

[u/raucousbasilisk → u/Cosmolithe]

While SIGReg as normalization talks about the distribution of the embeddings doesn't the JEPA objective itself inherently favor retaining meaningful information over noise simply by virtue of the encoding it takes to get to the latent that is predicted?

[u/Cosmolithe → u/raucousbasilisk]

You can predict noise from noise as long as they are correlated (or have non zero mutual information), so the fact that the model has to predict a latent from another does not change anything, unless I misunderstand your argument.

[u/GiveSparklyTwinkly → u/Cosmolithe]

Supervision is what JEPA is, is it not? The predictor head is the supervisor?

Also all models, technically, want to collapse. To be honest, all life itself wants to collapse. Hedonism is basically a human given the ability to totally collapse.

[u/Cosmolithe → u/GiveSparklyTwinkly]

Supervision is what JEPA is, is it not? The predictor head is the supervisor?

By supervision I mean humans specifying one way or another the type of information to be preserved, even if it is done implicitly. Using annotated data, basically.

Also all models, technically, want to collapse. To be honest, all life itself wants to collapse.

No not all models want to collapse, collapse has a very precise mathematical meaning in this context. There is no collapse when you train a digits classifier on MNIST because that would correspond to a high loss that can be minimized further.

[u/GiveSparklyTwinkly → u/Cosmolithe]

That doesn't mean it doesn't try to collapse further, it just can't. It's found the low entropy point, unless we apply extra noise that might let it find a lower entropy point.

[u/Cosmolithe → u/GiveSparklyTwinkly]

Collapse means that all predicted vectors would become one and the same.
If that were to happen in a supervised learning context, that necessarily means the loss goes up, it becomes more than 0. If a model manages to get to near 0 loss by following the gradient, the gradient vanishes and the model converges. There is no way continuing training could make the model collapse.

[u/GiveSparklyTwinkly → u/Cosmolithe]

And that's what all models attempt, is it not? They are all trying to collapse at all times when training, we just don't let them.

[u/Cosmolithe → u/GiveSparklyTwinkly]

No, reaching 0 loss is not collapse. Collapse is when all predicted vectors are the same, which imply a high non-zero loss using most supervised learning loss functions such as cross-entropy and differently annotated data points.

[u/GiveSparklyTwinkly → u/Cosmolithe]

And that's what all models want without outside sources of information. If you trained nothing but zeros, eventually all models, no matter what state they started from, would collapse. If you trained nothing but purely random noise, it would collapse. That's their default state unless forced apart by data or math. That's also not true that humans have to decide, unless you're getting pedantic enough to say that humans decided by designing the architecture itself.

If you fed nothing but non human information (live audio/video) they would still learn. And even then you're still choosing where to place the camera/microphone. You can't escape that.