Something I think that’s overlooked with many latent thinking options is
1. the final representation is not so in domain for what input expects, even adaptation can be learned it sort of conflicts with the main usage of embedding tokens and sampling them as usual, they seem to have different adapting directions
2. From stages of inference paper and many other we know early layers act as extended feature engineering, late layers act as pruning and preparation for loss, middle layers are where representations are rich
A few things come to mind;
1. Wanting to take the output from say, the 80% or so layer, and input it back around 50% or so. Possibly even initializing a zero initialized few transformer layers to adapt the output feature we take, sort it better fits towards the layer we aim to input back into?
2. Possibly a skip connection of earlier features or initial embedding, sort of like what looped transformers do
3. A bit of a different formulation, and one that might be preferrable for backprop with a target rather than RL strategies, is further leaning into fixed point iteration from looped transformers, incorporating some ideas from diffusion/ddpm of routine noise injection to help jitter out of poor minima and keep minimizing energy. I’d wonder too if Gaussian noise aligned with data covariance may be interesting if that better reflects replicating natural error/better treating exposure bias?
4. Might we want to consider different sampling operations at the end, or deterministic? Energy transformer, small diffusion head decoder do well for high dimensions, normalizing flow or Gaussian mixture models also may be possible?
It took me a long time to build an intuition for why CoT works. My thinking was always.. if the model can predict it downstream of 10k thinking tokens, it should have been able to predict it from the outset too.
My intuition now is:
- During inference, the correct paths are indeed somewhere in the hidden states, represented purely as probabilities
- However, in the process of sampling, we're forced to materialize just one path. This is destructive -- a 30% chance of ending up at the answer can become 0 if we sample the wrong token.
- The constant backtracking reasoning models do protect against this. Every "wait" or "but" is another chance for a shot on target.
- By the time models exhaust their reasoning budget, they've already seen a bunch of possible answers
- And since these models are also generally better at verifying answers than generating them, the chances of choosing the correct path, conditioned on this prefix, are much higher than it was at the start.