How does one implemented a parametrized meta-learner in Pytorch's higher library?

pinocchio: <p>I wanted to implement the meta-lstm meta-learner in the paper <a href="https://openreview.net/pdf?id=rJY0-Kcll" rel="noopener nofollow ugc">OPTIMIZATION AS A MODEL FOR FEW-SHOT LEARNING</a> using higher but I found problems. I found that I cannot make it work without removing (what seems to be this crucial line):</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/facebookresearch/higher/blob/8f0716fb1663218324c02dabdba26b639959cfb6/higher/optim.py#L101">
  <header class="source">

      <a href="https://github.com/facebookresearch/higher/blob/8f0716fb1663218324c02dabdba26b639959cfb6/higher/optim.py#L101" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/facebookresearch/higher/blob/8f0716fb1663218324c02dabdba26b639959cfb6/higher/optim.py#L101" target="_blank" rel="noopener nofollow ugc">facebookresearch/higher/blob/8f0716fb1663218324c02dabdba26b639959cfb6/higher/optim.py#L101</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="91" style="counter-reset: li-counter 90 ;">
          <li>        permit backpropagation through the optimization process. Setting</li>
          <li>        this to False allows the differentiable optimizer to be used in</li>
          <li>        "test mode", without potentially tracking higher order</li>
          <li>        gradients. This can be useful when running the training loop at</li>
          <li>        test time, e.g. in k-shot learning experiments, without</li>
          <li>        incurring a significant memory overhead.</li>
          <li>"""</li>
          <li>reference_params = list(reference_params)</li>
          <li></li>
          <li># Copy param groups and set up structures for copy state</li>
          <li class="selected">self.param_groups = _copy.deepcopy(other.param_groups)</li>
          <li>self._group_to_param_list: _typing.List[_typing.List[int]] = []</li>
          <li>self.state: _StateType = [</li>
          <li>    _collections.defaultdict(dict)</li>
          <li>    for _ in range(len(self.param_groups))</li>
          <li>]</li>
          <li></li>
          <li># Deal with override</li>
          <li>if override is not None:</li>
          <li>    self._apply_override(override)</li>
          <li></li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>to:</p>
<pre><code class="lang-auto">#self.param_groups = _copy.deepcopy(other.param_groups)
self.param_groups = other.param_groups
</code></pre>
<p>I provide an extremely simplified self-contained implementation of something similar here:</p>
<aside class="onebox githubgist" data-onebox-src="https://gist.github.com/renesax14/8499e0314351ea4199a17e494bff5c4d">
  <header class="source">

      <a href="https://gist.github.com/renesax14/8499e0314351ea4199a17e494bff5c4d" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/renesax14/8499e0314351ea4199a17e494bff5c4d" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/renesax14/8499e0314351ea4199a17e494bff5c4d</a></h4>

  <h5>simple_meta_lstm_meta_learner_parametrized_optimizer.py</h5>
  <pre><code class="Python"># base on the paper "OPTIMIZATION AS A MODEL FOR FEW-SHOT LEARNING": https://openreview.net/pdf?id=rJY0-Kcll

class EmptySimpleMetaLstm(Optimizer):

    def __init__(self, params, trainable_opt_model, trainable_opt_state, *args, **kwargs):
        defaults = {
            'trainable_opt_model':trainable_opt_model, 
            'trainable_opt_state':trainable_opt_state, 
            'args':args, 
            'kwargs':kwargs</code></pre>
    This file has been truncated. <a href="https://gist.github.com/renesax14/8499e0314351ea4199a17e494bff5c4d" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>but I will copy paste here to keep the discussion in one place:</p>
<pre><code class="lang-auto"># base on the paper "OPTIMIZATION AS A MODEL FOR FEW-SHOT LEARNING": https://openreview.net/pdf?id=rJY0-Kcll

class EmptySimpleMetaLstm(Optimizer):

    def __init__(self, params, trainable_opt_model, trainable_opt_state, *args, **kwargs):
        defaults = {
            'trainable_opt_model':trainable_opt_model, 
            'trainable_opt_state':trainable_opt_state, 
            'args':args, 
            'kwargs':kwargs
        }
        super().__init__(params, defaults)

class SimpleMetaLstm(DifferentiableOptimizer):

    def _update(self, grouped_grads, **kwargs):
        prev_lr = self.param_groups[0]['trainable_opt_state']['prev_lr']
        eta = self.param_groups[0]['trainable_opt_model']['eta']
        # start differentiable &amp; trainable update
        zipped = zip(self.param_groups, grouped_grads)
        for group_idx, (group, grads) in enumerate(zipped):
            for p_idx, (p, g) in enumerate(zip(group['params'], grads)):
                if g is None:
                    continue
                # get gradient as "data"
                g = g.detach() # gradients of gradients are not used (no hessians)
                ## very simplified version of meta-lstm meta-learner
                input_metalstm = torch.stack([p, g, prev_lr.view(1,1)]).view(1,3) # [p, g, prev_lr] note it's missing loss, normalization etc. see original paper
                lr = eta(input_metalstm).view(1)
                fg = 1 - lr # learnable forget rate
                ## update suggested by meta-lstm meta-learner
                p_new = fg*p - lr*g
                group['params'][p_idx] = p_new
        # fake returns
        self.param_groups[0]['trainable_opt_state']['prev_lr'] = lr

higher.register_optim(EmptySimpleMetaLstm, SimpleMetaLstm)

def test_parametrized_inner_optimizer():
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from collections import OrderedDict

    ## training config
    #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    track_higher_grads = True # if True, during unrolled optimization the graph be retained, and the fast weights will bear grad funcs, so as to permit backpropagation through the optimization process. False during test time for efficiency reasons
    copy_initial_weights = False # if False then we train the base models initial weights (i.e. the base model's initialization)
    episodes = 5
    nb_inner_train_steps = 5
    ## get base model
    base_mdl = nn.Sequential(OrderedDict([
        ('fc', nn.Linear(1,1, bias=False)),
        ('relu', nn.ReLU())
        ]))
    ## parametrization/mdl for the inner optimizer
    opt_mdl = nn.Sequential(OrderedDict([
        ('fc', nn.Linear(3,1, bias=False)), # 3 inputs 1 for parameter, 1 for gradient, 1 for previous lr
        ('sigmoid', nn.Sigmoid())
        ]))
    ## get outer optimizer (not differentiable nor trainable)
    outer_opt = optim.Adam([{'params': base_mdl.parameters()},{'params': opt_mdl.parameters()}], lr=0.01)
    for episode in range(episodes):
        ## get fake support &amp; query data (from a single task and 1 data point)
        spt_x, spt_y, qry_x, qry_y = torch.randn(1), torch.randn(1), torch.randn(1), torch.randn(1)
        ## get differentiable &amp; trainable (parametrized) inner optimizer
        inner_opt = EmptySimpleMetaLstm(base_mdl.parameters(), trainable_opt_model={'eta': opt_mdl}, trainable_opt_state={'prev_lr': 0.9*torch.randn(1)})
        with higher.innerloop_ctx(base_mdl, inner_opt, copy_initial_weights=copy_initial_weights, track_higher_grads=track_higher_grads) as (fmodel, diffopt):
            for i_inner in range(nb_inner_train_steps): # this current version implements full gradient descent on k_shot examples (which is usually small  5)
                fmodel.train()
                # base/child model forward pass
                inner_loss = 0.5*((fmodel(spt_x) - spt_y))**2
                # inner-opt update
                diffopt.step(inner_loss)
            ## Evaluate on query set for current task
            qry_loss = 0.5*((fmodel(qry_x) - qry_y))**2
            qry_loss.backward() # for memory efficient computation
        ## outer update
        print(f'episode = {episode}')
        print(f'base_mdl.grad = {base_mdl.fc.weight.grad}')
        print(f'opt_mdl.grad = {opt_mdl.fc.weight.grad}')
        outer_opt.step()
        outer_opt.zero_grad()
        
if __name__ == '__main__':
    test_parametrized_inner_optimizer()
    print('Done \a')

"""
output when deep copy is uncommented (parametrized optimizer trains properly):
episode = 0
base_mdl.grad = tensor([[-0.0351]])
opt_mdl.grad = tensor([[0.0085, 0.0000, 0.0204]])
episode = 1
base_mdl.grad = tensor([[0.0311]])
opt_mdl.grad = tensor([[-0.0086, -0.0100,  0.0358]])
episode = 2
base_mdl.grad = tensor([[0.]])
opt_mdl.grad = tensor([[0., 0., 0.]])
episode = 3
base_mdl.grad = tensor([[0.0066]])
opt_mdl.grad = tensor([[-0.0016,  0.0000, -0.0032]])
episode = 4
base_mdl.grad = tensor([[-0.0311]])
opt_mdl.grad = tensor([[0.0077, 0.0000, 0.0130]])
Done 
when deep copy is on (paremeters of inner optimizer are not train, sad!):
episode = 0
base_mdl.grad = tensor([[0.]])
opt_mdl.grad = None
episode = 1
base_mdl.grad = tensor([[0.]])
opt_mdl.grad = None
episode = 2
base_mdl.grad = tensor([[0.0069]])
opt_mdl.grad = None
episode = 3
base_mdl.grad = tensor([[0.]])
opt_mdl.grad = None
episode = 4
base_mdl.grad = tensor([[0.]])
opt_mdl.grad = None
Done
The deep copy line in higher I am referencing:
        self.param_groups = _copy.deepcopy(other.param_groups)
        #self.param_groups = other.param_groups
