PaperlayerView exact revision on arXiv

Source-linked AI summary

Pretraining Latent Information Feedback Transformers with Teacher Supervision

Dor Tirosh, Ido Amos, Mor Geva

arXiv:2609.38149v1cs.CL

TL;DR

Transformers pass information downward only through decoded tokens. LIFT trains in parallel on teacher-provided next-token-distribution states and feeds back its own predicted states at inference. Across evaluations, it outperformed token-matched Transformers and feedback baselines, matched or beat compute-matched Transformers, and solved 12-step S5 sequences at 100% accuracy.

  • Problem

    Transformers lack deep-to-shallow representation feedback: downward information passes only through decoded tokens, a narrow channel associated with recomputation or dropped continuations.

    Source: §1 · p. 1, §2 · p. 3

  • Method

    LIFT pairs each token with a pretrained teacher’s next-token-distribution state and trains parallel next-token and next-state prediction, then feeds back its own predicted states at inference.

    Source: §1 · p. 2

  • Results

    Across evaluations, LIFT outperformed token-matched Transformers and distillation/Jacobi baselines and matched or beat compute-matched Transformers; it solved 12-step S5 sequences at 100% accuracy.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The controlled task shows that a fixed-depth LIFT can solve longer sequences than a same-depth Transformer, even when trained on states from a teacher that fails at that length.

    Source: §1 · p. 2

    The study focuses on teachers sharing the student architecture; effects of other teacher sizes or vocabularies and training far beyond the teacher’s budget remain open.

    Source: §7 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Standard Transformers lack a deep-to-shallow feedback path during generation, while prior approaches to latent feedback face costly sequential training or iteration trade-offs. LIFT uses teacher-forced states for parallel pretraining and feeds back its own predicted states at inference, with gains across language modeling and downstream tasks.

1 INTRODUCTION

  • Prior post-training latent-feedback methods rely on expensive sequential rollouts, while Jacobi-style recurrence approximations trade improved quality for more iterations and training cost.LIFT addresses the pretraining question of exploiting deep-to-shallow feedback without sequential processing.Source: §1 · pp. 1–2
  • LIFT turns recurrent-state learning into teacher-forced prediction: it receives a token and state at each position, predicts both next outputs, then feeds back its own states at inference.Teacher states come from next-token distributions of an external model and are produced independently, preserving parallelism across positions.Source: §1 · p. 2
  • At 12 steps, two-layer LIFT achieves 100% accuracy with Transformer-teacher states, while the vanilla Transformer fails and the teacher scores 3%.This shows fixed-depth LIFT can solve longer sequences even when its teacher fails at that length.Source: §1 · p. 2
  • Across 135M–1B-scale evaluations, LIFT outperforms token-matched Transformers and feedback baselines, and is level with or ahead of compute-matched Transformers.Its largest gains are on procedural tasks, where it matches or beats Transformers trained on 2× more tokens at every scale.Source: §1 · p. 2
  • LIFT introduces a teacher-forced latent feedback channel for Transformer pretraining, with a pretrained LM supplying training states while the model generates its own states at inference.Its states are next-token distributions, aligning state supervision with the model’s native prediction task.Source: §1 · pp. 2–3

2 INFORMATION FLOW BOTTLENECK IN TRANSFORMER LMS

In a standard Transformer, attention lets later positions read earlier representations, but information moves downward across generation steps only through the decoded token. This narrow channel can force recomputation, discard continuations, and limit fixed-depth models on long state-tracking tasks.

2 INFORMATION FLOW BOTTLENECK IN TRANSFORMER LMS

  • Downward feedback is limited to a single decoded token, carrying at most log2 |V| bits—about 17 bits when |V| = 100,000.Attention mediates upward communication, allowing later positions to read earlier representations.Source: §2 · p. 3
  • Allowing lower layers to access earlier higher-layer representations could let models reuse deep computations instead of recomputing them or failing.The section also describes carrying several continuations forward rather than committing to one sampled token.Source: §2 · p. 3
  • State-tracking tasks whose sequential steps grow with input length are a challenge for fixed-depth Transformers.The section identifies recurrent computation, whose depth grows with sequence length, as relevant to such tasks.Source: §2 · p. 3

3 LATENT INFORMATION FEEDBACK TRANSFORMERS (LIFT)

