Controlling Gradient Propagation To Shared Layer From Multi-task Model

nmuralid: <p>I have a standard multi-task DNN model with 2 shared layers followed by a single layer for each of two separate but related regression tasks (one main task, one auxiliary task). My question is regarding how I could manage to selectively backpropagate gradients from the auxiliary task based on some condition. A relevant paper in this space is titled: <a href="https://arxiv.org/abs/1812.02224#" rel="nofollow noopener">Adapting Auxiliary Losses Using Gradient Similarity</a>, wherein the authors backpropagate gradients from the auxiliary task only when they are in a similar direction to the main task.</p>
<p>Here is the implementation I have so far</p>
<h1>Multi-Task DNN Model</h1>
<pre><code class="lang-auto">class Net(nn.Module):
  def __init__(self,inputsize,hiddensize,outputsize_main,outputsize_aux):
    super(Net,self).__init__()
    self.fc_shared1 = nn.Linear(inputsize,hiddensize)
    self.fc_shared2 = nn.Linear(hiddensize,hiddensize)
    self.fc_main1 = nn.Linear(hiddensize,outputsize_main)
    self.fc_aux1 = nn.Linear(hiddensize,outputsize_aux)
    self.tanh1 = nn.Tanh()
    self.tanh2 = nn.Tanh()

  def forward(self,x):
    out = self.fc_shared1(x)
    out = self.tanh1(out)
    out = self.fc_shared2(out)
    out = self.tanh2(out)

    #Main
    outmain = self.fc_main1(out)
    #Aux
    outaux = self.fc_aux1(out)

    return outmain, outaux
</code></pre>
<h1>Simple Hook Script (borrowed from <a href="https://www.kaggle.com/sironghuang/understanding-pytorch-hooks" rel="nofollow noopener">here</a>)</h1>
<pre><code class="lang-auto"># A simple hook class that returns the input and output of a layer during forward/backward pass
class Hook():
    def __init__(self, module,name, backward=False):
        self.name = name  #Name of the layer. Useful for debugging.
        if backward==False:
            self.hook = module.register_forward_hook(self.hook_fn)
        else:
            self.hook = module.register_backward_hook(self.hook_fn)
    def hook_fn(self, module, input, output):
        self.input = input
        self.output = output
    def close(self):
        self.hook.remove()
</code></pre>
<h1>Hyperparameters</h1>
<pre><code class="lang-auto">INPUTSIZE=4
OUTPUTSIZE_MAIN=1
OUTPUTSIZE_AUX=1
HIDDENSIZE=16
BATCHSIZE=32
NUMEPOCHS=100
learningrate=0.001
model=Net(INPUTSIZE,HIDDENSIZE,OUTPUTSIZE_MAIN,OUTPUTSIZE_AUX)
optimizer=torch.optim.Adam(model.parameters(),lr=learningrate)
criterion=nn.MSELoss()

# trainX, testX are n X 4 matrices containing input features where `n` is the number of instances.
# trainY, testY are n X 1 arrays containing target values for the main task.
# trainAuxY , testAuxY are also n X 1 arrays containing the auxiliary variable target values.

train_dataset = SampleDataLoader(trainX,trainY,trainAuxY)
train_loader = DataLoader(train_dataset,batch_size=BATCHSIZE)
test_dataset = SampleDataLoader(testX,testY,testAuxY)
test_loader = DataLoader(test_dataset,batch_size=BATCHSIZE)
</code></pre>
<h1>Training Loop</h1>
<pre><code class="lang-auto">hookF = [Hook(layer,name) for name,layer in list(model._modules.items())]
hookB = [Hook(layer,name,backward=True) for name,layer in list(model._modules.items())]

losses=list()
for ep in range(NUMEPOCHS):
  losses_epoch=list()
  for batchX,batchY,batchAuxY in train_loader:
    optimizer.zero_grad()
    _X=torch.autograd.Variable(batchX.float()).to(device)
    _Y=torch.autograd.Variable(batchY.float()).to(device)
    _AuxY=torch.autograd.Variable(batchAuxY.float()).to(device)

    outMain,outAux=model(_X)
    lossMain = criterion(outMain,_Y)
    lossAux = criterion(outAux,_AuxY)
    loss = lossMain + lossAux
    losses_epoch.append(loss.item())
    loss.backward()
    optimizer.step()
  
  print('***'*3+'  Forward Hooks Inputs &amp; Outputs  '+'***'*3)
  for hook in hookF:
      print(hook.name)            
      print("Num Items Input Forward= {} ".format(len(hook.input)))
      print("Size of Each Input = {}".format([i.size() for i in hook.input]))          
      print("Num Items Output Forward= {}".format(len(hook.output)))         
      print("Size of Each Output = {}".format([o.size() for o in hook.output]))
      print('---'*17)
  
  print('\n')
  print('***'*3+'  Backward Hooks Inputs &amp; Outputs  '+'***'*3)
  for hook in hookB: 
      print(hook.name)            
      print("Num Items Input Backward= {} ".format(len(hook.input)))
      print("Size of Each Input = {}".format([i.size() for i in hook.input]))          
      print("Num Items Output Backward= {}".format(len(hook.output)))         
      print("Size of Each Output = {}".format([o.size() for o in hook.output]))
      print('---'*17)
      break
  print("=============================================\n\n")
  losses.append(np.mean(losses_epoch))

