Reproduce “Corrector Sampling in Language Models” on OpenResearch

About this paper

Itai Gat, Neta Shaul, Uriel Singer, Yaron Lipman

Meta FAIR researchers introduce Resample-Previous-Tokens (RPT), a corrector sampling method that enables autoregressive language models to iteratively revisit and potentially replace previously generated tokens within a local window during inference, achieved by fine-tuning pre-trained models on only 10% additional tokens to learn Previous-Token-Prediction (PTP) conditionals alongside standard Next-Token-Prediction, demonstrating consistent ~10% relative improvements on reasoning and coding benchmarks including HumanEval+ (28.6 vs 25.6), MBPP, and GSM8K through a lightweight training modification that introduces local permutation awareness within fixed-size windows while preserving original NTP performance, with theoretical analysis confirming that PTP achieves empirically lower cross-entropy losses than NTP due to additional future context availability, and practical implementation requiring only 1-2 iterations to achieve convergence with reduced total variation distance between predicted and ground truth distributions.

Key insights

• Previous-Token-Prediction achieves significantly lower cross-entropy losses than Next-Token-Prediction because future context aids in predicting past tokens
• Local permutation awareness within fixed windows preserves autoregressive efficiency while enabling targeted correction capabilities
• RPT sampling reduces total variation distance between predicted and ground truth distributions after just one iteration, with 64.5% of cases showing improved ground truth token probabilities

Summary by alphaXiv.