Meta-Learning Priors Using Unrolled Proximal Networks
Authors: Yilang Zhang
Georgios B. Giannakis
Venue: ICLR 2024 poster (2024)
Keywords: Meta-learning
prior extraction
algorithm unrolling
Primary area: transfer learning, meta learning, and lifelong learning

Abstract: Relying on prior knowledge accumulated from related tasks, meta-learning offers a powerful approach to learning a novel task from a limited number of training data. Recent approaches use a family of prior probability density functions or recurrent neural network models, whose parameters can be optimized by utilizing labeled data from the observed tasks. While these approaches have appealing empirical performance, expressiveness of their prior is relatively low, which limits generalization and interpretation of meta-learning. Aiming at expressive yet meaningful priors, this contribution puts forth a novel prior representation model that leverages the notion of algorithm unrolling.  The key idea is to unroll the proximal gradient descent steps, where learnable piecewise linear functions are developed to approximate the desired proximal operators within *tight* theoretical error bounds established for both smooth and non-smooth proximal functions. The resultant multi-block neural network not only broadens the scope of learnable priors, but also enhances interpretability from an optimization viewpoint. Numerical tests conducted on few-shot learning datasets demonstrate markedly improved performance with flexible, visualizable, and understandable priors.

Published as a conference paper at ICLR 2024
META-LEARNEL PRIORS USING UNROLLED PROXIMAL NETWORKS
Yilang Zhang, Georgios B. Giannakis
Department of Electric and Computer Engineering University of Minnesota
Minneapolis, MN 55414, USA {zhan7453, georgios @umn edu
ABSTRACT
Relying on prior knowledge accumulated from related tasks, meta-learning offers a powerful approach to learning novel task from limited training data. Recent approaches parameterize the prior with a family of probability density functions recurrent neural networks, whose parameters can be optimized by utilizing val- idation data from the observed tasks. While these approaches have appealing empirical performance, the expressiveness of their prior is relatively low, which limits the generalization and interpretation of meta-learning. Aiming at expres- sive yet meaningful priors, this contribution puts forth novel prior representation model that leverages the notion of algorithm unrolling. The key idea is to unroll the proximal gradient descent steps, where learnable piecewise linear functions are developed to approximate the desired proximal operators within tight theoreti- cal error bounds established for both smooth and non-smooth proximal functions. The resultant multi-block neural network not only broadens the scope of learnable priors, but also enhances interpretability from an optimization viewpoint. Numer- ical tests conducted on few-shot learning datasets demonstrate markedly improved performance with flexible, visualizable, and understandable priors.
INTRODUCTION
While deep learning has achieved documented success in a broad range of applications (Silver al,, 2016 He et al. 2016 Vaswani et al.] (017), it often requires huge data records to train large-scale and high-capacity models. In contrast, human intelligence is capable of identifying new objects or concepts from merely few samples. How incorporate this ability into "machine intelligence" has garnered great attention and interest in a number of domains, especially when data are scarce or costly to collect. Examples of such applications include drug molecule discovery (Altae- Tran al. 2017). low-resource machine translation (Gu et al.) 2018), and robotics (Clavera et al. 2019).
Motivated by the fact that humans acquire new knowledge efficiently from past experiences, prin- cipled framework has been investigated to mimic this ability of humans, known as learning-to- -treern or meta-learning (Thrun & Pratt 1998). Meta-learning aims to identify a task-invariant prior from a class of (partially) related tasks, which can be used to facilitate the learning of new tasks from the same class. The underlying assumption of meta-learning is that all tasks of interest are linked through their data distribution or latent problem structure. Thus, task-invariin edge can be acquired as an inductive bias, and thereby transferred to new tasks (Thrun& & Prati 1998). By doing so, even a couple of training data can suffice for learning new task.
common prior knowl-
Conventional meta-learning methods rely on prescribed criteria to extract the prior; see e.g., [Schmidhuber] 1993, Bengio et al.) 1993). With recent advances of deep learning, these hand- crafted approaches have been replaced by data-driven ones, where a meta-learner captures the prior information across tasks, while a base-learner utilizes this prior aid per-task learning. The desired prior encoded in the base-learner parameters shared across tasks, and can be learned by optimizing loss over the given tasks. Early attempts to this end utilize neural network (NN) to represent the prior (Santoro et al. 2016 Mishra et al. 2018)) Ravi & Larochelle 2017). The base-learner employs e.g., recurrent neural networks (RNNs) with input training data per task, and output parameters for

---

Published as a conference paper at ICLR 2024
the task-specific model. However, the choices of the NNs heavily depend on the task-specific model, and the blxck-box nature of NNs makes them susceptible to poor interpretability and reliability.
As opposed to model-based meta-learming, model-agnostic meta-learning (MAML) extracts the prior without presuming the task-specific model beforehand (Finn et al. 2017). MAML resorts to iterative optimizer obtain the per-task model parameters. The prior information is reflected in the initialization of the model parameters, which is shared across tasks. Building upon MAML, various optimization-based meta-learning algorithms have been investigated to further improve its performance; see e.g., (Li et al. 2017: Bertinetto et al. 2019; Lee et al. 2019). Convergence guar- antees have also been established gain insights about these methods (Fallah et al., 2020 Ji et al.) 2020 2022). Interestingly, Grant et al. (2018) pointed out that the initialization learned in MAML is approximately tantamount to the mean Gaussian prior probability density function (pdf) over the model parameters. This motivates well Bayesian formulations of meta-learning to further quantify the uncertainty in model parameters (Finn et al., 2018 Ravi & Beatson 2019) Nguyen et al. [020] Zhang et al. 2023). Nevertheless, the priors learned by these MAML- variants are confined to spe- citic pdfs, including the Gaussian and degenerate ones. As a result, generalizing optimization- meta-learning practical domains that may require sophisticated priors is challenging.
This work advocates a novel meeaa-lerrnng approach termed MetaProxNet that offers sufficient prior expressiveness, while maintaining the highly desirable interpretability. Our contribution is fourfold. i) A prior representation framework introduced using the algorithm unrolling technique. The novel framework overcomes the interpretability challenge and breaks the expressiveness bottle- neck, thus enabling one to meta-leam complicated yet interpretable priors.
ïi) Instead of employing fixed proximal operator induced by a certain prior pdf. piecewise linear functions (PLFs) are developed to learn further generalized priors.
iii) Theoretical analysis provides tight PGD eπror bounds between the learnable PLFs and the opti- mal proximal operators, which can be readily minimized under mild conditions.
iv) Numerical tests compare MetaProxNet with state-of-the- of-the-art methods having different priors, and confirm superiority of MetaProxNet. PLFs are visualized to depict the explainable prior.
2 PROBLEM SETUP
Meta-learning extracts task-invariant prior information from a collection of relevant tasks to aid the learning of new tasks, even if only a small number of training data are available. Formally, let 1,..., T index the aforementioned relevant tasks, each with corresponding dataset D, := {(x[, 9/ ) comprising N₁ input-output data pairs. Set Dr formed with a training subset D C Dr and a validation subset D := D2 Dim Likewise, a new task (with subscript *) will comprise a training subset Dtr", and a test input xtat, for which the corresponding output y!t is to be predicted. Typically, [Dtr"| rather small compared to what is required in supervised deep learning tasks. Due to the limited training data, directly learning the new task by optimizing its task-specific model over Dir is infeasible. However, since can considerably large, one prudent remedy is to leverage the cumulative prior knowledge across other related tasks.
Let 0, E Rd denote the task-specific model parameter for task t, and 0 E E the prior parameter shared across tasks. The prior can be learned via empirical risk minimization (ERM) alternating between i) base-learner optimization per t that estimates 8, using D and 0; and, ii) meta- optimization that updates the estimate of 0 using (Dyal)?" It=1- This nested structure can be intrinsi- cally characterized by a bilevel optimization problem
min
(();
(la)
6
s.to ( = argmin L(O,; D( + R( 0), Vt
(1b)
0.
where is the loss function assessing the performance of the model, and R is the regularizer that captures the task-invariant prior. From the Bayesian viewpoint, L(O.; Dira) and R(0,; 0) in (lb) are typically selected to be the negative log- likehooo (nll) - log p(y? 0 Xtr)., and negative log- prior (nlp) - log p(,,; 0), where matrix X!! is formed by all input vectors in Dirn and y is the vector collecting their corresponding outputs. Hence, (Ib) can be interpreted as the maximum a posteriori (MAP) estimator 0; (0) argmax p(O.ly!' trm. ttm 0) upon invoking Bayes rule.
2

---

Published as a conference paper at ICLR 2024
It is worth stressing that R(0;;;) is instrumental in learning task t, when !D is small. Without it, an over parameterizze model such as a deep NN could easily overfit Dt Moreover, it is gen- erally infeasible to reach the global minimum 07, especially with highly non-convex optimization involved in learning the task-specific model. Thus, a practical alternative to rely on suboptimal solution O obtained by a parameterized base-learner B. Then, problem
boils down to
T
min 0
((),
(2a)
s.oo 0 (0) B(D 0), Vt.
(2b)
Depending on the choices of B, meta-learning approaches can be either NN-based or optimization based ones. The former typically employ an RNN to learn the mapping from Di to 0; using the premise that the recurrent cells of an RNN correspond to the iterations for optimizing (Tb) Ravi & Larochelle| (017)) However, there is no analytical guarantee regarding the convergence of this "RNN-based optimization." " and it is also hard to specify what priors have been learned by these RNNs. In contrast, the optimization- -based approaches solve (Ib) through an iterative optimizer, with R being the nip term linked with preselected pdf. For example, it has been reported in (Grant at. 2018) that the optimization strategy adopted by MAML (Finn et al. 2017) corresponds up an implicit Gaussian pdf p(0,; 0) = (e, Q.). where Q₂ is associated with the hyperparameters of B. Besides implicit prior pdfs, their explicit counterparts have also been investigated; see e.g., isotropic
Gaussian (Rajeswaran et al. 2019), and diagonal Gaussian (Ravi & Beatson 2019) examples.
3 INTERPRETABLE AND GENERALIZED PRIORS USING UNROLLED NNS
Existing meta-learning algorithms rely on either a blackbox NN or a preselected pdf (such as a Gaussian one) parameterize the prior. However, the NN often lacks interpretability and the chosen pdf can have limited expressiveness. Consider for instance a preselected Gaussian prior pdf, which is inherently unimodal, symmetric, log-concave, and infinitely differentiable by definition. Such prior may not be well-suited for tasks with multimodal or asymmetric parametric pdf'; see App. !] for a case study. To enhance the prior expressiveness as well as offer the desired interpretability. our key idea is to learn a data-driven regularizer R, which dynamically adjusts its form to fit the provided tasks. This learnable effected by an unrolled NN, which drives our base-learner B.
3.1 PRIOR REPRESENTATION VIA ALGORITIM UNROLLING
Algorithm unrolling was introduced in (Gregor & LeCun 2010) to learn the optimal update rule for the reconstruction of sparse signals from their low-dimensional linear measurements. In particular, algorithm unrolling involves unfolding the iterations of an optimization algorithm to create repeating blocks of NN. In doing so, the desired prior is parameterized using learnable weights of the NN; see for brief introduction. Following this work, several unrolling methods have been reported learn interpretable priors for natural and medical signals, especially for images (Monga et al. 2021). Algorithm unrolling is also adopted here, but for a different purpose. While earlier efforts focus on learning the prior for single task in the (transformed) signal space X' C IR ddim(x"") ) here it employed for fask-invariant prior extraction in the model parameler space O C Rd, that is, the prior we aim to learn is p(O,), Vt rather than p(x,) ) for t given. The widely adopted convolutional (C)NNs, which exhibit remarkable effectiveness in representing priors for 2-dimensional images, may not fit well with the 1-dimensional 0,, A better alternative will be sought after the ensuing discussion that links prior representation with proximal function learning.
To solve the regularized problem (lb] we consider unrolling the proximal gradient descent (PGD) algorithm (Parikh et al. [014), which allows one to "divide and conquer" the objective function by separately optimizing L and R. Each PGD iteration indexed by k includes two steps: i) optimization of C(0) :Drm) wrt x: using GD, with the update represented by an auxiliary variable 2 E R4; and ii) optimization of R(e, :: O) using ž to update 0² An upshot of the PGD algorithm is that it only requires L(O,: -) to be differentiable wrt 0,, while R(0,; -) can be con-ifferrentable and even discontinuous. Thus, the expanded choices of R broaden the range of representable priors. The
3

---

Published as a conference paper at ICLR 2024
Algorithm 1: Vanilla PGD algorithm for solving (1b)
Input: Dir hyperparameters 0, step size o, and maximum iteration K.
Initialization: initialize 0g according to 0, and z 0%
for k I,...: do
Descend z t.- a 0 - )):
Update O = p (z) 1);
end
Output: 0, = =
DC
prox
8
0 -1
prox
B
block
block K
Figure I: Illustrating diagram of the multi-block NN by unrolling the PGD algorithm.
steps of PGD are summarized in Algorithm 11 where the so-termed proximal operator
proX (z) == argmin 2/ ||0, 211 O)
1
(3)
e, 2a
For a broad range of R. their corresponding prOXR.a has an analytical form. One well-known exam- ple is the indicator function R = Ils for some set S, which is discontinuous and non-differentiable However, corresponds to well-defined proXR,a. namely the projection operator Ps onto set S.
Using algorithm unrolling, our idea is to search for the unknown optimal regularizing function R* (i.e., the one minimizing (T) through learning its corresponding proximal operator proxæ .a with an unrolled NN. In particular, each PGD iteration indexed by l is replaced by block consisting of a data consistency (DC) module, and a learnable NN-based prox". While the former ensures
that the task-specific estimate 0, --1 of the unrolled NN is consistent with D (by minimizing _(O) :D!r)) wrt 0, dht-)), ), the latter looks for the optimal per-step prior that calibrates 0r The pipeline of this unrolled NN is illustrated in Fig. [ where the DC module can be either a naive GD as in line 4 of Algorithm T] or, data-driven rule such GD with a learnable a. Let us for simplicity adopt the naive GD as DC module, which aligns with MAML (Finn et al. (2017). and can be readily generalized other iterative descent rules (Li et al. 2017 Lee & Choil 2018 Park & Oliva 2019 Flennerhag et al. 2020). The typical choice for each prox* is an NN. Although p(O,; 0) may not be available since the NN mapping nonlinear, it can serve as generalized prior, properly scaled.
Unlike previous works (Mardani et al. 2018 Hosseini et al. 2020) that model [prox*]K with 2- dimensional convolutions, here the input and output of prox* are both 1-dimensional vectors in Rd; cf. 3 Our motivation comes from the two most widely-used priors in optimization based meta- learning. The first prior is the diagonal Gaussian one with R(0; 0) 1/2 (0. ginit)" diag (0,
giait), ), where giait 0% is the task-invariant initialization of (Tb). and = [iiuit
is the
vector parameterizing R (Ravi & Beatson 2019 Rajeswaran
2019: Nguyen al. 2020).
can be easily verified that prox (z) = (z 0 )/ + aλ) + 0 iit with / being the element-
wise division and J E IRd denoting the constant vector all I's. The second example is the shifted sparse prior that shares pre-defined portion of 0, across tasks (Raghu et 2020 Bertinetto al.
2019 Lee et al. 2019). Here, consider its variant R(0,; 0) A 0imity I|| that can be learned (Tan al. 2020b). This results in prox (z) SaA (z 0 ginit) + ginit, where SaA is the
element-wise shrinkage (a.k.a. soft-thresholding) operator such that its i-th element
2, + adi
2, < -a
SaA (z)); == Sax, =
0.
-a < < ad;
24 adi
N ad;
For notational simplicity, denote by shifted vectors 0¹ 0 - ginit うेद z ginit shifted loss
((::-) ) = L(O oi) and shifted proximal operator prox (2) = prox (z 0init) finit
4