LIFT adds a state alongside each token so information from deeper layers can propagate across generation steps. It trains the model to predict both next tokens and next states, using teacher-forced states to preserve parallel training.

3.1 LIFT ARCHITECTURE FOR NEXT TOKEN-STATE PREDICTION

  • LIFT predicts a state alongside each token, extending the information carried across generation beyond the sampled token.The state is derived from a top-k-renormalized output distribution, which retains ranked alternative continuations.Source: §3.1 · pp. 3–4
  • State distributions provide shared token-probability coordinates for teacher targets and the model’s own inference-time predictions.The paper favors distributions over hidden representations because hidden states need not align across models.Source: §3.1 · p. 4

3.2 LIFT TRAINING WITH TEACHER-FORCED STATES

  • A frozen pretrained teacher computes states in advance, letting token-state pairs train in parallel through the standard Transformer backbone.The teacher’s weights remain fixed, and gradients are not propagated into it.Source: §3.2 · p. 4
  • Teacher-forced states address the training–inference mismatch with state-alignment supervision, followed by a final 10% adaptation phase using the model’s own states.The adaptation phase drops both the teacher and alignment term.Source: §3.2 · p. 5

4 EVALUATION OF STATE TRACKING CAPABILITIES

On S5 permutation composition, LIFT’s state feedback helps a fixed-depth model handle longer sequences than a standard Transformer. Parallel LIFT training can benefit even from a teacher that fails at the tested lengths, but does not give LIFT unbounded depth beyond its training range.

The S5 task

  • S5 requires composing a sequence of permutations to recover the final state, making it a state-tracking task.For growing N, fixed-depth Transformers cannot solve the task unless TC0 = NC1, while recurrent models need only a constant number of layers.Source: §4 · p. 5

Experimental setup

  • The study trains two-layer, 2M-parameter models on sequences up to N=16 and evaluates accuracy on held-out sequences through N=64.Results average three seeds; the Transformer×8 baseline trains for 40,000 steps versus 5,000 for the token-matched Transformer.Source: §4 · p. 5

Results

  • When trained in parallel on states from a recurrent LIFT-RNN teacher using lengths up to N=16, LIFT extrapolates, achieving 96% at N=32 and 65% at N=64.Source: §4 · p. 6
  • With a depth-bounded Transformer teacher, LIFT solves lengths through N=12 and half the sequences at N=16, but fails beyond its training range.Its teacher, trained with 8× the budget, is at chance beyond N=8.Source: §4 · p. 6

5 EXPERIMENTS WITH PRETRAINED LMS

Across pretrained-LM experiments, LIFT improves on Transformer baselines under token-matched budgets and remains competitive when compute is matched. Its largest reported gains are on procedural tasks, and it also compares favorably with feedback Transformers that use iterative state approximation.

5.1 EXPERIMENTAL SETTING

  • Experiments cover 135M, 350M, and 1B OLMo 2 models trained for 5× their Chinchilla token budgets, with same-size teachers trained for 10×.The feedback-model comparison uses a shared SmolLM2-135M backbone and evaluates matched-token and matched-compute budgets.Source: §5.1 · pp. 6–7

5.2 EVALUATING LANGUAGE MODELING, KNOWLEDGE, AND REASONING CAPABILITIES

  • 5–5.5% lower perplexity and wins in 33 of 36 downstream comparisons: LIFT leads token-matched Transformers at every scale.At matched compute, where Transformers train on over 40% more data, LIFT leads at 350M and 1B and is level with or ahead at 135M.Source: §5.2 · p. 7, Table 1 · p. 8
  • The token-efficiency advantage over a vanilla Transformer grows during training at every scale.Figure 3b marks the additional tokens the Transformer needs to reach LIFT’s loss at matched checkpoints.Source: §5.2 · p. 7, Figure 3 · p. 8
  • LIFT’s largest gains are on procedural tasks: arithmetic results beat a teacher trained on 2× more tokens, and pattern-continuation BPB improves over token-matched Transformers by 6–10%.Arithmetic gains persist as operand counts increase; pattern-continuation gains over 2×-token Transformers are 4–5% at 350M and 1B, and a tie at 135M.Source: §5.2 · p. 7, Figure 3 · p. 8
  • LIFT w/o states performs comparably to token-matched Transformers, showing that the model retains strong capabilities for parallel prompt processing.Iterative prompt refinement recovers much of the gain of sequential processing within a few iterations.Source: §5.2 · pp. 7–8

