Goodhart in RL with KL: Appendix

[Thomas Kwa]

This is the appendix to the previous post on Goodhart’s Law and KL regularization, containing all of our proofs.

Theorem about distributions

Theorem 1: Given any heavy-tailed reference distribution $Q$ over $R$ with mean ${\mu }_Q$, and any $M,ϵ>0$, there is a distribution $P$ with mean ${\mu }_P>M$ and $D_{KL}(P\parallel Q)<ϵ$.

Proof: WLOG let ${\mu }_Q=0$. We construct a sequence of distributions $\left\{P_t\right\}$ such that ${lim}_{t\to \infty }{E}_{P_t}\left[X\right]\ge c$ for any constant $c$, and ${lim}_{t\to \infty }{D}_{KL}(P_t\parallel Q)=0$. We define $P_t$ for any $t>c$ thusly. Writing $F_{P_t}\left(x\right)$ for the CDF $P{r}_{X\sim P_t}(X\le x)$ and ${\overline{F}}_{P_t}\left(x\right)$ for $1-F_{P_t}\left(x\right)$, we let

$${\overline{F}}_{P_t}\left(x\right)=\{\begin{array}{cc}1-\frac{1-c/t}{F_Q\left(t\right)}{F}_Q\left(x\right)& x\le t\\ \frac{c/t}{{\overline{F}}_Q\left(t\right)}{\overline{F}}_Q\left(x\right)& x>t\end{array}$$

Intuitively, we rescale the part of the distribution to the right of $t$ evenly to have total probability $c/t$, which is less than 1 because $t>c$.

We must check that ${lim}_{t\to \infty }{E}_{P_t}\left[X\right]=c$. We can write

$$\begin{array}{cc}{E}_{P_t}\left[X\right]=F_{P_t}\left(t\right)E_{P_t}[X|X\le t]& +{\overline{F}}_{P_t}\left(t\right)E_{P_t}[X|X>t]\\ =F_{P_t}\left(t\right)E_Q[X|X\le t]& +{\overline{F}}_{P_t}\left(t\right)E_Q[X|X>t]\\ =F_Q\left(t\right)E_Q[X|X\le t]& +{\overline{F}}_Q\left(t\right)E_Q[X|X>t]+\\ \left(F_{P_t}\right(t)-F_Q(t\left)\right)& E_Q[X|X\le t]+\left({\overline{F}}_{P_t}\right(t)-{\overline{F}}_Q(t\left)\right)E_Q[X|X>t]\\ =E_Q\left[X\right]+\left({\overline{F}}_{P_t}\right(t)-& {\overline{F}}_Q\left(t\right)\left)\right(E_Q[X|X>t]-E_Q[X|X\le t])\end{array}$$

We know that $E_Q[X|X>t]>t$ because it is an integral of values strictly greater than t. Because $E_Q\left[X\right]=0$ is a weighted average of $E_Q[X|X>t]$ and $E_Q[X|X\le t]$, and $E_Q[X|X>t]>0$, we know $E_Q[X|X\le t]<0$. So $E_Q[X|X>t]-E_Q[X|X\le t]>t.$ We also know that for sufficiently large $t$, $\left(F_{P_t}\right(t)-F_Q(t\left)\right)>0$. Intuitively, starting from $Q$, which has mean 0, $P_t$ moves a probability mass approaching $\frac{c}{t}$ from mean <0 to mean >t.

Now we can say

$$\underset{t\to \infty }{lim}{E}_{P_t}\left[X\right]>\underset{t\to \infty }{lim}\left[E_Q\right[X]+({\overline{F}}_{P_t}\left(t\right)-{\overline{F}}_Q\left(t\right)\left)\right(t-0\left)\right]=\underset{t\to \infty }{lim}(\frac{c}{t}-{\overline{F}}_Q(t\left)\right)t=\underset{t\to \infty }{lim}c-t{\overline{F}}_Q\left(t\right)$$

Because $Q$ has a finite mean, ${lim}_{t\to \infty }t{\overline{F}}_Q\left(t\right)=0$, and so ${lim}_{t\to \infty }{E}_{P_t}\left[X\right]\ge c$.

Now we check that ${lim}_{t\to \infty }{D}_{KL}(P_t\parallel Q)=0$:

$$\begin{array}{cc}{D}_{KL}(P_t\parallel Q)& ={\int }_R\log \frac{P_t\left(dx\right)}{Q\left(dx\right)}{P}_t\left(dx\right)\\ & ={\int }_{x\le t}\log \frac{P_t\left(dx\right)}{Q\left(dx\right)}{P}_t\left(dx\right)+{\int }_{x>t}\log \frac{P_t\left(dx\right)}{Q\left(dx\right)}{P}_t\left(dx\right)\\ & =F_{P_t}\left(t\right)\log \frac{F_{P_t}\left(t\right)}{F_Q\left(t\right)}+{\overline{F}}_{P_t}\left(t\right)\log \frac{{\overline{F}}_{P_t}\left(t\right)}{{\overline{F}}_Q\left(t\right)}\text{quad since both ratios are constant}\\ & =F_{P_t}\left(t\right)\log \frac{1-c/t}{F_Q\left(t\right)}+{\overline{F}}_{P_t}\left(t\right)\log \frac{{\overline{F}}_{P_t}\left(t\right)}{{\overline{F}}_Q\left(t\right)}\end{array}$$

Since both $1-c/t$ and $F_Q\left(t\right)$ go to $1$ as $t\to \infty$, the left term goes to $0$, and so

$$\begin{array}{cc}\underset{t\to \infty }{lim}{D}_{KL}(P_t\parallel Q)& \le 0+\underset{t\to \infty }{lim}{\overline{F}}_{P_t}\left(t\right)\log \frac{{\overline{F}}_{P_t}\left(t\right)}{{\overline{F}}_Q\left(t\right)}\\ & =\underset{t\to \infty }{lim}\frac{c}{t}\log \frac{c}{t{\overline{F}}_Q\left(t\right)}\le \underset{t\to \infty }{lim}\frac{c}{t}\log \frac{1}{{\overline{F}}_Q\left(t\right)}\\ & =\underset{t\to \infty }{lim}-\frac{c}{t}\log {\overline{F}}_Q\left(t\right)\text{since t>c}\end{array}$$

$Q$ is heavy tailed, so by definition ${lim}_{t\to \infty }{e}^{at}{\overline{F}}_Q\left(t\right)=\infty \text{for all}a>0$. This implies that for every $a>0$ there is a sufficiently large $t_c$ so that for all $t>t_c$, ${\overline{F}}_Q\left(x\right)>e^{-at}$, which means that $\log {\overline{F}}_Q\left(t\right)>-at$.

Therefore for every $a>0$, ${lim}_{t\to \infty }{D}_{KL}(P_t\parallel Q)\le {lim}_{t\to \infty }-c/t\log {\overline{F}}_Q\left(t\right)<{lim}_{t\to \infty }-\frac{-act}{t}=ac$, which since KL divergence is nonnegative means that ${lim}_{t\to \infty }{D}_{KL}(P_t\parallel Q)=0$ as desired. $■$

Theorem about deterministic Markovian-return MDPs

Definition: A deterministic-transition MDP with Markovian returns (DMRMDP) is an MDP $(S,A,P,R)$ such that:

• The transition function $P:S\times A\to S$ is deterministic, i.e., for each state $s\in S$ and action $a\in A$, there exists a unique state $s^{\prime }\in S$ such that $P(s^{\prime }|s,a)=1$.

• There is a set of sink states $E\subseteq S$ that terminate a trajectory, which is disjoint with the set of start states.

• Returns are Markovian; that is, for any two trajectories $\tau =(s_1,a_1,\dots ,s_n),{\tau }^{\prime }=(s_1^{\prime },a_1^{\prime },\dots ,s_n^{\prime }),$ if $s_n=s_n^{\prime }$, then $\tau$ and ${\tau }^{\prime }$ have identical return distributions. Equivalently, for the trajectory random variable $T=(S_1,A_1,\dots )$ distributed according to any policy, with return $G$, $G\perp \perp (S_{<i},A_{<i})|S_i$ for any $i\ge 1$.

Note: Sampling from a language model and applying RLHF is well-modeled as a DMRMDP, since the state is a sequence of tokens (actions) which deterministically results from the last token and returns depend only on the final state.

Theorem 2: Let $W=(S,A,P,R)$ be a deterministic-transition MDP with Markovian returns. Given $W$ we define the function that takes policies to trajectories $Tr:(S\to \Delta A)\to \Delta (S\times A{)}^{\ast }$, and the average return function $g:(S\times A{)}^{\ast }\to R$ which induces a function $G:\Delta (S\times A{)}^{\ast }\to \Delta R$. Let ${\pi }_0:S\to \Delta A$ be some reference policy. If $G\circ Tr\left({\pi }_0\right)$ is heavy-tailed with finite mean ${\mu }_Q$, then for any $M,ϵ>0$, there is a policy $\pi$ with mean return $E[U|U\sim G\circ Tr(\pi \left)\right]>M$ and $E_{s\in T,T\sim Tr\left(\pi \right)}\left[D_{KL}\right(\pi \left(s\right)\parallel {\pi }_0\left(s\right)\left)\right]<ϵ$.

Proof: We will exhibit a distribution of trajectories $\rho$ such that $D_{KL}(\rho \parallel Tr({\pi }_0\left)\right)<ϵ$ and $E\left[G\right(\rho \left)\right]>M$, and then construct a policy $\pi$ with $Tr\left(\pi \right)=\rho$. Note that this proof applies for continuous action spaces if trajectories are replaced with measurable sets, but this would make it harder to read.

Let ${\rho }_{{\pi }_0}=Tr\left({\pi }_0\right)$. We have a heavy-tailed distribution of return $Q\triangleq G\left({\rho }_{{\pi }_0}\right)$ over $R$, so we can apply Theorem 1. But to define $\rho$, we can construct $P_t$ in the proof of Theorem 1 in a particular way. For any $t>c$, we need a $P_t$ that uniformly upweights values of mean return such that ${\overline{F}}_{P_t}\left(t\right)=c/t$. We can define ${\rho }_t$ such that any trajectory $\tau$ is upweighted by a factor depending only on its mean return:

$${\rho }_t\left(\tau \right)=\{\begin{array}{cc}\frac{1-c/t}{F_Q\left(t\right)}{\rho }_{{\pi }_0}\left(\tau \right)& g\left(\tau \right)\le t\\ \frac{c/t}{{\overline{F}}_Q\left(t\right)}{\rho }_{{\pi }_0}\left(\tau \right)& g\left(\tau \right)>t\end{array}$$

Then we can let $P_t\triangleq G\circ {\rho }_t$ and the rest of the proof of Theorem 1 applies. Therefore, applying the theorem, we can let $\rho ={\rho }_t$ for sufficiently large $t$, and then ${\mu }_{G\circ \rho }>M$ and $D_{KL}(G\circ \rho ,G\circ {\rho }_{{\pi }_0})<ϵ$. But by the chain rule for KL divergence, $D_{KL}(\rho ,{\rho }_{{\pi }_0})=D_{KL}(G\circ \rho ,G\circ {\rho }_{{\pi }_0})+E_{\gamma \sim G\circ \rho }\left[D_{KL}\right(\rho \left(T\right)|G\left(T\right)=\gamma \parallel {\rho }_{{\pi }_0}\left(T\right)|G\left(T\right)=\gamma \left)\right]$. Since we constructed $\rho$ so that the probabilities of each $\tau$ conditional on its return being $\gamma$ are equal, the second term is zero, and we also have $D_{KL}(\rho ,{\rho }_{{\pi }_0})<ϵ$.

Finally, since the KL divergence between trajectory distributions is the sum of KL divergence between policies at each action in the trajectory, and each trajectory has at least one action, $E_{s\in T,T\sim Tr\left(\pi \right)}\left[D_{KL}\right(\pi \left(s\right)\parallel {\pi }_0\left(s\right)\left)\right]\le E_{T\sim Tr\left(\pi \right)}{\sum }_{s\in T}\left[D_{KL}\right(\pi \left(s\right)\parallel {\pi }_0\left(s\right)\left)\right]=D_{KL}(\rho \parallel {\rho }_{{\pi }_0})<ϵ$ as desired.

To define $\pi$ such that $Tr\left(\pi \right)=\rho$, we let $\pi (s,a)=Pr(a_i=a|\tau =(...,s,a_i,...)\sim \rho )$.

Then the probability that any trajectory $\tau =(s_1,a_1,\dots ,a_n)$ is sampled is:

$$\begin{array}{cc}Tr\left(\pi \right)\left(\tau \right)& =\prod _{i=1}^n\pi (s_i,a_i)\\ & =\prod _{i=1}^{n}Pr(a_i=a_i^{\prime }|{\tau }^{\prime }=(...,s,a_i^{\prime },...)\sim \rho )\\ & =\prod _{i=1}^{n}Pr(a_i=a_i^{\prime }|{\tau }^{\prime }=(s_1^{\prime },a_1^{\prime },...,s,a_i^{\prime },...)\sim \rho ,s_{<i}=s^{\prime }<i,a<i=a_{<i}^{\prime })\&=\rho \left(\tau \right)\end{array}$$

In (2), returns are Markovian, so all trajectory prefixes ending in state $s$ have the same distribution of returns under any policy. In the construction of $\rho$, all trajectories with the same mean return have equal measure. Therefore, conditioning on earlier states and actions of $\tau$ does not change the measure, so we can write (3). So $Tr\left(\pi \right)=\rho$ as desired. $■$

Lagrange multipliers to minimize KL divergence

Theorem 3: If $V$ is light-tailed, $E_Q\left[V\right]$ is finite, and $d=D_{KL}(P\parallel Q)$ is bounded, then $E_P\left[V\right]$ is bounded, and $E_P\left[V\right]\to 0$ as $d\to 0$.

Using Lagrange multipliers, we find that when KL divergence is minimized, we have $P\left(V\right)[{\lambda }_1\log \frac{P\left(V\right)}{Q\left(V\right)}+{\lambda }_2-X]=0$ for some constants ${\lambda }_1,{\lambda }_2$, so

$$\log \frac{P\left(V\right)}{Q\left(V\right)}=\frac{V-{\lambda }_2}{{\lambda }_1}$$$$P\left(V\right)=Q\left(V\right)\exp \left(\frac{V-{\lambda }_2}{{\lambda }_1}\right)=Q\left(V\right)e^{V/{\lambda }_1}{e}^{-{\lambda }_2/{\lambda }_1}=CQ\left(V\right)e^{V/{\lambda }_1}$$

That is, the new PDF is an exponential tilting of the old PDF. Now what is $E_P\left[V\right]$? It’s just ${\int }_{-\infty }^{\infty }CV{e}^{V/{\lambda }_1}Q\left(X\right)dV$. If the distribution of V is heavy-tailed distribution, this is $\infty$; if it is light-tailed, this is some finite value.

When $d=0$, $P$ and $Q$ are identical and $E\left[V\right]=0$. So by a continuity argument, $E_P\left[V\right]\to 0$ as $d\to 0$. $■$

Light tails + independence imply $EV\to \infty$

Theorem 4: If $U=X+V$with $X$ and $V$ both light-tailed, and the distribution of U is continuous, and ${\pi }^{\ast }\left(\beta \right)\stackrel{△}{=}\arg {max}_{\pi }E\left[U\right(\pi \left)\right]-\beta D_{KL}(\pi ,{\pi }_0)$, then ${lim}_{\beta \to 0^{+}}E\left[V\right({\pi }^{\ast }\left(\beta \right)\left)\right]=\infty$.

Proof: Fix some $\beta$. Using Lagrange multipliers, we find that for any event $S$, ${Pr}_{\pi }\left(S\right)={Pr}_{{\pi }_0}\left(S\right)e^{\lambda U\left(S\right)}$. Let $c\left(\beta \right)$ be the median value of $U$ under the policy ${\pi }^{\ast }\left(\beta \right)$; that is, $Pr(U>c(\beta )|U\sim G\circ Tr({\pi }^{\ast }\left(\beta \right)\left)\right)=\frac{1}{2}.$ This exists because $U$ has a continuous distribution. Then:

$$\begin{array}{cc}E\left[V|\pi \right]& =\frac{1}{2}E[V|\pi ,U<c]+\frac{1}{2}E[V|\pi ,U\ge c]\\ & \ge \frac{1}{2}E[V|\pi ,U<c]+\frac{1}{2}E\left[V|\pi \right]\\ \underset{\beta \to 0^{+}}{lim}E\left[V|\pi \right]& \ge \underset{\beta \to 0^{+}}{lim}\frac{1}{2}E[V|\pi ,U<c]+\underset{\beta \to 0^{+}}{lim}\frac{1}{2}E\left[V|\pi \right]\end{array}$$

The left term is $c$, while the right term is $\infty$, so the overall limit is $\infty$. $■$