What does the batch norm layer for MAML (model-agnostic meta-learning) do for during inference?

Question

I was mainly wondering if we should use the running statistics we used during meta-training or the batch statistics for the current task (during meta-evaluation).

---

Detailed thoughts (from git issue here https://github.com/tristandeleu/pytorch-maml/issues/19):

I was thinking that one would do it as follows:

1. During meta-training (fitting):

• inner loop (support set) it does have the mdl.train() (because we want to collect the running average accross tasks)
• query set, it has the same mdl.train() (to use the same params)

which is what your doing here: https://github.com/tristandeleu/pytorch-meta/blob/d487ad0a1268bd6e6a7290b8780c6b62c7bed688/examples/maml-higher/train.py#L93

The real question is what to do during evaluation (since at meta-eval, the tasks are completely different e.g. image classes we've never seen). There really are 3 options (call them a b c)

2.a. During meta-eval (inference e.g. validation, testing): 2.a. - use .train() for both the support (inner loop) and query set. Here the issue is the model would (accidently) cheat since it would use the stats of the eval set

2.b. - use .eval() for both the support (inner loop) and query set. Here the model would use the stats from training and would not cheat. The pro is that the model was trained with those stats so perhaps thats good - but the true stats of the eval set is something completely different (most likely since the classes have not been seen)

2.c. - use eval() AND set track_running_stats = False. This would use batch statistics. Which would mean the model uses "the right stats" but it was not trained on them...so, who knows if that is better. Plus idk what the BN layer would do for 1-shot learning...probably crash unless it uses layer norm LN.

I am basically curious what the standard maml does. From your code here: https://github.com/tristandeleu/pytorch-maml/blob/44104272a0140b35e2223ba68750e7e715315653/maml/metalearners/maml.py#L231 I infer that you choose option 2.b. So during the inner loop (support set) and the query set your model has eval and uses stats from training.

Is that right?

---

my implementation currently:

        # inner_opt = torch.optim.SGD(self.base_model.parameters(), lr=self.lr_inner)
        inner_opt = NonDiffMAML(self.base_model.parameters(), lr=self.lr_inner)
        # inner_opt = torch.optim.Adam(self.base_model.parameters(), lr=self.lr_inner)
        self.args.inner_opt_name = str(inner_opt)

        # Accumulate gradient of meta-loss wrt fmodel.param(t=0)
        meta_batch_size = spt_x.size(0)
        meta_losses, meta_accs = [], []
        for t in range(meta_batch_size):
            spt_x_t, spt_y_t, qry_x_t, qry_y_t = spt_x[t], spt_y[t], qry_x[t], qry_y[t]
            # if torch.cuda.is_available():
            #     spt_x_t, spt_y_t, qry_x_t, qry_y_t = spt_x_t.cuda(), spt_y_t.cuda(), qry_x_t.cuda(), qry_y_t.cuda()
            # Inner Loop Adaptation
            with higher.innerloop_ctx(self.base_model, inner_opt, copy_initial_weights=self.args.copy_initial_weights,
                                      track_higher_grads=self.args.track_higher_grads) as (fmodel, diffopt):
                diffopt.fo = self.fo
                for i_inner in range(self.args.nb_inner_train_steps):
                    fmodel.train()

                    # base/child model forward pass
                    spt_logits_t = fmodel(spt_x_t)
                    inner_loss = self.args.criterion(spt_logits_t, spt_y_t)
                    # inner_train_err = calc_error(mdl=fmodel, X=S_x, Y=S_y)  # for more advanced learners like meta-lstm

                    # inner-opt update
                    diffopt.step(inner_loss)

            fmodel.train() if self.args.split == 'train' else fmodel.eval()
            # Evaluate on query set for current task
            qry_logits_t = fmodel(qry_x_t)
            qry_loss_t = self.args.criterion(qry_logits_t, qry_y_t)

            # Accumulate gradients wrt meta-params for each task: https://github.com/facebookresearch/higher/issues/104
            # qry_loss_t.backward()  # note this is more memory efficient (as it removes intermediate data that used to be needed since backward has already been called)
            (qry_loss_t / meta_batch_size).backward()  # note this is more memory efficient (as it removes intermediate data that used to be needed since backward has already been called)

            # get accuracy
            if self.target_type == 'classification':
                qry_acc_t = calc_accuracy_from_logits(y_logits=qry_logits_t, y=qry_y_t)  #
            else:
                qry_acc_t = r2_score_from_torch(qry_y_t, qry_logits_t)
                # qry_acc_t = compressed_r2_score(y_true=qry_y_t.detach().numpy(), y_pred=qry_logits_t.detach().numpy())

            # collect losses & accs for logging/debugging
            meta_losses.append(qry_loss_t.item())
            meta_accs.append(qry_acc_t)

Answer 1

TLDR: Use mdl.train() since that uses batch statistics from the current task (but inference will not be deterministic anymore). You probably won't want to use mdl.eval() in meta-learning since that uses stats collected from (meta) training from different tasks.

---

I believe the following is correct:

BN intended behaviour:

• During inference (eval/testing) BN uses the running_mean, running_std collected during training (in meta-learning during meta-learning). Therefore, the (meta) eval performance might be worse if you use mdl.eval() since that would use batch statistics from different tasks (since during meta-learning the tasks at train, val and tests are different e.g. in Mini-Imagenet by having different classes available)
• During training the batch statistics is used but a population statistic is (attempted to be) estimated with running averages. However, the tasks are different for each split during meta-learning so you aren't using the "population statistics" you usually want. As a side note, I assume the reason batch_stats is used during training is to introduce noise that regularizes training (training under noise forces the model to be robust, since it has to perform the task even with the noise)
• in meta-learning I recommend using batch statistics all the time. Therefore always using mdl.train() -- even during testing/evaluation. This avoids using the running means from different task (e.g. different distributions) -- we are supposed to be seeing new tasks/distribution anyway. The price for increased accuracy/performance is loss of determinism -- but precision can be increased with a bigger (meta) batch (i.e. more tasks). This error (confidence interval) shrinks pretty quickly -- even if the sqrt(N) makes it seem it would not. Note the CI formula is z(significance_level) * std / sqrt(N) where N is the meta batch size. For accuracy std is maximally 1.0 and N can be increased to 500, which takes longer to run but the estimates become more precise. z(alpha=0.95) ~ 1.96 or z(alpha=0.99) ~ 2.58 which are fine with a bigger meta-batch.

This is likely why I don't see divergence in my testing with the mdl.train().

So just make sure you use mdl.train(), since that uses batch statistics, reference: https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html#torch.nn.BatchNorm2d) but that either the new running stats that cheat aren't saved or used later.

This likely collects "cheating" statistics but it won't matter for us because we never run inference with .eval() in meta-learning.

For more details see comments on question: https://stackoverflow.com/questions/69845469/when-should-one-call-eval-and-train-when-doing-maml-with-the-pytorch-highe/69858252#69858252