</code></pre>
<p>I have the following questions that I haven’t found a clear / convincing answer to online:</p>
<ol>
<li>
<p>What is the difference between using <code>loss.backward()</code> vs. <code>out.backward()</code>?  <code>out</code> (ofcourse in the above model it would be <code>outmain</code>, <code>outaux</code>) is the output of the call to <code>model.forward()</code>. I have read (and maybe misunderstood) that <code>loss.backward()</code> doesn’t work with hooks (although I didn’t see any errors being thrown by the hooks when I ran the aforementioned code with <code>loss.backward()</code>. So is it okay to use <code>loss.backward()</code> with hooks?</p>
</li>
<li>
<p>How can I isolate the two specific sets of gradients to be compared, i.e in this case the gradients are <code>d(lossMain)/d(fc_shared1)</code> , <code>d(lossAux)/d(fc_shared1)</code> and  <code>d(lossMain)/d(fc_shared2)</code> , <code>d(lossAux)/d(fc_shared2)</code>. The way I see it, my current implementation will only return <code>d(loss)/d(fc_shared1)</code> and <code>d(loss)/d(fc_shared2)</code>, should I keep the losses separate and invoke <code>lossMain.backward()</code> and <code>lossAux.backward()</code> separately or is there some other solution? In the case of separate invocation of <code>backward()</code> on losses, would the way losses are backpropagated through the network change in anyway w.r.t the single loss case i.e <code>loss = lossMain + lossAux</code> ?</p>
</li>
<li>
<p>Once gradients are isolated and I am able to compare them, will zeroing out <code>d(lossAux)/d(fc_shared1)</code>, <code>d(lossAux)/d(fc_shared2)</code> be enough to ensure that the gradient from the auxiliary loss won’t propagate to the shared layer(s)? The way I envision doing this is, zeroing out <code>fc_shared2.autograd.grad</code> or <code>fc_shared1.autograd.grad</code> before the call to <code>optimizer.step()</code> at times when <code>fc_shared2.autograd.grad</code> and <code>fc_shared1.autograd.grad</code> satisfies some condition?</p>
</li>
<li>
<p>For now, I am just trying to print out, the inputs and outputs to the forward and backward hook functions to better understand them but I am having a hard time understanding them.</p>
</li>
</ol>
<p>For example, here is the result of printing both the inputs and outputs of the forward hook function for each layer:</p>
<pre><code class="lang-auto">*********  Forward Hooks Inputs &amp; Outputs  *********
*********  Forward Hooks Inputs &amp; Outputs  *********
fc_shared1
Num Items Input Forward= 1 
Size of Each Input = [torch.Size([14, 4])]
Num Items Output Forward= 14
Size of Each Output = [torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16])]
---------------------------------------------------
fc_shared2
Num Items Input Forward= 1 
Size of Each Input = [torch.Size([14, 16])]
Num Items Output Forward= 14
Size of Each Output = [torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16])]
---------------------------------------------------
fc_main1
Num Items Input Forward= 1 
Size of Each Input = [torch.Size([14, 16])]
Num Items Output Forward= 14
Size of Each Output = [torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1])]
---------------------------------------------------
fc_aux1
Num Items Input Forward= 1 
Size of Each Input = [torch.Size([14, 16])]
Num Items Output Forward= 14
Size of Each Output = [torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1]), torch.Size([1])]
---------------------------------------------------
tanh1
Num Items Input Forward= 1 
Size of Each Input = [torch.Size([14, 16])]
Num Items Output Forward= 14
Size of Each Output = [torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16])]
---------------------------------------------------
tanh2
Num Items Input Forward= 1 
Size of Each Input = [torch.Size([14, 16])]
Num Items Output Forward= 14
Size of Each Output = [torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16]), torch.Size([16])]
---------------------------------------------------


*********  Backward Hooks Inputs &amp; Outputs  *********
fc_shared1
Num Items Input Backward= 3 
Size of Each Input = [torch.Size([16]), -1, torch.Size([4, 16])]
Num Items Output Backward= 1
Size of Each Output = [torch.Size([14, 16])]
---------------------------------------------------
fc_shared2
Num Items Input Backward= 3 
Size of Each Input = [torch.Size([16]), torch.Size([14, 16]), torch.Size([16, 16])]
Num Items Output Backward= 1
Size of Each Output = [torch.Size([14, 16])]
---------------------------------------------------
fc_main1
Num Items Input Backward= 3 
Size of Each Input = [torch.Size([1]), torch.Size([14, 16]), torch.Size([16, 1])]
Num Items Output Backward= 1
Size of Each Output = [torch.Size([14, 1])]
---------------------------------------------------
fc_aux1
Num Items Input Backward= 3 
Size of Each Input = [torch.Size([1]), torch.Size([14, 16]), torch.Size([16, 1])]
Num Items Output Backward= 1
Size of Each Output = [torch.Size([14, 1])]
---------------------------------------------------
tanh1
Num Items Input Backward= 1 
Size of Each Input = [torch.Size([14, 16])]
Num Items Output Backward= 1
Size of Each Output = [torch.Size([14, 16])]
---------------------------------------------------
tanh2
Num Items Input Backward= 1 
Size of Each Input = [torch.Size([14, 16])]
Num Items Output Backward= 1
Size of Each Output = [torch.Size([14, 16])]
---------------------------------------------------
</code></pre>
<p>The <code>Num Items Output Forward</code> and <code>Num Items Input Forward</code> i.e the size of the tuples stored in <code>input</code> and <code>output</code> seem to be some function of the <code>BATCHSIZE</code> i.e the number changes when I change with <code>BATCHSIZE</code> but I wasn’t able to find how this is calculated.</p>
<p>Also, if someone could clearly elucidate exactly what the <code>input</code>,<code>output</code> and <code>grad_input</code>, <code>grad_output</code> parameters do in the case of the forward and backward hooks respectively, it would help me greatly in my task of working with <code>hooks</code> in pytorch.</p>