Online Learning 2: Bandit learning with catastrophes

[RyanCarey]

Note: This describes an idea of Jessica Taylor’s.

The usual training procedures for machine learning models are not always well-equipped to avoid rare catastrophes. In order to maintain the safety of powerful AI systems, it will be important to have training procedures that can efficiently learn from such events. [1]

We can model this situation with the problem of exploration-only online bandit learning. In this scenario, we grant the AI system an exploration phase, in which it is allowed to select catastrophic arms and view their consequences. (We can imagine that such catastrophic selections are acceptable because they are simulated or just evaluated by human overseers). Then, the AI system is switched into the deployment phase, in which it must select an arm that almost always avoids catastrophes.

---

##Setup In outline, the learner will receive a series of randomly selected examples, and will select an expert (modeled as a bandit arm) at each time step. The challenge is to find a high-performing expert in as few time steps as possible. We give some definitions:

• Let $X$ be some finite set of possible inputs.

• Let $A:=\{1,2,...,K\}$ be the set of available experts (i.e. bandit arms).

• Let $R:X\times A\mapsto [0,b]$ be the reward function.$R(x_t,i)$ is the reward for following expert $i$ on example $x_t$.

• Let $C:X\times A\mapsto [0,1]$ be the catastrophic risk function.$C(x_t,i)$ is the catastrophic risk incurred by following expert $i$ on example $x_t$.

• Let $M(x,i):=R(x,i)-\frac{C(x,i)}{\tau }$ be the mixed payoff that the learner is to optimize, where $\tau :(0,1]$ is the risk-tolerance.$\tau$ can be very small, on the order of ${10}^{-20}$.

• Let $p:X\mapsto [0,1]$ be the input distribution from which examples are drawn in the deployment phase.

• Let $q_i\left(x\right):=\frac{p\left(x\right)R(x,i)}{{\sum }_{x^{\prime }}p\left(x^{\prime }\right)R(x^{\prime },i)}$ be the risk distribution. This is an alternative distribution that assigns higher weight to inputs that are more likely to be catastrophic for expert $i$.

• Let ${\hat{q}}_i\left(x\right)$ be the learner’s guess at $q_i$. This guess is assumed to contain a grain of truth: $$\forall i\in A\forall x\in X{\hat{q}}_i\left(x\right)\ge \frac{q_i\left(x\right)}{a}$$ where $a\ge 1$ is some known constant. This assumption is similar to the one in Paul Christiano’s post on active learning.

• Let ${\mu }_i:=E_{x\sim p}\left[M\right(x,i\left)\right]$ be the expected overall payoff of expert $i$ in the deployment phase.

The learner’s exploration phase consists of time steps $1,2,...,T$. At each time step $t$, the learner chooses an expert $i\in A$. Then, the learner samples from $p$, $q_i$ or both and observes a tuple containing a risk and a reward. i.e. it observes either i) $\left(R\right(x_t,i),C(x_t,i\left)\right)$, ii) $\left(R\right(x_t,i),C(x_t^{\prime },i\left)\right)$, iii) $\left(R\right(x_t^{\prime },i),C(x_t,i\left)\right)$, or iv) $\left(R\right(x_t,i),C(x_t^{\prime },i\left)\right)$ where $x_t\sim p,x_t^{\prime }\sim q_i$. After $T$ steps, the learner selects some final arm $j$.

The deployment phase just computes the performance of $j$ on the deployment distribution $p$. The mean payoff of the final expert is denoted ${\mu }^{\prime }:=E_{x\sim p}\left[M\right(x,j\left)\right]$. The highest mean payoff that any expert(s) can achieve is denoted ${\mu }^{\ast }$. Then, the simple regret of the learner is ${\mu }^{\ast }-{\mu }^{\prime }$.

The aim for the learner is to select an expert that achieves low regret with high probability: $P({\mu }^{\ast }-{\mu }_f\le ϵ)<1-\delta$, using as few exploration steps as possible.$ϵ$ is some regret tolerance, which is less than the reciprocal of the (very small) risk-tolerance: $ϵ<\frac{1}{\tau }$.

We assume that the agent knows the deployment distribution $p$ and knows the estimate ${\hat{q}}_i$ of the risk distribution but does not know the actual risk distribution $q$. Additionally, we assume that there exists some expert $j$ whose recommendations incur no risk of catastrophe, $$\exists j\in As.t.\forall x\in XC(x,j)=0$$

##Bounds on the number of time steps

The standard approach

A natural way to get this PAC bound is to use a bandit algorithm like Median Elimination. The overall payoff of each expert is defined by the random variable: $$\begin{array}{cc}{X}_i& :=R(x,i)-\frac{C(x,i)}{\tau }x\sim p\\ & =M(x,i)x\sim p\end{array}$$

