Novel embedding conditioning for a free improvement - free to use

[kooshi]

I don't want to fight for a winning run, but I want to contribute a helpful technique.
I've been independently testing a lot of stuff, even before this competition, but one of the most simple things I found was also one of my first findings.
Embeddings have a little bit of drift to them that make them less than ideal, and we can directly condition them with a loss penalty.

Here's a graph of the penalty: https://www.desmos.com/calculator/padllb0rwj
Each activation is individually penalized by $f(r)=\max(0,\sqrt{2d/(r^{2}+0.1)}-1)$ to prevent collapse from the second penalty,
acting on the mean of all activations: $m^{2}/\sqrt{d}$
and the whole thing is scaled by lambda. I found 0.008 to be reasonable.

This effectively conditions the embedding to stay outside of a shell of size ~ $\sqrt{d-0.1}$, while conditioning the mass of that shell to be nearly perfectly centered. Embeddings can cluster, but cannot form a cone.

class ShellCenteringPenalty(nn.Module):
    """Prevents activation collapse and centroid drift from the origin"""
    def __init__(self, d_model, lam: float = 0.008):
        super().__init__()
        self.lam = lam
        self.d_model = d_model
        self.penalty = None

    def forward(self, x, step: int | None = None):
        # x: [B, T, d_model]
        # Keep individual activations from collapsing to the origin
        norms = x.norm(dim=-1, p=2)  # [B, T]
        self.penalty = self.lam * (((2 * self.d_model/((norms**2)+0.1))**0.5)-1).clamp(min=0).mean()

        # Force the centroid of the activations to stay near the origin
        norm_of_mean = x.flatten(0, 1).mean(dim=0).norm(p=2)
        self.penalty += self.lam * ((norm_of_mean ** 2)/(self.d_model**0.5))

        return x  # pass-through, graph intact

usage:

  x = self.tok_emb(input_ids)
+ x = self.shell_centering_penalty(x)

  logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap)
- return F.cross_entropy(
+ loss = F.cross_entropy(
      logits.reshape(-1, logits.size(-1)).float(),
      flat_targets,
      reduction="mean",
  )
+ return loss + self.shell_centering_penalty.penalty.to(loss.dtype)

The conditioned embeddings should result in an immediately improved final loss, basically for free.

If this ends up being helpful and anyone wants to help optimize the shape or orherwise collaborate, let me know. This is just a hobby for me, but I have other things worthy of formalizing too.

[CiprianFlorin-Ifrim]

This looks interesting and there have been a few papers that looked at this as a post-training "fix". Would be interesting to see how well it actually improves things (especially given modern architectures employ techniques like RMSNorm to help with scaling factors, directly related to this work) and also how it affects models with weight tying between emb and lm head. Will do some work on this in 1 week from now, outside this challenge, to see how it behaves.

Alos is sqrt(2d / (r² + 0.1)) - 1 correct? Because when it fires, only r < √(2d - 0.1) ≈ √(2d), would then push the norms towards √(2d) nor √(d).

[kooshi]

Kindof, it was a last minute edit. The original was without the 2x, but I did a couple more quick tests and found that the majority of the improvement came from the model learning to scale up the embeddings, though the centering is also important as a secondary signal.

Following that, I added the 2x, which led to a significant improvement in my model. I then jumped to 20x, which reduced performance slightly (and the embeddings never got higher than $\sqrt{2d}$ despite the loss penalty pushing them higher, meaning the rest of the model pressured them to be low even stronger).

The centering term will cause the activations to climb up the 1/x penalty slightly, so in the original, the mean activation never makes it all the way to $\sqrt{d}$. It seems the model strongly prefers the embedding magnitude to be somewhere between $\sqrt{d}$ and $\sqrt{2d}$. Adding that 2x puts the mean magnitude in that range.

I have a lot of evidence for significant improvements in the 2.7M to 10M parameter range, which is what I can train locally in a reasonable amount of time. I did modify PR #1908 with it as well, but also only ran it locally at small scale.

I have a handful of other novel techniques I want to explore if this one pans out. I was targeting unlimited compute bpb per parameter, so some of them don't really help with the speedrun paradigm, which is why I wasn't really interested in competing.

My latest 2.7M param architecture hits 1.4839 bpb, trained in 1 hour on a 5090 and 2 3090s, and I really want to explore it further.

another edit: also I use this with tied embedding weights, so if anything I'm not sure how they'd act on non-tied weights

[CiprianFlorin-Ifrim]

Interesting, was the lambda also chosen in a similar fashion? Trial and error? 0.008 is very specific.
Will give it a go with some models that are around 20M, as well as depth recurrence and looped, and let you know if I find anything interesting.

[kooshi]

Indeed. Most of my AI findings have been the result of intuition, trial, and error. I spend hours just watching graphs as I train each iteration. Hence why I'd love help analyzing and refining more rigorously. I'm an enterprise software dev by trade.

I landed on 0.008 just by watching where the penalty landed with respect to the overall loss. I think I only tested 0.01 before that, and I'm not even certain that the move to 0.008 changed much. It was weeks ago.

I've been doing most of my tests at d=128. My gut says lambda might need to increase slightly at higher d to push the embeddings out far enough, but higher lambda might make the initial penalty too high relative to the loss from the model itself.

The shape of the penalty might be able to be refined as well, to make it more agnostic to the dimension, but I'm confident you'll see an improvement as-is for now.

[kooshi]

In case anyone does use this, turns out I've been accidentally using it with lam=0.1 for a while, so... that works too I guess

[CiprianFlorin-Ifrim]

Ok so I tried this at d=768 on a small LM trained on TinyStories, three lambda values:

lambda	native val delta vs baseline (512 ctx)	extended-ctx (4096) val delta
0.001	+0.002	+0.001
0.005	+0.009	+0.006
0.010	+0.011	+0.009

Baseline (no penalty) was native val 1.3494, ctx-4096 val 1.2832. Every lambda I tested hurt, with the regression scaling with penalty strength, clean monotonic.

My best guess on why it didn't transfer to this setup: I use the Vaswani-style x = tok_emb(idx) * sqrt(d) input scaling. At d=768 the trained raw tok_emb settles at a per-row norm around 2.8, so the post-scale residual is around 77. The shell target sqrt(2d) ~=39 on the raw embedding is roughly a 14x stretch from where training puts it, which then cascades into a circa 14x-larger residual than the downstream RMSNorm gain init and subsequent layer scales are calibrated for. The penalty seemed to fight the input scaling.

Separately I have an auxiliary memory module that gates against the embedding stream, and its usage rose monotonically with lambda which is consistent with tok_emb capacity being eroded and the model leaning harder on the auxiliary path to compensate.

The clean ablation I haven't run yet is dropping the sqrt(d) input scaling and re-testing. That could separate the penalty not helping. If you tested it on an arch that doesn't pre-scale tok_emb at the input, I'd be interested in what you got.