Meta-Learning Priors Using Unrolled Proximal Networks

Thank you for the interest in this work, and for the valuable questions. The two concerns regarding extended numerical tests and generalization to Bayesian meta-learning approaches will be answered below.

1. Extended numerical tests
While our current numerical results are evaluated on miniImageNet and tieredImageNet, we are writing codes to implement MetaProx on the Caltech-UCSD Birds (CUB)-200-2011 dataset and Meta-Dataset. In addition, experiments on reinforcement learning tasks has also been prioritized in our research agenda.

2. Generalization to Bayesian meta-learning
The objective of Bayesian meta-learning takes into account the uncertainties in the model parameter $\boldsymbol{\theta}_t$, which incurs a similar bilevel objective

$$\max_{\boldsymbol{\theta}} \prod_{t=1}^T p (\mathbf{y}_t^\mathrm{val} | \mathcal{D}_t^{\mathrm{trn}}; \mathbf{X}_t^\mathrm{val}, \boldsymbol{\theta}) = \prod_t \int p (\mathbf{y}_t^\mathrm{val} | \boldsymbol{\theta}_t; \mathbf{X}_t^\mathrm{val}) p(\boldsymbol{\theta}_t | \mathbf{y}_t^{\mathrm{trn}}; \mathbf{X}_t^{\mathrm{trn}}, \boldsymbol{\theta}) d\boldsymbol{\theta}_t$$ $$\text{s.t.} ~p(\boldsymbol{\theta}_t | \mathbf{y}_t^{\mathrm{trn}}; \mathbf{X}_t^{\mathrm{trn}}, \boldsymbol{\theta}) \propto p(\mathbf{y}_t^\mathrm{trn} | \boldsymbol{\theta}_t; \mathbf{X}_t^\mathrm{trn}) p(\boldsymbol{\theta}_t; \boldsymbol{\theta}), ~\forall t.$$

Akin to its deterministic counterpart, the exact posterior pdf $p(\boldsymbol{\theta}_t | \mathbf{y}_t^{\mathrm{trn}}; \mathbf{X}_t^{\mathrm{trn}}, \boldsymbol{\theta})$ is generally intractable. A practical alternative is to rely on approximate inference techniques. In Bayesian MAML (BMAML), the inference is carried out through particle sampling, where the complexity scales in linear with the number of particles. In other Bayesian meta-learning approaches such as Probabilistic MAML, Amortized Bayesian Meta-Learning (ABML) [1] and VAMPIRE [2], the posterior is approximated through variational inference (VI). In particular, VI seeks for a tractable surrogate pdf $p(\boldsymbol{\theta}_t;\mathbf{v})$ that best estimates $p(\boldsymbol{\theta}_t | \mathbf{y}_t^{\mathrm{trn}}; \mathbf{X}_t^{\mathrm{trn}}, \boldsymbol{\theta})$ by minimizing their KL divergence

$$\mathbf{v}_t^* = \mathop{\arg\min}_v~\mathrm{KL} \big(p (\boldsymbol{\theta}_t; \mathbf{v}) || p(\boldsymbol{\theta}_t | \mathbf{y}_t^{\mathrm{trn}}; \mathbf{X}_t^{\mathrm{trn}}, \boldsymbol{\theta})\big)$$

which is equivalent to the maximization of the so-termed evidence lower bound (ELBO)

$$\mathbf{v}_t^*=\mathop{\arg\max}_v~\mathbb{E}_p{\scriptstyle (\boldsymbol{\theta}_t; \mathbf{v})} [ \log p(\mathbf{y}_t^{\mathrm{trn}} | \boldsymbol{\theta}_t;\mathbf{X}_t^{\mathrm{trn}})] - \mathrm{KL} \big( p(\boldsymbol{\theta}_t; \mathbf{v}) || p(\boldsymbol{\theta}_t; \boldsymbol{\theta})\big).$$

