Paper: Variance-Reduced Meta-Learning via Laplace Approximation (ICLR-2024)

Reviewer comment (Experiments):
Task overlapping needs more examples. Because this is the core to this work, it is beneficial to provide more real-world examples. Current sinusoidal example appears to be oversimplified and not convincing enough.

Author response:
We thank the reviewer for their comments and wish to address their concerns regarding real-world experiments below.

We have included experiments on mini-ImageNet in the Appendix, and outline that real-world experiments such as classification fall outside the scope of the paper as they induce a discrete distribution of tasks. We agree with the reviewer that additional examples of task-overlap would help clarify the motivation of the paper.

Reviewer comment (Writing):
The notation keeps changing, which is very unclear. Please make notation clear and consistent.

Author response:
We agree with the reviewer that the presentation can be clarified. We have updated the notation and type of references in a revised version of the paper and hope the clarity has improved. We have additionally edited the objective described in Equation 1.

Reviewer comment (Writing):
Eq 1. is not clear. It is better to rewrite the objective.

Author response:
We agree with the reviewer that the presentation can be clarified. We have updated the notation and type of references in a revised version of the paper and hope the clarity has improved. We have additionally edited the objective described in Equation 1.

Reviewer comment (Theory):
Eq 3. is unclear, as $D_{\tau}^{Q}$ is just one realisation of $D_{\tau}$. It seems that the authors miss another expectation there given that the authors claim that the uncertainty comes from data points.

Author response:
The dataset is fixed with a finite collection of tasks and samples per task. In Equation 3, we define our objective as the maximum likelihood estimate of the parameters $\theta$. This is similarly expressed in Equation 2 in [1]. The uncertainty arises in the estimate of the posterior p(\theta_\tau|x,y,\theta) precisely because the number of samples for each task are finite. We introduce a better estimate of this that reduces this uncertainty. 

[1]: Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient-based meta-learning as hierarchical bayes. International Conference on Learning Representations (ICLR), 2018.

Reviewer comment (Theory):
The authors claim that $D$ is a collection of all tasks, but is it finite or countable infinite? Are there any uncertainties in the generating process of tasks themselves?

Author response:
The dataset is fixed with a finite collection of tasks and samples per task. In Equation 3, we define our objective as the maximum likelihood estimate of the parameters $\theta$. This is similarly expressed in Equation 2 in [1]. The uncertainty arises in the estimate of the posterior p(\theta_\tau|x,y,\theta) precisely because the number of samples for each task are finite. We introduce a better estimate of this that reduces this uncertainty. 

[1]: Erin Grant, Chelsea Finn, Sergey Levine, Trevor Darrell, and Thomas Griffiths. Recasting gradient-based meta-learning as hierarchical bayes. International Conference on Learning Representations (ICLR), 2018.

Reviewer comment (Theory):
In theory, it is even difficult to have a local solution of MAML's objective. To ensure it, we often need additional assumptions. The assumption 1 and the paragraph underneath do not make sense from a theoretical aspect though the authors try to provide some theory.

Author response:
We agree with the reviewer that MAMLs objective is a non-convex optimization problem and thus lacks a guarantee of convergence. Nevertheless, the effectiveness of gradient-based meta-learning has been proven in many scenarios thus we include the optimality as an assumption. We include this assumption to imply that there exists a local minimum for the task as this yields a non-zero Hessian. We agree with the reviewer that the assumption of global minimum

Reviewer comment (Theory):
The definition of $h_\psi$ on page 5 is problematic as covariance needs to be positive semi-definite.

Author response:
The covariance does indeed have to be positive semi-definite. Enforcing such a condition is possible e.g. using the Cholesky decomposition similar to as defined in [2]. We have updated the definition of the map to be $h_\psi : \mathcal{X} \times \mathcal{Y} \to \mathbb{S}^{d}_+$ (the space of positive semi-definite matrices)-

[2] Xu, Kailai, Daniel Z. Huang, and Eric Darve. "Learning constitutive relations using symmetric positive definite neural networks." Journal of Computational Physics 428 (2021): 110072.

Reviewer comment (Novelty):
I do not understand why the authors need to have Proposition 1 if it is a well-known result.

Author response:
We include Proposition 1 to draw a connection between the MAP estimate of the task-adapted parameters and the problem of variance reduction. To the best of our knowledge, we have not found any direct references for this result, which motivates us to include it in our presentation.

Reviewer comment (Evaluation):
Why does the proposed method perform worse than MAML on mini-ImageNet?

Author response:
We wish to clarify our assumptions about the task distribution. An additional assumption we make is that the distribution of tasks p(\theta|x, y) is an unimodal and continuous distribution, as the induced map from task to parameters space is a continuous function. Image classification introduces a discrete, multi-modal distribution over the task-manifold, and as such is not suitable to be approximated with a normal distribution. This assumption will be clarified in Section 3 on posterior estimation.

