Standard Transformer language models operate as pure feed-forward networks: information flows upward through layers, and the only mechanism for passing information across generation steps is the decoded token itself. This design forces the model to recompute intermediate representations at every position and to discard alternative continuations that might be useful for future reasoning. The paper under review introduces LIFT (Latent Information Feedback Transformer), an architecture and training method that enables deep-layer representations to feed back into shallower layers during pretraining. The key insight is to frame recurrent-state learning as a teacher-forced prediction problem, where an off-the-shelf pretrained LM supplies information-dense states that the model learns to predict alongside the next token. Because these teacher states are precomputed, training remains fully parallel across sequence positions. At inference, the model's own predicted states are fed back, adding only minor overhead that diminishes with model size.

The feed-forward bottleneck in standard Transformers

In a conventional Transformer LM, each layer processes the output of the layer below, and the final layer produces the next-token distribution. There is no pathway for information to travel from deep layers back to shallow layers within the same forward pass, and there is no persistent state carried across generation steps other than the discrete token that gets sampled and re-embedded. This means that any computation performed in deep layers — for instance, recognizing a syntactic pattern, tracking a variable assignment, or planning a multi-step reasoning chain — must be recomputed from scratch at the next position. The model cannot "remember" its own intermediate conclusions except by encoding them into the single token it emits, which is a severely constrained channel.

This limitation has been recognized in prior work on recurrent memory Transformers, state-space models, and various forms of latent recurrence. However, most approaches either break the parallelizability of pretraining (by requiring sequential state updates during training) or add substantial parameter overhead. LIFT addresses both concerns: it preserves parallel training through teacher forcing, and it adds only a small number of parameters for the state prediction head and the feedback injection mechanism.

Teacher-forced state prediction during pretraining

The LIFT training procedure begins with an off-the-shelf pretrained Transformer, called the teacher. For each token in the training corpus, the teacher runs a forward pass and records a latent state vector — specifically, a projection of the deep-layer hidden activations that captures the model's next-token distribution information. This state is computed once per token and stored. The student model, which has the same base architecture as the teacher but with two modifications, is then trained on the same corpus with an augmented objective: predict the next token and predict the next state.

The two modifications to the base Transformer are (1) a state prediction head attached to the final layer, which outputs a vector of the same dimensionality as the teacher state, and (2) a feedback injection module that adds the predicted state to the shallow-layer activations at the next position. During training, the input state for each position comes from the teacher's precomputed states, not from the student's own predictions. This means the student sees a clean, information-dense signal at every step, and the training objective is a simple sum of cross-entropy for the token and mean-squared error for the state. Because the input states are fixed and known in advance, the entire sequence can be processed in parallel, just like standard LM pretraining.

The teacher state is derived from the next-token logits of the pretrained LM. The authors use a linear projection of the logits to a lower-dimensional space (e.g., 512 dimensions for a 1B parameter model), which preserves the information relevant to future predictions while keeping the state compact. This projection is learned as part of the teacher's forward pass and frozen during student training. The result is that the student learns to mimic the teacher's internal "thought process" — the deep-layer representation that anticipates the next token — and to make that representation available to its own shallow layers at the next step.

Inference with self-generated feedback

At inference time, no teacher is available. The model must generate its own states and feed them back. The process is autoregressive: for the first token, a zero or learned initial state is injected into the shallow layers. The model produces the first token prediction and a predicted state. For the second token, the predicted state from the first step is injected, and so on. The feedback injection is a simple additive merge: the state vector is projected to the hidden dimension of the target shallow layer and added to the layer's activations before the attention and feed-forward sublayers.

The computational overhead of this feedback loop is minimal. The state prediction head is a single linear layer, and the injection is a single element-wise addition. The paper reports that the overhead decreases as a fraction of total compute with larger models, because the state dimension is fixed while the model's hidden dimension grows. For a 1B parameter model, the overhead is under 3% of forward-pass time; for a 135M model, it is around 8%. This is a favorable trade-off compared with the quality gains.

Experimental results across scales and tasks

The authors train LIFT models at four scales: 135M, 360M, 760M, and 1B parameters. Each model is trained on the same token budget as a standard Transformer baseline, using the same data mixture and optimization hyperparameters. The teacher for each scale is a pretrained Transformer of the same size. Evaluation covers language modeling perplexity, downstream reasoning benchmarks (including GSM8K for math reasoning, ARC for commonsense reasoning, and MMLU for multi-task knowledge), and procedural tasks that require tracking state across many steps (such as algorithmic reasoning and instruction following).

Under token-matched budget, LIFT consistently outperforms the standard Transformer baseline at every scale. The perplexity improvement on the validation set ranges from 0.8 to 1.5 points depending on scale. On GSM8K, the 1B LIFT model achieves 42% accuracy versus 36% for the baseline Transformer. On MMLU, the 1B LIFT reaches 58% versus 54%. These gains are achieved without additional training tokens — the same data yields better models because the feedback mechanism allows more effective use of each token.

