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

Autoregressive Large Language Models (LLMs) often fail in complex reasoning because earlystage 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. Our code can be found at github: https://github.com/Zmfei/hedit