I am trying to implement a loss function that takes as input 3 embeddings and output a value that is proportional to the collinearity of the embeddings. This is to shape the latent space of a convolutional autoencoder for embedding interpolation as described in this paper : Alon Oring Et al.- 2020.
I currently have code to get the embeddings during training as tensors and use a basic mse loss function. I've had multiple attempt to implement this but none worked, my training loss would always get stuck at the start.
Would you have suggestions to implement this collinearity loss function and the sum of it with the other loss in pytorch ?