Steepish Descent

A great divergence … or … the greatest divergence?

In my last post I derived reward weighted regression (RWR) via optimization of reverse KL divergence to the Gibbs policy. It looks promising, but RWR is not popular amongst LLM practitioners, perhaps for good reason: the Advantage Policy Alignment (APA) paper compares RWR and APA to PPO and consistently finds RWR to be inferior to APA. Turns out APA is using log-squared divergence instead of reverse KL divergence; this is the same divergence used in the IPO paper to improve over DPO. $$
\begin{align*}
\ell_{\log^2}(p, q) &= \frac{1}{2} \mathbb{E}_{z \sim \mu}\left[ \left( \log \frac{p(z)}{q(z)} \right)^2 \right].
\end{align*}
$$ This divergence is well positioned to tame importance weights. First, this can be evaluated using any reference measure $\mu$ which dominates $p$ and $q$, so no more pesky importance weights from change of measure! It’s also working in log-space, which is numerically beneficial.

Deriving APA

Using this divergence and the Gibbs policy we can derive Advantage Policy Alignment. Recall the Gibbs policy results from pretending to run the following optimization $$
\begin{align*}
\pi^* &= \sup_{\pi}\ \mathbb{E}_x\left[ \left. \mathbb{E}_{\substack{y \sim \pi(\cdot|x) \\ r \sim P(r|x,y)}}\left[r\right] – \lambda^{-1} \text{KL}(\pi(\cdot|x) \| h(\cdot|x)) \right| x\right] ,
\end{align*}
$$ then you end up with $$
\begin{align*}
\pi^*(y|x) &= Z(x)^{-1} h(y|x) \exp\left( \lambda \mathbb{E}\left[r | x, y\right]\right), \\
Z(x) &= \mathbb{E}_{y \sim h(\cdot|x)}\left[ \exp\left( \lambda \mathbb{E}\left[r | x, y\right]\right) \right].
\end{align*}
$$ Optimizing our current policy via the $\log^2$ divergence yields $$
\begin{align*}
\ell_{\log^2}(x, \theta) &= \frac{1}{2} \mathbb{E}_{y \sim \mu}\left[ \left. \left( \log \frac{\pi(y | x; \theta)}{\pi^*(y|x)} \right)^2 \right| x \right], \\
-\nabla_\theta \ell_{\log^2}(x, \theta) &= \mathbb{E}_{y \sim \mu}\left[ \left. \left( \log \frac{\pi^*(y|x)}{\pi(y | x; \theta)} \right) \nabla_\theta \log \pi(y|x; \theta) \right| x \right].
\end{align*}
$$ That ratio can be simplified to $$
\begin{align*}
\log \frac{\pi^*(y|x)}{\pi(y | x; \theta)} &= \log \frac{\pi^*(y | x; \theta)}{h(y|x)} – \log \frac{\pi(y | x; \theta)}{h(y|x)} \\
&= \lambda \mathbb{E}\left[r|x,y\right] – \log Z(x) – \log \frac{\pi(y | x; \theta)}{h(y|x)}
\end{align*}
$$ and the reward expectation is linear so $$
\begin{align*}
&-\nabla_\theta \ell_{\log^2}(x, \theta) \\
&\quad = \mathbb{E}_{\substack{y \sim \mu\\ r \sim P(r|y,x)}}\left[ \left. \left( \lambda r – \log \frac{\pi(y | x; \theta)}{h(y|x)} – \log Z(x)\right)\nabla_\theta \log \pi(y|x; \theta) \right| x \right].
\end{align*}
$$ Note: unless we sample from $\mu = \pi$, the expected gradient of the log likelhood is not zero. This indicates we should use the empirical estimator of $\log Z(x)$ for centering.

The gradient has some nice properties. First, this is risk-neutral because the exponential gets linearized, unlike RWR which is risk-seeking. Second, we are also free to sample from any distribution $\mu$ for learning. But what distribution to use?

Aside: When do importance weights get big?

Suppose we sample a (potentially infinite) sequence of tokens autoregressively from $h(y|x)$ and compute the running importance weight $$
\begin{align*}
M_t &= \prod_{s<t} \frac{\pi(y_s|y_{<s},x;\theta)}{h(y_s|y_{<s})}.
\end{align*}
$$ This is a non-negative martingale under $h$ and therefore from Ville’s inequality it rarely gets big: the chance that it exceeds a million is literally less than one-in-a-million. So $\log M_t$ never gets big when sampling from $h$ but it can get very small. Analogously $1/M_t$ is a martingale under $\pi$, so $-\log M_t$ can never get big when sampling from $\pi$ but it can get very small.

Thus, if we sample from $h$, the log density ratio will have a lot of negative skew; and if we sample from $\pi$, the log density ratio will have a lot of positive skew (which looks like the puke). Maybe if we sample from $\mu = (h + \pi)/2$ the skew balances out? The APA paper is silent on this issue.