Loss of BERT mask language modeling and next sentence prediction

pascal_notsawo: <p>I train bert to do mask language modeling (MLM) of next sentence prediction (NSP) tasks. But I am confused about the loss function.</p>
<pre><code class="lang-auto">import torch
from torch import tensor
import torch.nn as nn
</code></pre>
<p>Let’s start with NSP.<br>
Indeed, let’s suppose that I have three pairs of sentences (ie batch_size=3) and that for these three sentences the labels are the following (0 = noNext, 1=isNext) :</p>
<pre><code class="lang-auto">is_next = tensor([0, 1, 1]) # batch_size
</code></pre>
<p>Since in my dataset some sentences are too short for the NSP task, I also have a tensor that informs me whether this was done (1) or not (0): i.e. whether the entry was in the form <code>[CLS] segment_A [SEP] segment_B [SEP]</code> (1) or just <code>[CLS] segment_A [SEP]</code> (0).</p>
<pre><code class="lang-auto">is_next_weight = tensor([1, 1, 0]) # batch_size, no SNP for the last sentence
</code></pre>
<p>Let’s suppose that my model returns the following logits:</p>
<pre><code class="lang-auto">logits_clsf = tensor([[0.3797, 0.6203], [0.4363, 0.5637], [0.3797, 0.6203]]) # batch_size x 2
</code></pre>
<p>I have a loss function :</p>
<pre><code class="lang-auto">criterion2 = nn.CrossEntropyLoss()
</code></pre>
<p>At this level I have two ways to have my final loss.</p>
<ul>
<li>either I calculate its value without excluding the unwanted entries, then I average the desired entries.</li>
</ul>
<pre><code class="lang-auto">loss_clsf = criterion2(logits_clsf, is_next)
loss_clsf = (loss_clsf*is_next_weight.float()).mean()
&gt;&gt;&gt; tensor(0.4516)
</code></pre>
<ul>
<li>or I select only the desired entries and calculate the loss with them.</li>
</ul>
<pre><code class="lang-auto">is_next = is_next [is_next_weight == 1]
logits_clsf = logits_clsf [is_next_weight == 1]
loss_clsf = criterion2(logits_clsf, is_next)
&gt;&gt;&gt; tensor(0.7261)
</code></pre>
<p>I also have this concern at the MLM level. As the number of masked tokens depends on each sentence, I sometimes complete, with <a href="https://pytorch.org/docs/stable/generated/torch.nn.utils.rnn.pad_sequence.html" rel="noopener nofollow ugc">torch.nn.utils.rnn.pad_sequence</a>, masked_ids, which contains here the ids of the masked tokens, so that they have the same length for the same batch (pytorch constraint).<br>
Let us take an example with a batch_size of 3 and a vocabulary of size <code>n_vocab=5</code> (it is just to make simple: 0 = pad_index, one ignores the ids of the special tokens like <code>[UNK]</code>, <code>[CLS]</code> and <code>[SEP]</code>): 1 token of id <code>[2]</code> is masked in the first sentence of the batch, 2 tokens of ids <code>[1, 2]</code> are masked in the second and 3 in the last, of ids <code>[1,3,4]</code>.<br>
Then, <code>max_len = 3</code> (number of masked tokens in the last one) and :</p>
<pre><code class="lang-auto">masked_ids = tensor([[2, 0, 0], [1, 2, 0], [1, 3, 4]]) # batch_size x max_len
masked_weights = tensor([[1, 0, 0], [1, 1, 0], [1, 1, 1]]) # batch_size x max_len
</code></pre>
<p>masked_weights is 1 for real masked tokens and 0 for padding tokens.</p>
<p>Let’s suppose that our model returns the following logits:</p>
<pre><code class="lang-auto">batch_size, max_len, n_vocab = 3, 3, 5
torch.manual_seed(0)
logits_lm = torch.rand(batch_size, max_len, n_vocab) # batch_size x max_len x n_vocab
&gt;&gt;&gt; tensor([[[0.4963, 0.7682, 0.0885, 0.1320, 0.3074],
	         [0.6341, 0.4901, 0.8964, 0.4556, 0.6323],
	         [0.3489, 0.4017, 0.0223, 0.1689, 0.2939]],

	        [[0.5185, 0.6977, 0.8000, 0.1610, 0.2823],
	         [0.6816, 0.9152, 0.3971, 0.8742, 0.4194],
	         [0.5529, 0.9527, 0.0362, 0.1852, 0.3734]],

	        [[0.3051, 0.9320, 0.1759, 0.2698, 0.1507],
	         [0.0317, 0.2081, 0.9298, 0.7231, 0.7423],
	         [0.5263, 0.2437, 0.5846, 0.0332, 0.1387]]])
</code></pre>
<p>I also have a loss function:</p>
<pre><code class="lang-auto">criterion1 = nn.CrossEntropyLoss(reduction='none')
</code></pre>
<p>As with NSP, I have two choices:</p>
<ul>
<li>calculate the loss before reducing it:</li>
</ul>
<pre><code class="lang-auto">loss_lm = criterion1(logits_lm.transpose(1, 2), masked_ids)
loss_lm = (loss_lm*masked_weights.float()).mean()
&gt;&gt;&gt; tensor(1.0664)
</code></pre>
<ul>
<li>or reduce the entries before compute the loss.</li>
</ul>
<pre><code class="lang-auto">masked_ids = masked_ids [masked_weights == 1]
logits_lm = logits_lm [masked_weights == 1]

loss_lm = criterion1(logits_lm, masked_ids)
loss_lm = loss_lm.mean()
&gt;&gt;&gt; tensor(1.5996)
</code></pre>
<p>I would like to know which choices are better and why? Because they return different values.</p>