skipgram word embeddings results in nan values

medmsion690: <p>After training a word embedding model on a large-ish corpus, my embeddings converge to <code>nan</code> values.</p>
<p>The model is very simple (skipgram with negative sampling)</p>
<p>full model:</p>
<pre><code class="lang-python">class NEG_loss(nn.Module):
    def __init__(self, vocab_size, embed_size, neg_sampling_table=None):
        """
        :param vocab_size: An int. The number of possible classes.
        :param embed_size: An int. EmbeddingLockup size
        :param num_sampled: An int. The number of sampled from noise examples
        :param neg_sampling_table: A list of non negative floats. Class neg_sampling_table. None if
            using uniform sampling. The neg_sampling_table are calculated prior to
            estimation and can be of any form, e.g equation (5) in [1]
        """

        super(NEG_loss, self).__init__()

        self.device = t.device("cuda:0" if t.cuda.is_available() else "cpu")

        self.vocab_size = vocab_size
        self.embed_size = embed_size

        self.out_embed.weight = nn.Parameter(
            t.cat(
                [
                    t.zeros(1, self.embed_size),
                    FT(self.vocab_size - 1, self.embed_size).uniform_(
                        -0.5 / self.embed_size, 0.5 / self.embed_size
                    ),
                ]
            )
        )


        self.in_embed.weight = nn.Parameter(
            t.cat(
                [
                    t.zeros(1, self.embed_size),
                    FT(self.vocab_size - 1, self.embed_size).uniform_(
                        -0.5 / self.embed_size, 0.5 / self.embed_size
                    ),
                ]
            )
        )

        self.neg_sampling_table = neg_sampling_table
        if self.neg_sampling_table is not None:
            assert min(self.neg_sampling_table) &gt;= 0, "Each weight should be &gt;= 0"

            self.neg_sampling_table = Variable(t.from_numpy(neg_sampling_table)).float()

    # TODO this is bad - find more elegant solution
    def sample(self, num_sample):
        """
        draws a sample from classes based on neg_sampling_table
        """

        return self.neg_sampling_table[
            t.randint(0, len(self.neg_sampling_table), (num_sample,))
        ]

    def forward(self, input_labels, out_labels, num_sampled):
        """
        :param input_labels: Tensor with shape of [batch_size] of Long type
        :param out_labels: Tensor with shape of [batch_size, window_size] of Long type
        :param num_sampled: An int. The number of sampled from noise examples
        :return: Loss estimation with shape of [1]
            loss defined in Mikolov et al. Distributed Representations of Words and Phrases and their Compositionality
            papers.nips.cc/paper/5021-distributed-representations-of-words-and-phrases-and-their-compositionality.pdf
        """

        input_labels = input_labels.to("cuda:0")
        out_labels = out_labels.to("cuda:0")

        batch_size = out_labels.size()[0]

 
        input_ = self.in_embed(input_labels.view(-1))
        output = self.out_embed(out_labels.view(-1))

        if self.neg_sampling_table is not None:
           
            noise_sample_count = batch_size * num_sampled
            draw = self.sample(noise_sample_count)
            noise = draw.view(batch_size, num_sampled).long()
            
        else:


            noise = Variable(
                t.Tensor(batch_size, num_sampled)
                .uniform_(0, self.vocab_size - 1)
                .long()
            )
        # if use_cuda:
        noise = noise.to(self.device)
        noise = self.out_embed(noise).neg()

        log_target = (input_ * output).sum(1).squeeze().sigmoid().log()

        """ ∑[batch_size * window_size, num_sampled, embed_size] * [batch_size * window_size, embed_size, 1] -&gt;
            ∑[batch_size, num_sampled, 1] -&gt; [batch_size] """
        sum_log_sampled = (
            t.bmm(noise, input_.unsqueeze(2)).sigmoid().log().sum(1).squeeze()
        )

        loss = log_target + sum_log_sampled

        return -loss.mean()

    def input_embeddings(self):
        return self.in_embed.weight.detach().cpu().numpy()
</code></pre>
<p>and training loop and optimizer code</p>
<pre><code class="lang-python">    # NEG loss and optim
    neg = NEG_loss(vocab_size, dim, neg_sampling_table=neg_sampling_dist)
    neg.to("cuda:0")
    optimizer = Adam(neg.parameters(), 0.01)

    sys.stdout.write("BEGINNING TRAINING\n")
    for i in range(epochs):
        sys.stdout.write("-" * 35 + "\n")
        sys.stdout.write(f"EPOCH {i+1}\n")
        for i, batch in enumerate(tqdm(dataloader)):
            input_, output_ = batch
            optimizer.zero_grad()
            loss = neg(input_, output_, neg_samples)
            loss.backward()
            optimizer.step()
        sys.stdout.write("-" * 35 + "\n")
</code></pre>
<p>This model is about as simple as it gets so I am a bit surprised I am having this issue. Everything is fine on a smaller corpus (7k batches of size 128) but gets <code>nan</code>-y with a larger corpus (200k batches of size 128).</p>
<p>Anyone see anything immediately wrong or have any tips to figure out what is going?</p>
<p>Any help is appreciated. Also pretty new to pytorch - so if you see anything that is a bit dumb feel free to point it out <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>Could you check, if all inputs contain valid values via <code>torch.isfinite(input).all()</code> and <code>torch.isfinite(output).all()</code>?<br>
If the data is alright, I would recommend to set <code>torch.autograd.detect_anomaly(True)</code> at the beginning of the script and post the stack trace here. This should hopefully point to the operation, which created the NaN outputs.</p>

medmsion690: <p>Thanks for your response <a class="mention" href="/u/ptrblck">@ptrblck</a>. I do not believe my inputs are the culprit - they are just indices that then look up the embedding within the <code>in_embed</code> layer. Do you mean to check this for at model init for all of <code>in_embed</code>.</p>
<p>In the meantime, I am currently running the model using the <code>torch.autograd.detect_anomaly(True)</code>.</p>
<p>Thanks again</p>

medmsion690: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I included the <code>torch.autograd.detect_anomaly(True)</code> before my training loop and my script ran to completion - still yielding embedding weights of <code>nan</code>. Any other recommendations?</p>

ptrblck: <p>Are you seeing a NaN output after your embedding weights are getting NaN values?<br>
Anomaly detection should raise an error, so I’m a bit confused why your script finished.</p>