How is cross-entropy used in seq2seq models?

kirk86: <p>Hi folks,<br>
I have a bit of clarifying (technical) question regarding using cross-entropy with seq2seq models?</p>
<pre><code class="lang-auto">B = batch_size
S = sequence_length (subscripts, i=input, o=output)
E = embedding_size
V = vocabulary_size
C = num_classes/labels

X = [B, S_i, V]  # our input data

Y = [B, S_o, C] # our targets, consider them one-hot encoded

X_enc = encoder(X) # X_enc = [B, S_i, E]

X_dec = decoder(X_enc, Y) # X_dec = [B, S_o, E]

loss = nn.CrossEntropyLoss()(X_dec, Y) # this shoudn't work due to shape mismatch
</code></pre>
<p>How should cross entropy work in this case? There’s a shape mismatch between <code>X_dec != Y</code>?</p>
<p>Plus, to my understanding cross entropy accepts inputs of shape <code>[B, C]</code> and targets of shape <code>[B, 1]</code>, what am I missing here?</p>
<p>Is there another way how is this computed?</p>
<p>One way I was thinking about going at it would have been the following, but I’m not sure if its correct?</p>
<pre><code class="lang-auto">torch.mean(torch.sum((-Y * F.log_softmax(X_dec, dim=-1)), dim=-1))
</code></pre>

ptrblck: <aside class="quote no-group" data-username="kirk86" data-post="1" data-topic="196871">
<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/kirk86/48/151_2.png" class="avatar"> kirk86:</div>
<blockquote>
<p>Plus, to my understanding cross entropy accepts inputs of shape <code>[B, C]</code> and targets of shape <code>[C, 1]</code>, what am I missing here?</p>
</blockquote>
</aside>
<p>That’s not the case as described in the <a href="https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html">docs</a>.<br>
If you are using class indices the target should have the shape <code>[B]</code> given your model output shape. If the target contains probabilities it should match the model output shape, so <code>[B, C]</code>.</p>

kirk86: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="196871">
<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>the target should have the shape <code>[B]</code> given your model output shape</p>
</blockquote>
</aside>
<p>Yes correct, my mistake, I’ve corrected it in the post.</p>
<p>But, do you have any idea how is cross entropy computed when both <code>outputs = [B, S_o, E]</code> and <code>targets_one_hot = [B, S_o, C]</code> are 3-dim? This seems quite common loss in seq2seq tasks but since this case is not covered in the docs (unless I’m missing something), is there another reference or MWE?</p>

ptrblck: <p>The docs also cover this use case and will treat the output as <code>[batch_size, nb_classes, additional_dimension]</code>. Since the target is also 3-dimensional is must have the same shape, so <code>C==E</code> in your case.<br>
Here is an example:</p>
<pre data-code-wrap="python"><code class="lang-python">batch_size = 2
nb_classes = 3
seq_len = 4
output = torch.randn(batch_size, nb_classes, seq_len, requires_grad=True)
target = torch.randint(0, nb_classes, (batch_size, seq_len))
target = torch.nn.functional.one_hot(target, num_classes=nb_classes).permute(0, 2, 1).float()
print(target.shape)
# torch.Size([2, 3, 4])

criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)
</code></pre>
<p>The criterion will compute the cross-entropy loss for each sample in <code>seq_len</code> in the same way it’s doing so for each sample in the batch dimension.</p>

kirk86: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="196871">
<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>criterion will compute the cross-entropy loss for each sample in <code>seq_len</code> i</p>
</blockquote>
</aside>
<p>Awesome, thanks a lot for clarifying this!</p>
<p>As I was searching around I saw that in many cases another approached that was used was to flatten the 3-dim tensors from <code>[B, S, E]</code> to <code>[B*S, E]</code> and then compute the loss, which I presume is similar to what you described.</p>

ptrblck: <p>Yes, but be careful with the dimensions.<br>
Using my example you would need to <code>permute</code> the target before flattening:</p>
<pre data-code-wrap="python"><code class="lang-python">batch_size = 2
nb_classes = 3
seq_len = 4
output = torch.randn(batch_size, nb_classes, seq_len, requires_grad=True)
target = torch.randint(0, nb_classes, (batch_size, seq_len))
target = torch.nn.functional.one_hot(target, num_classes=nb_classes).permute(0, 2, 1).float()
print(target.shape)
# torch.Size([2, 3, 4])

criterion = nn.CrossEntropyLoss()
loss = criterion(output, target)

# with flattened tensors treating the sequence dimension as separate samples in the batch dimension
print(output.shape)
# torch.Size([2, 3, 4]) = [batch_size, nb_classes, seq_len]
output = output.permute(0, 2, 1).contiguous().view(-1, nb_classes)
print(output.shape)
# torch.Size([8, 3]) = [batch_size * seq_len, nb_classes]

print(target.shape)
# torch.Size([2, 3, 4]) = [batch_size, nb_classes, seq_len]
target = target.permute(0, 2, 1).contiguous().view(-1, nb_classes)
print(target.shape)
# torch.Size([8, 3]) = [batch_size * seq_len, nb_classes]

loss2 = criterion(output, target)

print(loss - loss2)
# tensor(0., grad_fn=&lt;SubBackward0&gt;)
</code></pre>

kirk86: <p>Thanks a lot.</p>
<p>As I understood from your MWE there are 2 <code>key</code> points here.</p>
<ol>
<li>
<p>In order for cross-entropy to work with 3-dim tensors we should have <code>nb_classes</code> as dim=1 and let cross-entropy compute the loss over the <code>nb_classes</code>, e.g., both <code>output</code> and <code>target</code> to have shape <code>[batch_size, nb_classes, seq_len]</code></p>
</li>
<li>
<p>Otherwise if we have to flatten dimensions we should flatten across <code>batch_size*seq_len</code>, therefore our tensors should have shape <code>[batch_size, seq_len, nb_classes]</code>, i.e., my case.</p>
</li>
</ol>