Reviewer comment (Experiments):
Given the worse computation requirement of LAVA compared to CAVIA and MAML, it is necessary for the authors to further discuss and investigate possible techniques for computation savings. For example, instead of computing and storing the Hessian, techniques such as 1) only performing Hessian vector products (which only takes constant back props to compute) and 2) using conjugate gradients to approximately find the (inverse Hessian, vector) product could be used. These explorations are necessary for the scaling of the proposed method, as purely scaling up the non-adaptable parameters while keeping the number of adaptable parameters to a very small number (e.g. $<10$) may not provide sufficient performance for LAVA on real world problems.

Author response:
The Hessian computation indeed requires O(d) more computational time. However, this calculation can be done in parallel, which by assuming d^2 processes yields a Hessian computation that is only twice that of the Jacobian, trading computational time for memory complexity. 
We agree with the reviewer in that with a sensible increase in the dimensionality of the task space the computational cost would be infeasible. Most real-world problems have an intrinsic dimensionality low enough to allow these memory optimizations. In the mini-ImageNet experiments, we experiment with a context size of 100 within a reasonable computational time. 
As the reviewer points out, additional measures can be taken to reduce the time and memory complexity of the problem. This is an important avenue for future work and we will append Section 4.3 with an additional discussion on alternative methods to compute the Hessian.

Reviewer comment (Experiments):
**Evaluating on problems with higher inherent dimensions**. The three few-shot experiments are all internally parameterized by a small number of parameters ($<10$), which might make it appropriate to consider LAVA which could adapt a correspondingly small number of context parameters. However, for real-world problems, the inherent problem dimensions could be much higher or unknown. One experiment setup that more closely captures this property is the image completion task used in CAVIA. The current experiments in the paper haven't demonstrated the ability of LAVA to handle such problems.

Author response:
While it is true that in the experiments described in the main text of the paper, both the model and the context size are relatively small, we do however provide results for the Mini-ImageNet experiments in the Appendix. For this experiment, the model used is a CNN architecture of similar size to the other baselines (4 CNN layers and 1 or 2 MLP layers) and a context size of 100 dimensions. LAVA does not achieve state-of-the-art performances on this experiment but it provides evidence for the usability of LAVA on more high-dimensional experiments.

Reviewer comment (Experiments):
**Comparing gradient-based meta-learning methods under the same number of inner-problem backprops used**. In the experiment section, the performances of different baseline gradient-based meta-learning methods are reported. However, it is unclear whether all these methods are restricted to using 1 inner step (the same as LAVA). If so, it should be noted that LAVA’s single inner adaptation update step requires a number of backprops proportional to the adaptable parameter dimension, while the other methods only need 1 backprop per single inner adaptation step. For fairness, the other baseline methods should also be allowed to use the same number of backprops by allowing them to take more inner adaptation steps (preferably during meta-training, at least during meta-test). This might improve the performances of some of the baseline methods.

Author response:
Equating the models under the same computational complexity is an interesting proposition, however, we would like to point out some potential pitfalls: it introduces overfitting in the GBML methods, especially in higher dimensions, which may not yield an accurate comparison. Secondly, it neglects the tradeoff between computational and memory complexity that our method can exploit, by parallelizing the Hessian computation which has no direct equivalent in GBML methods, as introducing more gradient steps can only be done sequentially. We thank the reviewer for pointing out that a further analysis of the complexity of the method is warranted, and we will add a discussion on this in a future version of the manuscript.

Reviewer comment (Theory):
**Choice of technical language – Assumption 1**. It is unclear how to define the optimal $\theta_{\tau}^*$. The right-hand side of Equation (6) should not necessarily be the best possible $\theta$ in the space $\Theta$ for the task $\tau$. Instead, it makes more sense to use $\theta_{\tau}^*$ to denote the one-step gradient adapted model parameters over an infinite number of support examples sampled from task $\tau$ (as shown in Appendix A.2). However, if more than one gradient step is taken, the unbiasedness (and thus Assumption 1) would no longer hold.

Author response:
The expectation on the left-hand side of Equation 6 is over different samples x from the support dataset. This implies that the number of support samples is finite, but the expectation is taken over different samples of a set of x. Assumption 1 then implies that each $\hat{\theta}_i$ is the same as the optimal on expectation. This justifies the variance reduction as each Z_i in Equation 12 is then assumed to have the same mean.

Reviewer comment (Writing):
**Choice of technical language – Assumption 2**. The notation $f_\tau(x)$ is defined without exactly specifying what its output means until Assumption 2. In the interpretation of Assumption 2, the authors mention that an alternative formulation is that the conditional task-distribution has a non-zero “covariance”. However, the terminology should be changed to “having a non-singleton distribution support”.

