Explorative Modeling: Unlocking a Third Pretraining Axis and E2E Generation

Author summary thread on twitter https://x.com/AlexiGlad/status/2083230922196107288

The idea seems mind-bogglingly simple, instead of

  y    = model(sample_latent())   # generate one output (from noise, a mask, …)
  loss = recon_loss(y, x)         # score it against the data target x
  loss.backward()

you do

  losses = []
  for _ in range(K):                   # explore K candidate outputs
      y = model(sample_latent())       # generate one candidate
      losses.append(recon_loss(y, x))  # score each against x
  min(losses).backward()               # train only the closest candidate

I guess the intuition is that if you just generate one output and score it, your loss function forces it to split the difference between all samples and converge to the average (even if it produces _a_ valid output, if it's not the exact target being trained on it gets penalized). Whereas if you do basically best-of-n, it's not penalized for generating other valid samples as well.