One example is a transformer model I’m training to generate token sequences with specific CLIP features. Initially, I trained it using teacher forcing for the first few epochs to bootstrap it—essentially matching specific tokens to guide it early on. However, that approach doesn’t fully align with the model’s true purpose. Instead, I want to take the token outputs, feed them into a CLIP model, and calculate the cosine similarity between the resulting features and a target.

The challenge is that the embedding layer, which converts discrete tokens into inputs for the encoder, acts as a lookup operation that kills gradients. Additionally, since my output is logits over the vocabulary, taking the argmax to get tokens also breaks gradient flow.

One potential solution is multiplying the logits with the embeddings matrix. However, that involves multiplying a tensor with the shape (batch_size, 77, 49807) by a (49807, 768) matrix on every batch. This not only explodes memory usage but also makes gradient calculations extremely complicated.

Fortunately, I found a loss function with help from LLMs that avoids the massive memory overhead while still being differentiable, making it a more practical solution for my training process.

Another experiment I’m running involves taking two sets of latent features—one representing “original” features and the other representing “theme” features—and combining them in a specific way to produce new latent features.

I’m doing this by generating a third set of features, which start as the mean of the original and theme features that serve as parameters for an optimizer. The goal is to minimize a loss function designed to position the new latent features in a part of latent space that has the desired properties related to the original and theme features in ways that closed-form arithmetic operations can't achieve.

I also have a large set of static “reference” features, and I calculate the similarity between these references and the original, theme, and optimized features. The loss function operates on these similarities, using them to create a gradient that shifts the optimized features based on how the original and theme reference similarities relate to each other.

The loss function breaks down into several discrete cases based on these relationships:

Case 1: The theme similarity is roughly orthogonal to the reference features. The optimized features should lean toward the original, with more weight if the theme is on the same side of the orthogonal line as the original.

Case 2: The original similarity is roughly orthogonal. The optimized feature should be pushed toward matching the theme’s reference similarity but with moderate weight, regardless of the original’s position relative to the orthogonal line.

Case 3: Neither is orthogonal, and the original and theme are in opposite directions. The optimized feature should strongly match the theme reference similarity, making it very close to the theme for these reference features.

Case 4: Neither is orthogonal, but both lie on the same side of the orthogonal line. Here, the optimized feature should be pushed moderately toward matching the average of the original and theme reference similarities.

These four cases need to blend smoothly without any discontinuities, which is tricky. Using torch.where isn’t sufficient because the loss changes in a complex way across multiple reference similarities per original-theme pair.

The loss function is critical to my experiment for reasons a bit too complicated to explain here. It's challenging to design loss algorithms that are acceptably efficient and ensuring smooth loss curves with respect to optimized similarities for all possible original-theme similarity pairs.

LLMs also noted ways to improve stability in edge cases that wouldn't have been obvious to me. I'd have needed to notice the subtle problems those stabilities issues sometimes cause and diagnosis it after the fact--a significant time sink.