L Rudolf L comments on Deconfusing Direct vs Amortised Optimization

[L Rudolf L]

I was very happy to find this post—it clarifies & names a concept I’ve been thinking about for a long time. However, I have confusions about the maths here:

Mathematically, direct optimization is your standard AIXI-like optimization process. For instance, suppose we are doing direct variational inference optimization to find a Bayesian posterior parameter $\theta$ from a data-point $x$, the mathematical representation of this is:

$\begin{array}{c}{\theta }_{\text{direct}}^{\ast }={\text{argmin}}_{\theta }KL\left[q\right(\theta ;x\left)||p\right(x,\theta \left)\right]\end{array}$

By contrast, the amortized objective optimizes some other set of parameters $\phi$ over a function approximator $\hat{\theta }=f_{\varphi }\left(x\right)$ which directly maps from the data-point to an estimate of the posterior parameters $\hat{\theta }.$ We then optimize the parameters of the function approximator $\varphi$ across a whole dataset $D=\left\{\right(x_1,{\theta }_1^{\ast }),(x_2,{\theta }_2^{\ast })\dots \}$ of data-point and parameter examples.

$\begin{array}{c}{\theta }_{\text{amortized}}^{\ast }={\text{argmin}}_{\varphi }{E}_{p\left(D\right)}\left[L\right({\theta }^{\ast },f_{\varphi }\left(x\right)\left)\right]\end{array}$

First of all, I don’t see how the given equation for direct optimization makes sense.$KL\left[q\right(\theta ;x\left)||p\right(x,\theta \left)\right]$ is comparing a distribution over $\theta$ over a joint distribution over $(x,\theta )$. Should this be $KL\left[q_{\psi }\right(\theta \left)||p\right(\theta |x\left)\right]$for variational inference (where $\psi$ is whatever we’re using to parametrize the variational family), and $KL\left[q\right(\theta |x\left)||p\right(\theta |x\left)\right]$ in general?

Secondly, why the focus on variational inference for defining direct optimization in the first place? Direct optimization is introduced as (emphasis mine):

Direct optimization occurs when optimization power is applied immediately and directly when engaged with a new situation to explicitly compute an on-the-fly optimal response – for instance, when directly optimizing against some kind of reward function. The classic example of this is planning and Monte-Carlo-Tree-Search (MCTS) algorithms [...]

This does not sound like we’re talking about algorithms that update parameters. If I had to put the above in maths, it just sounds like an argmin:

$g\left(x\right)={argmin}_{y\in A}L\left(y\right)$

where $g$ is your AI system, $A$ is whatever action space it can explore (you can make $A$ vary based on how much compute you’re wiling to spend, like with MCTS depth), $L$ is some loss function (it could be a reward function with a flipped sign, but I’m trying to keep it comparable to the direct optimization equation.

Also, the amortized optimization equation RHS is about defining a $\varphi$, i.e. the parameters in your function approximator $f$, but then the LHS calls it ${\theta }_{amortized}^{\ast }$, which is confusing to me. I also don’t understand why the loss function is taking in parameters ${\theta }^{\ast }$, or why the dataset contains parameters (is $\theta$ being used throughout to stand for outputs rather than model parameters?).

To me, the natural way to phrase this concept would instead be as

$g\left(x\right)=f_{\hat{\varphi }}\left(x\right)$

where $g$ is your AI system, and $\hat{\varphi }={argmin}_{\varphi }{E}_{(x,y)\sim D}\left[L\right(y,f_{\varphi }\left(x\right)\left)\right]$, with the dataset $D=\left\{\right(x_1,y_1),(x_2,y_2)\dots )\}$.

I’d be curious to hear any expansion of the motivation behind the exact maths in the post, or any way in which my version is misleading.