5.3 COMPARISON WITH JACOBI-ITERATION FEEDBACK TRANSFORMERS

  • LIFT is equal to or better than Jacobi-iteration feedback methods on every metric under both token-matched and compute-matched budgets.Iterative refinement trades training efficiency for approximation quality; LIFT also uses less peak memory than both feedback baselines at the measured batch sizes.Source: §5.3 · pp. 8–9, Table 2 · p. 9

6 RELATED WORK

Related work addresses latent feedback in both post-training and pretraining, but often relies on sequential rollouts, iterative approximations, or annotated reasoning. LIFT’s distinction is teacher-forcing pretrained-LM states during pretraining.

Latent feedback in post-training

Latent feedback in pretraining

  • Pretraining approaches use exact recurrence or approximate Jacobi iterations, whereas LIFT teacher-forces states from a trained LM.Jacobi-based methods trade training cost for state-approximation quality.Source: §6 · p. 9
  • Other teacher-forced approaches differ in their supervision source: Kumar & Isola (2026) use a jointly trained encoder, while Tack et al. (2026) use a pretrained teacher for targets.Tack et al. (2026) mix predicted concepts into hidden states.Source: §6 · p. 9

7 CONCLUSION

LIFT uses teacher-provided states to train Transformers for deep-to-shallow feedback in parallel, then predicts its own states at inference. Across the reported scales and settings, it outperforms standard Transformers on varied tasks; the study’s teacher choice and training regime leave questions open.

7 CONCLUSION

  • Across 135M–1B scales, LIFT outperforms standard Transformers on varied tasks when matched in tokens and compute.Against alternative feedback methods, it is comparable to or better than all baselines in both settings.Source: §7 · pp. 9–10
  • LIFT uses an external teacher’s next-token output distribution to provide states for deep-to-shallow feedback while maintaining parallel training.During pretraining, LIFT learns to predict its own states, so it does not rely on the teacher at inference.Source: §7 · p. 9
  • The study focuses on teachers that share the student’s architecture; effects of other teacher sizes or vocabularies remain future work.Performance and training dynamics when LIFT is trained far beyond its teacher’s training budget also remain open questions.Source: §7 · p. 10

A ADDITIONAL METHOD DETAILS

These appendices detail state-dropout and alignment-loss choices, the adaptation procedure, the S5 setup, model-training configurations, and the feedback-transformer comparison. The procedures are designed to expose models to missing or self-generated states while retaining parallel teacher-forced training.

A.1 PREFIX STATE DROPOUT

  • Prefix state dropout randomly selects a prefix to process without states, and can expose the model to an entire sequence without states.The first position always keeps its learned initial state; the augmentation also applies to the adaptation phase’s loss-carrying pass.Source: §A.1 · p. 19

A.2 THE STATE-ALIGNMENT LOSS

  • The state-alignment loss compares teacher and student distributions on the state’s top-k tokens and a combined tail outcome.Including the tail lets the student match the teacher’s mass outside the top-k set instead of being pushed to sharpen its distribution; the coarsened KL is a lower bound on full-vocabulary KL.Source: §A.2 · p. 19

A.3 ADAPTATION PHASE

  • In the last 10% of training, adaptation uses one or two no-gradient passes to generate states, then a loss-carrying pass fed by prior-pass states.The final pass applies prefix state dropout and uses the language-modeling loss, with alignment turned off; no-gradient passes add 1.5 forward passes per step on average, and the teacher is not needed.Source: §A.3 · p. 20

B.1 S5 STATE TRACKING: SETUP

  • The S5 experiments train two-layer, width-256 models on sequences from 12 fixed generators and evaluate accuracy on held-out sequences at each length.LIFT uses a linear fusion layer and top-k states; teacher-fed models switch to their own sequentially computed states for the final 10% of training.Source: §B.1 · p. 20

B.2 IMPLEMENTATION DETAILS

  • The language-model experiments use OLMo 2 blocks, three model scales, shared LIFT hyperparameters, and a stratified corpus drawn from five sources.Training budgets are 13.4B, 34.9B, and 107.4B tokens; the 135M, 350M, and 1B models use pretrained vanilla teachers at their own or larger sizes.Source: §B.2 · pp. 20–21