Author response:
In the beginning of section 2, we define $f_\tau : X \to Y$ as a map between two spaces X, Y for a specific task $\tau$. Its output defines different tasks as $\tau$ is varied. We agree with the reviewer that a non-singleton distribution support is a more correct terminology, and will update the end of section 3 in a revised version of the manuscript.

Reviewer comment (Theory):
**Equation 9 is incorrect without additional assumptions**. The single-example posterior terms on the right-hand side are each proportional to the product of the example likelihood and the prior $\propto p(y_i \mid x, \theta) p(\theta)$. Hence, the product on the right-hand side of Equation 9 would be proportional to the product of $N$ prior probabilities $p(\theta)^N$. In contrast, the left-hand side of Equation 9 is proportional to only $p(\theta)$, meaning the prior is weighted more heavily on the right-hand side. To fix this, a uniform prior assumption should be introduced earlier than Equation (9) instead of after.

Author response:
We thank the reviewer for their suggestion. The prior $p(\theta)$ is indeed assumed to be uniform. We will clarify these assumptions in Section 2.1.

Reviewer comment (Theory):
**The explanation of the variance reduction is a bit lacking**. Equation (12) formulates a variance reduction problem, but the connection to its use of Assumption (1), which would imply that the estimators $Z$ under different values of $\{W_i\}$ have the same expectation, is missing. Without clearly specifying this, the variance reduction problem seems out of place and unmotivated.

Author response:
The observation that the expectation of each Z_i is equal to the optimal as is stated in Assumption 1. That the Z_i have an equal mean is indeed an assumption in the variance reduction problem. We will add this assumption in proposition 1.

Reviewer comment (Experiments):
The experiments do not contain ablations in certain axes. The problem dimensionality is varied, but the model dimension (the network architecture) is not.

Author response:
We thank the reviewer for their suggestion and agree that a study of the network architecture would be a valuable addition to the paper. On page 5, paragraph 1, we make the assumption that our network is expressive enough to model the task distribution as a Gaussian in parameter space. This raises an interesting question of what architectural choices affect the expressivity in the parameter space. [1] Suggests that the number of layers is more important than the network width for gradient-based meta-learning. It would be interesting to extend this analysis to what architectures allow for arbitrary shaping of the parameter space. We will append this analysis to the experimental section in a revised version of the manuscript. 

[1] When MAML Can Adapt Fast and How to Assist When It Cannot, Arnold et. al. (2019)

Reviewer comment (Theory):
The assumptions of the algorithm is not clear. The variance calculation implicitly assumes that $\theta_0$ is deterministic, but is that reasonable given that it depends on previous parameter updates (which are not deterministic)?

Author response:
The reviewer is correct that the variance calculation only takes into account the variance of the sampled support data. The variance of the task-adapted parameters arises from both the finite sampling of the support and the distribution of $\theta_0$, which as the reviewer points out arises during stochastic gradient descent. In this work, we study the variance of the task-adapted parameters when $\theta_0$ is given and the support data is varied. In the paper, we show that GBML methods might still suffer from a high-variance estimation even when having access to the optimal $\theta_0$. We can make this assumption, and focus on the adaptation-induced variance only, as $\theta_0$ is not dependent on the current sampling of the support, which affects only the task-adapted parameters. We will further clarify this in Section 3.

Reviewer comment (Experiments):
How do the results vary when the model architecture varies, e.g. when the number of parameters in the neural network increases?

Author response:
We thank the reviewer for their suggestion and agree that a study of the network architecture would be a valuable addition to the paper. On page 5, paragraph 1, we make the assumption that our network is expressive enough to model the task distribution as a Gaussian in parameter space. This raises an interesting question of what architectural choices affect the expressivity in the parameter space. [1] Suggests that the number of layers is more important than the network width for gradient-based meta-learning. It would be interesting to extend this analysis to what architectures allow for arbitrary shaping of the parameter space. We will append this analysis to the experimental section in a revised version of the manuscript. 

[1] When MAML Can Adapt Fast and How to Assist When It Cannot, Arnold et. al. (2019)

Reviewer comment (Theory):
Is there a proof of the claim made in the first paragraph of page 5?

Author response:
The adaptation process essentially induces a continuous map $M : T \to \theta$ from which we model the image $M(T)$ as a Normal distribution using the Laplace approximation. For this to be possible, we assume there exists a diffeomorphism between the true task distribution and the normal. This might not be true in general, for example if the topology of the space of tasks $T$ is not simply-connected or doesn’t contain only one connected component.

If these assumptions are satisfied, $f$ will find a solution given that it is a universal function approximator. This is reasonable to assume given $f$ is of a sufficient size. As the reviewer has suggested a study of the effect of parameter count and performance would be a valuable addition to validate this claim. We will append the experimental section with further results on different architectures and add it to a revised version of the manuscript.