Theoretical analysis of the proposed prior reweighting technique and proof
Several reviewers (8Er5, DFgi, n6Yf) are concerned with the theoretical property of our method. Here, we present the **formal theoretical analysis of the prior reweighting technique**.
Consider a pair $(u,w)$ sampled using the prior reweighting technique with hyperparameter $\gamma$: $(u,w)\sim p_{\gamma}(u,w) = p(u|w) p^{\gamma}(w)/C_{\gamma}$.
Denote $J^*$ as the global minimum of the control objective $J$. Define $Q(\varepsilon)$ as the "$\epsilon$-optimal" solution set of $(u, w)$ such that $J(u, w) - J^* \leq (\epsilon) $, whose complement set is $Q(\varepsilon)^c$, and denote $\mathbb{I}\_{Q(\varepsilon)}(u,w)$ as its indicator function, i.e. $\mathbb{I}\_{Q(\varepsilon)}(u,w)=1$ if $(u,w)\in Q(\varepsilon)$; otherwise 0.
Define $Y$ to be the random variable of "whether use $J$ as a guidance for sampling", namely,
$
p(Y|u,w)=\begin{cases}
e^{-J(u,w)} / Z, &Y=1 \\\\
1 - e^{-J(u,w)} / Z, & Y=0
\end{cases}
$
Consider $E(\gamma)=\mathbb{E}\_{(u,w)\sim p_{\gamma}(u,w)}[\mathbb{I}\_{Q(\varepsilon)}(u,w)|Y=1]$, which indicates the expectation of getting an $\epsilon$-optimal solution by using the prior reweighting technique with $\gamma$ under the guidance of $J$.
Define
$
F(\gamma) = \frac{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w) \ln(p(w))]}{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)]} - \frac{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w) \ln(p(w))]}{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)]}
$. Then, we have the following theorem:
**Theorem**: Assume $E(\gamma)$ is a smooth function, then:
- (i): If $F(1) < 0$, there exists $\gamma_{-} < 1$, s.t., $E(\gamma_{-})> E(1)$;
- (ii): If $F(1) > 0$, there exists $\gamma_{+} > 1$, s.t., $E(\gamma_{+})> E(1)$.
**Proof**
\begin{align}
E(\gamma)= &\int \mathbb{I}\_{Q(\varepsilon)}(u,w) p\_{\gamma}(u,w|Y=1) \mathrm{d}(u,w) \\\\
= &\int \mathbb{I}\_{Q(\varepsilon)}(u,w) \frac{p\(Y=1|u,w) p\_{\gamma}(u,w)}{p\(Y=1)} \mathrm{d}(u,w) \\\\
= &\frac{\mathbb{E}\_{(u,w)}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)]}{\mathbb{E}\_{(u,w)}[p(Y=1|u,w)]} \\\\
= &\frac{\mathbb{E}\_{(u,w)\}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p\(Y=1|u,w)]}{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p\(Y=1|u,w)] + \mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)]} \\\\
= &\frac{1}{1 + \frac{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p\(Y=1|u,w)]}{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p\(Y=1|u,w)]}}
\end{align}
Define
\begin{align}
G(\gamma) = &\frac{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)]}{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)]} \\\\
= &\frac{\mathbb{E}\_w[\mathbb{E}\_u[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)|w]]}{\mathbb{E}\_w[\mathbb{E}\_u[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)|w]]} \\\\
= &\frac{\int \mathbb{E}\_u[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)|w] p^{\gamma}(w) \mathrm{d}w}{\int \mathbb{E}\_u[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)|w] p^{\gamma}(w) \mathrm{d}w}
\end{align}
Then
$E(\gamma) = \frac{1}{1 + \frac{1}{G(\gamma)}}$. Since $G(\gamma)>0$, $E(\gamma)$ and $G(\gamma)$ have the same monotonicity.
\begin{align}
G'(\gamma) = &\frac{\int \mathbb{E}\_{u}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)|w] p^{\gamma}(w) \ln(p(w)) \mathrm{d}w \mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)]}{(\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)])^2 } \\\\
& -\frac{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)] \int \mathbb{E}\_{u}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)|w] p^{\gamma}(w) \ln(p(w)) \mathrm{d}w}{(\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)])^2} \\\\
= & \frac{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w) \ln(p(w))] \mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)]}{(\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)])^2} \\\\
& - \frac{\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)}(u,w) p(Y=1|u,w)] \mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w) \ln(p(w))]}{(\mathbb{E}\_{u,w}[\mathbb{I}\_{Q(\varepsilon)^c}(u,w) p(Y=1|u,w)])^2}
\end{align}
By definition, $F(\gamma)$ is a positive multiple of $G'(\gamma)$, which implies our conclusion:
- If $F(1) < 0$, then $G'(1) < 0$ and thus $E(\gamma)$ decreases around 1. Hence, there exists $\gamma_{-} < 1$, s.t., $E(\gamma_{-})> E(1)$, thus (i) holds;
- otherwise, for similar reason, (ii) holds.
**Remark**: Here $F(\gamma)$ can be interpreted as some kind of difference between "entropies" in $Q(\varepsilon)^c$ and $Q(\varepsilon)$. When $F(1) < 0$, it means that $Q(\varepsilon)^c$ has higher "entropies", implying that the training trajectories are far from optimal. As a result, we may need to flatten the distribution of training trajectories, which corresponds to using the prior reweighting technique with $\gamma<1$. Since this is the most common case, we usually set $\gamma<1$.