BERT-BiLSTM integration: BiLSTM layer is returning 0 gradient value after first few iterations

Buzz: <p>Hi,</p>
<p>I have been trying to integrate BERT and BiLSTM as a part of a <em>one class classification</em> experiment and I noticed the gradient values of  LSTM are returning 0 after few iteration. I am trying to play with my own defined loss function. Could you please help me finding out the hidden bug in this code? Or is it the usual scenario for such loss function?</p>
<p>This is my model’s structure:</p>
<pre><code class="lang-auto">class Experiment3(nn.Module):
    
    def __init__(self, bert):
        super(Experiment3, self).__init__()
        self.bert = BertModel.from_pretrained('bert-base-uncased')        
        self.lstm = torch.nn.LSTM(768, 128,2,batch_first=True,bidirectional=True)
        self.fc = nn.Linear(128*2, 1)
    def forward(self, ids, mask):
          sequence_output, pooled_output = self.bert(ids,attention_mask=mask)
          lstm_output, (h,c) = self.lstm(sequence_output[:, 0, :])
          weights_=self.fc.weight
          weights=(torch.transpose(weights_, 0, 1))
          biases=nn.Parameter(torch.tensor([0.01])) 
          fc_output=torch.matmul(lstm_output,weights)-biases 
          return weights,biases,fc_output
</code></pre>
<p>The sample dataset has 100 data with labels(0 or 1). I have used a dataloader object with random sampling.</p>
<pre><code class="lang-auto">batch_size=1
train_data = TensorDataset(train_seq, train_mask, train_y)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)
</code></pre>
<p>The train method is following:</p>
<pre><code class="lang-auto">def train():
    model.train()
    for i, batch in enumerate(train_dataloader):
        batch = [b.to(device) for b in batch]
        sent_id, mask, labels = batch
        model.zero_grad()
        weights,biases,output_= model(sent_id, mask)
        loss = customLoss(weights,biases,output_)
        loss.backward()
        for name, param in model.named_parameters():
            print(name, param.grad)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        
    gc.collect()
def customLoss(weight,bias,output):
    a=0.5*torch.square(torch.linalg.norm((weight)))
    temp=bias-output
    z = torch.zeros_like((temp))    
    b=[torch.mean(torch.maximum(z,temp))]
    l2 = a+(1/(0.01))*b[0]-bias
    return l2
</code></pre>
<p>The driver’s code is following:</p>
<pre><code class="lang-auto">model = Experiment3(bert)
model = model.to(device)
optimizer = AdamW(model.parameters(), lr=2e-5)
epochs = 1
current = 1
while current &lt;= epochs:
    print(f'\nEpoch {current} / {epochs}:')
    train()
    current = current + 1
</code></pre>
<p>I would appreciate if you help me understanding this NN and its tricky gradient nature.</p>

ptrblck: <p>In <code>Experiment3.forward</code> you are recreating the <code>biases</code> parameter so note that it won’t be trained and I assume that’s on purpose.</p>
<p>I don’t fully understand your custom loss function. It seems you are adding weight decay to the <code>b</code> tensor. If <code>temp</code> has a negative value (i.e. if <code>output&gt;0.01</code>) your model would create zero gradients, as <code>torch.maximum</code> would select <code>z</code> and would not backpropagate through the <code>output</code> to the model. (Please correct me if I’m missing something)</p>
<p>Also,</p>
<aside class="quote no-group" data-username="Buzz" data-post="1" data-topic="157378">
<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/b/ee59a6/48.png" class="avatar"> Buzz:</div>
<blockquote>
<p><em>one class classification</em> experiment</p>
</blockquote>
</aside>
<p>sounds strange, since the model would just try to learn the single class. Wouldn’t <code>return 0</code> be a perfect classifier in this case?</p>

Buzz: <p>ptrblck Hi</p>
<p>Thanks for your reply. According to your query on the loss function, I have tried to implement the following loss function:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb277e72ed218f71b7633711bb9f570c40319c81.jpeg" alt="s1" data-base62-sha1="qHDTWDTqCeHlngswX1p38F084q5" width="476" height="73"><br>
In my code, I extract the <code>w</code> from FC layer and set a random value to <code>bias</code>. Here I assume the <code>r=bias</code>. What do you think on my loss function’s implementation? Am I on the right track?</p>
<p>I have noted your observation regarding the <code>temp</code> value. I have increased the value of the <code>bias</code> to make the <code>temp</code> value positive which in turn shows the <code>lstm</code> layers returning the gradient values other than zero.</p>
<p>Here my intention is <code>w</code> and <code>bias</code> must be trainable parameters. In my implementation, I found <em>FC layer’s <code>bias</code> parameter is not learning(it returns None).</em></p>
<p>Regarding your last query, my intention is just to learn the features of only a single class.</p>
<p>Hope, these explanations will help you to advise me further.</p>

ptrblck: <aside class="quote no-group" data-username="Buzz" data-post="3" data-topic="157378">
<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/b/ee59a6/48.png" class="avatar"> Buzz:</div>
<blockquote>
<p>In my implementation, I found <em>FC layer’s <code>bias</code> parameter is not learning(it returns None).</em></p>
</blockquote>
</aside>
<p>Yes, that’s expected as mentioned in my previous post:</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="157378">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>In <code>Experiment3.forward</code> you are recreating the <code>biases</code> parameter so note that it won’t be trained and I assume that’s on purpose.</p>
</blockquote>
</aside>
<p>If you want to reuse the <code>bias</code> from the linear layer and train it, you should access it and not create a new <code>nn.Parameter</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">bias = self.fc.bias
# use bias
</code></pre>
<aside class="quote no-group" data-username="Buzz" data-post="3" data-topic="157378">
<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/b/ee59a6/48.png" class="avatar"> Buzz:</div>
<blockquote>
<p>Regarding your last query, my intention is just to learn the features of only a single class.</p>
</blockquote>
</aside>
<p>I don’t know if and how it would work. Assuming your model indeed tries to predict a single class, the model would only need to make sure to predict the largest logits for the single class which would not depend of any inputs, so I’m not sure how much the model would “learn” the dataset.</p>

Buzz: <p>Hi,</p>
<p>I have experimented the model with 2 cases:<br>
<strong>Case 1</strong>: <code>bias=self.fc.bias</code><br>
In this case, I notice either the gradient value of lstm layers’ are showing <code>0</code> for each batch(if the bias value is negative) or it starts to show <code>0</code>(if the bias value is positive) after certain iterations.<br>
<strong>Case 2:</strong> a big value is initialized to bias such as <code>biases=nn.Parameter(torch.tensor([10]))</code><br>
In such scenario, I observed the gradient of the FC layers bias is returning a constant value in each iteration, <code>fc.bias tensor([199.])</code>.</p>
<p>What might be the reasons behind such behavior?</p>