Regret Pre-training: Bridging Prior and Posterior Views for Enhanced Knowledge Grounding
Abstract
Causal language models factorize sequence probabilities using only preceding context, leaving future information unexploited during training despite its availability in the training data. This paper introduces Regret Pre-training, a self-supervised framework grounded in the Learning Using Privileged Information (LUPI) paradigm. The framework employs a dual-view architecture in which a single model generates both a causal Student distribution and a future-conditioned Teacher distribution. The training objective augments standard language modeling with a regret loss that minimizes the KL divergence from teacher to student, transferring future-aware signals to the causal representations. We investigate two teacher configurations on the OLMoE-1B-7B architecture:LocalRegret, which extends attention by one future token, andGlobalRegret, which conditions on bidirectional context with the target position masked. Experiments on nine downstream tasks following 4 billion tokens of training demonstrate that both configurations consistently outperform the baseline. On average,GlobalRegret andLocalRegret achieve 33.9% and 32.2% accuracy respectively, surpassing the baseline's 30.2%. Most notably,GlobalRegret improves BoolQ performance by 18.1 percentage points (61.0% vs 42.9%). The framework introduces no additional parameters and requires only one extra inference-mode forward pass per training step.
Lay Summary
When today's AI language models — the kind that power chatbots and writing assistants — learn from text, they read it the way a person would read a book with a sliding window: only the words that come before the current word are visible, while everything that comes after is hidden. This left-to-right reading style is what allows the model to generate text one word at a time later on. But it also wastes information. During training, the full text is actually available; the model just isn't allowed to peek ahead. As a result, the model misses out on hints from later sentences that could help it understand the earlier ones better — much like trying to solve a crossword clue without ever glancing at the intersecting answers.We propose a training method called Regret Pre-training that lets the model briefly "peek ahead" during training, but only to teach itself. The same model reads the text twice: once in the normal restricted way, and once with access to the surrounding context. It then compares the two readings and learns to close the gap. At test time, the peeking ability is removed, so the model still generates text in the usual left-to-right manner — but it has internalized the lessons learned from the broader view.In our experiments, this simple change consistently improved the model's performance across nine common reasoning and comprehension tests, with the largest improvement reaching 18 percentage points on a reading comprehension task. The method adds no extra parameters and requires only a modest increase in training time, making it a practical way to get more out of the same training data.