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

Here we provide background knowledge on PID control, which is an essential building block of our proposed framework. A PID (Proportional, Integral, Derivative) controller is a feedback-based control system extensively utilized in industrial settings and numerous other domains where continuous adjustment is essential. Specifically, given a continuous dynamic system, typically described as an ordinary differential equation,

where $\dot{\mathbf{x}}_{t}$ represents the derivative with respect to time, while $\Psi:\mathbb{R}^{d}\times\mathbb{R}^{m}\rightarrow\mathbb{R}^{d}$ defines a function or vector field. The term $\mathbf{x}_{0}$ specifies the initial state. Moreover, $\mathbf{u}_{t}\in\mathbb{R}^{m}$ denotes the control input exerted on the dynamic system (e.g. in this work, the control input is applied linearly to the state $\mathbf{x}_{t}$). The PID control continuously computes an error term, $e_{t}$, by calculating the difference between a target reference $\mathbf{r}_{t}$ and the actual state variable $\mathbf{x}_{t}$,

The construction for the target reference generally depends on the application. In this work, the target reference is constructed by a sequence of embedding manifolds (see Section 2.1), and the process variable represents the hidden state of a deep neural network during forward propagation.

Then, a PID controller applies a control $\mathbf{u}_{t}$ based on proportional, integral, and derivative terms to correct the measured error. In the continuous case,

where $K_{p}$, $K_{i}$, and $K_{d}$, all non-negative, denote the coefficients for the proportional, integral, and derivative terms respectively. In the PID control design,

• •

  The output from proportional control directly correlates with the current value of the error, $e_{t}$. Thus, a larger error will yield a proportionally larger control output, adjusted by the gain factor $K_{p}$. However, employing only proportional control leads to a persistent difference between the desired target and the actual process variable, since the generation of the proportional response necessitates the presence of an error.

• •

  The output from integral control considers the accumulation of past error values over time. This means that when a residual error remains after proportional control is applied, the integral control works to correct this residual error by leveraging the historical total error. This will result in the proportional effect diminishing as the error decreases, but this is compensated for by the growing integral effect.

• •

  The output from derivative control provides an estimate of the trend of the error, using its current rate of change as a basis. It effectively seeks to reduce the effect of the error by exerting a control influence generated by the rate of error change. Hence, the more rapid the error’s progression, the more intense the applied correction.

Although a PID controller has three control terms, some applications need only one or two terms to provide appropriate control. This is achieved by setting the unused parameters to zero and is called a PI, PD, P, or I controller in the absence of the other control actions. PD controllers are fairly common in applications where integral action would be sensitive to measurement noise, but the derivative term is often needed for the system to reach its target reference.

In analyzing an LLM through the lens of discrete dynamical systems, we observe that its state trajectory, governed by the composition of transformations, forms a lower-dimensional structure embedded in the ambient state space, also known as the “manifold hypothesis" (Fefferman et al., 2016) (empirical evidence is shown in Table 7). This can be conceptualized as a sequence of embedding manifolds. We consider a r-dimensional smooth manifold embedded in $\mathbb{R}^{d}$ as $\{\mathbf{x}:f(\mathbf{x})=\mathbf{0}\}$, where $f:\mathbb{R}^{d}\rightarrow\mathbb{R}^{(d-r)}$ is a surjective mapping that can be used to measure the distance between a state $\mathbf{x}_{t}$ to the embedding manifold. For instance, $\lVert f(\mathbf{x})\rVert=0$ if $\mathbf{x}$ belongs to the embedding manifold, and $\lVert f(\mathbf{x})\rVert>0$ if $\mathbf{x}$ is outside the embedding manifold.

At each time step $t$ (e.g., the ${\rm t^{th}}$ hidden states), we construct three embedding manifolds with three distinct surjective functions $f_{t}^{P}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{(d-r)}$, $f_{t}^{I}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{(d-r)}$, and $f_{t}^{D}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{(d-r)}$, that represent the embedding manifolds of the state, the integration of past states, and the derivative of the state, respectively. In a discrete setting, integration corresponds to the accumulation of past states, while the derivative is approximated by the difference between two successive states. Under this setting, $f_{t}^{I}$ denotes the embedded manifold of past states, and $f_{t}^{D}$ represents the embedded manifold derived from the difference between two consecutive states. Given these embedding functions, we propose the following running loss to evaluate the controlled state at time step $t$,

where the layer-dependent regularization term $c_{t}$ prevents using large controls. The running loss consists of three components, each assessing the error in the controlled state through distinct embedding functions: proportional, integration, and derivative. In this construction of running loss, the optimal controller results in a controlled state $(\mathbf{x}_{t}+\pi_{t}(\mathbf{x}_{t}))$ which is expected to closely align with the state embedding manifold, evaluated by $f_{t}^{P}$. Additionally, the controller must ensure that past controlled states remain close to the embedding manifold of integrated states, as evaluated by $f_{t}^{I}$, and the state’s derivative should similarly align closely with the manifold of the state’s derivative embedding, as quantified by $f_{t}^{D}$.