---

Published as a conference paper at ICLR 2024
aa (Rt)
pros (2)
prinf (n))
Kb.bt Ch ved
(a)
(h)
(c)
Figure 2: Proximal operators for: (a) Gaussian prior; (b) sparse prior; (c) unrolling traselnn--based prior.
The PGD iteration can be thus reformulated as
a V 0. -00, k-1 :)!
(4a)
0 prox (z)), k k= L......,K
(4b)
with initialization 02 = O2 and output Ú, = 0x ginit Further, the prox (z) operator of the forgoing two examples reduces to 2/(1a a + aA) and (z) respectively.
Inspired by the fact that proX (z) of both examples belongs to the family of piecewise linear functions (PLFs), the fresh idea is to parameterize the shifted per-step přox*(z: ; 0) = prox + ginit) O of the unrolled NN using learnable PLFs. We first show that the wanted přox* Rd + Rd can be effectively decomposed and thus simplified under the following assumption that is widely adopted in meta-learning (Ravi & Beatson 22199) Rajeswaran et al. [019 Nguyen et al.) 2020).
Assumption 3.1. The optimal regularizer R factorizes across its input dimensions; that is,
R" (0, 0) Σd= R, ((01;;); 0).
With Assumption 3.1 in effect, an immediate result is the element-wise proximal operator
[proxR*, (z)], = argmin 1 0 =12 + [0c). 2a
R (())):0)
argmin 1 ((0,): 21 R 01]]; 0) = proxR; (2)), i= 1......,d d. (5)
[.).
2a
This observation suggests that we can alternatively model the dimension decomposition
přox = přox*] for each i 1,...,d, with a handy 1-dimensional PLF
w/ (tt, =t) (t.)
24 < ch
-4.00 -.) st/o-))
Ch S 24 ck and 2,... C -
(rox((2))
(6)
Ct Ct--ctto-1
-2) Wf o01 -t.c G, ct.
24 ck
where 2 1 is pre-selected constant indicating the total number of pieces, and ((C w Do are the learnable control points parametrizing přox.. To ensure přox, is valid function, we further require cto .... S ç for Vi, k. To this end, the problem of finding a proper task-invariant prior p0,,; 0) boils down to learning the parameters of PLFs that are shared across tasks. Comparison of baseb and PLF-based proximal operators can be visualized in Fig.
3.2 PRIOR LEARNING ALTERNATING OPTIMIZATION
Building upon the unrolling- bardel prior information representation, we are ready to elucidate how the prior can be learned by alternately optimizing the meta-learner and base-learner. We term the proposed method as meta-learning via proximal networks (MetaProxNet)..
Let r and denote iteration indices for Ta and lb respectively. For notational brevity, de- fine vectors 2h == Isi......, st.ccTT and & = [... with chh.c]] of the PLF control points, and
0" the concatenation of ginit,r -2. S K.r 2br 2 Kr in the r-th iteration of (T)). Given

---

