# Amortized Bayesian Meta-Learning with Accelerated Gradient Descent Steps
## 3 Meta-Learning with Amortized Bayesian Inference
### Probabilistic Model  
According to the classical hierarchical variational Bayes theory discussed in [10], the marginal likelihood of a hierarchical Bayes model is written in Equation (1) as follows:  
\[p_{f}(\mathcal{D})=\int_{\theta}p_{f}(\mathcal{D}|\theta)p(\theta)=\int_{\theta} \Biggl{[}\prod_{t=1}^{T}\int_{\psi}p_{f}(d_{t}|\psi_{t})p(\psi_{t}|\theta) \Biggr{]}p(\theta). \tag{1}\]  
Its generative process is shown as a directed graphical model in Figure 2, where \(\theta\) is denoted as the global latent variable. For each task \(t\), a task-specific parameter \(\psi_{t}\) is sampled from its priority \(p(\psi_{t}|\theta)\). Hierarchical variational inference can be used to estimate the lower bound of the likelihood in Equation (2) as follows:  
\[\log p_{f}(\mathcal{D})\] \[\geq\mathbb{E}_{q_{\theta}}\left[\sum_{t=1}^{T}\log\biggl{(}\int_ {\psi}p_{f}(d_{t}|\psi_{t})p(\psi_{t}|\theta)\biggr{)}\right]-\mathrm{KL}(q( \theta)||p(\theta))\] \[\geq\sum_{t=1}^{T}\Bigl{[}\mathbb{E}_{q_{\theta_{t}}}[\log p_{f} (d_{t}|\psi_{t})]-\mathrm{KL}(q_{\theta_{t}}(\psi_{t})||p(\psi_{t}|\theta)) \Bigr{]}-\mathrm{KL}(q(\theta)||p(\theta)), \tag{2}\]  
where \(\log p_{f}(\mathcal{D})\) is denoted as the _evidence lower bound_ (ELBO) on the log version. \(q_{\theta_{t}}(\psi_{t})\) is the introduced variational distribution with parameter \(\theta_{t}\). Given \(\theta_{t}(0\leq t\leq T)\), variational inference solves the following optimization problem:  
\[\operatorname*{argmin}_{\theta,f,\theta_{1},\cdots,\theta_{T}} \mathbb{E}_{q_{\theta}}\left[\sum_{t=1}^{T}-\mathbb{E}_{q_{\theta_{t}}}[\log p _{f}(d_{t}|\psi_{t})]+\mathrm{KL}(q_{\theta_{t}}(\psi_{t})||p(\psi_{t}|\theta ))\right]\] \[+\mathrm{KL}(q(\theta)||p(\theta))\] \[=\operatorname*{argmin}_{\theta,f,\theta_{1},\cdots,\theta_{T}} \frac{1}{T}\sum_{t=1}^{T}\Bigl{[}\mathbb{E}_{q_{\theta_{t}}}[-\log p_{f}(d_{t }|\psi_{t})]+\mathrm{KL}(q_{\theta_{t}}(\psi_{t})||p(\psi_{t}|\theta))\Bigr{]}\] \[+\mathrm{KL}(q(\theta)||p(\theta)). \tag{3}\]  
Note that the optimization process in (3) results in heavy computational costs since it has to optimize \(T\)-different variational distributions \(q_{\theta_{t}}(0\leq t\leq T)\) when \(T\) is large. To overcome this shortcoming, we introduce a _transductive_ amortization network \(\phi\) that takes \(x_{n}^{(t)}\in\mathcal{D}_{S}^{(t)}\) as input, and it outputs the variational distribution \(q_{\phi(d_{t}^{g})}\). Here, we use \(d_{t}^{S}\) to represent samples in \(\mathcal{D}_{S}^{(t)}\) for short. The idea of amortized variational inference (AVI) was first introduced in [26] as a trainable autoencoder that produces variational parameters  
Figure 2: Directed graphical model for the Bayesian meta-learning framework. The meta-learners learn the transductive amortization function \(f\) from the task-specific latent variables \(\psi\) and the global variables \(\theta\). for each data point, and has been extended to Bayesian variational inferences [27]. Assume that \(q_{\phi(d_{t}^{S})}\) follows the Dirac delta distribution (e.g., \(q(\theta)=\mathbb{I}\{\theta=\theta^{*}\}\)), where \(\phi(d_{t}^{S})=\theta_{t}^{0}\). For the \(k\)-th stochastic gradient descent iterate:  
\[\theta_{t}^{k+1}=\theta_{t}^{k}+\eta\
abla_{\theta_{t}}\mathrm{KL}(q_{\theta_{ t}^{k}}(\psi_{t})||p_{f,\theta}(\psi_{t}|d_{t}))\text{ with }\theta_{t}^{0}=\phi, \tag{4}\]  
where \(\eta\) denotes the descent step size. By conducting this optimization process up to the \(K\)-th step, such that \(\theta_{t}^{K}=\phi^{K}(d_{t}^{S})\), this guarantees that \(\theta_{t}^{K}\) is a good approximation to its optimal value \(\theta_{t}^{*}\).