LIFT Architecture Enables Transformers to Pass Hidden State to Next Token via Teacher Supervision

Standard Transformers operate as feed-forward networks where deep-layer representations are never fed back to shallower layers, forcing models to recompute intermediate results and discard alternative continuations. The LIFT (Latent Information Feedback Transformer) architecture removes this bottleneck during pretraining by enabling models to propagate state across generation steps.

How LIFT Works

The method turns recurrent-state learning into a prediction problem by pairing each input token with an information-dense state derived from an off-the-shelf pretrained LM. The model, extended with a small number of additional parameters, is trained to predict both the next token and the next state. Because these input states are precomputed, the pretraining process remains fully parallel across positions.

Performance and Efficiency

During inference, the model uses its own predicted states, introducing a minor computational overhead that decreases as model size increases. Experiments with models ranging from 135M to 1B parameters demonstrate that LIFT consistently outperforms standard Transformers and baselines on downstream reasoning tasks and language modeling under token-matched budgets, while matching or exceeding compute-matched Transformers. Additionally, a controlled study on a state-tracking task showed that a tiny LIFT model outperformed a standard Transformer trained on 8x more data.