[Paper Review] Adaptive Sampled Softmax with Kernel Based Sampling

Guy Blanc, Steffen Rendle|arXiv (Cornell University)|Dec 2, 2017

Domain Adaptation and Few-Shot Learning3 citations

TL;DR

This paper proposes kernel-based sampling for adaptive sampled softmax, using a quadratic kernel to approximate the softmax distribution efficiently. It enables low-bias training with significantly fewer samples—1–2 orders of magnitude less than uniform sampling—by adapting to model outputs via a divide-and-conquer algorithm in O(D log n) time, where D is the kernel space dimension.

ABSTRACT

Softmax is the most commonly used output function for multiclass problems and is widely used in areas such as vision, natural language processing, and recommendation. A softmax model has linear costs in the number of classes which makes it too expensive for many real-world problems. A common approach to speed up training involves sampling only some of the classes at each training step. It is known that this method is biased and that the bias increases the more the sampling distribution deviates from the output distribution. Nevertheless, almost any recent work uses simple sampling distributions that require a large sample size to mitigate the bias. In this work, we propose a new class of kernel based sampling methods and develop an efficient sampling algorithm. Kernel based sampling adapts to the model as it is trained, thus resulting in low bias. Kernel based sampling can be easily applied to many models because it relies only on the model's last hidden layer. We empirically study the trade-off of bias, sampling distribution and sample size and show that kernel based sampling results in low bias with few samples.

Motivation & Objective

• Address the high bias in sampled softmax when using non-softmax sampling distributions, especially with large class counts.
• Overcome the inefficiency of uniform or popularity-based sampling, which require large sample sizes to reduce bias.
• Develop a sampling method that adapts to model outputs, model structure, and parameters during training.
• Enable efficient sampling for large-scale multiclass problems by leveraging kernel methods and logarithmic-time algorithms.
• Demonstrate that kernel-based sampling achieves full softmax quality with far fewer samples than standard approaches.

Proposed method

• Propose a kernel-based sampling distribution that depends on the model’s last hidden layer output, enabling example-, model- and parameter-dependent sampling.
• Use a quadratic kernel to approximate the exponential softmax function, allowing efficient computation of partition functions in kernel space.
• Design a divide-and-conquer sampling algorithm that runs in O(D log n) time, where D is the dimension of the kernel space and n is the number of classes.
• Apply the kernel-based sampling only during training, preserving the full exponential softmax for loss computation to maintain model quality.
• Ensure the method is generic and applicable to any model with a dot-product final layer between hidden representations and class embeddings.
• Leverage kernel properties to enable fast sampling without requiring full softmax computation during training steps.

Experimental results

Research questions

• RQ1 Can kernel-based sampling reduce bias in sampled softmax training compared to uniform or popularity-based sampling?
• RQ2 How does the sample efficiency of kernel-based sampling compare to uniform sampling in terms of convergence quality and speed?
• RQ3 Can a kernel-based sampling method be efficiently computed for large-scale multiclass problems with n >> 10^4 classes?
• RQ4 Does adaptive sampling based on model outputs lead to faster convergence or better generalization than non-adaptive sampling?
• RQ5 To what extent can kernel-based sampling approximate the full softmax distribution while maintaining computational efficiency?

Key findings

• Kernel-based sampling, particularly with a quadratic kernel, reduces bias significantly compared to uniform sampling, achieving full softmax performance with one to two orders of magnitude fewer samples.
• The proposed divide-and-conquer sampling algorithm runs in O(D log n) time, making it scalable for large numbers of classes.
• Sampling distributions that adapt to model outputs—such as those based on kernels—dramatically improve sample efficiency compared to static distributions like uniform or global popularity.
• Even with minimal sampling (e.g., m ≈ 10–50), kernel-based sampling achieves convergence quality comparable to full softmax, while uniform sampling requires much larger m to reach similar performance.
• Once bias is sufficiently reduced, increasing the sample size beyond a certain point does not improve convergence speed, indicating diminishing returns for larger m.
• The method preserves the high-quality output of exponential softmax by using the kernel only for sampling, while computing the loss via full softmax, thus maintaining model accuracy.