Upon defining loss $\mathcal{L}(\mathbf{v};\mathcal{D}_t^\mathrm{trn}):=\mathbb{E}_p{\scriptstyle (\boldsymbol{\theta}_t; \mathbf{v})} [ -\log p(\mathbf{y}_t^{\mathrm{trn}} | \boldsymbol{\theta}_t;\mathbf{X}_t^{\mathrm{trn}})]$ to be the expected nll, and regularizer $\mathcal{R}(\mathbf{v};\boldsymbol{\theta}):=\mathrm{KL} \big( p(\boldsymbol{\theta}_t;\mathbf{v}) || p(\boldsymbol{\theta}_t; \boldsymbol{\theta})\big)$, the VI objective reduces to

$$\mathbf{v}_t^*=\mathop{\arg\min}_v~\mathcal{L}(\mathbf{v};\mathcal{D}_t^\mathrm{trn})+\mathcal{R}(\mathbf{v};\boldsymbol{\theta})$$

which is exactly (1b) of our manuscript. In fact, the regularizer in VI can be any function assessing the divergence between two pdf. Similar to the deterministic case where Gaussian prior is popular, most existing Bayesian meta-learning algorithms select KL divergence as the regularizer, which is known to be unreliable in certain cases (e.g., when two pdfs have disjoint supports). Moreover, since the ELBO is typically optimized using gradient ascent in Bayesian meta-learning, MetaProx can be readily incorporated into VI to learn a data-driven divergence function.

[1] S. Ravi, and A. Beatson, "Amortized bayesian meta-learning," ICLR 2018.

[2] C. Nguyen et al., “Uncertainty in Model-Agnostic Meta-Learning using Variational Inference,” WACV 2020.

We appreciate your interest in this work and the constructive feedback provided. The issues raised are addressed one-by-one next.

1. Response to Weakness 1 & Question 1
As a flexible plug-in module, our MetaProx can be readily extended to few-shot regression and reinforcement learning (RL) tasks. Remarkably, the performance improvement by MetaProx mainly attributes to the enhanced prior expressiveness. Hence, it is expected to be particularly effective for tasks that exhibit intricate data distributions and require complex priors. To support this claim, a numerical case study on few-shot regression tasks will be provided in Appendix I of the updated manuscript.
In addition, implementing MetaProx in RL applications such as few-shot policy learning is also an intriguing direction which has been prioritized in our research agenda. One noteworthy challenge towards this direction is known as “meta-overfitting”, wherein an overly complex prior can overfit when the number $T$ of given tasks is also limited. While this challenge can be alleviated by employing a small $C$, Theorems 3.2 and 3.3 indicate a potential increase of approximation error in such case. Therefore, a potential remedy is to explore a proximal operator with improved error bounds.

2. Response to Weakness 2 & Question 2
The evolution of the prior in each PGD step $k=1,\ldots,5$ has been visualized in Figure 4, marked in different colors. It is seen that the acquired PLFs keep shallow layer weights being sparse around the initial values (i.e., less updated with a soft-thresholding) when $k$ is small, while deep layers can be updated freely along its gradient directions. In other words, the learned prior freezes the weights of the embedding function and exclusively train the model head in the initial iterations. Once the head has attained sufficiently training with the increasing of $k$, the prior gradually unfreeze the shallow layers and proceed with optimizing the entire model. In fact, this learned update strategy closely resembles the widely adopted “gradual unfreezing” training approach for fine-tuning large-scale models, which has been proven effective in various practical applications of computer vision and natural language processing.

3. Response to Weakness 3
Thank you for highlighting the parallels and distinctions between MetaProx and [1]. The missing reference and comparison will be added to the revised manuscript.

4. Response to Question 3
The PLFs in MetaProx can be initialized according to the prior knowledge about the optimization process. In particular, we set initially $\psi_{i,c}^k = \zeta_{i,c}^k, ~\forall i,c,k$ in our experiments. This initialization renders the proximal operator an identical mapping, and thus PGD boils down to GD. Consequently, MetaProx exhibits no impact on the backbone method in the initial stage. As meta-training progresses, the PLFs dynamically adjust its parameter $\psi_{i,c}^k$ to learn a data-driven prior.
Additionally, we also observed that the training of the PLFs occasionally diverge with the Adam optimizer, but remains robust when trained by SGD with Nesterov momentum. Numerically, this divergence appears linked to the instability in second moment estimation of the gradient.

