Correcting in Hindsight: Editing Past Key-Value States for Robust LLM Reasoning

Autoregressive Large Language Models (LLMs) often fail in complex reasoning because early-stage errors remain uncorrectable in subsequent steps—a limitation fundamentally rooted in the inherent irreversibility of the Transformer architecture. In this paper, we propose HEdit, a lightweight reasoning enhancement paradigm that equips models with a "hindsight-like" capability for dynamic error correction during generation. Our core insight involves deconstructing reasoning failures into two pivotal stages: latent representational biases emerging at logical anchors, and the subsequent eruption of explicit cognitive dissonance at trigger points. Based on these observations, the HEdit framework detects internal inconsistency signals at trigger points in real-time, actively backtracks to critical anchors, and utilizes a lightweight trainable editor to precisely refine their Key-Value (KV) caches. This mechanism effectively breaks the unidirectional constraints of autoregressive inference. Empirical results demonstrate that HEdit significantly enhances the performance of various models on mathematical reasoning tasks—with average accuracy improvements ranging from 2.2\% to 10.8\%—while maintaining extremely low overhead (add parameters $<0.5\%$). HEdit provides a dynamic, pluggable and lightweight solution, making it particularly beneficial for users in low-resource environments.

When solving a complex puzzle, people often notice they took a wrong turn earlier and backtrack to fix it. Today's AI language models cannot do this: once they commit to an intermediate step, they treat it as unchangeable truth, even if it was incorrect. Consequently, a tiny early mistake—such as copying a wrong number—can quietly corrupt every subsequent step until the final answer fails. We present HEdit, a lightweight method that gives language models a kind of "hindsight." While the model reasons through a problem, it watches for internal signs of confusion or contradiction. When such a signal appears, HEdit pinpoints the specific earlier decision that started the trouble and makes a minimal correction to the model's memory of that moment. The model then resumes reasoning from that corrected point, as if the error never happened. This intervention adds less than 0.5% extra size to the model and only triggers when necessary. Experiments on difficult math benchmarks show that HEdit improves accuracy by 2% to over 10% across various AI models, with almost no added computational cost. By enabling AI to revise its past reasoning rather than stubbornly marching ahead, HEdit makes advanced problem-solving more robust and practical, even for smaller systems.