Introduction

Negative sampling is a training trick for problems with a very large output space, such as word prediction or item recommendation. Instead of comparing each example against every possible class on every step, TensorFlow lets you train against the true target plus a small sample of negatives, which makes training much cheaper while still learning useful embeddings.

Why negative sampling exists

Suppose you are training a skip-gram model with a vocabulary of 100,000 words. A full softmax update touches all 100,000 output classes for each example. That is expensive in both time and memory.

Negative sampling replaces the full comparison with a smaller problem:

• one or more true labels
• a limited number of sampled negatives

The model then learns to score the true label higher than the sampled negatives.

TensorFlow APIs used for sampled losses

TensorFlow exposes this idea through functions such as tf.nn.sampled_softmax_loss and tf.nn.nce_loss. Both expect a weight matrix for output classes, bias values, the true labels, and input activations.

A minimal example looks like this:

1import tensorflow as tf
2
3vocab_size = 10000
4embedding_dim = 64
5batch_size = 4
6num_negative = 20
7
8inputs = tf.random.normal((batch_size, embedding_dim))
9labels = tf.constant([[3], [17], [25], [999]], dtype=tf.int64)
10
11weights = tf.Variable(tf.random.normal((vocab_size, embedding_dim)))
12biases = tf.Variable(tf.zeros((vocab_size,)))
13
14loss = tf.nn.sampled_softmax_loss(
15    weights=weights,
16    biases=biases,
17    labels=labels,
18    inputs=inputs,
19    num_sampled=num_negative,
20    num_classes=vocab_size,
21)
22
23print(tf.reduce_mean(loss).numpy())

The important shapes are:

• 'weights: one row per class'
• 'inputs: one row per training example'
• 'labels: shape [batch_size, num_true], often one true class per example'

Understand what the sampled loss is approximating

sampled_softmax_loss is an approximation used during training. It is not the same as computing exact probabilities over all classes at inference time. That distinction matters when people expect the sampled loss itself to produce calibrated final probabilities.

In many embedding tasks, the training objective is mainly about learning useful vector geometry, not about exact normalized class probabilities on every step.

A small train step example

1import tensorflow as tf
2
3vocab_size = 5000
4embedding_dim = 32
5num_negative = 10
6
7class ToyModel(tf.keras.Model):
8    def __init__(self):
9        super().__init__()
10        self.encoder = tf.keras.layers.Embedding(vocab_size, embedding_dim)
11        self.output_weights = tf.Variable(tf.random.normal((vocab_size, embedding_dim)))
12        self.output_biases = tf.Variable(tf.zeros((vocab_size,)))
13
14    def call(self, inputs):
15        return self.encoder(inputs)
16
17model = ToyModel()
18optimizer = tf.keras.optimizers.Adam(1e-3)
19
20context_ids = tf.constant([1, 4, 9, 16], dtype=tf.int32)
21target_ids = tf.constant([[2], [5], [10], [20]], dtype=tf.int64)
22
23with tf.GradientTape() as tape:
24    embeddings = model(context_ids)
25    loss = tf.reduce_mean(
26        tf.nn.sampled_softmax_loss(
27            weights=model.output_weights,
28            biases=model.output_biases,
29            labels=target_ids,
30            inputs=embeddings,
31            num_sampled=num_negative,
32            num_classes=vocab_size,
33        )
34    )
35
36grads = tape.gradient(loss, model.trainable_variables)
37optimizer.apply_gradients(zip(grads, model.trainable_variables))
38print(float(loss))

This shows the usual workflow: encode the input, compute sampled loss against target classes, and update the model.

nce_loss versus sampled_softmax_loss

Both losses use sampled negatives, but they are not identical objectives. nce_loss comes from noise-contrastive estimation, while sampled_softmax_loss approximates a softmax-style objective. In practice, both are used for large-vocabulary training, and the better choice depends on the model and evaluation goal.

If you are following a paper or an existing TensorFlow example, use the loss that matches that training recipe instead of swapping them casually.

Sample quality matters

Negative sampling is only as good as the negatives you draw. If the sampled negatives are too easy, the model learns slowly. If the sampling distribution is badly mismatched to the task, the learned embeddings may be less useful.

That is why production recommender and NLP systems often spend real effort designing their sampling strategy instead of treating num_sampled as the only tuning knob.

Common Pitfalls

• Expecting sampled loss to behave exactly like a full softmax probability calculation.
• Passing labels with the wrong shape or dtype.
• Forgetting that the output weight matrix must have one row per class.
• Using too few negative samples and then wondering why training quality is weak.
• Swapping nce_loss and sampled_softmax_loss without understanding the objective difference.

Summary

• Negative sampling makes large-output training cheaper by comparing true labels against sampled negatives.
• TensorFlow supports this with APIs such as tf.nn.sampled_softmax_loss and tf.nn.nce_loss.
• The sampled objective is a training approximation, not a full probability computation.
• Correct tensor shapes and class counts are essential.
• Sampling strategy and the number of negatives affect the quality of the learned embeddings.