## Background Link to heading

In Lex Fridmans interview of Yann Lecun, Yann Lecun says there are a few things LLMs fundamentally can’t do in their current state that a human can.

Some of thee include:

- Planning
- Logical reasoning

He is an advocate of self-supervised learning, and in the interview, talks about JEPA and how self-supervised learning could allow models to construct a world model by just observing the world.

The joint embedding cuts out the irrelevant details of a feed (like cutting out specific leaf positions in a dashcam video) and keeps the important information (like which way the road is turning in a dashcam video).

Once this embedding is trained, a predictor can be trained on top of it to predict the next frames embedding from the previous frames embedding.

This is essentially trying to predict the next frame in a given video, but cutting out irrelevant details by doing the predictions in embedding space.

Yann Lecun mentions VICReg as an example embedding network, so I looked up the paper and tried it out on CIFAR-10.

## results Link to heading

The final network has the following architecture:

```
class VICReg(nn.Module):
def __init__(self, embedding_dimention, representation_dimention):
super(VICReg, self).__init__()
self.features = models.vgg16().features
for param in self.features.parameters():
param.requires_grad = False
self.embedding = nn.Sequential(
nn.Linear(512, 128),
nn.ReLU(),
nn.Dropout(),
nn.Linear(128, 32),
nn.ReLU(),
torch.nn.BatchNorm1d(32),
nn.Linear(32, embedding_dimention),
)
self.expand = nn.Sequential(
nn.Linear(embedding_dimention, 80),
torch.nn.BatchNorm1d(80),
nn.ReLU(),
nn.Linear(80, 80),
torch.nn.BatchNorm1d(80),
nn.ReLU(),
nn.Linear(80, representation_dimention),
)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
x = self.embedding(x)
x = self.expand(x)
return x
model = VICReg(50, 100)
```

If you want to try this yourself, I used this tutorial.

The coolest thing about VICReg is the loss function. With just a MSE loss, all embdding would collapse to a single constant point (you get 0 loss if you map everything to a single point).

I omitted the pushing forces in my loss function and got a collapse of the 2-dimentional embedding space like this:

There are contrastive and non-contrastive methods for preventing collapse.

Instead of using a contrastive method (having positive and negative examples to pull and push examples from eachother), VICReg has a loss function with a regularization term so only positive examples are needed.

The loss function is made of 3 parts…

### V loss Link to heading

The V stands for variance. This looks at the variance of the embeddings over the batch. The more variance there is between embeddings in a batch, the lower the loss is. This pushes apart representations and prevents collapse.

The formula looks like this

For example, if batch size was 2 and we had the embeddings [1,2,3] and [1,2,3] for both samples, the variance loss would be high. If we had [0,0,0] and [1,1,1] for both samples, the variance loss would be low.

### I loss Link to heading

The I stands for invariance. This is the normal MSE loss that pulls together similar samples.

### C loss Link to heading

The C stands for covariance. This is not necessary but nice to have. This prevents dimentionality collapse. This means if we allow a dimention 3 embedding space, we encourage the model to use all 3 dimentions for embedding and discourage correlations between dimentions.

### VICReg loss Link to heading

Combining it all together we get a pulling force from I loss, pushing force from V loss and C loss.

This is the code I used:

```
I_importance = 25
V_importance = 25
C_importance = 1
epsilon = 0.0001
relu = nn.ReLU()
def loss_var(image1):
std_z_a = torch.sqrt(image1.var(dim=0) + epsilon)
loss_v_a = torch.mean(relu(1 - std_z_a))
return loss_v_a
mse = torch.nn.MSELoss()
def loss_cov(image1):
N, D = image1.shape
image1 = image1 - image1.mean(dim=0)
cov_z_a = ((image1.T @ image1) / (N - 1)).square() # DxD
return (cov_z_a.sum() - cov_z_a.diagonal().sum()) / D
def VICReg_loss(image1, image2):
S_error = mse(image1, image2)
var_loss = loss_var(image1) + loss_var(image2)
cov_loss = loss_cov(image1) + loss_cov(image2)
return I_importance*S_error + C_importance*cov_loss + V_importance*var_loss
```

The embedding space I got looks a bit funky but that is probably because I am visualizing a 50 dimentional embedding space in 2 dimentions with T-SNE.

Loss averaged over 200 batches each:

The diagonal pattern is probably because the covariance error is too small, which encourages points to be along the diagonal instead of using the whole space.

I tried making the weights 1,1,1 and increased the embedding dimention to 50 and the points spread out a bit.