"""
</code></pre>
<hr>
<h1><a name="real-solution-1" class="anchor" href="#real-solution-1"></a>Real Solution</h1>
<p>The real solution is if I could pass an arbitrary dictionary to a differentiable optimizer and if I could do whatever I wanted with it.</p>
<hr>
<p><span class="hashtag-raw">#update:</span></p>
<p>perhaps this can be implemented with override:</p>
<blockquote>
<p>override (optional) – a dictionary mapping optimizer settings (i.e. those which would be passed to the optimizer constructor or provided within parameter groups) to either singleton lists of override values, or to a list of override values of length equal to the number of parameter groups. If a single override is provided for a keyword, it is used for all parameter groups. If a list is provided, the ith element of the list overrides the corresponding setting in the ith parameter group. This permits the passing of tensors requiring gradient to differentiable optimizers for use as optimizer settings.</p>
</blockquote>
<p>Didn’t work with override:</p>
<pre><code class="lang-auto">Exception has occurred: ValueError 
Mismatch between the number of override tensors for optimizer parameter trainable_opt_model and the number of parameter groups.
seems like it checks that these lengths match... 

def _apply_override(self, override: _OverrideType) -&gt; None: 
for k, v in override.items(): 
# Sanity check 
if (len(v) != 1) and (len(v) != len(self.param_groups)

</code></pre>
<hr>
<p>I think this is all I need:</p>
<pre><code class="lang-auto">inner_opt = EmptySimpleMetaLstm( base_mdl.parameters() )
with higher.innerloop_ctx(base_mdl, inner_opt, copy_initial_weights=copy_initial_weights, track_higher_grads=track_higher_grads) as (fmodel, diffopt):
     diffopt.override = {'trainable_opt_model': opt_mdl, 'trainable_opt_state': {'prev_lr': 0.9*torch.randn(1)} }
</code></pre>
<hr>
<p>cross-posted:</p>
<ul>
<li>git issue in higher: <a href="https://github.com/facebookresearch/higher/issues/62" class="inline-onebox" rel="noopener nofollow ugc">How does one implemented a parametrized meta-learner (like meta-lstm optimizer) in higher? · Issue #62 · facebookresearch/higher · GitHub</a></li>
<li><a href="https://www.reddit.com/r/pytorch/comments/hbp1n5/how_does_one_implemented_a_parametrized/" class="inline-onebox" rel="noopener nofollow ugc">Reddit - Dive into anything</a>?</li>
<li><a href="https://discuss.pytorch.org/t/how-does-one-implemented-a-parametrized-meta-learner-in-pytorchs-higher-library/85988" class="inline-onebox">How does one implemented a parametrized meta-learner in Pytorch's higher library?</a></li>
</ul>
<p>related:</p>
<ul>
<li><a href="https://stackoverflow.com/questions/62437960/why-does-higher-need-to-deep-copy-the-parameters-of-the-base-model-to-create-a-f/62438602?noredirect=1#comment110458561_62438602" class="inline-onebox" rel="noopener nofollow ugc">machine learning - Why does higher need to deep copy the parameters of the base model to create a functional model? - Stack Overflow</a></li>
<li><a href="https://stackoverflow.com/questions/60311183/what-does-the-copy-initial-weights-documentation-mean-in-the-higher-library-for/62378264?noredirect=1#comment110457791_62378264" class="inline-onebox" rel="noopener nofollow ugc">machine learning - What does the copy_initial_weights documentation mean in the higher library for Pytorch? - Stack Overflow</a></li>
</ul>