ST$^2$U: Stateful Test-Time Unlearning via Restricted Knowledge Boundary Control
2026-08-24 • Machine Learning
Machine LearningComputation and Language
AI summaryⓘ
The authors address the problem of safely removing restricted knowledge from large language models during their use, without retraining. They show that previous methods only fix small parts of the model's memory at a time, which lets the unwanted knowledge sneak back in later. Their new method, called ST²U, watches the whole generation process and adjusts the model's hidden states more globally and carefully to prevent this re-entry. Tests on different models and tasks show that ST²U forgets the restricted knowledge better while keeping the model useful and fast.
large language modelstest-time unlearningactivation editingautoregressive generationhidden statesmodel alignmentrestricted knowledgetrajectory controlinvertible coordinatesknowledge forgetting
Authors
Xunlei Chen, Qinghui Gong, Ruini Xue, Yaodong Hu, Tian Lan, Wenhong Tian
Abstract
Controlling restricted knowledge in large language models is essential for model alignment and safe deployment. Test-time unlearning avoids costly retraining and parameter updates by intervening only during inference. However, existing activation-editing methods apply isolated pointwise corrections, overlooking how autoregressive generation continually reconstructs hidden states from the prompt, cache, and generated prefix. Consequently, later states may return to restricted knowledge regions after a locally successful correction, causing restricted knowledge re-entry. In this work, we propose Stateful Test-Time Unlearning via restricted knowledge boundary control (ST$^2$U), which formulates test-time unlearning as trajectory-wide boundary control. ST$^2$U first models restricted knowledge boundaries in low-dimensional invertible coordinates while leaving orthogonal non-target components unchanged. During inference, ST$^2$U monitors risk along the trajectory, applies minimal boundary corrections with contextual anchoring, and propagates historical correction states across tokens to mitigate knowledge re-entry. This trajectory-wide control enables more persistent forgetting while preserving non-target capabilities and limiting inference overhead. Across three benchmarks and three model families, ST$^2$U delivers the strongest overall balance, combining best or second-best retention with competitive forgetting and substantially less restricted-knowledge re-entry than test-time baselines (13.76%-19.84% versus 46.50%-59.10%).