Latent Lookahead Training for Transformers
Latent lookahead gives an autoregressive transformer extra internal computation before it commits to a difficult next token. The additional steps happen in continuous hidden-state space rather than as visible generated tokens.
Transformer training
Imagine writing the moves for a maze one at a time. At the first junction, left and right both look fine. Left ends three turns later at a wall. If you have to write your first move before tracing those turns, you have to make the choice with less information.
Latent lookahead gives a transformer a few extra passes at selected positions before it emits a token. During those passes it feeds back hidden vectors and gets more hidden vectors. During training, putting a prediction target on each vector gives those extra passes a specific job.
What the model feeds back
A language model already produces a vector before it produces a word. Normally it converts that vector into vocabulary scores, picks a token, and feeds the corresponding token embedding into the next step. At a lookahead position, the method takes the final hidden vector and feeds it back directly. It repeats this process for a chosen number of steps, which the paper calls τ.
Those intermediate vectors add positions to the model’s internal context. Each is a list of numbers. The model can wait until the end of the rollout to pick the next visible token.
Show the targets on a short route
Suppose the correct route begins RIGHT, UP, UP, and each move is a token. Set τ to 3 and put a lookahead position right before the answer. On a training example, the first latent state is scored against RIGHT, the second against UP, and the third against the next UP. The answers provide loss targets. They are unavailable as input to this rollout at inference time.
When the latent vectors are ready, the method lets the vectors within that group attend to one another. The first vector can be revised using information in the later ones. The system then uses that first revised state to choose the next move. After RIGHT is emitted, the model still has to generate the rest of the route.
Training target: RIGHT UP UP
Latent supervision: z₁ z₂ z₃
First visible prediction: read the refined z₁What extra time buys you
One way to check whether the future targets are doing any work is to give another model the same number of extra positions as pause tokens. Both models get room to compute. The lookahead model also gets separate supervision for the upcoming tokens. In the paper’s 4×4 Sudoku experiment, the two-layer model trained from scratch got 78% with standard next-token prediction, 86% with pause tokens, and 93.5% with latent lookahead. Those are results from the paper’s specific task setup.
The fine-tuning experiments on pretrained models were much less consistent. The training and inference cost also matters. The latent rollouts add context, and the steps within each rollout have sequential dependencies. Spending those passes at a junction is easier to justify than spending them on every trivial continuation. The paper places the first rollout after the question in its planning tasks. Learning where to pause is a separate question.