B.3 FEEDBACK-TRANSFORMER COMPARISON: SETUP

  • The 135M feedback comparison uses a shared 10B-token FineWeb-Edu setup, with LIFT trained without the KL objective and evaluated in sequential recurrence.T2MLR uses 2.28× vanilla-step FLOPs per token; LIFT reaches that compute-matched budget after 16.2B tokens and the vanilla model after 23.0B.Source: §B.3 · p. 22

B.4 EVALUATION BENCHMARKS

The benchmark appendix specifies how downstream suites are scored and selected for the main table, and describes the arithmetic and source datasets. It also identifies excluded suites and clarifies why arithmetic is reported in bits per byte.

B.4 EVALUATION BENCHMARKS

  • The twenty OLMES suites use each model’s trained context window; cloze options are scored by length-normalized log-likelihood, while bits per byte normalizes answer log probability by UTF-8 bytes.At 135M, generation is capped at 256 tokens across models, so those generative scores are comparable within, not across, model sizes.Source: §B.4 · p. 22
  • The main table omits suites at chance or floor for every model and BoolQ, while all excluded suites are reported in appendix tables.BoolQ is a yes/no cloze with a 62% majority class, and its scores vary by ten points across training seeds for some comparisons.Source: §B.4 · p. 22
  • Arithmetic prompts contain three worked examples followed by a left-to-right expression using random integers and addition or subtraction.Bits per byte is reported because teacher-forced argmax exact match is substantial only for two operands and only a few percent from three operands onward for every model.Source: §B.4 · pp. 22–23

C.1 TRANSFORMER TRAINING BUDGET ON S5

The appendix tests whether extra training raises the two-layer Transformer’s S5 ceiling and reports additional breakdowns of perplexity and downstream results. It also specifies seed, arithmetic, and pattern-continuation analyses.

C.1 TRANSFORMER TRAINING BUDGET ON S5

  • The two-layer Transformer gains nothing beyond 100,000 steps and remains near 10% at N=12 and at chance from N=16 onward across budgets.The 20× and 40× curves coincide in every seed.Source: §C.1 · p. 24, Figure 4 · p. 24

C.2 PERPLEXITY BY SOURCE AND EVALUATION UNCERTAINTY

  • LIFT leads token-matched, compute-matched, and distillation baselines on all eleven perplexity sources at every size.The paired bootstrap intervals measure dependence on held-out text at a fixed training run, not seed variance.Source: §C.2 · p. 24, Table 5 · p. 25, Table 6 · p. 25

C.3 DOWNSTREAM RESULTS PER SUITE

  • Tables 7–9 report OLMES results at 135M, 350M, and 1B, including suites omitted from the main table.The tables distinguish accuracy for multiple-choice and generative tasks from bits per byte for the BPB rows.Source: §C.3 · p. 24, Table 7 · p. 26, Table 8 · p. 26

C.4 SEED VARIANCE AT 135M

  • The 135M main-table results are also reported across three training seeds as mean ± standard deviation.Differences are treated as beyond seed noise when the mean ± sd ranges do not overlap.Source: §C.4 · p. 24, Table 10 · p. 27

C.5 ARITHMETIC BY NUMBER OF OPERANDS

  • Tables 11 and 12 break arithmetic performance down by operand count using bits per byte and teacher-forced argmax exact-match accuracy.The bits-per-byte table also includes a vanilla Transformer trained on twice the tokens.Source: §C.5 · p. 24, Table 11 · p. 28, Table 12 · p. 28

C.6 PATTERN CONTINUATION

  • The pattern-continuation results cover 534 five-shot questions and report cloze accuracy and answer bits per byte.The table includes the paper’s training-seed rows and a Transformer trained on twice the tokens.Source: §C.6 · p. 29, Table 13 · p. 29

C.7 FEEDBACK-TRANSFORMER COMPARISON: RESULTS

The feedback-Transformer comparison reports seed variability and suite-level results, while the ablation framing states that LIFT leads compute-matched Transformers at every scale. The ablations then examine where that lead comes from.

C.7 FEEDBACK-TRANSFORMER COMPARISON: RESULTS

  • Table 14 reports three-seed means and standard deviations for most Table 2 models; Table 15 lists seed-42 results for every OLMES suite.T2MLR (HF) has no seed statistics because only one released checkpoint is available.Source: Table 14 · p. 29, §C.7 · p. 29, Table 15 · p. 30

