Why not accumulate query loss and then take derivative in MAML with Pytorch and Higher?
when doing MAML (Model agnostic meta-learning) there are two ways to do the inner loop:


def inner_loop1():
    n_inner_iter = 5
    inner_opt = torch.optim.SGD(net.parameters(), lr=1e-1)

    qry_losses = []
    qry_accs = []
    meta_opt.zero_grad()
    for i in range(task_num):
        with higher.innerloop_ctx(
            net, inner_opt, copy_initial_weights=False
        ) as (fnet, diffopt):
            # Optimize the likelihood of the support set by taking
            # gradient steps w.r.t. the model's parameters.
            # This adapts the model's meta-parameters to the task.
            # higher is able to automatically keep copies of
            # your network's parameters as they are being updated.
            for _ in range(n_inner_iter):
                spt_logits = fnet(x_spt[i])
                spt_loss = F.cross_entropy(spt_logits, y_spt[i])
                diffopt.step(spt_loss)

            # The final set of adapted parameters will induce some
            # final loss and accuracy on the query dataset.
            # These will be used to update the model's meta-parameters.
            qry_logits = fnet(x_qry[i])
            qry_loss = F.cross_entropy(qry_logits, y_qry[i])
            qry_losses.append(qry_loss.detach())
            qry_acc = (qry_logits.argmax(
                dim=1) == y_qry[i]).sum().item() / querysz
            qry_accs.append(qry_acc)

            # Update the model's meta-parameters to optimize the query
            # losses across all of the tasks sampled in this batch.
            # This unrolls through the gradient steps.
            qry_loss.backward()

    meta_opt.step()
    qry_losses = sum(qry_losses) / task_num
    qry_accs = 100. * sum(qry_accs) / task_num
    i = epoch + float(batch_idx) / n_train_iter
    iter_time = time.time() - start_time

def inner_loop2():
    n_inner_iter = 5
    inner_opt = torch.optim.SGD(net.parameters(), lr=1e-1)

    qry_losses = []
    qry_accs = []
    meta_opt.zero_grad()
    meta_loss = 0
    for i in range(task_num):
        with higher.innerloop_ctx(
            net, inner_opt, copy_initial_weights=False
        ) as (fnet, diffopt):
            # Optimize the likelihood of the support set by taking
            # gradient steps w.r.t. the model's parameters.
            # This adapts the model's meta-parameters to the task.
            # higher is able to automatically keep copies of
            # your network's parameters as they are being updated.
            for _ in range(n_inner_iter):
                spt_logits = fnet(x_spt[i])
                spt_loss = F.cross_entropy(spt_logits, y_spt[i])
                diffopt.step(spt_loss)

            # The final set of adapted parameters will induce some
            # final loss and accuracy on the query dataset.
            # These will be used to update the model's meta-parameters.
            qry_logits = fnet(x_qry[i])
            qry_loss = F.cross_entropy(qry_logits, y_qry[i])
            qry_losses.append(qry_loss.detach())
            qry_acc = (qry_logits.argmax(
                dim=1) == y_qry[i]).sum().item() / querysz
            qry_accs.append(qry_acc)

            # Update the model's meta-parameters to optimize the query
            # losses across all of the tasks sampled in this batch.
            # This unrolls through the gradient steps.
            #qry_loss.backward()
            meta_loss += qry_loss

    meta_loss.backward()
    meta_opt.step()
    qry_accs = 100. * sum(qry_accs) / task_num
    i = epoch + float(batch_idx) / n_train_iter
    iter_time = time.time() - start_time


are they truly equivalent?



cross-posted:


git issue: https://github.com/facebookresearch/higher/issues/60


The only difference is that in second approach you'll have to keep much more stuff in memory - until you call backward you'll have all unrolled parameters fnet.parameters(time=T) (along with intermediate computation tensors) for each of task_num iterations as part of the graph for the aggregated meta_loss. If you call backward on every task then you only need to keep full set of unrolled parameters (and other pieces of the graph) for one task.
So to answer your question's title: because in this case the memory footprint is task_num times bigger.
In a nutshell what you're doing is similar to comparing loopA(N) and loopB(N) in the following code. Here loopA will get as much memory as it can and OOM with sufficiently large N, while loopB will use about same amount of memory for any large N:
import torch
import numpy as np
a = 0
np.random.seed(1)
v = torch.tensor(np.random.randn(1000000))
y = torch.tensor(np.random.randn(1000000))
x = torch.zeros(1000000, requires_grad=True)

def loopA(N=1000):
    a = 0
    for i in range(N):
        a += ((x * v - y)**2).sum()
    a.backward()

def loopB(N=1000):
    for i in range(N):
        a = ((x * v - y)**2).sum()
        a.backward()

Regarding the normalization - two approaches are equivalent (up to numerical precision maybe): if you first sum up individual losses, then divide by task_num, then finally call backward then you'll effectively compute d((Loss_1 + ... + Loss_{task_num})/task_num) / dw (where w is one of the weights meta-optimizer is fitting). On the other hand if you call backward for each loss divided by task_num you'll get d(Loss_1/task_num)/dw + ... + d(Loss_{task_num}/task_num)/dw which is the same because taking gradient operation is linear. So in both cases your meta-optimizer step will start with pretty much same gradients.