5. Response to Question 4
The proposed MetaProx is expected to be particularly effective on more challenging tasks, which naturally necessitate sophisticated priors. As an example, in the ablation test (cf. Table 2), MetaProx improves MAML by $(53.58 - 48.70)/48.70 = 10.0\\%$ in the 5-way 1-shot setup, while the relative performance improvement is $11.3\\%$ in the 10-way 1-shot case.

6. Response to Question 5
The tradeoff is hidden in the regularization term of (1b). Consider for instance an isotropic Gaussian prior pdf $p(\boldsymbol{\theta_t}; \boldsymbol{\theta}) = \mathcal{N}(\boldsymbol{\mu}, \lambda^{-1} \mathbf{I}_d), ~\boldsymbol{\theta}:=[\boldsymbol{\mu}^\top, \lambda]^\top$, the corresponding regularizer can be analytically written as $\mathcal{R}(\boldsymbol{\theta_t}; \boldsymbol{\theta}) = -\log p(\boldsymbol{\theta_t}; \boldsymbol{\theta}) = \frac{\lambda}{2} || \boldsymbol{\theta}_t - \boldsymbol{\mu} ||_2^2 + \text{constant}$. In this example, the isotropic precision $\lambda$ serves as the (learnable) tradeoff weight between loss and regularization.

Thank you for your interest in this work and for providing the insightful suggestions. Next, the issues raised are addressed one-by-one below.

1. Regarding Weakness 1
As suggested by the reviewer, a toy numerical test will be provided next to demonstrate the tightness of the bound in Theorem 3.2. Consider tasks defined by linear relationship $y_t^n =\mathbf{w}_t^{\*\top} \mathbf{x}_t^n + e_t^n$, and linear prediction model $\hat{y}_t^n = f(\mathbf{x}_t^n;\boldsymbol{\theta}_t) := \boldsymbol{\theta}_t^\top \mathbf{x}_t^n$ with squared $\ell_2$ loss $\mathcal{L} (\boldsymbol{\theta}_t; \mathcal{D}_t^\text{trn}) := \frac{1}{2} \sum_n || y_t^n - \boldsymbol{\theta}_t^\top{\bf x}_t^n ||_2^2$, where the unknown oracle $\mathbf{w}_t^* \sim \text{Uniform}([-3,3]^d)$ and $\mathbf{x}_t^n \sim {\cal N}({\bf0}_d,\mathbf{I}_d)$. In the test, we set $K = 5, \alpha=0.01, d=64, |\mathcal{D}_t^\text{trn}|=8, \boldsymbol{\theta}^\text{init}=\mathbf{0}_d, A = 3$ and $C$ varying from 2 to 20. The target optimal proximal operator to be approximated is defined in eq.(27) with $G_1 = 1$. Additionally, the Lipschitz constant $G_2$ of $\nabla \mathcal{L}(\boldsymbol{\theta}_t; \mathcal{D}_t^\mathrm{trn})$ is numerically computed from the generated data $\\{ \mathbf{x}_t^n \\}_n$.

The plot comparing the numerical PGD error and its upper bound can be found via this anonymous URL. It is observed that the numerical error aligns with the theoretical bound up to a small constant in the log-scale. This discrepancy arises because the upper bound considers the worst-case scenario, where the largest error between ${\rm prox}_{{\cal R}^*, \alpha}$ and ${\rm prox}^k$ is reached at each PGD step. Furthermore, this constant gap suggests that the numerical error is in the same order with the bound (notice that both axes are in log-scale); that is, ${\cal O}(1/C^2)$. This empirically corroborates our theoretical proofs.

2. Regarding Weakness 2
We are in the process of revising the introduction to improve its clarity. In particular, an additional paragraph has been added to offer a comprehensive description of our methodology, before listing the contributions. Owing to space limitation, related work and an introduction to algorithm unrolling have been provided in the Appendices K and J, respectively.

