DEVOURED
A new transformer passes its hidden state to the next token instead of recomputing it
Researchers introduced LIFT, a transformer architecture that passes internal states between generation steps to outperform standard transformers on reasoning and procedural tasks.
What: The Latent Information Feedback Transformer (LIFT) feeds its previous internal state into the next token generation step. During pretraining, it is trained to predict both the next token and the next state using teacher-forced supervision.
Why it matters: This approach addresses the 'narrow channel' problem where standard transformers discard deep-layer representations between generation steps, forcing them to redundanty recompute state.
Deep dive
- LIFT models (135M to 1B parameters) consistently beat token-matched standard transformers.
- Pretraining remains parallel because state inputs are precomputed from an existing model's predictions.
- Inference requires feeding back predicted states with minor compute overhead.
- Exhibits strong performance on state-tracking and procedural benchmarks.
- Demonstrated effectiveness using supervision from models that were otherwise ineffective at the task.
Decoder
- Teacher-forced: A training technique where a model is trained to predict the next step using the ground-truth previous state rather than its own (possibly flawed) previous output.
- Hidden State: The intermediate, high-dimensional representation of input data within a neural network.
Original article
Transformer language models (LMs) are feed-forward: deep-layer representations are never fed back to shallower layers, and the only pathway for information to flow downward across generation steps is the decoded token. This narrow channel forces models to recompute intermediate results and to discard alternative continuations. In this work, we remove this bottleneck during pretraining, introducing the LIFT (Latent Information Feedback Transformer) architecture and training method which enable LMs to propagate state across generation. We achieve this by turning recurrent-state learning into a teacher-forced prediction problem: each input token is paired with an information-dense state, derived from the next-token distribution of an off-the-shelf pretrained LM. The model, extended with a small number of additional parameters, is then trained to predict both the next token and the next state. As the input states are precomputed, pretraining remains fully parallel across positions. At inference, the model's own predicted states are fed back, with a minor computational overhead that decreases with model size. Experiments with pretrained models ranging from 135M to 1B parameters show that LIFT consistently outperforms standard Transformers and baselines on language modeling, downstream reasoning tasks, and procedural tasks under token-matched budget, while being on par with or ahead of compute-matched Transformers. Moreover, a controlled study on a state-tracking task shows that a tiny LIFT outperforms same-size Transformers trained on 8x more data, even when trained with the states of a Transformer that fails the task. Overall, we show that LMs can learn to exploit deep-to-shallow feedback during pretraining via scalable teacher supervision.
Explore all 58 articles in today's digest
Fresh dev news, AI-digested daily. Never miss a trend at devoured.fyi