Published as a conference paper at ICLR 2024
Algorithm 2: MetaProxNet algorithm
Input: (D,) step sizes a and B, batch size B, and maximum iterations K and R.
Initialization: randomly initialize 0".
for r Il,,....-l do
Randomly sample mini-batch T C c(,,..., T) of cardinality B; for / E T" do
Initialize 0% = 2 0a;
for k = l.... do Descend (0 } = 0 10--(00 10 - 0: 10-1:(0))
Update 0 (0--1) přox* * (2) (0 ): thr ubt:r;;
end
0KK Shift ()(0-=1)- = (0--1) + ginit,r:
10
end
Update 0" = 0-- -0-11-81 -B EIET. L(
end
Output: 0
(Drra) 1 the goal of (lb) is to learn the task-specific estimate (0") that depends on 0 per task This can leverage the current base-learner estimate B(;; 0"), which the unrolled multi-block
NN of our MetaProxNet. In the k-th block, its DC module and PLFs optimize (Tb) through
(0") 0. 20") V 1 2(0) - (0"):
(7a)
0t (0") = přox" (z) (0"): k.r s (rr) I....,... K.
(7b)
where äे and 0t denote the shifted iterative variables of the unrolled NN in (4)
After obtaining (6)) = 0. (0") ginit.r, the next step is to optimize La by updating 0" pop-
ular strategy is the mini-batch stochastic GD (SGD). Specifically, a subset T {L...., ., T) of tasks are randomly selected assess the performance of 0" on Diai, which yields a loss L(O, Dial) for Vt e T". Then, + is reached by descending the averaged loss with step size B, that is
r-+1 = 0"
V ((=(
(8)
T
TET
The step-by-step pseudo-codes for our novel MetaProxNet approach are listed under Algorithm 2
incurs cumbersome gradient computations due to the entangled structure of (1) To relieve this burden, we fix the former by uniformly partitioning closed interval [-A, A), while optimizing only the latter. In other words, we let ch = (2 1)A, Vc, i, k. where A > 0 is pre-selected constant that is sufficiently large; see Assumption A.3 In fact, this setup can be viewed as a uniform discretization of the continuous variable Ç E R on AA, A]. Non-uniform discretization can be alternatively sought, if p(c?) or its estimate is available a priori.
In practice however, simultaneously optimizing both { and { { Jk=
Jk=1
3.3 ERROR BOUNDS FOR PLF-BASER PROXIMAL OPERATOR
Having introduced how to model and learn priors using unrolled NNs, this subsection analyzes the performance by bounding the approximation error on 0, induced by replacing the unknown optimal proxæ-.0 . with the learned PI.F-aseed přox*. Sharp bounds will be separately established for smooth and non-smooth prox aa operators under mild conditions. Utilizing these bounds, quantitative criterion will be provided for choosing the hyperparameter C, All proofs and technical assumptions can be found in Apps. C Smooth přoxæ... E C² (I([[A,, A, A]d) will be first considered.
The following theorem offers upper bound for the normalized enror on (shifted) 0..
Theorem 3.2 (Finite-step PGD error for smooth proximal operators). Consider přox* defined [ab] ..... =b denote the matrix parameteriz-
by (6) with fixed ct. (2g 1)A and let T
ing (prox* ) Let 0x and 0% be the K K-step PGD outputs using proxR=,a E CC A, A)d) and

---

Published as a conference paper at ICLR 2024
přox*, respectively. Under mild assumptions, it holds t 1...., T that
min 1 1100 (v) -0 (g))
This bound is tight when v prox A-) and v - prOXR; aa (A), 4k,
Theorem 3.2 asserts that by optimizing over T of the PLFs, přox* approximate any smooth prox .a with Ktstep PGD error in the order O (zzz). ). In other words, an e-approximant 0 of 0, can be obtained upon choosing C' = 9 )/) and optimizing V. The tightness of the bound implies that there exists least proxR- that reaches the upper bound when enforcing the first and last control points of each PLF to align with the desired proxR;.0 operator.
Unfortunately, directly optimizing the left-hand side of ( is impossible, because the optimal proxR*, corresponding the oracle prior p(0,;; 0*) is unknown. A feasible alternative is to per- form the ERM in by leveraging the datasets {D.)== generated with 0r ~ p(0r; 0*). As a result, the (unknown) optimal PLF parameters & argming 110, 0 (I))l, and the sub-optimal estimate obtained by solving (T). satisfy the inequality
- Fall -1 0. (V)12 S Vd o0o 0. (VV))ll
Val
1 th (i)) 0 (V (V²))l2-
(10)
The extra error V 00 (i) 0 (vr can be further bounded in linear order O (+tal v of the normalized ERM error; see App. for further elaboration.
Aside from smooth ones, non-smooth proXR. has gained attention in various PGD- (GD-guided appli- cations. The next theorem forgoes the smooth assumption to yield a more generic but looser bound.
Theorem 3.3 (Finite-step PGD error for continuous proximal operators). Consider the notational conventions of Theorem 3.2 with continuous proxæ. a ([-A, A)d). Under mild assumptions, it holds for t= 1..... T that
1 1 0 L 111 [1
min
(11)
This bound is tight when oh prox aa (-A) A) and vi.c proXR; aa (A), Vk, i.
Compared to the smooth case, the error bound in Theorem 3.3 has an order of O This implies that by selecting C 2 operator přox* can approximate any continuous prox with nor- malized K-step PGD error no larger than e. This increased order implies that one can easily expand the range of learnable priors with larger C. Moreover, the discussion following (TO) regarding the
sub-optimality L,, applies to Theorem 3.3 too, and is deferred to App. C
4 NUMERICAL TESTS
In this section, numerical tests are presented on several meta-learning benchmark datasets to evaluate the empirical performance of MetaProxNet. Hyperparameters and datasets described in App. L All experiments are run on server with RTX A5000 GPU, and our codes are available online at https://github. con/ http://haagyyllaa//aaaaaaaooo..a
4.1 COMPARISON OF META-LFARRING METHODS HAVING DIFFERENT PRIORS
The first test is on fho- classification datasets minilmageNNt (Vinyals et al. 2016) and Tiered- ImageNet (Ren et al., 2018), where "shot" signifies the per-class number of labeled training data for each t. The default model is a standard 4-layer CNN (Vinyals et al. (0166), each layer comprising a 3 x 3 convolution operation of 64 channels, a batch normalization, ReLU activation, and 2 max pooling. linear regressor with softmax is appended to perform classification.
To demonstrate the superiority of unrolling-basee
priors over the RNN-based and handerafted ones, we first compare MetaProxNet against several state-of-the -hee----tthh-tttt meta- meta-learning methods. As discussed in Sec. 3.1 our NetaProxNet can be readily integrated with other optimization-baser meta-learning
7

---

Published as a conference paper at ICLR 2024
Table 1: Comparison of MetaProxNet against meta-learning methods with different priors. The
highest accuracy as well as mean accuracies within its 95% confidence interval are bolded.
5-class minilmageNet
5-class Tiered
Method
Prior
1-shot (%)
5-shot (%)
1-shot (%) 5-shot (%)
LSTM (Ravi & Larochelle
2017)
RNN-based
43. 13.4440.77
60.604..71
MAML Finn al. 2017)
implicit Gaussian
48.70
63. LL
51.67 53.31
70.30
ProtoNets (Snell et al. 22017
shifted sparse
49.422
68.20440.66
72.6920.74
R2D2
(Bertnetto el al 20199
shifted sparse
51.8
68.4
MC (Par & Oliva 2019
block-diag. Gaussian
54.08
67.9940.73
L2F
Bark ct al. (2020)
implicit Gaussian
52.10.00.50
69.38
54.40
73.34
KML.
Abdollahzadeh al.
2021
shifted sparse
54.10.20611
68.07
54.67=0.39
72.09
MeTAL
[Bark al. 2021
implicit Gaussian
52.63
70.52.0.22
54.3420.31
70.40=0.21
MinimaxMMAML (Wang el al. 2023)
inverted nlp
51.70+0.42
1..4
MctaProxNet+MMML
unrolling- brasling-bbscd
53.70+1.4u
70.08
54.56+1.44
71.80+a.78
MetaProxNet+MC
unrolling-based
55.9411.39
71.97 =0.97
57.34 42 73.38
Table 2: Ablation tests of MetaProxNet using minilmageNet dataset with a 4-layer 32-channel CNN.
5-class
10-class
Method
Presct prior?
DC
1-shot (%)
5-shot (%)
1-shot (%)
5-shot (%)
MAML
Yes
GD
48.70 18.70
63.1 63.11
31.27=1.15
46.92+1
PGD-Gaussie
Yes
PGD
48. 18.5841.40
64.56
30.04_0.89
47. 47.30.0.49
MetaProxNe++MAML
No
PGD
53.5841144
67.88 7.
34.80=0.91
51.03_0.51
methods through simple substitution of the DC module. Tab. lists the performance of MetaProx- Net assessed using 1, 000 random new tasks, with MAML (Finn et al.) 2017) and MetaCurvature (MC) (Park & Oliva 2019) serving backbones. For apples-to-apples comparison, methods that use different models (e.g., residual networks) or pretrained feature extractors are not included the table. It is seen that our MetaProxNNtt performs competitively in terms of classification accuracy when compared to state-of-the-art meta-learning methods. This empirically confirms the effective- ness of MetaProxNet. Additional discussions regarding the efficiency of MetaProxNet and extra tests with tied weights can be found in the Apps. and G
4.2 ABLATION TESTS
0.55
Ablation tests are also carried out to investigate the essential reason for the performance gain of MetaProxNet. Evidently, MetaProxNet+MAMI. differs from its backbone MAML in two key as- pects: task-level optimization algorithm (PGD vs.
0.54
Accuracy
0.52
0.51
GD) and prior (unrolled-NN based vs. Gaussian). To assess which of the two contributes more to the performance gain of MetaProxNet, the ablation tests compare three methods: i) MAML that employs GD and Gaussian prior; ii) a variant with PGD and Gaus- sian prior; and, iii) MetaProxNNtt+MAR lizes PGD and an unrolled- -NN based prior. To avoid overfitting in MAML. the models for all methods are fixed to a 4-layer 32-channel CNN. Tab. 2 lists the performance of the three methods. It is performance, while MetaProxNet outperforms both in all tests. This reveals that the key factor contributing to MetaProxNet's success is the more expressive prior relative to PGD.
0.49
0.48
that uti-
Number C of pleces in PLFe
Figure 3: Classification accuracy against the
number C of PLF pieces.
that the PGD baseline and MAML exhibit comparable
4.3 IMPACT OF HYPEREPARAMETER C'
Numerical tests are also carried out verify the theoretical analysis in Sec. 3.3 which upper bounds the 8-2 error between two PGD optimization outputs: one using the optimal prior and the other using PLF-induced prior. Specifically, Theorems 3.2 and 3.3 state that this 6-2 error bounds will reduce C increases, thus offering a better calibrated 0,. To examine the qualities of Ô with different
8

---

Published as a conference paper at ICLR 2024
0.5
-0.5
k=1
-0.5
-1
km3
-1
k=4
K=5
-1.5
-1.5
-0.5
0
0.5
-0.5
0
0.5
(a)
(b)
Figure 4: Visualization of the learned PL.Fs averaged across CNN layers; (a) first layer; (b) last layer.
Fig. depicts the test accuracies of MetaProxNet-MAML. function of C. It can be observed that the accuracy improves with C' increasing, which corroborates with theories. Moreover, C' 5 suffices achieve satisfactory performance, while larger values of C' only have a minor impact on MetaProxNer's empirical performance. This suggests that the constants hidden within the error bounds O and ((zz) can be small enough in practice. To avoid potential overfitting of priors. we set C = in all the tests.
5-class 1-shot minilmageNet
4.4 INTERPRETING UNROLLING- BASED PRIORS BY VISUALIZING THE LEARNED PLFs
From an optimization viewpoint, the learned PLFs correspond an implicit prior pdf that gener- ally comes with no analytical expression. These PLFs can be visualized to further understand the behavior of the unrolled NN. Figs. 4a and 4b respectively depict the averaged přox. for i's that correspond to the first and last CNN layers. The visualization showcases that the averaged PLF for the first layer is similar to the soft shrinkage function S,,,, of the sparse prior mentioned Sec. 3.1 while the last layer tends to have a linear PLF, which resembles that of a Gaussian prior.
In practice, the visualization of the PLFs can be utilized to examine the impact of the prior when updating model parameters, thus guiding the model training process. In Fig. 4 the acquired PLFs keep shallow layer weights being sparse around the initial value Oiuit (that is, less updated) when small, while deep layers can be updated freely along its gradient directions. This suggests, when fine-tuning pre-trained large-scale model on 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 sufficient training, gradually unfreeze the shallow layers and proceed with fine-tuning the entire model. 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; see e.g.. (Howard & Ruder 2018
then
CONCLUSIONS AND OUTLOOK
A novel prior information representation approach was pursued in this work using algorithm un- rolling to learn more flexible and generalized priors. Under this framework, meta-leaming method termed MetaProxNet was developed with learnable PLFs effecting implicit prior. The learned prior enjoys interpretability from an optimization vantage point, and can be well explained by vi- sualizing its PLFs. Further, performance analysis established that the PLFs are capable of fitting smooth/continuous proximal functions with a proper selection C. Numerical tests further cor- roborated empirically the superiority of MetaProxNet relative meta-learning alternatives in prior representation and learning.
Our future research agenda includes exciting themes on i) investigating various optimizers besides PGD: ii) implementing MetaProxNet with more complicated backbones and DC modules; and, iii) establishing bilevel convergence guarantees for MetaProxNet.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by NSF grants 2102312, 2103256, 2128593, 2126052, and 2212318.
REFERENCES
Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yo Yu Chen, and Tianyi Chen. Sharp-MAML: Sharpness- aware model-agnostic meta learning. In Proc. Int. Conf. Machine Learn., volume 162, pp. 10-32, 17-23 Jul 2022.
Milad Abdollahzadeh, Touba Malekzadeh, and Ngai-Man (Man) Cheung. Revisit multimodal meta- learning through the lens of multi-task learning. In Proc. Adv. Neural Info. Process. Systems, volume 34, pp. 14632-14644, 2021.
Han Altae- Tran, Bharath Ramsundar, Aneesh S. Pappu, and Vijay Pande. Low data drug discovery with one-shot learning. ACS Central Science, 3(4):283-292, 2017.
Marcin Andrychowicz, Misha Denil, Sergio Gómez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. In Proc. Adv. Neural Info. Process. Systems, volume 29, 2016.
Sungyong Baik, Scokil Hong, and Kyoung Mu Lee. Learning to forget for meta-learning. In Proc. Conf. Computer Vision and Pattern Recognition, June 2020.
Sungyong Baik, Janghoon Choi, Heewon Kim, Dohee Cho, Jaesik Min, and Kyoung Mu Lee. Meta- learning with task- adsk-idae loss function for few-shot learning. In Proc. Int. Conf. Computer Vision, pp. 9465-9474
October 2021.
Samy Bengio, Yoshua Bengio, and Jocelyn Cloutier. On the search for new learning rules for anns. Neural Processing Letters, 2(4):26-30, 1995.
Luca Bertinetto, Joao F. Henriques, Philip Torr, and Andrea Vedaldi. Meta-learning with differen- tiable closed-form solvers. In Proc. Int. Conf. Learn. Represention, 2019.
Ignasi Clavera, Anusha Nagabandi, Simin Liu, Ronald S. Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning adapt in dynamic, real-world environments through meta- reinforcement learning. In Proc. Int. Conf. Learn. Represention, 2019.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. On the convergence theory of gradient- based model-agnostic meta-learning algorithms. In Proc. Inf. Conf. Artif. Intel. and Stats., volume 108, pp. 1082-1092, 26-28 Aug 2020.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Proc. Int. Conf. Machine Learn., volume 70, pp. 1126-1135, 06-11 Aug 2017.
Chelsea Finn, Kelvin Xu, and Sergey Levine. Probabilistic model-agnostic meta-learning. In Proc. Adv. Neural Info. Process. Systems, volume 31, 2018.
Sebastian Flennerhag, Andrei A. Rusu, Razvan Pascanu, Francesco Visin, Hujun Yin, and Raia Hadsell. Meta-learning with warped gradient descent. In Proc. Int. Conf. Learn. Represention, 2020.
Jonathan Gordon, John Bronskill, Matthias Bauer, Sebastian Nowozin, and Richard Turner. Meta- learning probabilistic inference for prediction. In Proc. Int. Conf. Learn. Represention, 2019.
Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient- based meta-learning as hierarchical bayes. In Proc. Inf. Conf. Learn. Represention, 2018
Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In Proc. Int. Conf. Machine Learn., pp. 399-406, 2010. ISBN 605588977.
Rémi Gribonval and Mila Nikolova. A characterization of proximity operators. J. Mathematical Imaging and Vision, 62:773-789, July 2020.
10

---

Published as a conference paper at ICLR 2024
Jiatao Gu, Yong Wang, Yun Chen, Victor O. K. Li, and Kyunghyun Cho. Meta-learning for low- resource neural machine translation. In Proc. Empirical Methods in Natural Language Process., pp. 3622-3631, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proc. Conf. Computer Vision and Pattern Recognition, June 2016.
Sepp Hochreiter, A. Steven Younger, and Peter R. Conwell. Learning to learn using gradient descent. In Prof. Intl. Conf. Artif. Neural Networks, pp. 87-94, 2001.
Seyed Amir Hossein Hosseini, Burhaneddin Yaman, Steen Moeller, Mingyi Hong, and Mehmet Akçakaya. Dense recurrent neural networks for accelerated mri: History-cognizant unrolling of optimization algorithms. IEEE J. Sel. Topics Sig. Process., 14(6): 1280-1291, 2020.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. arXiv preprint 1.XXv:8801.061444 2018.
Samuel Hurault, Arthur Leclaire, and Nicolas Papadakis. Proximal denoiser for convergent plug- and-play optimization with nonconvex regularization. Proc. Int. Conf. Machine Learn., volume 162, pp. 9483-9505, 17-23 Jul 2022.
Kaiyi Ji, Jason D Lee, Yingbin Liang, and H. Vincent Poor. Convergence of meta-learning with task-specific adaptation over partial parameters. Proc. Adv. Neural Info. Process. Systems, volume 33, pp. 11490-11500 2020.
Kaiyi Ji, Junjie Yang, and Yingbin Liang. Theoretical convergence of multi-step model-agnostic meta-learning. Mach. Learn. Res., 23(1), jan 2022. ISSN 1532-4435.
Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. Meta-learning with differentiable convex optimization. In Proc. Conf. Computer Vision and Pattern Recognition, June 2019.
Yoonho Lee and Seungjin Choi. Gradient-basee meta-learning with learned layerwise metric and subspace. In Proc. Inf. Conf. Machine Learn., volume 80, pp. 2927-2936, 10-15 Jul 2018.
Ke Li and Jitendra Malik. Learning optimize. In Proc. Int. Conf. Learn. Represention, 2017.
Mao Li, Yingyi Ma, and Xinhua Zhang. Proximal mapping for deep regularization. In Proc. Adv. Neural Info. Process. Systems, volume 33, pp. 11623-11636, 2020a.
Yuelong Li, Mohammad Tofighi, Junyi Geng, Vishal Monga, and Yonina C. Eldar. Efficient and interpretable deep blind image deblurring via algorithm unrolling. IEEE Comput. Imaging, 6: 666-66 681, 2020b.
Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li. Meta-sgd: Learning to learn quickly for few- shot learning. arXiv preprint 1oXiv:/707.000 09835, 2017.
Morteza Mardani, Qingyun Sun, David Donoho, Vardan Papyan, Hatef Monajemi, Shreyas Vasanawala, and John Pauly. Neural proximal gradient descent for compressive imaging. In Proc. Ady. Neural Info. Process. Systems, volume 31, 2018.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta- learner. In Proc. Int. Conf. Learn. Represention, 2018.
Vishal Monga, Yuelong Li, and Yonina C. Eldar. Algorithm unrolling: Interpretable, efficient deep learning for signal and image processing. IEEE Sig. Process. Mag., 38(2):18 44, 2021.
Cuong Nguyen, Thanb-Toan Do, and Gustavo Carneiro. Uncertainty in model-agnostic meta- learning using variational inference. In Proc. Winter Conf. App. Computer Vision, March 2020.
Neal Parikh, Stephen Boyd, al. Proximal algorithms. Foundations and trends@ in Optimization, 1(3):127-239, 2014.
Eunbyung Park and Junier B Oliva. Meta-curvature. In Proc. Adv. Neural Info. Process. Systems, volume 32, 2019.
11

---

Published as a conference paper at ICLR 2024
Aniruddh Raghu, Maithra Raghu, Samy Bengio, and Oriol Vinyals. Rapid learning feature reuse?
towards understanding the effectiveness of maml. In Proc. Int. Conf. Learn. Represention, 2020.
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine Meta-learning with im-
plicit gradients. In Proc. Adv. Neural Info. Process. Systems, volume 32, 2019.
Sachin Ravi and Alex Beatson. Amortized bayesian meta-learning. In Proc. Int. Conf. Learn. Represention, 2019.
Sachin Ravi and Hugo Larochelle. Optimization model for few-shot learning. In Proc. Int. Conf. Learn. Represention, 2017.
Mengye Ren, Sachin Ravi, Eleni Triantafillou, Jake Snell, Kevin Swersky, Josh B. Tenenbaum, Hugo Larochelle, and Richard S. Zemel. Meta-learning for semi-supervised few-shot classification. In Proc. Int. Conf. Learn. Represention, 2018.
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta- learning with ume 48, pp. 1842-1850, 20-22 Jun 2016.
neural networks. Proc. Int. Conf. Machine Learn., vol-
J. Schmidhuber. A neural network that embeds its meta-levels. In IEEE Insl. Conf. on Neural Nenworks, pp. 407-412 vol.1, 1993.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):444488, 2016.
Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Proc. Adv. Neural Info. Process. Systems, volume 30, 2017.
Flood Sung, Yongxin Yang, Li Zhang, Tao Xiang. Philip H.S. Torr, and Timothy M. Hospedales. Learning compare: Relation network for few-shot learning. In Proc. Conf. Computer Vision and Pattern Recognition, June 2018.
Sebastian Thrun and Lorien Pratt. Learning to Learn. Springer, 1998.
Hongduan Tian, Bo Liu, Xiao-Tome Tong Yuan, and Qingshan Liu. Meta-learning with network pruning. In Proc. European Conf. Computer Vision, pp. 675-700, 2020a.
Hongduan Tian, Bo Liu, Xiao-Tong Yuan, and Qingshan Liu. Meta-learning with network pruning. In Proc. European Conf. Computer Vision, pp. 675-700, 2020b. ISBN 978-3-038 978-3-030-58529-7.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Adv. Neural Info. Process. Systems, volume 30, 2017.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, koray kavukcuoglu, and Daan Wierstra. Match- ing networks for one shot learning. In Proc. Ady. Neural Info. Process. Systems, volume 29, 2016.
Lianzhe Wang, Shiji Zhou, Shanghang Zhang, Xu Chu, Heng Chang, and Wenwu Zhu. Improv- ing generalization of meta-learning with inverted regularization at inner-level. In Proc. Conf. Computer Vision and Pattern Recognition, pp. 7826-7835, June 2023.
Yilang Zhang. Bingcong Li, Shijian Gao, and Georgios B. Giannakis. Scalable bayesian meta- learning through generalized implicit gradients. In Proc. AAAI Conf. Artif. Intel., volume 37(9), pp. 11298-1111 1306, 2023.
12

---

Published as a conference paper at ICLR 2024
A PROOF OF THEOREM 3
Smooth proxR*,o E ((), will be first considered under the following four technical as- sumptions.
Assumption A.1. proxR- a E ((=,(- A19) has Gi-Lipschitz gradient on [-A, AJ4.
Assumption A.2. [ has Gy-Lipschitz gradient on [-A, Al
Assumption A.3. Constant A is sufficiently large so that ž 2/- AJr, Vt, k, where z!_- is the PGD auxiliary variable generated with proXR* ,a'
A Remark A.5 (Mild assumptions). Assumption A.1 and A.2 the optimal proxg-,0 aa and the loss C are only assumed to be Lipschitz smooth on the compact subset [-A, A]" c R2, without imposing any strong premise regarding their convexity or Lipschitz continuity.
Assumption A.4. Operator prox--,0 aa A| is L-Liscchitz on |-A,
For Assumption A.3 the existence of such an A can be easily guaranteed when e.g., task-level step size a S 2/G.. and level sets {0, Z(0,: D) S C(0,: )] {0, R) s R(0,; 0)} and {0, (0,) < (0,) are bounded.
In addition, Assumption A.4 can be readily satisfied as well. For example, when R has Gr-
Lipschitz gradient on |- A, Ald and a < 1/GR. it follows from the stationary condition of (3) that (proxx (z) z) Prropr (a (z) O. Hence, it holds for Vz, z' E [-A, that
VR (proxx- a (z2) = (1 - aGR) Il prax .a (z) pr al (2)) In other words, the Lipschitz constant in this case is
|zzzz'll 2 prox aa (z) pr a az(z)( ll2 - VR (proxx- a (z))
upper bounded by L - 11(1 - aGR).
To prove Theorem 3.2 we first show lemma that is important for bounding the error p -prox all on
A, A].
Lemma A.6. Let I E C² (R) R + I be function with G-Lipschitz gradient. For VCI, S2 E I and (i + S2. define
f(z) (G2 - ==)=(((()) + (=-G))(EE)
(12)
C2
It then holds for Vy E [0, that
G
|f((1 [f((1 - 7) 752) f((117)C - + x52) (52 - (1)2.
(13)
Proof. For notational convenience, let (())(() =
+ 752) f(( 7] + 7 Using
the definition of f and g, it can be easily verified for Vyy E (0, that E (IR and
g(7) 2 g(0) = 9(1 0.
(14)
Therefore, there exits at least one maximizer 7 argmax,Eto.1) g(7) inside the open interval (0, For brevity, define the corresponding c==(( (1 >-) + + Through Fermat's 's stationary point theorem (a.k.a. interior extremum theorem), it turns that g (T*) which implies
(G1 - -a)(()((()))) = -C1 (C)
(15)
Since Çi & (2, we obtain
-((*)=(((T)
(16)
Next, we discuss the following two possible cases of 7
Caes i)
7" E E0,
It follows from 12 16 and the Lipschitzness of f' that
f ((1 7) + 7(2) f'((1 7)1- + 752)
= | + - f(()))
= ||'((1 -
- -((C))
s G| -7) +12---000 S"1
Gly -7'11< -1. G/.
(17)
13

---

Published as a conference paper at ICLR 2024
As result, it holds for byy [0.1] that
L.¹ [62 (1) 752) (52 -fif((( -7) +G2) dy
$ 162 sil L/¹ [(1 -766 + 752) - ((1 - (((1-))6 + 752) d7
3 G|62 c/12 l-'
66" -7ddy
G(2 (1) 2
(0) G (C2 - {i)
(18)
8
where (a) uses the fact that f(G) = ((G)) (b) from 17 and (c) is due to y" $ 1/2.
Case ii) 7" E [1/2. 1) Likewise, we can also have
(f + (1 n)(2) 8 (n1i + (1 n)2) = f + (1)) ((5))
(' - f(()
S G| + +1 -51
- - Sı/-
(19)
It then holds for by [0, 1] that
= [ 7'
(<>) S 9(7")
(62
-C2 (1) () xx52) dn
$ 162 Gl [[]
((11 7/s1 + 7(2) (11 + 762) dr
162 61 1-3"
(( + (1 (1- 1-1))) - ( + (1 m)(2 dn
S GK22 - 2/2 /
(1 (1-1) 11-1-3
- -000~1000 )dn
= G((2 - (1) (1-7²)2
2
G
S g (C2 21
(20)
where (a) follows by the substitution integral variable y = 1 - 17. and (b) uses 19
Combining these two cases with [14 yields the desired conclusion.
The next theorem bounds the per-step error | - -prxxxx.. utilizing Lemma A.6
Theorem A.7 (Per-step error for smooth proximal operator). Consider přox* defined by 6 with
fixed Sh. (2 1)A. Define J [ | ... v IT the vector parameterizing přox Then under Assumptions 3.7 and A. it holds for i -,..., d and k 1...., K that
p (2) přox (2: bb) A1
min max 3 =E[-A.A)
(21)
202
This bound is tight with the additional constraints that v proxR; (-A) and wh. proxR; (A), Vk, i.
14

---

Published as a conference paper at ICLR 2024
Proof. Define bb. = [proxR; .a( (-o)) a (-A + ?A.,.... ........... , (A)] be the vector
collecting the proximal function values at the partition points Sh = - - 1)A. It then follows that p (2) - přox (z; (22)
min max |prox (())) přox (2: $1 < =E|-A.A]
max
4 =E[-A.A]
Next, applying Lemma A. each piece of přox holds for V/y E [0, 1] and 1,..., C that
GIA²
p. ((1-)))t--1 - ((1 (4 21)1 - 8
ch - -i,- 1)²
20
(23)
By noticing that U -= that
E [0, = [c 6 = | A, A], we obtain from
23
G
p a (z) - ( 2a)1 -
Vz E I-A.A
(24)
20
Relating
22 to 24
leads to the desired error bound (2I).
For later use, we define distance dist ((stt(,, [a, 6]; vb,)) = max | (2) - přox (2; v5,) sE[a.b]
(25)
To illustrate this bound is tight with the additional constraints stated in Theorem A.7 a specific ex- ample will be constructed to show that the upper bound can be actually attained. To be more specific, and
will be shown that for any given C 2 1, there exists a proxx; that satisfies Assumption A. reaches the right side of (1)), with the minimizer exactly being
argmin Va.0 popoxpoox p (.!A) prc==poox
dist(l- Al;
(26)
13 15-1] -14 in the sequel. Consider the For simplicity, we drop the superseript & to write Si.c following proximal function
A
proXR; (2)
G - 2 + 20
2
C² e
G, S x < (i,041, c= 00,,, 2= Gi. S 2 < Si,tt1, c= 2,
2 (z (i,e-12 + 2,,42 (c+1)
C+1
C²
2
AGI - G (])( A) + 20
(27)
can be verified that this function satisfies Assumption A.I
by showing that
0, Gi(z - Si,),,
2 < -A
S,e S x < (1,011, c= 0., 21
prox;; (n (2)
(28)
Gi(Si.c-1 - 2).
St. S 2 < St,ct1, 1.3, 1.3....,2 C+1
-1
AGI 1g))
A
continuous, and the second-order derivate
GG,
-A otherwise
[prox"k;
(29)
0,
is bounded by G1.
In such case, the e-th element of bb, is Hie- proxR;,alae (Sr.c) = 20, 20, that
6a c. It then follows from (
přax p (b) = GIA
0 (z (A)
(30)
As result, one can have
A1A22
dist([-A,A): = dist((Gi.e-1, Srcc: b,)))
(31)
2
where the maximum (hidden inside dist; cf. (5 is attained at 2 S.e 1 +G.g for each c = 1,..., C.
15

---

Published as a conference paper at ICLR 2024
What remains is to prove (26 which relies the mathematical induction of C. The proof
starts with the base case that C 1. With the two extra constraints in effect, we already have
argmin #1. #Lo=proxx... m... (-A)
dist([- A, A); p .a( (- proxR; (A)) ab, (32)
and the minimum value dist([-A, A|: +2) G14² 20
Now, with the inductive hypothesis that (26) holds for C' = L....,
C' (C" 2 1), now prove that (26) is also true for C = C² Without loss of generality, assume C' is even so that C' + is odd. similar analysis can be readily carried out when C' is odd.
Next, we discuss the following three possible cases to determine the optimal J, that minimizes dist([- A, A); vb.). In particular, will be proved that the minimum distance of G, only in the first case, where the induction hypothesis for C' = C" + holds. 2C can be reached
Case i) b = lbbll = prox pa (Si,or) With this additional condition, it holds that
argmin
dist(1- Al;
wi ppoooro (-AA
argmnin
max (dist([-A, Al; (33)
vy =poo=ro (-A)
c= (A)
#1,0" proxxp (Si,er)
where the last equality follows from the definition (55
By applying the base inductive case (i.e., C = 1) on the interval [S. A] that contains one piece of přox (note that G,0'+1 = Si, = A by definition). it follows that
G1(A - ).c')²
G1
min
dis
(34)
# o - prxx R7 (c) Wie proxmpr (Gi.e)
20
where the second equality utilizes that ( 6,0++1- -0.0 24 C
1C uniformly partition [-A, A] so that A Gi,0"
Moreover, using the inductive hypothesis for C C' on the interval [-A, containing C" pieces of přoxt, obtain
= GilGio + +A)
G1A²
min
(35)
1. vh oro= (-A) or prrox pa (s..e))
80
20
and that the first C' elements of the minimizer of 355 are equal to [bb.... Therefore, combining
33
35
we arrive at
| prox;; . (A)] V.
(36)
which indicates that the inductive hypothesis 26) also holds for C = C" +
Next, we show that the minimum distance of 26 in the following two cases is larger than 4.4² 20
Case ii) Vi... > J. = proxR; (St,or) According to easily verify that
27
30 and that C' even, one can
prox (2) - přoxf (2: xb,) 2 0, z (.
(37a)
prox (2) - přox/ (z; (b; & 0, z A
(37b)
16

---

S Published as a conference paper at ICLR 2024
Since vi, vi,c" and vi,c Vi.cc. it follows from the definition (6 that
přox (2; (zi +/)) > přoxf v.). 2 (s),
(38)
Then, it holds that
4
= (ist( 1);
|prox přox (2; 2b))
2C
max
(a)
Si.0 + A
(5i.0r +
proxR;
přox
:255)
2
2
(b)
S1,C' +
A
S1,0' +
přox;
(b:)
proXR;,a
2
(0)
Si,C' +
A
Si,C + A
přox k
:bb,)
pi
2
2
max 26 AA
proXR a (3) přoxf (z; bb)
(39)
where (a) uses that 31 is achieved at Sie 2
tie (b) follows from (37b), and (c) is due to
38
In other words, if o' Vi it must hold that
-
dist((C' (3:))))
8C2
which is larger than that case i). Therefore, the optimal of must satisfy Vi, S Vi.c..
Case iii) oi < Vi. = prox:; (a,,0
Again with 31 one can easily get
A
202 = (i:t(((s.cr-11:::c)::)))
(1.00-1+ (iC") přox! (4 :b..)
S..c"
= proxo;; a
(40)
Recall from the definition c that (2; 2 (1. is defined as the line seg- ment connecting points (G, 11,10 and (G,,r, v.)))). To ensure
G 20 a necessary condition is přox (5.0r-1thh ab!) 2 přox (G..rr-tts +5,00 /b. cf. (37a).
Since we have fi ehi,0" in this case, it must hold that 1 - > - By ap- plying this analysis recursively, one can proceed to obtain a series of necessary conditions of dist(I-A, < G.4² 20 which are (recall that O' is presumed even)
1h > bi,0-11, Vi-c--2 < 22.0'-2
(41)
(42)
1 V1,1,
v0 Vi.0-
(43)
This contradicts with the constraint that v.. = prox a (-A) = D.u; ef. (26). That is to say, G,4²
requiring Vi..r < o.. will lead to dist( (lst([-A,,, Al; xb)) >
202 which is not optimal.
To the end, through the three cases we conclude that the minimizer tp must satisfy vi, 4/,.0 which implies (26 holds for C" + 1; see (36). The proof is thus completed.
Building upon the per-step error bound established in Theorem A.7 the K K-step cumulative error bound will next be proved. In particular, the following theorem offers an upper bound for the normalized error on (shifted) 0,.
Theorem A.8 (Formal statement: finite-step PGD enror for smooth proximal operators). Consider přox* defined by (6) with fixed Ct: = - 1)A. Define V == (ab'....,,, ...., 1 the matrix param- eterizing {přox*)K J x-1- Let 0. and 0K be the K K-step PGD outputs using prox and přox*,
17

---

Published as a conference paper at ICLR 2024
respectively. With Assumptions 3.7 and A.7
A.4 A.7
in effect, it holds for 1,...,I that
min
0 v - 0 (
(44)
This bound is tight with the additional constraints that v (alxR; and vf.cc proxR; (A), Vk,
Proof. For notational compactness, define $ argming max-A1 [| ,a (z) přox (z; *b*))= and v [ ..... 3 J. Since both proxR aa and přox* factorable across the dimensions of their inputs, we know that &O is the concatenation of the minimizers
R = argmin max |ppoox;;
přoxf (z: 112122, i = 1,...,d
# =E[-A.A]
for each dimension. It then holds for h: 1,...,, K that
min |00 05( (1)l2
10, 0 ( 1111 = |poOOx aa (zt) - přox, (ž;; أ²>"" )l1
p aa (z+) proxx. ,a (z.) + ||prOxR.... .a (Et) přoxf (zt; 30)12
| (,)))) proXR-,. (a(22, 112 + V
202
2 11 (1)1 + V@AAG1
202
( 1.100=1 ا
vdA²Gi
1|2 + +aT.| 80 0.-1 ;
2C2
(4)
vdA²G
< I(1 + 0 - 0: (
(45)
2C
where (a) follows from Theorem A.7 and Assumptions A.1 A.3 uses Assumption A.4 from and 7 and (d) is due to Assumption A.2
(c) is
Using this recursive relationship between |0 line), together with the boundary condition 0 01112 = 00, - Oallz = 0, we arrive at the solution
0 1 l| (line 2) and |0 k ي
|0)l (the last
(1+G.) L(+4GG)
vaAzG. 20
if L(1 + oG2) + 0
Vd C
(46)
LACA 20
otherwise
Dividing by vd and minimizing v on both side of 110, - 0. v 112 - 100¹
05 ||
lead to
min 1
Vd o. o. S
1 V Io 0 0 c2
1
(47)
Plugging in (46) with A: K gives 9
In addition, the tightness 9 follows from the tightness of Theorem A.7
B PROOF OF THEOREM 3.3
The proof of Theorem 3.3 relies on a more generic lemma which can be applied to non-smooth (but still Lipschitz) proxR-
Lemma B.1. Let f E (R): R + R be an L-Lipschitzz function. For VC1, Ca E IR and G # Ca define
f(z) = (Ce ==))f((())) + (2 -f)f((C)
(48)
C2 S1
18

---

Published as a conference paper at ICLR 2024
It then holds for vy [0, 1] that
| (1(1-))61 + 752) ((1--))
+ + - 2 62 Sil-
(49)
Proof. Define g(7) = [
+ + Following the same step (4)
of Lemma A.6. it can be shown that there also exists at least one maximizer > E (0, 1) of g(7).
Thus obtain
5(7) S ,(>")
(4)
([(11)))(++'' -10-++++))-111111100 (1100-7/(+110)001--133---=1---0-
< (1 (1- - - +1)) + 7 - f((2)
(6)
2 2y 4S2 Gi
- L
2 K<2 -cl
(50)
where (a) follows from the definition 488 of f. (b) exploits the Lipschitzness of f. and (c) due to that y* (1 -²) 5 - 1/4 for y 1).
With Lemma B.I at hand, Theorem 3.3 can be proved using similar techniques as Theorem 3.2
Theorem B.2 (Formal statement: sinite-step PGD error for continuous proximal operators), Con- With Assumptions 3.7 and A.2 A.4
sider the notations defined in Theorem 3.2 t == -...., T that
4.41o in effect, it holds for
min V
o (w)122
(51)
This bound is tight with the additional constraints that v proxR; a (A), Vk, i.
(roxx;aa(-A0 and whe
=
Proof. Following the same steps of Theorem A.7 it can be shown that the per-step error bound for continuous prox is
AL
min max S A.A]
proX (2) (x)))
(b))
<
Vi.
(52)
C
Also, this bound is tight provided with the additional constraints that w. = prox (-A) and d/.cc proxR; (A).
Likewise, we define $² == argming max ALs |proXr ,a (z) přox (z; 26 ( || and L
1 ....,.b" K Then, it follows from the first two lines of 45) that
min 10, (v) < ||poOx (z!) proxR aa (2)) l12 + proXR aa (2t)) přox* prox"(2t;e
(a) T.
VaAL
b
क (
|11 +
C
(6) L(1 + -G2)
k-1 - + VdAL
(53)
where (a) is from Assumption A.4 and
52 and (b) utilizes and Assumption A.2
To the end, this recursive relationship, combined with the condition 110, 0 l2 0, results in
1-4" (=2"()+))))* VeAL
if + +G2) + o Vg.
Vd
min 10000 0. (+) 5
1-L11+
(54)
C
K 2
otherwise
Dividing both sides by vdd completes the proof.
19

---

Published as a conference paper at ICLR 2024
C UPPER BOUND OF 10
3.3 As discussed Sec. feasible approximation the sub-optimal & obtained from the ERM T. which brings about an extra error term Va 0 (T) 0. (I )12: (10). This section derives its upper bound, based on the following extra assumption.
the optimal & that minimizes 9 and is typically unavailable.
Assumption C.1. přox* (z; w²* (1 is L1- and Lz-Lipschitz w.r.t z and ab"h for AAla 3 2 3 Ala, respectively.
Denoting by Φ and whth the k-th columns of L and & that parametrize přox*, it holds for ,,....... K that
00 ) (.((1)
(
p (2) (1); *²²) přox" (2) (VV)) (b *12
p (Z) ());; "²) přox* (z) ()); Il2 + přox (2) (Z) (/*)) přox (Z)((()) *b>))
(
L2lle" : + 112 + L1 1 (i) 전 112
1/2111 x - 16" k 1/2 + Tn (10) (ي) - 0, (Tr") 1|2+
V0 2(0) (((--1(1) D! V L ) )l2)
L2 26 L1 aG2) 0 ( 0:
)|l
(8)
- (1 + oG22)l 11)) 0: I|le + Ln
||
(55)
where (a) follows from (b). (b) uses Assumption C.1 (c) from
(7 and Assumption A.2
and
(d) is due to that |
*k 2 S max
2 nk
12 = 1
Solving the recursive relationship
55 using (0)( 0: v
")2 gives
1- Li (1+aG2)* x
Lz L2Il v Tl if L aG2) +
0 K
1 (i(taaaa))
(56)
K KL2 v
11,
otherwise
which concludes that
-0) =0 -
& 1'1l1).
(57)
D ADDITIONAL REMARKS REGARDING THE THEORETICAL RESULTS
Next, three important remarks regarding the derived error bounds will be provided. Remark D. (Difference with convergence rate analysis). It is worth stressing these approximation bounds are different from the convergence rate analysis. Essentially, it quantifies the impact of using parametric přox to approximate the optimal yet unknown proxæ ,ar Moreover, although the bounds increase with K, it important to note that K is a sufficiently small constant (typically 1 < K < 5) in the context of meta-learning (Finn et al., 2017). as the overall complexity for solving (2) scales linearly with K. Furthermore, the learning rate must satisfy a E (0. 2/Gz) to guarantee gradient-related descent direction, with consequence of this choice is that (1 + aG2) K E [2, 32]. which ensures that the constant in the upper bounds will not diverge.
1/G₂ being the optimal choice. A
Remark D.2 (Factorability and scalability). Assumption 3.1 ensures that the prior dimension D scales with the task-specific parameter dimension d. As d in practice can be extremely large (e.g., 9(10a)), )), complete prior such as full Gaussian pdf would incur prohibitively high complexity; that is, D = ((d)) = 10) A feasible simplification is to approximate the prior in Rd using the multiplication of d pdfs in R. This assumption essentially considers each dimension of 0, to be mutually independent, leading to D = (d)). Such an independence assumption is prevalent not only in meta-learning, but also in high-dimensional statistics when dealing with deep NNs.
20

---

Published as a conference paper at ICLR 2024
Remark D.3 (Validity of learned proximal operator). The learnable prox in this paper mains a proximal operator, even when the corresponding regularizer is non-convex. Indeed, let prox (x: 0) == arg ming [|z]]] prox((;; 0) = x} for {prox(z: |), 3 A1,) = X, and Xo = (2; 0) E X. The stationary point condition of the proximal operator indi- cates prox (xo: 0) Xo E ax(: I); thus, one of the regularizers satisfying this condition is R(0;; 0) = J xxxxxxx00 (prox- (x; 0) x)dx. Compared to non-expansive operators, proximal operators induced by non-convex regularizers have gained popularity in recent years thanks to their enhanced expressiveness, and their convergence guarantees (Hurault al., 2022). Our method fails precisely within this category. Additionally, the PLFs should be monotone to qualify as a valid proximal operator; see e.g., (Gribonval & Nikolova [2000 Theorem 1). In practice, the sought monotonicity can be established by enforcing V. S Wi,er Vcc S c'. This can be readily achieved upon defining v = v/c exp(At/i Vii c, ki and then learn the log-increment (A Interestingly, we have observed that the learned PLFs are exactly monotone functions (see e.g., Fig. even without an explicit constraint. This observation suggests an inherent preference for monotonic PL.Fs by the data.
(xo:
E DETAILED SETUPS OF NUMERICAL TESTS
In this section, we introduce the dataset and elaborate the detailed setups of the numerical tests.
The minilmage dataset (Vinyals et al. 2016) consists of 60, 000 natural images sampled from the full ImageNet (ILSVRC-12) 12) dataset. These images are categorized into 100 classes, each with 600 labeled samples. As suggested by (Ravi & Larochelle| 2017). all images are cropped and resized to size 84 x 84. The dataset is split into disjoint groups containing 64, 16 and 20 classes, which can be respectively accessed during the training, validation, and testing phases of meta- meta-learning. The experimental setups follow from the standard M-class --shot few-shot learning protocol (Ravi & Larochelle, 2017] Finn et al. 2017). Specifically, D!nm per task includes M classes randomly drawn from the dataset, each containing N labeled data. As a result, it is clear that |Dirm MN for each t.
The TieredImageNet (Ren et al. 2018) dataset is a larger subset of the ImageNet dataset, composed of 779, 165 images from 608 classes. Likewise, all the images are preprocessed have size 84 84. Instead of using random split, classes are partitioned into 34 categories according to the hierarchy of ImageNet dataset. Each category contains 10 to 30 classes. These categories are further grouped into 3 different sets: 20 for training, 6 for validation, and 4 for testing.
We utilized the group of hyperparameters described in MAML (Finn et al. 2017 consistently throughout all the tests. To be specific, the maximum number K of PGD steps is 5, and the total number R of mini-batch SGD iterations (8) is 60, 000. The number of convolutional channels is 64 for MetaProxNet+MAMM. and 128 for MetaProxNet+MC The learning rates for PGD and SGD are a = 0.01 and 3 0.001, with batch size tieredImageNett while SGD with Nesterov momentum of 0.9 and weight decay of 10 is used for minilmageNet.
Adam optimizer employed for
The interval A-A, A] and number C' of pieces are determined through grid search leveraging the validation tasks. For both minilmageNet and TieredImgeNet datasets, A = 0.02 and C' = 5. We found that C suffices to reach a satisfactory performance, while larger C only contributes marginally to MetaProxNet's empirical performance. This suggests the constants hidden inside the error bounds (11/C) and (11C22) can be sufficiently small in practice.
F COMPLEXITY ANALYSIS OF METAPROXNET
It can be observed from and (b) that the per-step piecewise linear function (PLF) prox* applies a dimension-wiin affine transformation its input. Although the per-dimension přox/ consists of
C ++ parameters, the affine transformation merely relies on a single piece [S), <t of the PLF ........d. complexity dKd)). While the K K-step GD of (7a) also exhibits forward and backward complexities of O (Kd), its constant hidden within O can be much larger, as the convolutional operations in the
Thus, each computation involves only two control points vhth v/f..-1. v6 for every l = 1, 1...., K and i = 1......d d. As a result, the forward calculation and backward differentiation of PLFs both incur
21

---

Published as a conference paper at ICLR 2024
Table 3: Comparison of (normalized) running time on the 5-class 5-shot minilmageNeet dataset.
Method
Forward
Backward
Total
MAMI,
x0.182
x0.818
x1 (reference)
MctaProxNet+MAML
x0.202
x0,836
x1.038
MetaCurvature
x0.188
x0.834
x1.022
NetaProxNee+-MeaCCuvauuue
x0.208
x0.863
x1.071
Table 4: Ablation test regarding weight untying using 4-layer 32-channel CNN.
5-class
10-class
Method
1-shot (%)
5-shot (%)
1-shot (%)
5-shot (%)
MAMI
48.70 11.4
63.11 10.92
31.27
46.92 11.25
MetaProxNet (shared přox) + MAML
51. 51.1611.45
67.081.0 10.71
34.62 10.93
50.26 10.48
MetaProxNet (per-step pÊox;) + MAML
53.5841.43
67. 67.88+0.72
34.80=0.91
51.03
CNN are more time-consuming than the affine ones. Consequently, PLFs contribute only marginally (< 5%) to the overall complexity compared to their backbone. This is also evidenced numerically in Tab. gl Moreover, Assumption (3.1) indeed ensures that the prior dimension D scales with d, therehy mitigating any substantial complexity increase. Specifically, in practical scenarios where d extremely large (e.g., 121, 093 or 463, 365 in our experiments), employing complete prior, such full Gaussian pdf, would yield D (df)). A feasible simplification is to approximate the prior in R& using the multiplication of d pdfs in I, which leads to D O(d).
G EXTRA ABLATION TESTS REGARDING WEIGHT UNTYING
The next test examines the effectiveness of the weight-untying technique; i.e., the per-step přox. The experiment is conducted on the minilmageNet dataset with 4-layer 32-channel CNN, and the corresponding results are summarized in Tab. 4 It is seen that the per-step proximal operator consistently outperforms the shared one in all four tests. In fact, the per-step přox,, inherently cor- responds to an adaptive prior, which evolves with the optimization process. The same technique was originally provided by the renowned LISTA algorithm (Gregor & LeCun 2010), which pioneered algorithm unrolling, and has since been widely adopted by the community on inverse problems.
H NUMERICAL VERIFICATION OF ERROR BOUNDS
Next, a toy numerical test is carried out verify the derived PGD error bounds. For simplicity. we will exclusively focus on Theorem 3.2 while similar analysis can be readily applied to Theorem 3.3 Consider tasks defined by the linear relationship ye w x + eg. and a linear prediction model
Ur = fx;; ; 0.) 0. x² with squared Eg loss L(O,; Dr = 1 Σ Nara
llv!"
0, x;" 12 where the
2
Numerical Theoretical bound
K-step PGD eroo
100
10
10
10!
Figure 5: Comparison of theoretical error bound
9 and numerical error in linear regression.
22

---

Published as a conference paper at ICLR 2024
unknown oracle w ~ Uniform (1-3, 3| and x ~ (N Lz). In the test, we set K = 5, 0.01, 64, |D 8, 0 0a, A = 3 and C varying from to 20. The target optimal proximal operator to be approximated is defined in (27 with G1 = 1. Additionally, the Lipschitz constant G2
of ( mi is numerically computed from the randomly generated training data matrix X!!"M.
The plot comparing the numerical PGD error and its upper bound can be found in Fig. 5 It is
observed that the numerical error aligns with the theoretical bound up to small constant in the log- scale. This discrepancy arises because the upper bound considers the worst-case scenario, where the largest error between proxR',0 .a and prox* is reached at each PGD step. Furthermore, this constant gap suggests that the numerical eπror is in the same order with the bound (notice that both axes are in log-scale); that is, (g2)). This empirically corroborates our theoretical proofs.
I CASE STUDY: FEW-SHOT REGRESSION
Next, a straightforward yet illuminating numerical case study is provided to show the claimed su- perior prior expressiveness. Consider few-shot regression tasks defined by the linear data model 3/" = w x + ef, where the unknown per-task weights {w] are ii.d. samples from the oracle pdf p(w"() ) := 1/2 [nnifrm([[111.. 10111-110000 - Uniform [10, 11)"). and et ~ -V(00 o²), Vt, is the ad- ditive white Gaussian noise. Further, consider linear prediction model o7 = f(x" 0x) == 0 x²- Since p(w") ) is symmetric and isotropic, the optimal Gaussian prior of 0, for this case must have a mean of 0₂ an isotropic covariance. In other words. if the prior pdf is chosen to be Gaussian p(0,; 0) Ni,, Σ), == [n. T, vec(E)]]] its optimal parameter 0* must consist of Ju' 0₂ and E = I-I for some A E R. As a result, the corresponding regularizer for this prior is R(0,; = +110,11: which prevents O, deviating far from O2. However, the ground-truth task model implies that the optimal 0; should belong to the set S ularizer is thus a barrier for optimizing 0, In contrast, if the prior pdf is allowed to be non-Gaussian the optimal prior will be the ground-truth one, i.e., p(0;; 0¹)
11]1 11j6, and the reg-
Uniform(S). The corresponding proximal operator in this case is the projection proxæ- (z) = I' (z). which is exactly a piecewise linear function driving Og to the oracle set. In summary, the Gaussian pdf fails to match the underly- ing prior due to its inherent unimodality, while the more expressive PLF-inducee prior can perfectly align with the groundtruth prior to enhance the lask- learning.
For visualization purpose, numerical test is carried out with d = 1. The remaining parameters are [D!'i = 2, |DY
5, de 1, K = 5, R 10, 000, a = 0.1, / 0.01, A = 15, C 30, B- = 4 and x. ~ (0, 1). In Fig. 6 the linear function learned by MAML is inclined to have a slope close to O, while the PLFs in MetaProxNet quickly refine 0g into the set S = i-11, -10 U [10, 11j. This empirical observation substantiates the advocated superior prior expressiveness.
In practical tasks such as drug discovery and robotic manipulations, the oracle model parameters can have 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.
(+)
peont]
10
10
-0.b
-10
Groundruth MANI
-MMa-rraeaa Training dats
(a)
(b)
Figure 6: Visualization of 1-d few-shot regression case study; (a) comparison of MAML and MAML+MetuProxXee (b) the learned PLF-based proximal operator averaged across PGD steps.
23

---

Published as a conference paper at ICLR 2024
J A BRIEF INTRODUCTION TO ALGORITHM UNROLLING
Algorithm unrolling was first introduced in (Greegor & LeCun 2010) to solve the inverse problem. In particular, it aims to recover (transformed) signal x E I from its compressed measurements
y Ax e
(58)
where A E RRmx is a given matrix with m < n. and e is additive white Gaussian noise. Since the system (58) is under-determined, it has infinitely many solutions. To ensure the uniqueness of the solution, a prudent remedy is to rely on the prior p(x), which yields
& argmin lly
Ax 2 + R(x).
(59)
In the above, lyy - Ax and R(x) correspond to the nll log p(y: x) and nlp log p(x), spectively. As nature signals are inherently sparse in certain transform domains such as Fourier and wavelet ones, a popular choice is the sparse prior with R(x) X|x|||1. With such prior,
the resultant optimization problem (59) can be efficiently solved by the well-documented iterative soft-thresholding algorithm (ISTA), which involves a two-step update rule
z x aA (Ax*
x y) = (Ln aA A) A)x + AAAy
(60a)
x argmin
||
x1 ||||| = (z*", Ao
K
(60b)
2a
Here, S&A is the soft-thresholding operator shown in Fig. 2b
When given dataset {(Xm, (n ))) it is possible enhance the accuracy of x by learning more efficient optimization rule. In (Gregor& & LeCun 2010), the two steps (60) of ISTA for each
k l..... K are replaces by two learnable NN blocks
z" W * + W y.y
(61a)
x = S5: (x*)
(61b)
with 0 == ((W W ) being the learnable weights of the NN. Denoting by fly; 0) this multi-block NN mapping, the NN-based update rule can be learned via
x, y.
min Σ xx. - fly::
(62)
This method of unfolding and substituting the iterations of an optimization algorithm to form multi-block NN is known as algorithm unrolling, and the resultant NN is termed an unrolled NN.
K RELATED WORK
NN-based meta-learnin:: Recurrent neural network (RNN) has been introduced in (Hochreiter et al. 200T) to learn the update rule of the task-specific model parameters O,, with prior encoded in the RNN's weights. Following this work, different RNN architechtures have been explored to
enhance the learning of the update rules. On the one gradient information has been leveraged in (Andrychowicz et al. 2016: Li & Malik 2017 Ravi & Larochelle 2017) to mimic the gradient-
based optimization. On the other hand, temporal convolutions and soft attention have been utilized to aggregate and pinpoint information from past experiences (Mishra al.] 2018). More recently, this paradigm of NN-based optimization has been extended to Beyesian meta-learning. aiming to infer the posterior pdf p(0, p(O,|y!'n, X!t)) (Gordon et al., 2019). Due to the blackbox nature of the RNNs however, it is hard to interpret the impact of learned prior from the NN-based update rules.
Optimization-base based meta-learning: To empower fast adaptation to a new task, MAML capital- ized on learning task-invariant initialization (Finn et al. (017). with task-level learning defined by cascade of few GD steps on the task-specific parameter 0,. Intuitively, by descending a small number of steps, 0, should not deviate too far away from its initial value 9. In fact, it has been pointed out in (Grant et al. 2018) that MAML's GD solver satisfies (0) (,( argming, L(0,; Dirm) + 1, 0T2,. where A, is determined by hyperparamters of GD. This observation indicates that MAML's optimization strategy approximates an implicit Gaussian prior 0) (00 1) with initialization ginit serving the mean vector. Following MAMI.,
24

---

Published as a conference paper at ICLR 2024
spectrum of algorithms have been developed to encode different priors. For instance, Meta- SGD (Li et al., 2017) augments MAML by meta-learning a dimension-wise step size, which es- sentially corresponds to per-step diagonal Gaussian prior. Other examples of the induced priors include isotropic Gaussian (Rajeswaran et al. 2019 Abbas et et. al. diagonal Gaussian (Ravi &
Beatson 2019 Nguyen et 2020), per-step block-diagonal Gaussian (Park & Oliva 2019 Flen-
nerhag et al. (020), and implicit Gaussian (Baik et al., 20200 (021).
Another line of research termed metric-based meta-learning (which can be either NN-based or optimization- based) splits the model into an embedding "body" and classifierreegresoor "head," and learn their priors independently (Vinyals et al., 2016 [nell et al., 2017] [0n7/Sung Sung al.) [2018] Li al. 2020a). In particular, with o and 0bead denoting the corresponding partitions of 0,, the prior is presumed factorable as p(0,; 0) p[Ohody, boody )p(Oread, @head),
On the one hand, the head typically has nontrivial prior such the Gaussian one (Bertinetto al. 2019 Lee et al. (019)) On the other hand, the body's 's prior is intentionally restricted to a degenerate pdf p(O,ady bobdy
[0 [@body T @headd T
5(0-ody , gbody), ). where 5(.) is the Dirac delta function. Although freez- ing the body in task-level optimization remarkably reduces its complexity, it often leads to degraded performance compared to the full GD update (Raghu et al. 2020). In additional to degenerate priors, sparse priors have also been recently investigated to selectively update a subset of parameters (Lee & Choi 2018 Tian al. 2020a)).
Compared to these preset prior pdfs of fixed shapes, the focus of this work is to learn a data-driven prior pdf that can dynamically adjust itself to fit the given tasks.
Algorithm unrolling: The advocated MetaProxNNt pertains to the algorithm unrolling cate- gory (Gregor & LeCun [2000] Monga et al. [202] Li et et 2020b). Closely related our work is the deep regularization approach introduced in (Li al. 2020a). This method shares similar high-level idea of incorporating priors into PGD optimization iterations through algorithm unrolling. In (Li al. 2020a)), the hidden representations are transformed into domain conducive easy regular- ization by a predefined prior pdf (c.g.. isotropic Gaussian). In contrast, our MetaProxNet approach involves the direct learning of the proximal operator within the parametric space.
25


Reviews:
[{"review_id": "KvuLVOYsKw", "reviewer": "Reviewer_5m7L", "summary": "The authors present an approach for meta-learning that estimates per-task model\nparameters with an unrolled neural network. The neural network is derived from\nproximal gradient descent on a regularized ERM problem (per-task); the\nregularization term is learned by representing its proximal operator, which is\nparameterized as a piecewise linear function that separates over its\ncoordinates. A theorem is given that bounds approximation error between\npiecewise linear operators of this type and general $C^1$ maps asymptotically in\nterms of the number of discretization points. Experimental results demonstrate\nimproved performance over certain meta-learning benchmarks on few-shot image\nclassification, and visualize the proximal operators learned in these tasks.", "questions": "- The authors motivate their method at the start of section 3.1 by contrasting\n  their method, which learns a prior over model parameters\n  $\\boldsymbol{\\theta}_t$, with prior work on algorithm unrolling, where the focus\n  is to learn a prior over signals $\\mathbf{x}_t$ for one fixed task. It would\n  be helpful in supporting this point if the authors presented a concrete\n  example of a class of tasks where this difference can be appreciated. For\n  example, what form does this distinction take in a family of Gaussian\n  denoising tasks -- a typical setting for algorithm unrolling? I think this\n  would be helpful, as the presentation of the computational methodology in the\n  rest of the section tends towards abstraction.\n\n- The notion of an \"optimal\" regularizer $\\mathcal{R}^*$ is introduced on page\n  4, and plays a role in the authors' methodology. In what sense is this\n  regularizer \"optimal\"? I could not find any discussion of this notion in the\n  background section.\n\n- Is the content of Theorem 3.2 mainly an application of results on\n  one-dimensional piecewise linear function approximation (with uniform\n  discretization) of $C^1$ functions?\n\n- What is the \"apples-to-apples\" comparison mentioned when discussing the\n  experimental results on page 8? Are there SotA meta-learning methods that have\n  been excluded from this table? It seems natural that the MetaProx approach\n  might perform better than methods that use less computation, since Algorithm 2\n  could be interpreted for natural losses $\\mathcal{L}$ as a multilayer residual\n  network (with some per-task conditioning) with learnable nonlinearities, if\n  I'm not mistaken.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a clear and robust overview of background and prior work on\n  meta-learning, making the authors' contributions easy to understand and\n  appreciate in-context.\n\n- The paper features a mixture of theory, justifying the approach, and\n  experiment.\n\n- Experimental results demonstrate the approach performs favorably on few-shot\n  image classification with standard meta-learning datasets against standard\n  benchmarks.\n\n- The results on learned nonlinearities in section 4.4 are an interesting\n  consequence of the authors' method. In interpreting these results, it would be\n  helpful to know how the $\\zeta$ parameters are initialized.", "weaknesses": "- The conceptual contribution is not exactly made clear. Maybe the clearest\n  manifestation of this is that the presentation of the\n  methodology does not make it clear why the authors' approach should be\n  superior to other approaches to meta learning (say, to solving the formulation\n  in equation (1)). The methodological discussion emphasizes that several\n  commonly-used meta learning priors (under a common assumption on the \"optimal\n  regularizer\") can be associated to proximal operators that can be represented\n  as coordinatewise-applied piecewise linear functions. However, it does not\n  connect back to specific tasks of interest in practice mentioned in the intro\n  (robotics, say) and argue that the parametrized piecewise-linear proximal\n  operators (essentially learnable nonlinearities in a neural network that\n  implements the base-learner) are useful given specific structures the model\n  parameters $\\boldsymbol{\\theta}_t$ have in these tasks. I think this is\n  important to be able to support the authors claims of \"expressivity\" and\n  \"interpretability\" made in (i) of the claimed contributions.\n\n- Labeling the algorithm \"MetaProx\" seems like it might be a slight misnomer --\n  based on the discussion on page 5 (after Assumption 3.1), it seems that what\n  the authors are implementing is rather an amalgamation of one-dimensional\n  picewise linear functions, which need not correspond to the prox of any\n  function. In one dimension, it should be sufficient and necessary here to\n  constrain the piecewise linear function to be monotone.\n\n- It is not completely clear how the authors' eventual MetaProx algorithm\n  differs from other meta learning approaches; most of the authors' discussion\n  emphasizing the novelty of their approach seems to focus on contrasting with\n  prior works on algorithm unrolling. For example, is this the first work that\n  has made a connection between the problem in equation (1) and proximal\n  gradient descent? A small section discussing this after the method has been\n  presented would be helpful.\n\n\n### Minor Issues\n\n- Page 3: \"...the optimization-based approaches solve (2b) ...\" -- should this\n  be (1b)? The rest of this setence discusses $\\mathcal{R}$, which does not\n  appear in equation (2).\n- Text before equation (11) on page 7: seems to be a dangling clause, \"...has\n  gained attention in various PGD-guided a [sic?].\"\n- Notation for continuously differentiable functions: as I understand it, using\n  $\\mathbb{C}^k$ for this class is nonstandard (and clashes with the usual\n  notation for $k$-tuples of complex numbers in a jarring way). I guess this is\n  an attempt to avoid a conflict with the parameter $C$ in the authors'\n  algorithm; what about using $\\mathcal{C}^k$ or $\\mathsf{C}^k$ for this class?\n- Is it correct that Theorem 3.2 involves Assumption 3.1? If so, it would be\n  good style to reference this in the theorem statement.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OvYaWuUQ8I", "reviewer": "Reviewer_Da9s", "summary": "This paper presents a meta learning algorithm called MetaProx. When addressing meta learning algorithm with so-called proximal gradient descent, a central idea of this paper is to apply algorithmic unrolling. Algorithmic unrolling solves iterative optimization method by learning a deep neural network by cascading the steps of the iterative process. By applying algorithmic unrolling to proximal gradient descent, the paper attempts to enhance expressiveness and interpretability of the prior. Piecewise linear functions are chosen as proximal operators, for which, the paper provides error bounds theortically. In few-shot learning benchmarks, the paper shows advantages of the proposed method.", "questions": "Questions and suggestions have been addressed in the section above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I noted the following strenghts:\n\n1. The idea of bringing algorithmic unrolling to proximal gradient descent is new and interesting.\n\n2. The paper derives theoretic error bounds for both smooth and non-smooth proximal operator, which makes the approach theoretically sound with certain level of technical depth.\n\n3. The proposed method seems to perform well in the few shot benchmark. Improvements are consistent.", "weaknesses": "I thought following points could improve the paper.\n\n1. I wonder if the dervied theoretic bounds can be validated empirically and the results appear in appendix.\n\nIn 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.\n\n2. There are rooms to improve the presentation.\n\nIntroduction: Only few sentences are devoted to describe the contribution of the paper, while four paragraphs are devoted to motivation for this work. Within the descriptions of the contributions, most of the words are spent in giving benefits of the approach. I think this should appear after explaining what the actual approach is, e.g., explain what algorithm unrolling technique is, etc.\n\nRelated work: The paper misses related work section. Therefore, it was difficult to locate the contributions of the paper within the state of the art. \n\nSection 3.1: I had to read other papers to understand algorithmic unrolling. Looking back, the provided explanations seem rather difficult and this paragraph might be throughly revised.\n\n- The part on explainability may not be clear.\n\nYes, algorithmic unrolling can add interpretability by examining each layer mapping. But, the results look rather vague and I wasnt sure how this explainable prior can be useful in practice.\n\n----------------------------------------------\n\nOverall, the paper has many ideas and I am leaning towards recommending accept. I hope that my concerns are addressed during rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "40Yj77v0Mm", "reviewer": "Reviewer_W1mq", "summary": "This work presents a novel approach to meta-learning dubbed MetaProx, which leverages unrolled proximal neural networks to learn flexible and interpretable priors. The core concepts of MetaProx involve decompose the proximal gradient descent algorithm into a multi-block neural network, with each block comprising a data consistency module and a trainable piecewise linear function that approximates the proximal operator. \n\nTheoretical analysis of this approximation demonstrates its capability to effectively model smooth and non-smooth proximal functions. Empirical evaluations of MetaProx in the context of few-shot classification tasks reveal its superior performance compared to other SOTA meta-learning techniques employing either RNN-base or handcrafted priors.", "questions": "1. How do you anticipate the performance of MetaProx on few-shot regression or reinforcement learning tasks? Is it ready to be extend to these domains? Are there any challenges to be addressed?\n2. Do you have any ideas to help interpret the high-dimensional learned prior, beyond visualizing the 1D components? Would it be possible to visual the evolution of learned prior after each unrolled PGD steps?\n3. Does MetaProx require careful initialization or scheduling of $\\lambda$ to avoid instability during training? \n4. The method mainly evaluated on 5-way 1/5-shot classification benchmarks. How do you expect MetaProx to perform on more challenging few-shot settings, like 20-way 1/5-shot?\n5. Is there trade-off between the regularization term and the loss function in eq (1b)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea of learning priors through unrolled PGD is a novel and intriguing approach. This method facilitate the model to learn a broader range of complex and adaptable priors in a data-driven manner, thanks to the inclusion of non-smooth and even discontinuous functions in the regularization term of the proximal operator. Furthermore, it enhances both model expressiveness and interpretability when compared to traditional meta-learning approaches that rely on predefined priors. For example, in optimization-based meta-learning methods like LEO and ANIL, an explicit Gaussian or Laplace prior is imposed on model parameters during inner-loop adaptation. To maintain computational efficiency and feasibility, the authors employ piecewise linear functions to approximate the proximal operator. Additionally, they further reduce costs by employing fixed discretization points for the piecewise linear functions instead of simultaneous learning them. \n2. The paper provides a rigorous theoretical analysis bounding the approximation errors, offering valuable practical guidelines for hyper-parameter tuning, like C.\n3. Comprehensive empirical evaluation on benchmark datasets and comparisons with state-of-the-art meta-learning methods. The consistent performance gains validate the effectiveness of MetaProx.", "weaknesses": "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.\n\nThe 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.\n\nFurthermore, 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.\n\n[1]. Mao Li, Yingyi Ma, and Xinhua Zhang. Proximal mapping for deep regularization. Advances in Neural Information Processing Systems 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eQFW4FNlgg", "reviewer": "Reviewer_YbkU", "summary": "The paper considers optimization-based meta-learning algorithms, i.e., algorithms that employ a bi-level optimization strategy to optimize task-specific and task-global parameters. This approach can be interpreted as the task-specific parameters being regularized towards a task-global prior optimized in the outer loop. The authors propose a novel method to induce more general/expressive priors in comparison to related work, by learning proximal operators for proximal gradient descent using unrolled NNs. The authors motivate a piecewise linear parametrization of this proximal operator, derive an algorithm to optimize the corresponding parameters, and analyze error bounds. They compare their approach against a range of optimization-based meta-learning algorithms.", "questions": "cf. weaknesses", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written and both the theoretical exposition as well as the experimental evaluation seem to be well fleshed out. Algorithmic choices such as using piecewise linear approximations for the proximal operator appear well-motivated. The experiments contain comparisons to a wide range of optimization-based meta-learning algorithms, demonstrating superior performance of the proposed method on the miniImageNet benchmark. Further ablations underline the effectiveness of algorithmic design choices. In summary, the paper appears as an interesting and effective approach to derive priors for optimization-based meta-learning. \n\nTherefore, I vote for acceptance. However, my recommendation is with low confidence, so I might decrease my score during the rebuttal period.", "weaknesses": "The paper could be further strengthend by providing experimental results on benchmarks other than miniImageNet. Furthermore, I would appreciate a discussion of how the method relates to and/or could be extended to fully Bayesian approaches like Probabilistic MAML [1] or Bayesian MAML [2].\n\n[1] Kim et al., Bayesian Model-Agnostic Meta-Learning, NeurIPS 2018\n\n[2] Finn et al., Probabilistic Model-Agnostic Meta-Learning, NeurIPS 2018", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)