This algorithm will find an $ϵ$-optimal expert with $1-\delta$ probability. However, $X_i$ have a very large range, $[-\frac{1}{\tau },b]$, and so if the risk-tolerance is low, then the number of time steps required will be very large. In fact, if $N$ is the number of time steps required, then [2]:

$$N=\Theta \left(\frac{K(b+\frac{1}{\tau }{)}^2\ln (1/\delta )}{{ϵ}^2}\right)$$

Intuitively, the problem is that if some risky scenarios only occur rarely in $p$ then many samples will be required to identify these.

##Using a surrogate objective to reduce the required number of time steps Instead, the learner can sample using its guess ${\hat{q}}_i$ of the risk distribution. If we cap the extent to which we penalize risky arms, then we arrive at a surrogate bandit problem, in which the payoff of each expert is [2]:

$$S_i:=R(x,i)-min\left(\right(ϵ+b)a,\frac{p\left(x^{\prime }\right)C(x^{\prime },i)}{\hat{q}\left(x^{\prime }\right)\tau })x\sim p,x^{\prime }\sim {\hat{q}}_i$$

Since $S_i$ is bounded by $[-(ϵ+b)a,b]$, the required number of time steps for this problem is:

$$N=\Theta \left(\frac{K(b+a(ϵ+b\left){)}^2\ln \right(1/\delta )}{{ϵ}^2}\right)$$

This improves the speed of the standard approach by a factor of $\frac{a(ϵ+b)}{b+\frac{1}{\tau }}$, which may be a large speedup indeed if catastrophes are rare.

We can prove that any $ϵ$-optimal expert in this surrogate problem will also be $ϵ$-optimal in the original problem.

Theorem 1. Let $S^{\ast }$ be the overall payoff of the optimal expert.

For an expert $j$, if $E\left[S^{\ast }\right]-E\left[S_j\right]<ϵ<\frac{1}{\tau }$ then $E\left[S_j\right]=E\left[X_j\right]$.

To prove this, first we show that:

Lemma 2. $$E\left[S_i\right]=E_{x\sim p}\left[R\right(x,i\left)\right]-E_{x\sim q_i}[min(\frac{(ϵ+b)a{\hat{q}}_i\left(x\right)}{q_i\left(x\right)},\frac{1}{\tau }{E}_{x^{\prime }\sim p}[C(x^{\prime },i)\left]\right)]$$

Proof of Lemma 2. $$\begin{array}{ccc}E\left[S_i\right]& =E_{x\sim p}\left[R\right(x,i\left)\right]-E_{x\sim {\hat{q}}_i}[min\left(\right(ϵ+b)a,\frac{p\left(x\right)C(x,i)}{{\hat{q}}_i\left(x\right)\tau })]& \\ & =E_{x\sim p}\left[R\right(x,i\left)\right]-E_{x\sim q_i}[min(\frac{(ϵ+b)a{\hat{q}}_i\left(x\right)}{q_i\left(x\right)},\frac{p\left(x\right)C(x,i)}{q_i\left(x\right)\tau })]& \\ & =E_{x\sim p}\left[R\right(x,i\left)\right]-E_{x\sim q_i}[min(\frac{(ϵ+b)a{\hat{q}}_i\left(x\right)}{q_i\left(x\right)},\frac{1}{\tau }{E}_{x^{\prime }\sim p}[C(x^{\prime },i)\left]\right)]& \text{by the definition of}{q}_i\end{array}$$

Proof of Theorem 1. We have assumed that for the expert $j$: $$\begin{array}{ccc}E\left[S^{\ast }\right]-E\left[S_j\right]& <ϵ& \\ E\left[S_j\right]& >-ϵ& \text{(}\exists iC(x,i)=0\text{, so}E\left[S^{\ast }\right]\ge 0\text{)}\\ E_{x\sim p}\left[R\right(x,j\left)\right]-E_{x\sim q_j}[min(\frac{(ϵ+b)a{\hat{q}}_j\left(x\right)}{q_j\left(x\right)},\frac{1}{\tau }{E}_{x^{\prime }\sim p}[C(x^{\prime },j)\left]\right)]& >-ϵ& \left(\text{Lemma 2}\right)\\ E_{x\sim q_j}[min(\frac{(ϵ+b)a{\hat{q}}_j\left(x\right)}{q_j\left(x\right)},\frac{1}{\tau }{E}_{x^{\prime }\sim p}[C(x^{\prime },j)\left]\right)]& <(ϵ+b)& \text{(}R(x,i)\le b\text{)}\\ \therefore \frac{1}{\tau }{E}_{x^{\prime }\sim p}\left[C\right(x^{\prime },j\left)\right]& <(ϵ+b)& \left({\hat{q}}_i\right(x)\ge \frac{q_i\left(x\right)}{a})\\ \therefore \forall i\frac{1}{\tau }{E}_{x^{\prime }\sim p}\left[C\right(x^{\prime },j\left)\right]& <\frac{(ϵ+b)a{\hat{q}}_i\left(x\right)}{q_i\left(x\right)}& (\frac{a{\hat{q}}_i\left(x\right)}{q_i\left(x\right)}\ge 1)\end{array}$$