Under compute-matched budget, where the baseline Transformer is allowed more tokens to compensate for LIFT's small inference overhead, LIFT remains on par with or ahead of the baseline. This indicates that the quality improvement is not simply due to the model seeing more effective context; the feedback architecture itself provides a representational advantage.

State-tracking study with a failing teacher

A controlled experiment on a synthetic state-tracking task provides strong evidence that LIFT learns a general feedback capability rather than merely imitating the teacher. The task requires the model to track a hidden state variable through a sequence of operations (e.g., "add 3, multiply by 2, subtract 1") and output the final value. A small standard Transformer (50M parameters) trained on this task fails to generalize beyond the training distribution. The authors use this failing Transformer as the teacher for a LIFT student of the same size.

Remarkably, the LIFT student succeeds on the task, achieving near-perfect accuracy, while the teacher fails. Moreover, the LIFT student trained on 1x data outperforms a standard Transformer of the same size trained on 8x more data. This demonstrates that the feedback architecture enables the model to discover the correct state-tracking algorithm from the teacher's noisy states, even when the teacher itself cannot execute the algorithm reliably. The teacher's states contain enough information about the correct next-token distribution for the student to learn the underlying state dynamics, and the feedback loop allows the student to apply those dynamics recursively.

Architectural details of the feedback injection

The feedback injection module is designed to be lightweight and compatible with existing Transformer implementations. The predicted state vector is first projected from the state dimension (e.g., 512) to the model's hidden dimension (e.g., 2048 for the 1B model) using a learned linear layer. This projected vector is then added to the input embeddings at the first layer, or alternatively to the output of the first attention block. The authors experiment with both injection points and find that injecting at the first layer works slightly better for smaller models, while injecting after the first attention block works better for larger models. The difference is small (under 0.2 perplexity points).

The state prediction head is a two-layer MLP with a GeLU nonlinearity, projecting from the final hidden state to the state dimension. This head is trained jointly with the language modeling head. The total parameter overhead of the state head, projection, and injection is approximately 0.5% of the base model parameters for the 1B model, and around 2% for the 135M model. The relative overhead decreases with scale because the state dimension is held constant while the model grows.

Limitations and trade-offs

The LIFT approach has several limitations that the authors acknowledge. First, the method requires a pretrained teacher model to generate the training states. This means LIFT cannot be used for from-scratch pretraining without an existing LM; it is a distillation-and-enhancement technique, not a standalone pretraining method. However, the authors note that any competent LM can serve as the teacher, and the student can exceed the teacher's capabilities on specific tasks, as shown in the state-tracking experiment.

Second, the state dimension is a hyperparameter that must be chosen. Too small a dimension loses information; too large increases the prediction head size and the injection overhead. The authors find that 512 dimensions works well across scales from 135M to 1B, but optimal values may differ for much larger models or different architectures.

Third, the feedback injection is currently a simple additive merge. More sophisticated integration mechanisms (such as gating or attention-based fusion) could potentially yield further gains, but they would increase complexity and parameter count. The authors chose the simplest effective design to isolate the contribution of the feedback idea itself.

Fourth, at inference, errors in the predicted state can compound across steps. The teacher-forced training does not expose the model to its own prediction errors. The authors experiment with scheduled sampling (mixing teacher states with student states during training) but find it does not consistently improve results and sometimes harms convergence. They suggest that future work could explore more robust inference-time correction mechanisms.

Practical implications for LLM development

For practitioners building language models, LIFT offers a concrete path to improve model quality without increasing the training token budget. The method can be applied as a second-phase pretraining step: take an existing pretrained model, generate teacher states for your training corpus, and continue training with the LIFT objective. The parallel training property means this phase is as fast as standard continued pretraining. The resulting model has the feedback capability baked in, which improves reasoning and procedural tasks at minimal inference cost.

The state-tracking experiment suggests a particularly valuable use case: training small models for specialized tasks that require multi-step reasoning or state maintenance. A tiny LIFT model can achieve performance that would otherwise require a much larger standard Transformer or much more training data. This is relevant for edge deployment, on-device inference, and scenarios where model size is constrained.

From an engineering perspective, implementing LIFT requires adding the state prediction head, the feedback projection, and the injection point to the model architecture. The training loop needs to load precomputed teacher states alongside tokens, which is a straightforward data pipeline change. The authors provide reference implementations for Hugging Face Transformers and report that integration takes a few hours for a team familiar with the codebase.

Broader significance for recurrent language modeling

LIFT demonstrates that the feed-forward bottleneck in Transformers is not a fundamental limitation but a design choice that can be overcome with a relatively simple architectural change and a scalable training procedure. The teacher-forced state prediction framework provides a general recipe: any source of informative latent states can be used to teach a model to predict and use its own feedback. This opens the door to using states from stronger teachers (larger models, models trained on more data, models with specialized capabilities) to bootstrap smaller or more efficient student architectures.

The fact that LIFT works even when the teacher fails the target task is particularly provocative. It suggests that the information content of a model's internal representations can exceed the model's own behavioral competence, and that a feedback architecture can unlock that latent capability. This has implications for model distillation, self-improvement loops, and the design of architectures that separate knowledge representation from computation.

Read the paper on arXiv