To Limit Impact, Limit KL-Divergence

[J Bostock]

TL;DR

Run a potentially-harmful model alongside a known-harmless model, such that their action-spaces (e.g. output token sets) are equivalent. Combine the output probabilities so as to limit the KL-divergence between the resulting token probabilities and the harmless model’s probabilities. This provides a mathematical ceiling on the impact of the resulting combined policy.

Unfortunately, the probability of catastrophe ($q$) scales linearly with the allowed $D_{KL}$, whereas depending on the distribution, reward may only scale with $\sqrt{D_{KL}}$ (though for unbounded rewards it often scales with $D_{KL}$. Bounds on probability of catastrophe compared to the probability of the weaker model causing a catastrophic outcome ($q_0$) are also very weak. Overall the probability of catastrophe scales as $q=\frac{D_{KL}}{\ln 1/q_0}$.

Long Version

We will use the following model of AI interactions:

In this case, $P$ represents the past, $A_0$ the output of an AI (which we will represent as having policy ${\pi }_0$) and $F_0$ the future after this AI’s actions have occcurred. For a “base” case, we will consider using an AI which is both relatively powerful, but that we think has a probability $p_0=ϵ$ of being catastrophic (like GPT-4) to calculate $P\stackrel{{\pi }_0}{\to }{A}_0$. We must consider the model as running at a temperature $T\ge 0$ for this to make sense.

Now imagine using a powerful AI with policy ${\pi }_1$, which is of unknown safety. If we do the same thing here, the results could be catastrophic!

What we want is a way to use some of ${\pi }_1$’s power, but not all. Let’s consider a modified case, where we only change the relation $P\stackrel{{\pi }_{\Pi }}{\to }{A}_{\Pi }$. We will generate $A_{\Pi }$ by interpolating between $A_0$ and $A_1$ according to some rule $\Pi$.

For example we might use the rule $A\left(b\right)=softmax(logits\left(A_0\right)+b\times logits\left(A_1\right))$, which smoothly interpolates between $A\left(b\right)=A_0$ and $A\left(\infty \right)=argmax\left(A_1\right)$. What is important is that we choose the maximum value of $b=b_{max}$ such that:

$\underset{p\in P}{max}\left[D_{KL}\right(A\left(b\right)|P=p\parallel A|P=p\left)\right]\le {\Pi }^2$

Then we can define $A_{\Pi }=A\left(b_{max}\right)$ for some allowed “power level” $\Pi$. This then limits the impact of the combined system on the future:

$D_{KL}(F_{\Pi }\parallel F_0)\le {\Pi }^2$

This lemma is proved at the end of this post.

Limiting Takeover Probability

If, as stated, our model ${\pi }_0$ has a probability $ϵ$ of taking over, then what’s the probability $q_{\Pi }$ of the model ${\pi }_{\Pi }$ being able to execute a successful takeover?

For a binary success variable, which has a probability $q$ of occurring in $F_0$, and $q_{\Pi }$ of occurring in $F_{\Pi }$, we can make the following relation:

$q_{\Pi }\ln \frac{q_{\pi }}{q_0}+(1-q_{\Pi })\ln \frac{1-q_{\Pi }}{1-q_0}\le D_{KL}(F_{\Pi }\parallel F_0)\le {\Pi }^2$

For $q_0=ϵ\approx 0$, the relation approaches:

$q_{\Pi }\le \frac{{\Pi }^2}{\ln 1/ϵ}$

Unfortunately, this scales with ${\Pi }^2$, and also only inversely with $\ln \left(1/ϵ\right)$. This means we cannot guarantee that security measures against a weaker model will work against a stronger one.

Experimental Validation in Toy Model

For this I used a sequence relation similar to my previous work on the topic:

$s_t=f(...,s_{t-1})=\{\begin{array}{cc}(\stackrel{7}{\prod _{i=1}}(s_{t-i}+1))\mathrm{mod}33\mathrm{mod}32& \text{With probabilty 0.98}\\ s\in \{0,...,31\}& \text{With probability 0.02}\end{array}$

${\pi }_0$ was trained as a next-token predictor for this sequence.

${\pi }_1$ was trained using RL, so $s_t$ is chosen such that $s_{t+1}=f(...,s_{t-1},s_t)$ maximizes the reward function $\rho \left(s\right)$:

$\rho :S\ni s\to {argmax}_n(n\in \{0,...,5\}|s\mathrm{mod}{2}^n=0)=r\in R$ which varies from 0 to 5.

${\pi }_1$ in this case was trained “normally” (i.e. fully online and unboxed). In theory we could train it using examples generated by $\pi \left(b\right)$.

I generate the following data in a slightly different way to the $\Pi$-limiting method described above above:

For each of a series of values of $b$, I calculate $A_b=softplus\left(logits\right(A_0)+b\times logits(A_1\left)\right)$ then calculate the expected value$E_P(D_{KL}(A_b|P=p\parallel A_0|P=p))$, then plot the square root of that value.

Using a fundamental result from statistical utility mechanics, we know that the distribution $P(R_b=r)\propto P(R_0=r)\exp \left(br\right)$ for a variable $b$ gives the maximum $E\left(R_b\right)$ for a given $D_{KL}(R_b\parallel R_0)\le {\Pi }^2$, so for a (different) set of $b$ values, I calculate the max possible RL score given a value of $D_{KL}^{1/2}$:

We can see that the maximum possible $D_{KL}(R_b\parallel R_0{)}^{1/2}$ approaches a value significantly less than the maximum $D_{KL}\left(A\right(b)|P\parallel A_0|P{)}^{1/2}$ value that $\pi \left(b\right)$ does. This is because our base model in this case is (in a sense) worse than useless! Not only does it not perform well on the RL task, it confidently predicts a “random” (from the perspective of the RL task) answer!

If we use a randomly-initialized (and untrained model) as the base policy ${\pi }_0$, the maximum $\sqrt{D_{KL}}$ value for our model approaches the maximum $\sqrt{D_{KL}}$ value for our RL score:

Interestingly, the pretrained ${\pi }_0$ does better at lower $\sqrt{D_{KL}}$ values than does the randomly-initialized ${\pi }_0$. I’m not really sure why this is!

Thoughts and Existing Work

Why It Might Be Useful

For box-like strategies, we want to limit the bit-rate of an AI’s output channel in order to limit that AI’s impact on the world. But the information density of even a superintelligent AI’s text output is unlikely to be precisely equal to the bitrate of the output channel. This means that we will struggle to actually put a numerical limit a powerful AI’s optimizing power without completely hamstringing its ability to do anything.

For a concrete illustration: imagine we have a doctor who asks the system “Which drug should I give to this sick person?”. For a disease with only two available drugs, only one bit of information is required, but it seems obvious that this is much easier to transmit smeared over a dozen output tokens than compressed into one.

This is especially true if the sick person in question asks the system “Which drug should I take?”. If the intermediary AI only needs one bit of information to figure it out (as GPT-4 might well do) then the information can be transmitted from our boxed AI to the user with only one bit of information leaving the box.

By leveraging a less powerful but “safe” AI, we are essentially having the intermediary AI “interpret” the more powerful one, spreading its optimizing power out across an output of reasonable space.

Further Thoughts

It’s a shame that the probability of catastrophe scales as ${\Pi }^2/\ln \left(1/ϵ\right)$.

The reason for the power of two in ${\Pi }^2$ becomes apparent when we start considering multiple rounds of AI action.$D_{KL}$ is a generalization of squared distance, not distance, so if the AI wants to break up its budget into ${\Pi }_1,{\Pi }_2,\dots$ we can limit the total $D_{KL}$ more sensibly by limiting each observation-action cycle to ${\Pi }_1^2,{\Pi }_2^2,\dots$ than by working with a bunch of square-roots.

We might allow the system to vary the amount of $D_{KL}$ budget spent on each token. I will have to think carefully about how this affects the total $D_{KL}$, given the system will be able to see whether it stochasticly took an “unlikely” action in the past and plan accordingly.

In the case that this system runs out of $D_{KL}$ budget, as long as the intermediary AI is relatively capable, I expect it to “fail gracefully”, by just continuing as the intermediary AI would.

I would like to try this out with some real language models.

This is very similar to the concept of a quantilizer, but it uses a “power budget” ${\Pi }^2$ rather than a q value.

Proof Of Core Lemma

For variables $X,Y,Z$, and related variables $X,Y^{\prime },Z^{\prime }$

And we have that that $P(Z^{\prime }|X,Y^{\prime })=P(Z|X,Y)$

We can combine the nodes $X+Y⟺\left[XY\right]$

$D_{KL}\left(\right[X{Y}^{\prime }]\parallel [XY\left]\right)=\sum _{x,y}P(X=x,Y^{\prime }=y)\ln \frac{P(X=x,Y^{\prime }=y)}{P(X=x,Y=y}$

Since $X\to Y$:

$=\sum _{x,y}P(X=x)P(Y^{\prime }=y|X=x)\ln \frac{P(Y^{\prime }=y|X=x)}{P(Y=y|X=x)}$

$=\sum _{x}P(X=x)\sum _{y}P(Y^{\prime }=y|X=x)\ln \frac{P(Y^{\prime }=y|X=x)}{P(Y=y|X=x)}$

$=E_X(D_{KL}(Y^{\prime }|X=x\parallel Y|X=x))$

If we therefore enforce $D_{KL}(Y^{\prime }|X=x\parallel Y|X=x)\le {\Pi }^2$ for all $x$ and some value $\Pi$, we can bound the value of $D_{KL}\left(\right[X{Y}^{\prime }]\parallel [XY\left]\right)\le {\Pi }^2$.

By the data processing inequality, since we have $\left[XY\right]\to Z$, we get:

$D_{KL}\left(Z^{\prime }|Z\right)\le D_{KL}\left(\right[X{Y}^{\prime }]\parallel [XY\left]\right)\le {\Pi }^2$