Candidate Sampling

Do you want to train a multiclass or multilabel model with thousands or millions of output classes (for example, a language model with a large vocabulary)? Training with a full Softmax is slow in this case, since all of the classes are evaluated for every training example. Candidate Sampling training algorithms can speed up your step times by only considering a small randomly-chosen subset of contrastive classes (called candidates) for each batch of training examples.

See our [Candidate Sampling Algorithms Reference] (../../extras/candidate_sampling.pdf)

Sampled Loss Functions

TensorFlow provides the following sampled loss functions for faster training.

tf.nn.nce_loss(weights, biases, inputs, labels, num_sampled, num_classes, num_true=1, sampled_values=None, remove_accidental_hits=False, partition_strategy='mod', name='nce_loss')

Computes and returns the noise-contrastive estimation training loss.

See [Noise-contrastive estimation: A new estimation principle for unnormalized statistical models] (http://www.jmlr.org/proceedings/papers/v9/gutmann10a/gutmann10a.pdf). Also see our [Candidate Sampling Algorithms Reference] (../../extras/candidate_sampling.pdf)

Args:
  • weights: A Tensor of shape [num_classes, dim], or a list of Tensor objects whose concatenation along dimension 0 has shape [num_classes, dim]. The (possibly-partitioned) class embeddings.
  • biases: A Tensor of shape [num_classes]. The class biases.
  • inputs: A Tensor of shape [batch_size, dim]. The forward activations of the input network.
  • labels: A Tensor of type int64 and shape [batch_size, num_true]. The target classes.
  • num_sampled: An int. The number of classes to randomly sample per batch.
  • num_classes: An int. The number of possible classes.
  • num_true: An int. The number of target classes per training example.
  • sampled_values: a tuple of (sampled_candidates, true_expected_count, sampled_expected_count) returned by a *_candidate_sampler function. (if None, we default to log_uniform_candidate_sampler)
  • remove_accidental_hits: A bool. Whether to remove "accidental hits" where a sampled class equals one of the target classes. If set to True, this is a "Sampled Logistic" loss instead of NCE, and we are learning to generate log-odds instead of log probabilities. See our [Candidate Sampling Algorithms Reference] (../../extras/candidate_sampling.pdf). Default is False.
  • partition_strategy: A string specifying the partitioning strategy, relevant if len(weights) > 1. Currently "div" and "mod" are supported. Default is "mod". See tf.nn.embedding_lookup for more details.
  • name: A name for the operation (optional).
Returns:

A batch_size 1-D tensor of per-example NCE losses.


tf.nn.sampled_softmax_loss(weights, biases, inputs, labels, num_sampled, num_classes, num_true=1, sampled_values=None, remove_accidental_hits=True, partition_strategy='mod', name='sampled_softmax_loss')

Computes and returns the sampled softmax training loss.

This is a faster way to train a softmax classifier over a huge number of classes.

This operation is for training only. It is generally an underestimate of the full softmax loss.

At inference time, you can compute full softmax probabilities with the expression tf.nn.softmax(tf.matmul(inputs, tf.transpose(weights)) + biases).

See our [Candidate Sampling Algorithms Reference] (../../extras/candidate_sampling.pdf)

Also see Section 3 of Jean et al., 2014 (pdf) for the math.

Args:
  • weights: A Tensor of shape [num_classes, dim], or a list of Tensor objects whose concatenation along dimension 0 has shape [num_classes, dim]. The (possibly-sharded) class embeddings.
  • biases: A Tensor of shape [num_classes]. The class biases.
  • inputs: A Tensor of shape [batch_size, dim]. The forward activations of the input network.
  • labels: A Tensor of type int64 and shape [batch_size, num_true]. The target classes. Note that this format differs from the labels argument of nn.softmax_cross_entropy_with_logits.
  • num_sampled: An int. The number of classes to randomly sample per batch.
  • num_classes: An int. The number of possible classes.
  • num_true: An int. The number of target classes per training example.
  • sampled_values: a tuple of (sampled_candidates, true_expected_count, sampled_expected_count) returned by a *_candidate_sampler function. (if None, we default to log_uniform_candidate_sampler)
  • remove_accidental_hits: A bool. whether to remove "accidental hits" where a sampled class equals one of the target classes. Default is True.
  • partition_strategy: A string specifying the partitioning strategy, relevant if len(weights) > 1. Currently "div" and "mod" are supported. Default is "mod". See tf.nn.embedding_lookup for more details.
  • name: A name for the operation (optional).
Returns:

A batch_size 1-D tensor of per-example sampled softmax losses.

Candidate Samplers

TensorFlow provides the following samplers for randomly sampling candidate classes when using one of the sampled loss functions above.


tf.nn.uniform_candidate_sampler(true_classes, num_true, num_sampled, unique, range_max, seed=None, name=None)

Samples a set of classes using a uniform base distribution.

This operation randomly samples a tensor of sampled classes (sampled_candidates) from the range of integers [0, range_max).

The elements of sampled_candidates are drawn without replacement (if unique=True) or with replacement (if unique=False) from the base distribution.

The base distribution for this operation is the uniform distribution over the range of integers [0, range_max).

In addition, this operation returns tensors true_expected_count and sampled_expected_count representing the number of times each of the target classes (true_classes) and the sampled classes (sampled_candidates) is expected to occur in an average tensor of sampled classes. These values correspond to Q(y|x) defined in this document. If unique=True, then these are post-rejection probabilities and we compute them approximately.

Args:
  • true_classes: A Tensor of type int64 and shape [batch_size, num_true]. The target classes.
  • num_true: An int. The number of target classes per training example.
  • num_sampled: An int. The number of classes to randomly sample per batch.
  • unique: A bool. Determines whether all sampled classes in a batch are unique.
  • range_max: An int. The number of possible classes.
  • seed: An int. An operation-specific seed. Default is 0.
  • name: A name for the operation (optional).
Returns:
  • sampled_candidates: A tensor of type int64 and shape [num_sampled]. The sampled classes.
  • true_expected_count: A tensor of type float. Same shape as true_classes. The expected counts under the sampling distribution of each of true_classes.
  • sampled_expected_count: A tensor of type float. Same shape as sampled_candidates. The expected counts under the sampling distribution of each of sampled_candidates.

tf.nn.log_uniform_candidate_sampler(true_classes, num_true, num_sampled, unique, range_max, seed=None, name=None)

Samples a set of classes using a log-uniform (Zipfian) base distribution.

This operation randomly samples a tensor of sampled classes (sampled_candidates) from the range of integers [0, range_max).

The elements of sampled_candidates are drawn without replacement (if unique=True) or with replacement (if unique=False) from the base distribution.

The base distribution for this operation is an approximately log-uniform or Zipfian distribution:

P(class) = (log(class + 2) - log(class + 1)) / log(range_max + 1)

This sampler is useful when the target classes approximately follow such a distribution - for example, if the classes represent words in a lexicon sorted in decreasing order of frequency. If your classes are not ordered by decreasing frequency, do not use this op.

In addition, this operation returns tensors true_expected_count and sampled_expected_count representing the number of times each of the target classes (true_classes) and the sampled classes (sampled_candidates) is expected to occur in an average tensor of sampled classes. These values correspond to Q(y|x) defined in this document. If unique=True, then these are post-rejection probabilities and we compute them approximately.

Args:
  • true_classes: A Tensor of type int64 and shape [batch_size, num_true]. The target classes.
  • num_true: An int. The number of target classes per training example.
  • num_sampled: An int. The number of classes to randomly sample per batch.
  • unique: A bool. Determines whether all sampled classes in a batch are unique.
  • range_max: An int. The number of possible classes.
  • seed: An int. An operation-specific seed. Default is 0.
  • name: A name for the operation (optional).
Returns:
  • sampled_candidates: A tensor of type int64 and shape [num_sampled]. The sampled classes.
  • true_expected_count: A tensor of type float. Same shape as true_classes. The expected counts under the sampling distribution of each of true_classes.
  • sampled_expected_count: A tensor of type float. Same shape as sampled_candidates. The expected counts under the sampling distribution of each of sampled_candidates.

tf.nn.learned_unigram_candidate_sampler(true_classes, num_true, num_sampled, unique, range_max, seed=None, name=None)

Samples a set of classes from a distribution learned during training.

This operation randomly samples a tensor of sampled classes (sampled_candidates) from the range of integers [0, range_max).

The elements of sampled_candidates are drawn without replacement (if unique=True) or with replacement (if unique=False) from the base distribution.

The base distribution for this operation is constructed on the fly during training. It is a unigram distribution over the target classes seen so far during training. Every integer in [0, range_max) begins with a weight of 1, and is incremented by 1 each time it is seen as a target class. The base distribution is not saved to checkpoints, so it is reset when the model is reloaded.

In addition, this operation returns tensors true_expected_count and sampled_expected_count representing the number of times each of the target classes (true_classes) and the sampled classes (sampled_candidates) is expected to occur in an average tensor of sampled classes. These values correspond to Q(y|x) defined in this document. If unique=True, then these are post-rejection probabilities and we compute them approximately.

Args:
  • true_classes: A Tensor of type int64 and shape [batch_size, num_true]. The target classes.
  • num_true: An int. The number of target classes per training example.
  • num_sampled: An int. The number of classes to randomly sample per batch.
  • unique: A bool. Determines whether all sampled classes in a batch are unique.
  • range_max: An int. The number of possible classes.
  • seed: An int. An operation-specific seed. Default is 0.
  • name: A name for the operation (optional).
Returns:
  • sampled_candidates: A tensor of type int64 and shape [num_sampled]. The sampled classes.
  • true_expected_count: A tensor of type float. Same shape as true_classes. The expected counts under the sampling distribution of each of true_classes.
  • sampled_expected_count: A tensor of type float. Same shape as sampled_candidates. The expected counts under the sampling distribution of each of sampled_candidates.

tf.nn.fixed_unigram_candidate_sampler(true_classes, num_true, num_sampled, unique, range_max, vocab_file='', distortion=1.0, num_reserved_ids=0, num_shards=1, shard=0, unigrams=(), seed=None, name=None)

Samples a set of classes using the provided (fixed) base distribution.

This operation randomly samples a tensor of sampled classes (sampled_candidates) from the range of integers [0, range_max).

The elements of sampled_candidates are drawn without replacement (if unique=True) or with replacement (if unique=False) from the base distribution.

The base distribution is read from a file or passed in as an in-memory array. There is also an option to skew the distribution by applying a distortion power to the weights.

In addition, this operation returns tensors true_expected_count and sampled_expected_count representing the number of times each of the target classes (true_classes) and the sampled classes (sampled_candidates) is expected to occur in an average tensor of sampled classes. These values correspond to Q(y|x) defined in this document. If unique=True, then these are post-rejection probabilities and we compute them approximately.

Args:
  • true_classes: A Tensor of type int64 and shape [batch_size, num_true]. The target classes.
  • num_true: An int. The number of target classes per training example.
  • num_sampled: An int. The number of classes to randomly sample per batch.
  • unique: A bool. Determines whether all sampled classes in a batch are unique.
  • range_max: An int. The number of possible classes.
  • vocab_file: Each valid line in this file (which should have a CSV-like format) corresponds to a valid word ID. IDs are in sequential order, starting from num_reserved_ids. The last entry in each line is expected to be a value corresponding to the count or relative probability. Exactly one of vocab_file and unigrams needs to be passed to this operation.
  • distortion: The distortion is used to skew the unigram probability distribution. Each weight is first raised to the distortion's power before adding to the internal unigram distribution. As a result, distortion = 1.0 gives regular unigram sampling (as defined by the vocab file), and distortion = 0.0 gives a uniform distribution.
  • num_reserved_ids: Optionally some reserved IDs can be added in the range [0, num_reserved_ids] by the users. One use case is that a special unknown word token is used as ID 0. These IDs will have a sampling probability of 0.
  • num_shards: A sampler can be used to sample from a subset of the original range in order to speed up the whole computation through parallelism. This parameter (together with shard) indicates the number of partitions that are being used in the overall computation.
  • shard: A sampler can be used to sample from a subset of the original range in order to speed up the whole computation through parallelism. This parameter (together with num_shards) indicates the particular partition number of the operation, when partitioning is being used.
  • unigrams: A list of unigram counts or probabilities, one per ID in sequential order. Exactly one of vocab_file and unigrams should be passed to this operation.
  • seed: An int. An operation-specific seed. Default is 0.
  • name: A name for the operation (optional).
Returns:
  • sampled_candidates: A tensor of type int64 and shape [num_sampled]. The sampled classes.
  • true_expected_count: A tensor of type float. Same shape as true_classes. The expected counts under the sampling distribution of each of true_classes.
  • sampled_expected_count: A tensor of type float. Same shape as sampled_candidates. The expected counts under the sampling distribution of each of sampled_candidates.

Miscellaneous candidate sampling utilities


tf.nn.compute_accidental_hits(true_classes, sampled_candidates, num_true, seed=None, name=None)

Compute the position ids in sampled_candidates matching true_classes.

In Candidate Sampling, this operation facilitates virtually removing sampled classes which happen to match target classes. This is done in Sampled Softmax and Sampled Logistic.

See our Candidate Sampling Algorithms Reference.

We presuppose that the sampled_candidates are unique.

We call it an 'accidental hit' when one of the target classes matches one of the sampled classes. This operation reports accidental hits as triples (index, id, weight), where index represents the row number in true_classes, id represents the position in sampled_candidates, and weight is -FLOAT_MAX.

The result of this op should be passed through a sparse_to_dense operation, then added to the logits of the sampled classes. This removes the contradictory effect of accidentally sampling the true target classes as noise classes for the same example.

Args:
  • true_classes: A Tensor of type int64 and shape [batch_size, num_true]. The target classes.
  • sampled_candidates: A tensor of type int64 and shape [num_sampled]. The sampled_candidates output of CandidateSampler.
  • num_true: An int. The number of target classes per training example.
  • seed: An int. An operation-specific seed. Default is 0.
  • name: A name for the operation (optional).
Returns:
  • indices: A Tensor of type int32 and shape [num_accidental_hits]. Values indicate rows in true_classes.
  • ids: A Tensor of type int64 and shape [num_accidental_hits]. Values indicate positions in sampled_candidates.
  • weights: A Tensor of type float and shape [num_accidental_hits]. Each value is -FLOAT_MAX.