The objective function defined in equation 2 and the associated running loss detailed in equation 3 can be solved via the dynamical programming principle (Bellman, 1952). However, this method faces exponential complexity in terms of the dimension of the state. To overcome this “curse of dimensionality", we can reinterpret the optimal control problem through Pontryagin’s Maximum Principle and approximate it using the method of successive approximation (Chen et al., 2020). To begin with, we define the Hamiltonian $H(t,\{\mathbf{x}_{s}\}_{s=0}^{t},\mathbf{p}_{t+1},\bm{\theta}_{t},\mathbf{u}_{t})$ as

where $\mathbf{u}_{t}=\pi_{t}(\mathbf{x}_{t})$ is a control solution. Pontryagin’s maximum principle consists of a two-point boundary value problem,

plus a maximization condition of the Hamiltonian.

To obtain a numerical solution, one can consider iterating through the forward dynamic equation 4 to obtain all states $\{\mathbf{x}_{t}\}_{t=0}^{T-1}$, the backward dynamic equation 5 to compute the adjoint states $\{\mathbf{p}_{t}\}_{t=0}^{T-1}$, and maximizing the Hamiltonian defined in equation 6 with current states and adjoint states via gradient ascent. This iterative process is continued until convergence. Given an initial condition $\mathbf{x}_{0}$, Pontryagin’s Maximum Principle characterizes the optimal feedback control $\pi_{t}(\mathbf{x}_{t})$ with an open-loop control $\mathbf{u}_{t}$, this open-loop control necessitates both forward and backward propagation through a pre-trained deep neural network over several iterations during online inference.

However, implementing the above Pontryagin’s Maximum Principle is generally infeasible for LLMs. In the subsequent section, we construct an analytic solution with certain relaxation assumptions.

In this section, we develop an analytic solution for solving the objective function defined in equation 2, under certain assumptions. These assumptions are summarized in the following,

• •

  Assumption $1$: Both embedding manifold and layer-wise transformation are simplified as linear functions. In this case, the layer-wise transformation, denoted as $F_{t}(\cdot)$, is linearized through a matrix $\bm{\theta}_{t}\in\mathbb{R}^{d\times d}$. A smooth embedding manifold is represented by a linear embedding subspace. This linear embedding subspace is defined by a set of basis vectors, which are captured by the column space of a matrix $\mathbf{V}\in\mathbb{R}^{d\times r}$, corresponding to an $r$-dimensional embedding subspace.

• •

  Assumption $2$: Both embedding manifold and layer-wise transformation are orthogonal. In this case, layer-wise transformations are represented by orthogonal matrices, satisfying $\bm{\theta}_{t}^{\top}\bm{\theta}_{t}=\bm{\theta}_{t}\bm{\theta}_{t}^{\top}=\mathbf{I}$. Additionally, the basis vectors $\mathbf{V}_{t}^{P}$, $\mathbf{V}_{t}^{I}$, and $\mathbf{V}_{t}^{D}$ are considered to be mutually orthogonal,

  $$(\mathbf{V}_{t}^{P})^{\top}\mathbf{V}_{t}^{I}=\mathbf{0},\;\;\;\;(\mathbf{V}_{t}^{P})^{\top}\mathbf{V}_{t}^{D}=\mathbf{0},\;\;\;\;(\mathbf{V}_{t}^{I})^{\top}\mathbf{V}_{t}^{D}=\mathbf{0}.$$

Based on these assumptions, the computational costs of the proposed control algorithm are similar to performing forward propagation with the original model. This implies that the PID control approach introduces negligible computational cost. The negative impact of these assumptions are discussed in Section 3.4.

Under both assumptions $1$ and $2$ defined in Section 2.2, the controlled dynamics, under some perturbations, are formulated as follows:

where $\mathbf{z}$ represents an arbitrary perturbation decomposable into two mutually orthogonal components $\mathbf{z}=\mathbf{z}^{\parallel}\oplus\mathbf{z}^{\perp}$: $\mathbf{z}^{\parallel}$, aligning within the data embedding subspace, and $\mathbf{z}^{\perp}$, orthogonal to the data manifold. We represent the state trajectory in the absence of input perturbation and without any applied control as $\mathbf{x}_{t+1}=\bm{\theta}_{t}(\mathbf{x}_{t})$, the following theorem quantifies the error as $\lVert\overline{\mathbf{x}}_{t}-\mathbf{x}_{t}\rVert_{2}^{2}$, which evaluates the difference between the perturbed state after control correction and the original state from unperturbed input data.

The detailed derivation is provided in Appendix 9. This computation rigorously demonstrates that perturbations, specifically those spanning the orthogonal complement denoted by $\mathbf{z}^{\perp}$, exhibit a decay phenomenon during the process of forward propagation. Furthermore, in scenarios where control parameters are subject to regularization constraints, when $c>0$, our analysis reveals nuanced insights. We establish that the optimal control solution, which is derived by considering the intricate interplay among different transformation layers, adheres to these constraints while optimizing performance, which captures the complex dynamics between layers.

