Update model A on model B's loss

pgujjar: <p>I am trying to implement the Meta-Pseudo Label paper (<a href="https://arxiv.org/pdf/2003.10580v2.pdf" rel="nofollow noopener">https://arxiv.org/pdf/2003.10580v2.pdf</a>)  where a teacher model is updated based on a student model’s performance on a validation set. It would look something like this:</p>
<pre><code class="lang-auto">modelA_optim = optim.Adam(modelA.parameters(), lr=0.01)
modelB_optim = optim.Adam(modelB.parameters(), lr=0.01)

pseudo_labels = modelA(X_train)

# Train model B on pseudo labels
logits = modelB(X_train)
loss1 = CrossEntropy(logits, pseudo_labels)
loss1.backward()
modelB_optim.step()

# Update model A with model B's loss on a validation set
logits = modelB(X_val)
loss2 = CrossEntropy(logits, y_val)
loss2.backward()
modelA_optim.step()
</code></pre>
<p>This in iteself wouldn’t update model A’s weights as the gradient cannot backpropate all the way through model B and eventually to model A through the graient of loss1.</p>
<p>I have tried using the higher library to make model B a functional one, thereby being able to backpropate through its one updat (becuase of loss1, model B’s weights_t will now be weights_t+1). However, torch.autograd.grad(loss, modelA.parameters()) yields zero gradients for model A. I suspect this is because the backward graph is broken at the loss1 point of connection between model A and model B. Is there a way to address this?!</p>

Ouasfi: <p>Hi,</p>
<p>I think the problem lies in the optimization step. As you can see here :<br>
<a href="https://github.com/pytorch/pytorch/blob/cd9b27231b51633e76e28b6a34002ab83b0660fc/torch/optim/sgd.py#L63" rel="nofollow noopener">SGD step </a> the step function is operating on p.data instead p (parameter). So the computational graph is not keeping track of these operations ( Equation 1 of the paper for instance).<br>
You may also need to set  <code>create_graph</code> to True in your backward function in order to compute the second order derivatives of loss1 since you will need it at some point to backpropagate the gradients of loss2.</p>

pgujjar: <p>I agree that the optimization is where the problem is in the code snippet above (even after setting create_graph=True). I looked in to the higher library (<a href="https://github.com/facebookresearch/higher" rel="nofollow noopener">https://github.com/facebookresearch/higher</a>) which implements a differentiable optimizer. This would allow you to track the updates to the parameters through the many SGD updates. The code would look like this:</p>
<pre><code class="lang-auto">with higher.innerloop_ctx(modelB, modelB_optim, copy_initial_weights=False) as (fmodel, diffopt):
    pseudo_labels = modelA(X_train)

    # Train model B on pseudo labels
    logits = fmodel(X_train)
    loss1 = CrossEntropy(logits, pseudo_labels)
    diffopt.step(loss1)

    # Update model A with model B's loss on a validation set
    logits = fmodel(X_val)
    loss2 = CrossEntropy(logits, y_val)
    grad = torch.autograd.grad(loss2, modelA.parameters())
</code></pre>
<p>The grad is all zero. I suspect this is because higher takes a snapshot of model B only (to make it functional), the gradients cannot be backprop-ed all the way to model A.</p>
<p>If there was a way to extend the graph for fmodel leading all the way back to model A, this would probably work.</p>

Ouasfi: <aside class="quote no-group" data-username="pgujjar" data-post="3" data-topic="83180">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/e9bcb4/48.png" class="avatar"> pgujjar:</div>
<blockquote>
<p>I suspect this is because higher takes a snapshot of model B only</p>
</blockquote>
</aside>
<pre><code class="lang-auto">with higher.innerloop_ctx(student, student_optim, copy_initial_weights=False) as (fmodel, diffopt):
   pseudo_labels = teacher(x)

   logits = fmodel(x)
   loss1 = criterion(logits, pseudo_labels)
   diffopt.step(loss1)

   logits = fmodel(x)
   loss2 = criterion(logits, y_val)
   grad = torch.autograd.grad(loss2, teacher.parameters())
</code></pre>
<p>I agree. It seems to be working when replacing modelB with fmodel  in second the update.</p>

pgujjar: <p>My bad. It was supposed to be fmodel in the original change itself. However, even after using fmodel, you’ll notice that the torch.autograd.grad return all zeros. And this I think is because the backward graph doesn’t extend to model A. Trying to find out how to do this so that the grads are not zero!</p>

Ouasfi: <aside class="quote no-group" data-username="pgujjar" data-post="5" data-topic="83180">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/p/e9bcb4/48.png" class="avatar"> pgujjar:</div>
<blockquote>
<p>you’ll notice that the torch.autograd.grad return all zeros.</p>
</blockquote>
</aside>
<p>It will depends on the model I think.  Here is a simple case where torch.autograd.grad don’t return zeros : <a href="https://colab.research.google.com/drive/1EeRSKTexXS1ym2Q3Sn7S5RP2X66JzWZX?usp=sharing" rel="noopener nofollow ugc">Diff opt Colab.</a></p>

pgujjar: <p>You’re absolutely right! This depends on the model, mostly the depth of it. I was using Resnets as teachers and students. It appears that the second gradients are so small that there are effectively zero.</p>
<p>Thank you very much <a class="mention" href="/u/ouasfi">@Ouasfi</a> for taking the time out and running an example. As a sidenote, the grad is zero if you interchange your teacher and student model definitions. Time to get deeper into the gradient norms I guess!</p>

lukemshepherd: <p>Hey!</p>
<p>I’m also trying to implement meta pseudo labels - this is the training loop I have, but not sure if <code>teacher_loss.backward()</code> is updating the parameters in the right way.</p>
<p>Any advice would be much appreciated!</p>
<pre><code class="lang-auto">for i in range(epochs):
    print(i)
    with higher.innerloop_ctx(student, student_optim, copy_initial_weights=False) as (fstudnet, diffopt):
        
        pseudo_labels = teacher(x_train) # labels from teacher

        logits = fstudnet(x_train)
        student_loss = student_loss_func(logits, pseudo_labels)
        
        diffopt.step(student_loss) # zero grad done implicitly 

        logits = fstudnet(x_val)
        teacher_loss = teacher_loss_func(logits, y_val)
        
#         grad = torch.autograd.grad(teacher_loss, teacher.parameters()) 
#         not sure how to pass grad explicitly to model parameters

        teacher_loss.backward()        
        teacher_optim.step() 
        teacher_optim.zero_grad()
        
</code></pre>