How does MAML inner loop optimization works?

Question

I started to learn meta-learning, reading the MAML paper https://arxiv.org/pdf/1703.03400.pdf

In the inner loop, I am calculating adapted parameters for each task, I will be doing multiple steps of inner SGD. I will calculate adapted parameters after two or more gradient steps ($\theta{'}$, $\theta^{''}$ , $\theta^{n}$), then using testing parameters, I will have a loss in respect to the original $\theta$ (If I understand the derivation correctly). Now I am supposed to backpropagate through the gradient. Unfortunately, I am not sure how it is done...

Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks Page 3.
"The MAML meta-gradient update involves a gradient through a gradient"

To do this, I understand that I have to store each $\theta^{n}$, but don't know how the loss from $\theta^{n}$ to $\theta^{n-1}$ is transferred up to the original $\theta$.
I guess that the for the last series of parameters ($\theta^{n}$) loss is calculated in the standard way with the testing set, but then I somehow need to pass information saying how much the previous set of parameters was wrong... (Gradient of the gradient?)

I see Hessian and vector products popping up everywhere on the internet, but I cannot imagine how that works, and have no idea how it is calculated and passed/implemented...

Can someone explain to me how the inner loop [back-propagation trough meta-gradient] is working - how the derivations go and how loss is transferred/updated?

Answer 1

According to: "Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks" it is initialized in Class MAML:

class MAML:
    def __init__(self, dim_input=1, dim_output=1, test_num_updates=5):
        """ must call construct_model() after initializing MAML! """
        self.dim_input = dim_input
        self.dim_output = dim_output
        self.update_lr = FLAGS.update_lr
        self.meta_lr = tf.placeholder_with_default(FLAGS.meta_lr, ())
        self.classification = False
        self.test_num_updates = test_num_updates
        if FLAGS.datasource == 'sinusoid':
            self.dim_hidden = [40, 40]
            self.loss_func = mse
            self.forward = self.forward_fc
            self.construct_weights = self.construct_fc_weights
        elif FLAGS.datasource == 'omniglot' or FLAGS.datasource == 'miniimagenet':
            self.loss_func = xent
            self.classification = True
            if FLAGS.conv:
                self.dim_hidden = FLAGS.num_filters
                self.forward = self.forward_conv
                self.construct_weights = self.construct_conv_weights
            else:
                self.dim_hidden = [256, 128, 64, 64]
                self.forward=self.forward_fc
                self.construct_weights = self.construct_fc_weights
            if FLAGS.datasource == 'miniimagenet':
                self.channels = 3
            else:
                self.channels = 1
            self.img_size = int(np.sqrt(self.dim_input/self.channels))
        else:
            raise ValueError('Unrecognized data source.')

"How does MAML inner loop optimization work?"

Where Rectified Linear Unit (ReLU) neural networks are locally almost linear (Goodfellow et al., 2015), and second derivatives close to zero, using a first-order approximation removes the need for computing Hessian-vector products in an additional backward pass.

It is initialized by xavier_initializer_conv2d from TensorFlow, as explained in "Understanding the difficulty of training deep feedforward neural networks", Xavier Glorot and Yoshua Bengio, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, PMLR 9:249-256, 2010.

A method to alleviate the risk of overfitting for gradient-based meta-learning is explained in: "Regularizing Meta-Learning via Gradient Dropout" - April 13 2020, by Tseng, Chen, Tsai, Liu, Lin, and Yang:

"[they] aim to find model parameters that are sensitive to changes in the task, such that small changes in the parameters will produce large improvements on the loss function of any task drawn from $p(T)$, when altered in the direction of the gradient of that loss.".

Combined, those methods shorten the time taken to complete each pass.

In the line:

FLAGS.datasource == 'omniglot' or FLAGS.datasource == 'miniimagenet':
            self.loss_func = xent

See pages 6 and 7:

"A significant computational expense in MAML comes.from the use of second derivatives when backpropagating the meta-gradient through the gradient operator in the meta-objective (see Equation (1)). On MiniImagenet, we show a comparison to a first-order approximation of MAML, where these second derivatives are omitted. Note that the resulting method still computes the meta-gradient at the post-update parameter values $θ^ 0_i$, which provides for effective meta-learning. Surprisingly however, the performance of this method is nearly the same as that obtained with full second derivatives, suggesting that most of the improvement in MAML comes from the gradients of the objective at the post-update parameter values, rather than the second order updates from differentiating through the gradient update.".
$$\theta \leftarrow \theta \; - \beta \nabla_\theta \sum_{\mathcal{T}_i\thicksim p (\mathcal{T})} \mathcal{L_{T_i}}(\mathcal{f}_{\theta^{{'}}_i}) \tag{1} $$