3. Regarding Weakness 3
In practice, the visualization of the PLFs can be utilized to examine the impact of the prior when updating model parameters, thus providing guidance for practical model training process. In Figure 4, the acquired PLFs keep shallow layer weights being sparse around the initial value $\boldsymbol{\theta}^{\rm init}$ (that is, less updated) when $k$ is small, while deep layers can be updated freely along its gradient directions. This suggests, when fine-tuning a pre-trained large-scale model on a specific task, it is advisable to freeze the weights of the embedding function and exclusively train the last few layers with a relatively large step size in the initial epochs. Once these deep layers have attained sufficiently training, one can then gradually unfreeze the shallow layers and proceed with fine-tuning the entire model. In fact, this learned update strategy closely aligns with the widely adopted “gradual unfreezing” training approach for fine-tuning large-scale models, which has been proven effective in various practical applications of computer vision and natural language processing.

6. Response to Question 3: connection between Theorem 3.2 and 1-d PLF approximation.
The proof of Theorem 3.2 relies on the 1-d PLF approximation; cf. Lemmas A.5 and B.1. However, the proof is not a trivial application of the PLF approximation, particularly in the construction of the example for tightness; see Theorems A.6 and A.7 for more details.

7. Response to Question 4: elaboration of “apples-to-apples” comparison and the extra computational cost of MetaProx.
Here “apples-to-apples” comparison ensures the methods listed in the tables are evaluated using the same task-specific learning model. In other words, the dimension $d$ of model parameter $\theta_t$ and the model architecture are consistent among the methods. For fairness, results obtained with larger models have been excluded from the table. Moreover, complexity analysis has been provided in Appendix F both theoretically and numerically, revealing that the learning of PLFs contributes only marginally ($< 5\\%$) to the overall complexity compared to its backbone. As a result, our MetaProx can serve as a flexible plug-in module to enhance the performance of backbone meta-learning algorithms without markedly adding to its complexity. Notably, this benefit is realized through the acquisition of a more expressive prior.

Thanks for the thorough review of our manuscript and the valuable feedback provided. The issues raised have been carefully addressed one-by-one as follows.

1. Response to Weakness 1 & Question 1: example of claimed superior expressiveness.
As a straightforward yet illuminating example, consider regression tasks defined by the 1-d linear data model $y_t^n = w_t^* x_t^n + e_t^n$, where the unknown per-task weights $\\{w_t^*\\}_{t=1}^T$ are i.i.d. samples from the oracle pdf $p(w^*):=\frac{1}{2}\text{Uniform}(-11,-10)+\frac{1}{2}\text{Uniform}(10,11)$, and $e_t^n\sim{\cal N}(0,\sigma_e^2),\forall t$ is the additive white Gaussian noise. Further, consider a linear prediction model $\hat{y}_t^n=f(x_t^n;\theta_t):=\theta_tx_t^n$. Since $p(w^*)$ is symmetric, the optimal Gaussian prior of $\theta_t$ for this case must have a mean of 0. In other words, if the prior pdf is chosen to be Gaussian $p(\theta_t;\theta)={\cal N}(\mu,\lambda^{-1}),\theta:=[\mu,\lambda]^\top$, the optimal parameter $\theta^*$ must consist of $\mu^*=0$ and some $\lambda^* \in \mathbb{R}$. As a result, the corresponding regularizer for this prior is ${\cal R}(\theta_t;\theta^*)=\frac{\lambda^*}{2}\theta_t^2$, which prevents $\theta_t$ deviating far from $0$. However, the ground-truth task model implies that the optimal $\theta_t^*$ should belong to the set $S:=[-11,-10]\cup[10,11]$, and the regularizer is thus a barrier for optimizing $\theta_t$. In contrast, if the prior pdf is allowed to be non-Gaussian, the optimal prior will be the ground-truth one, i.e. $p(\theta_t; \theta^*)={\rm Uniform}(S)$. The corresponding proximal operator in this case is the projection $\text{prox}(z)=\mathbb{P}_S(z)$, which is exactly a piecewise linear function (PLF) driving $\theta_t$ to the oracle set. In summary, the Gaussian pdf fails to match the underlying prior due to its inherent unimodality, while the more expressive PLF-induced prior can perfectly align with the groundtruth prior to enhance the task-level learning.
In practical tasks such as drug discovery and robotic manipulations, the ground-truth model parameters can exhibit similar multi-modal pdf defined on a bounded set. In drug discovery for instance, the efficacy of a drug might only manifest when one component accounts for a specific portion.
To enhance the clarity of the paper, this example and a corresponding experiment have been included in Appendix I of the revised manuscript. We express our gratitude for the insightful suggestion.