$$\begin{array}{ccc}\therefore E\left[S_j\right]=& E_{x\sim p}\left[R\right(x,j\left)\right]-\frac{1}{\tau }{E}_{x^{\prime }\sim p}\left[C\right(x^{\prime },j\left)\right]\left)\right]& \left(\text{Lemma 2}\right)\\ =& E\left[X_j\right]& \end{array}$$

Discussion

This result suggests that it might be possible to avoid catastrophes by training machine learning systems on inputs that induce catastrophic behavior. We have given a clear set of circumstances in which an agent can learn to avoid catastrophes. However, this set of requirements does seem very tricky to meet. Finding a hypothesis about the location of catastrophes that has a grain of truth for all experts with all examples may be very difficult. The assumption of separate training and execution phases may also be difficult to instantiate. For example, if the AI system is able to escape its training simulation, then it will no-longer be valid to allow failures during the exploration phase, and so this training procedure would not work.

Another problem is that it could be difficult to know the input distribution $p$. In future work we plan on exploring learning with catastrophes in an online active learning setting instead of the supervised learning setting of this post, so that this requirement can be relaxed.

Footnotes

1. Paul Christiano has discussed approaches to learning with catastrophes and their application to AI Control at Learning with catastrophes and red teams. The idea explored in this post can be seen as a form of red teaming (with the distribution ${\hat{q}}_i$ representing the red team).

2. The PAC bound for Median Elimination is proved in Theorem 4. Even-Dar, Eyal, Shie Mannor, and Yishay Mansour. “PAC bounds for multi-armed bandit and Markov decision processes.” International Conference on Computational Learning Theory. Springer Berlin Heidelberg, 2002. To get the exact bound, we Hoeffding’s inequality for a random variable with range $[1/\tau ,b]$ rather than on $[0,1]$ as in the original paper.

[paulfchristiano]

It seems like the qualitative conclusion here is similar to what you’d get by verifying that an arm is non-catastrophic by pulling it over and over again. You have a quadratic dependence on $a$ rather than linear, but I think that is just because you aren’t exploiting the fact that the target catastrophe probability is $0$? If that’s right, it might be better to stick with the simplest algorithm that exhibits the behavior of interest. (If the optimal dependence is actually quadratic in this situation that surprising.)

In the definition of $q_i$, presumably you want $C$ rather than $R$. Note that the risk distribution is just the posterior conditioned on catastrophe.

(I think today we don’t usually call arms experts—there may be both experts and arms, in which case an expert might advise you about which arm to pull, but I think it’s a bit confusing to talk about experts with partial feedback.)

Re footnote 1: I would call this an example of adversarial training. I’d suggest usage like: a red team is a group of humans (or in general a powerful hopefully-aligned AI) which is acting as an adversary for the purposes of adversarial training or development. I think the original version of my post may have overreached a bit on the definition and not given adequate credit to the adversarial training authors (who do seem to consider this kind of thing an example of adversarial training).

[RyanCarey]

Thanks!!

Additionally to Jessica’s comments: uniformly calling the selections ‘arms’ seems good, as does clarifying what is meant by ‘red teams’. I’ve corrected these, and likewise the definition of $q_i$.

[jessicata]

A couple notes:

• The quadratic dependence on $a$ is almost certainly unnecessary; we didn’t try too hard to reduce it. The way to reduce the bound is probably by observing that, for the best arms, the average importance-sampled estimate of the risk has a low mean; showing that as a result the estimate is subgaussian; and then applying a stochastic bandit algorithm that assumes subgaussian tails.

• If we just pulled an arm repeatedly to ensure it’s non catastrophic we’d get dependence on $1/\tau$ which is huge; the main idea of the post is that we can get dependence on $a$ instead of $1/\tau$.

[paulfchristiano]

(I meant sampling $x$ repeatedly from the distribution ${\hat{q}}_i$, I agree that sampling $x$ at random won’t help identify rare catastrophes.)

[jessicata]

The main qualitative difference from sampling from ${\hat{q}}_i$ is that we’re targeting a specific tradeoff between catastrophes and reward, rather than zero probability of catastrophe. I agree that when $\tau =0$ we’re just sampling from ${\hat{q}}_i$.