# PID Control-Based Self-Healing to Improve the Robustness of Large Language Models

arXiv: https://arxiv.org/abs/2404.00828 · HTML: https://mlsanity.com/paper/2404.00828/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 59 references that are in the map.
- First posted on arXiv 31 Mar 2024.

## Builds on (most cited first, 30 of 35)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Towards Deep Learning Models Resistant to Adversarial Attacks](https://mlsanity.com/paper/1706.06083/) (2017)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [A large annotated corpus for learning natural language inference](https://mlsanity.com/paper/1508.05326/) (2015)
- [Adversarial NLI: A New Benchmark for Natural Language Understanding](https://mlsanity.com/paper/1910.14599/) (2019)
- [Universal Adversarial Triggers for Attacking and Analyzing NLP](https://mlsanity.com/paper/1908.07125/) (2019)
- [Is BERT Really Robust? A Strong Baseline for Natural Language Attack on Text Classification and Entailment](https://mlsanity.com/paper/1907.11932/) (2019)
- [Adversarial Training for Free!](https://mlsanity.com/paper/1904.12843/) (2019)
- [ClinicalBERT: Modeling Clinical Notes and Predicting Hospital Readmission](https://mlsanity.com/paper/1904.05342/) (2019)
- [TextBugger: Generating Adversarial Text Against Real-world Applications](https://mlsanity.com/paper/1812.05271/) (2018)
- [Deep Speech: Scaling up end-to-end speech recognition](https://mlsanity.com/paper/1412.5567/) (2014)
- [Stable Architectures for Deep Neural Networks](https://mlsanity.com/paper/1705.03341/) (2017)
- [FreeLB: Enhanced Adversarial Training for Natural Language Understanding](https://mlsanity.com/paper/1909.11764/) (2019)
- [SMART: Robust and Efficient Fine-Tuning for Pre-trained Natural Language Models through Principled Regularized Optimization](https://mlsanity.com/paper/1911.03437/) (2019)
- [Adversarial Training and Robustness for Multiple Perturbations](https://mlsanity.com/paper/1904.13000/) (2019)
- [You Only Propagate Once: Accelerating Adversarial Training via Maximal Principle](https://mlsanity.com/paper/1905.00877/) (2019)
- [Towards Improving Adversarial Training of NLP Models](https://mlsanity.com/paper/2109.00544/) (2021)
- [Maximum Principle Based Algorithms for Deep Learning](https://mlsanity.com/paper/1710.09513/) (2017)
- [A Mean-Field Optimal Control Formulation of Deep Learning](https://mlsanity.com/paper/1807.01083/) (2018)
- [CAT-Gen: Improving Robustness in NLP Models via Controlled Adversarial Text Generation](https://mlsanity.com/paper/2010.02338/) (2020)
- [Grey-box Adversarial Attack And Defence For Sentiment Classification](https://mlsanity.com/paper/2103.11576/) (2021)
- [Deep Learning Theory Review: An Optimal Control and Dynamical Systems Perspective](https://mlsanity.com/paper/1908.10920/) (2019)
- [An Optimal Control Approach to Deep Learning and Applications to Discrete-Weight Neural Networks](https://mlsanity.com/paper/1803.01299/) (2018)
- [Bayesian Inference with Certifiable Adversarial Robustness](https://mlsanity.com/paper/2102.05289/) (2021)
- [Towards Robust Neural Networks via Close-loop Control](https://mlsanity.com/paper/2102.01862/) (2021)
- [The King is Naked: on the Notion of Robustness for Natural Language Processing](https://mlsanity.com/paper/2112.07605/) (2021)

## Built on by (most cited first, 3 of 3)

- [Perturbation-Restrained Sequential Model Editing](https://mlsanity.com/paper/2405.16821/) (2024)
- [Test-time reward-guided alignment of language models by importance sampling on pre-logit space](https://mlsanity.com/paper/2510.26219/) (2025)
- [MARR: Module-Adaptive Residual Reconstruction for Low-Bit Post-Training Quantization](https://mlsanity.com/paper/2605.17997/) (2026)

## Abstract (arXiv)

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.