PID Control-Based Self-Healing to Improve the Robustness of Large Language Models
Zhuotong Chen, Zihu Wang, Yifan Yang, Qianxiao Li, Zheng Zhang
TL;DR
This work introduces a PID control-based self-healing scheme to enhance the robustness of pre-trained large language models against perturbations encountered during online inference. By modeling an LLM as a discrete dynamical system and designing running losses via embedding manifolds for state, accumulation, and derivative, the approach yields corrective controls at each layer. It offers analytic solutions under linearity/orthogonality assumptions (with Assumptions 1 and 2) and a theoretical error bound showing perturbations in the orthogonal complement decay, while enabling Tucker-based construction of embedding subspaces for practical deployment. Empirically, the method improves robustness across SNLI, MNLI, and ANLI datasets with modest inference-time overhead and favorable comparisons to baseline adversarial training, demonstrating the practical viability of a low-cost, feedback-driven defense for trustworthy NLP.
Abstract
Despite the effectiveness of deep neural networks in numerous natural language processing applications, recent findings have exposed the vulnerability of these language models when minor perturbations are introduced. While appearing semantically indistinguishable to humans, these perturbations can significantly reduce the performance of well-trained language models, raising concerns about the reliability of deploying them in safe-critical situations. In this work, we construct a computationally efficient self-healing process to correct undesired model behavior during online inference when perturbations are applied to input data. This is formulated as a trajectory optimization problem in which the internal states of the neural network layers are automatically corrected using a PID (Proportional-Integral-Derivative) control mechanism. The P controller targets immediate state adjustments, while the I and D controllers consider past states and future dynamical trends, respectively. We leverage the geometrical properties of the training data to design effective linear PID controllers. This approach reduces the computational cost to that of using just the P controller, instead of the full PID control. Further, we introduce an analytical method for approximating the optimal control solutions, enhancing the real-time inference capabilities of this controlled system. Moreover, we conduct a theoretical error analysis of the analytic solution in a simplified setting. The proposed PID control-based self-healing is a low cost framework that improves the robustness of pre-trained large language models, whether standard or robustly trained, against a wide range of perturbations. A detailed implementation can be found in:https://github.com/zhuotongchen/PID-Control-Based-Self-Healing-to-Improve-the-Robustness-of-Large-Language-Models.