D ABLATIONS

  • LIFT is ahead of compute-matched Transformers at every scale, motivating ablations of how its extra forward pass is used.The compute-matched baseline spends the same compute on over 40% more tokens.Source: D ABLATIONS · p. 30

D.1 TEACHER FORCING: INPUTS, TARGETS, OR NO TEACHER

The ablations test how teacher states, channel width, the adaptation phase, teacher strength, and prompt-cache construction affect LIFT. They find that teacher-forced states and a sufficiently wide channel are central to the gains, while one parallel cache-refinement pass nearly recovers sequential prefill.

D.1 TEACHER FORCING: INPUTS, TARGETS, OR NO TEACHER

  • At 135M, only teacher states used as both inputs and targets beat the compute-matched Transformer (+0.47±0.08 perplexity); target-only distillation trails by 0.77±0.04.Training on the model’s own states ties the compute-matched baseline, while the alignment term helps transfer the teacher-forced channel to inference.Source: Figure 5 · p. 31, §D.1 · pp. 31–32

D.2 CHANNEL WIDTH

  • Perplexity improves monotonically with channel width: k = 1 is 1.16 behind the compute-matched Transformer, while k = 1024 is 0.38 ahead.The gain comes from what wider states carry at inference, not from training a better model without states; returns diminish at larger k.Source: Figure 6 · p. 32, §D.2 · pp. 32–33

D.3 TRUNK VS. ADAPTATION PHASE

  • The teacher-forced trunk builds most of LIFT’s gain; adaptation alone reaches 32.53 perplexity at 135M, slightly worse than the 32.43 Transformer.At 1B, training with the teacher until the end retains 73% of the gain, and a 5% adaptation phase recovers the rest.Source: §D.3 · p. 33

D.4 TEACHER STRENGTH

  • Results barely change across teacher choices, including larger teachers or teachers trained much longer, on every reported metric.Table 16 compares same-size and stronger teachers against the token-matched vanilla Transformer.Source: Table 16 · p. 33, §D.4 · p. 33

E PREFILL MODES: REFINING THE PROMPT’S KEY-VALUE CACHE

  • One parallel refinement pass recovers 93–104% of LIFT’s relative gain over sequential prefill, with three or five passes changing scores by at most 0.2 points.Without refinement, LIFT retains 37–51% of its gain on OLMES; one refinement costs an extra parallel pass rather than sequential prompt processing.Source: Figure 7 · p. 34, E PREFILL MODES: REFINING THE PROMPT’S… · pp. 34–35

F.1 TRAINING AND INFERENCE FLOPS

LIFT adds less than 10% to vanilla training FLOPs over a full run, and the feedback-model comparison attributes memory differences to the activations and caches each method retains. The related-work discussion distinguishes LIFT’s training-time teacher inputs from other feedback and reasoning approaches.

F.1 TRAINING AND INFERENCE FLOPS

  • Across a full run, LIFT’s additions total under 10% of vanilla training FLOPs, excluding the teacher’s forward pass from the itemized compute accounting.Compute-matched Transformers are budgeted using the whole LIFT run, including the teacher pass and adaptation passes.Source: §F.1 · pp. 35–36, Table 18 · p. 36

F.2 MEMORY FOOTPRINT

  • LIFT’s extra memory is transient teacher logits and a soft-token matrix, while recurrent caches and retained pass activations contribute to the feedback baselines’ memory use.The Multi-pass Transformer has the largest peak because its three-pass batches retain activations from all three passes.Source: §F.2 · p. 36, Table 19 · p. 37

G EXTENDED RELATED WORK

  • Unlike looped depth or chain-of-thought computation, LIFT’s depth grows with sequence length at fixed per-token cost, complementing both approaches.Looped-depth inference cost grows with iterations; chain of thought adds serial generated tokens, each passing through the decoded-token channel.Source: G EXTENDED RELATED WORK · p. 37
  • LIFT alone among the closest feedback methods uses an external pretrained LM to supply training states, allowing a single parallel pretraining pass.The student learns to replace those teacher distributions with its own states; the distillation baseline isolates the teacher’s target role.Source: G EXTENDED RELATED WORK · p. 37, Table 20 · p. 38

Related papers

Loading 2609.38149v1…