Training a single embedding using masking

amm90: <p>I am trying to train a one embedding layer using masking.</p>
<p>It takes a masked sentence of 10 tokens and predict the masked tokens.</p>
<p>The values are the ids of the tokens in a vocabulary.<br>
E.g:. [223, 444, 1, 11, 53, 232, 1, 435, 12, 43]</p>
<p>The target is the same sequence with all non-masked tokens are replaced by empty token 0:<br>
E.g.: [0, 0, 33, 0, 0, 0, 23, 0, 0, 0] During loss computation index 0 is ignored.</p>
<p>How to modle the classification problem. In layman terms for each entry I am computing<br>
the class of 10 tokens. How this should be reflected in the network.</p>
<p>The network is as follow:</p>
<pre><code class="lang-auto">        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.linear = nn.Linear(max_length * embedding_dim, max_length * vocab_size)

    def forward(self, x):
        embedded = self.embedding(x) # (batch_size, max_length, embedding_size)
        embedded = embedded.view(x.size()[0], -1) # (batch_size, max_length * embedding_size)
        output = self.linear(embedded) # (batch_size, max_length * vocab_size)
        output = output.view(?, ?, ?)
        return output
</code></pre>
<p>In particular what should be the shae of the output to be suitable for a<br>
crossEntropyLoss computation.</p>
<p>(A): (N, C, 10)<br>
or<br>
(B): (N, 10, C)</p>
<p>Where C is the vocab_size or n_classes<br>
Where N is batch size</p>
<p>I actually tried both.<br>
(A) is working fine.<br>
(B giving me an error: expected output of size [N, C] got [N, 10])</p>

ptrblck: <p>I’m not sure I entirely understand you use case, but given the mentioned output shapes I would say (A) is the expected shape as generally <code>nn.CrossEntropyLoss</code> expects a model output in the shape <code>[batch_size, nb_classes, *additional_dimensions]</code>.</p>

amm90: <p>To make the problem clearer:</p>
<p>Say, in the context of natural language processing (NLP) I have<br>
a model that predicts the next word. My vocabulary size is 60,000.<br>
Hence the model output (logits) should be <em>60,000</em>.<br>
In this case the Cross Entropy module will get the argument of<br>
the maximum activated unit using SoftMax and then apply entropy loss.</p>
<p>Now I want to extend the model to predict 10 tokens rather than only one.<br>
The output size is (10x60,000) for each token we 60,000 units that we extract the<br>
maximum activated unit from.</p>
<p>In the latter case how, we should reshape the module so that it works with cross entropy.<br>
Is this multidimensional loss?? Or I should manually calculate softmax and argmax.</p>

ptrblck: <p><code>60,000</code> would refer to the class count in your example and should thus be in <code>dim1</code>. I’m unsure how tokens are used and if these refer to independent samples, in which case they should land in <code>dim0</code> (the batch dimension), or if these tokens represent a temporal dependency they would land in <code>dim2</code> where <code>dim0</code> would still represent the batch dimension and could be <code>1</code> if a single sample with 60,000 tokens is used.</p>

amm90: <p>True. Solution as in <a href="https://discuss.pytorch.org/t/nn-crossentropyloss-for-text-with-multiple-dimension/90662">nn.CrossEntropyLoss() for text with multiple dimension</a></p>
<p>However it is misleading to have the number of classes as a dimension.<br>
It make sense to have the output as (batch_size, dimension1, nb_classes) which will be<br>
reduced to (batch_size, dimension1).</p>
<p>It is misleading because from application perspective the output should be (batch_size, dimension1, nb_classes). You hve to add another functional operation to clear the difference.</p>
<p>For example check this <a href="https://neptune.ai/blog/how-to-code-bert-using-pytorch-tutorial" rel="noopener nofollow ugc">bert training code</a>. They added “output.transpose(1, 2)” operation before computing the loss.</p>

ptrblck: <p>I don’t think it’s generally misleading as the needed permutations would depend on the use case. E.g. segmentation models using a <code>nn.Conv2d</code> layer as the output define the output channels as the number of classes and would thus not need any transpose. It could of course differ for your language model, but changing the default would force transposes for other models.</p>