Theorem 4 outlines how errors in state computations at any given time step are influenced by input perturbations represented by $\mathbf{z}$, despite these perturbations existing within the continuous domain of $\mathbb{R}^{d}$, this setting fits real-world adversarial attacks on LLMs, which involve modifying discrete elements, such as tokens or characters, in the input text. The act of modifying a word or substring through an adversarial attack leads to a discrepancy between the embedding sequences of the original and modified input tokens, manifesting as the perturbation vector $\mathbf{z}$ within the input embedding space. Specifically, the embedding manifolds, derived from unperturbed training data, capture the structure of this data in a lower-dimensional subspace. Adversarial examples, meanwhile, are designed to be semantically similar to the original input yet induce a marked divergence in the embedding space during the model’s forward propagation. Under these circumstances, the difference between the embedding sequences of the original input and the adversarial example can be quantified and adjusted within the PID control framework. We provide empirical error computation of Theorem 4 in Section 3.4.

Here, we provide details on implementing the proposed PID control method. We begin with constructing the linear embedding basis $\mathbf{V}_{t}^{P}$, $\mathbf{V}_{t}^{I}$, and $\mathbf{V}_{t}^{D}$ from training dataset. In NLP tasks, the hidden states are generally represented as $2$-dimensional matrix (sequence of embedding vectors), $\mathbf{X}_{t}\in\mathbb{R}^{l\times d}$, where $l$ denotes the temporal length. Given $N$ pieces of data sampled from the data distribution $\mathcal{D}$ ($N$ training data), we can concatenate the hidden states as a $3$-way tensor $\mathbfcal{X}_{t}\in\mathbb{R}^{N\times l\times d}$, and apply Tucker decomposition (De Lathauwer et al., 2000b; a; Kolda & Bader, 2009) (known as high-order singular value decomposition) to generate linear embedding basis along both temporal and state embedding dimensions.

Tucker Decomposition is an extension of the traditional singular value decomposition to higher-order tensors. Mathematically, Tucker decomposition represents an $I$-way tensor as $\mathbfcal{X}\approx\mathbfcal{G}\times_{1}\mathbf{V}^{(1)}\times_{2}\mathbf{V}^{(2)}\times_{3}\cdots\times_{I}\mathbf{V}^{(I)}$, where $\mathbfcal{G}$ is the core tensor, which governs the interaction between different modes, $\mathbf{V}^{i}$ are orthogonal bases corresponding to the principal components in each tensor mode, $\times_{n}$ is the mode-$n$ tensor product. The mode-$n$ product of a tensor $\mathbfcal{A}\in\mathbb{R}^{I_{1}\times\dots\times I_{n}\times\dots\times I_{d}}$ with a matrix $\mathbf{U}\in\mathbb{R}^{J\times I_{n}}$ is defined as

where the $(i_{1},i_{2},\cdots,i_{d})$-th elements of $\mathbfcal{A}$ and $\mathbfcal{B}$ are denoted as $a_{i_{1}i_{2}\cdots i_{d}}$ and $b_{i_{1}i_{2}\cdots i_{d}}$, respectively.

An implementation of Tucker decomposition is detailed in Algorithm 1. Along each of the $I$ modes, the concatenated high-dimensional states $\mathbfcal{X}$ are reshaped along the $\rm i^{th}$ dimension, which is used to compute the orthogonal basis from singular value decomposition. The core tensor $\mathbfcal{G}$ is computed by multiplying the states $\mathbfcal{X}$ with each of the $I$ basis along each mode. The low-rank reconstruction of concatenated states $\mathbfcal{X}$ can be obtained by $\mathbfcal{G}\times_{1}\mathbf{V}^{1}\times_{2}\mathbf{V}^{2}\times_{3}\cdots\times_{I}\mathbf{V}^{I}$.

Given a pre-trained LLM (naively trained or robustly trained), we collect the concatenated states from training data, which results in a set of $3$-way tensors $\{\mathbfcal{X}_{t}\}_{t=0}^{T-1}$. Then Tucker decomposition is applied at every $\mathbfcal{X}_{t}$ (refer Algorithm 1). Extending this to integral and derivative controls is straightforward, as one can substitute the concatenated states $\mathbfcal{X}$ by either the summation of past states $\mathbfcal{X}=\sum_{s=0}^{t}\mathbfcal{X}_{s}$ or the subtraction of two consequential states $\mathbfcal{X}=\mathbfcal{X}_{t}-\mathbfcal{X}_{t-1}$. Using the linear embedding bases $\mathbf{V}_{t}^{P}$, $\mathbf{V}_{t}^{I}$, and $\mathbf{V}_{t}^{D}$ obtained from Tucker decomposition, the construction of the feedback controller is achieved by adhering to the methodology outlined in Proposition 3.