2. Response to Weakness 2: misnomer and monotonicity of PLF.
In fact, a sufficient condition for the proximal operator to be monotone is the convexity of the regularizer $\cal R$. Specifically, denoting by $x=\text{prox}_{{\cal R},\alpha}(z)$, it follows from stationary point condition that $x\in z-\alpha\partial{\cal R}(x)$, which can be non-monotone when $\cal R$ is non-convex. In addition, the PLF in this case still induces a valid regularizer, although it may not have a closed-form solution. Please check Remark C.3 in Appendix C for more details.

3. Response to Weakness 3: comparison between MetaProx with other meta-learning algorithms.
Existing meta-learning algorithms rely on either a blackbox NN or a preselected pdf (such as a Gaussian one) to parameterize the prior. However, the NN often lacks interpretability and the chosen pdf can have limited expressiveness; that is, it may have insufficient ability to offer an accurate fit. Consider for instance a preselected Gaussian prior pdf, which is inherently unimodal, symmetric, log-concave, and infinitely differentiable by definition. Such a prior may not be well-suited for tasks with multimodal or asymmetric parametric pdfs, as illustrated in the aforementioned example. In contrast, our MetaProx introduces data-driven prior (induced by the learnable PLFs) which dynamically adjust its form to fit the provided tasks and offers the desired interpretability. Our theory and numerical tests are also centered around this claim of prior expressiveness and interpretability.
In response to the suggestion from the reviewer, we have incorporated this comparison into Sec.3 of the revised manuscript.

4. Response to Minor Issues.
We extend our appreciation for pointing out the typos and improper notations in the manuscript, which will be corrected accordingly. Regarding the last minor issue, a more formal version of Theorem 3.2 and 3.3 can be found in Appendices A and B, with Assumption 3.1 clearly involved in the statement.

5. Response to Question 2: definition of the optimal regularizer.
The optimal regularizer is defined to be the function $\cal R$ (possibly restricted to a space such as ${\cal C}^0$ or ${\cal C}^1$) that minimizes the meta-learning objective (1a). This missing definition has been added to the revised manuscript.

Dear authors,

Thank you for the thorough response to my review. I am happy with the clarifications you have provided, and will increase my score. I found the intuitive example you provided (first point response) to be helpful in clarifying the utility of your method.

I want to push back slightly on the convexity and prox issue, though -- I am not sure this is correct. An excellent characterization of this issue has been provided recently by Gribonval and Nikolova in https://centralesupelec.hal.science/hal-01835101/. One sees (see Theorem 1) that the characterization of when a map is a proximal operator amounts to it being the gradient of a convex function. For a 1d example, and in the presence of differentiability, I believe this is equivalent to the map in question being monotone (for example, imagine the twice-differentiable case, and note that the map is equal to the integral of its second derivative (up to a constant) -- and by convexity, this second derivative is positive).

I believe the issue you have in mind about convexity of the regularizer is different -- Gribonval and Nikolova clarify that the classical Moreau's characterization of proximal operators is that the map is the gradient of a convex function and nonexpansive. The characterization I discussed above (Theorem 1 of Gribonval and Nikolova) shows that in the nonconvex case, it is necessary and sufficient to discard this nonexpansivity.

This issue is not affecting my rating of the paper, but I do hope we can come to an agreement on this and clarify it in the work.

Hi Weisen,

To avoid potential confusion, the name of our method will be changed to MetaProxNet in our revised manuscript. In addition, the comparison will be included in the Appendix.

Thanks,
Authors of paper #4443