Paper: Meta-Learning Priors Using Unrolled Proximal Networks (ICLR-2024)

Reviewer comment (Experiments):
While this paper presents a solid contribution with rigorous analysis and experimental validation on few-shot classification, the scope of applications could be expanded to further demonstrate the generality of the proposed MetaProx method. In particular, the current empirical evaluation is limited to few-shot classification tasks. To reveal the full merits of MetaProx and its ability to learn flexible priors, additional experiments on other few-shot learning domains would make the work stronger. For example, further testing MetaProx on few-shot regression tasks, few-shot policy learning in reinforcement learning, or few-shot time series forecasting would verify its effectiveness beyond classification.

Author response:
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.

Reviewer comment (Experiments):
The paper currently visualizes the learned proximal operators averaged across iterations. An insightful extension would be to visualize the evolution of the learned priors at each unrolled PGD step, if possible. Analyzing how the proximal functions adapt over the course of optimization could provide deeper understanding into the dynamics of MetaProx and the induced priors.

Author response:
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](https://www.tensorflow.org/guide/keras/transfer_learning) and [natural language processing](https://aclanthology.org/P18-1031.pdf).

Reviewer comment (Novelty):
The work by [1] on proximal layers for deep regularization is relevant, as it shares the high-level idea of encoding priors into optimization-based neural network training. A key difference is that [1] transforms hidden representations to comply with preset priors, while this work learns proximal operators over model parameters. Comparing and contrasting with [1] could potentially help situate the contributions of this paper.

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

Reviewer comment (Experiments):
I wonder if the derived theoretic bounds can be validated empirically and the results appear in appendix. In toy settings, it is often possible to compute the error bounds. Numerical examination could validate the bounds and further show how tight the error bound can be. Specifically, the paper claims that the derived bounds are tight. It could help if this claim is verified through experiments.

Author response:
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](https://docs.google.com/presentation/d/e/2PACX-1vTtYYSMKjOnX8nlyNFWe3Q0980cG2MsoxkTPwFRY2joTSxkD4DwiQaxZFyZqYwwWWVu8AJao0r7pmWH/pub?start=false&loop=false&delayms=3000). 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.

Reviewer comment (Novelty):
The conceptual contribution is not exactly made clear. Maybe the clearest manifestation of this is that the presentation of the methodology does not make it clear why the authors' approach should be superior to other approaches to meta learning (say, to solving the formulation in equation (1)). The methodological discussion emphasizes that several commonly-used meta learning priors (under a common assumption on the "optimal regularizer") can be associated to proximal operators that can be represented as coordinatewise-applied piecewise linear functions. However, it does not connect back to specific tasks of interest in practice mentioned in the intro (robotics, say) and argue that the parametrized piecewise-linear proximal operators (essentially learnable nonlinearities in a neural network that implements the base-learner) are useful given specific structures the model parameters $\boldsymbol{\theta}_t$ have in these tasks. This is important to be able to support the authors' claims of "expressivity" and "interpretability" made in (i) of the claimed contributions.

Author response:
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.

Reviewer comment (Theory):
Labeling the algorithm "MetaProx" seems like it might be a slight misnomer -- based on the discussion on page 5 (after Assumption 3.1), it seems that what the authors are implementing is rather an amalgamation of one-dimensional piecewise linear functions, which need not correspond to the prox of any function. In one dimension, it should be sufficient and necessary here to constrain the piecewise linear function to be monotone.

Author response:
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.

Reviewer comment (Novelty):
It is not completely clear how the authors' eventual MetaProx algorithm differs from other meta learning approaches; most of the authors' discussion emphasizing the novelty of their approach seems to focus on contrasting with prior works on algorithm unrolling. For example, is this the first work that has made a connection between the problem in equation (1) and proximal gradient descent? A small section discussing this after the method has been presented would be helpful.

Author response:
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.

Reviewer comment (Writing):
Page 3: "...the optimization-based approaches solve (2b) ..." -- should this be (1b)? The rest of this sentence discusses $\mathcal{R}$, which does not appear in equation (2).

Author response:
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.

Reviewer comment (Writing):
Text before equation (11) on page 7: seems to be a dangling clause, "...has gained attention in various PGD-guided a [sic?]."

Author response:
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.

Reviewer comment (Writing):
Notation for continuously differentiable functions: as I understand it, using $\mathbb{C}^k$ for this class is nonstandard (and clashes with the usual notation for $k$-tuples of complex numbers in a jarring way). I guess this is an attempt to avoid a conflict with the parameter $C$ in the authors' algorithm; what about using $\mathcal{C}^k$ or $\mathsf{C}^k$ for this class?

Author response:
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.

Reviewer comment (Writing):
Is it correct that Theorem 3.2 involves Assumption 3.1? If so, it would be good style to reference this in the theorem statement.

Author response:
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.

Reviewer comment (Theory):
The notion of an "optimal" regularizer $\mathcal{R}^*$ is introduced on page 4, and plays a role in the authors' methodology. In what sense is this regularizer "optimal"? I could not find any discussion of this notion in the background section.

Author response:
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.

Reviewer comment (Novelty):
Is the content of Theorem 3.2 mainly an application of results on one-dimensional piecewise linear function approximation (with uniform discretization) of $C^1$ functions?

Author response:
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.

Reviewer comment (Experiments):
What is the "apples-to-apples" comparison mentioned when discussing the experimental results on page 8? Are there state-of-the-art meta-learning methods that have been excluded from this table? It seems natural that the MetaProx approach might perform better than methods that use less computation, since Algorithm 2 could be interpreted for natural losses $\mathcal{L}$ as a multilayer residual network (with some per-task conditioning) with learnable nonlinearities, if I'm not mistaken.

Author response:
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.