Sampled Softmax with Random Fourier Features

By Ankit Singh Rawat, Jiecao Chen, Felix X. Yu, Ananda Theertha Suresh, Sanjiv Kumar. 2019

61 citations recorded by the source corpus.

The computational cost of training with softmax cross entropy loss grows linearly with the number of classes. For the settings where a large number of classes are involved, a common method to speed up training is to sample a subset of classes and utilize an estimate of the loss gradient based on these classes, known as the sampled softmax method. However, the sampled softmax provides a biased estimate of the gradient unless the samples are drawn from the exact softmax distribution, which is again expensive to compute. Therefore, a widely employed practical approach involves sampling from a simpler distribution in the hope of approximating the exact softmax distribution. In this paper, we develop the first theoretical understanding of the role that different sampling distributions play in determining the quality of sampled softmax. Motivated by our analysis and the work on kernel-based sampling, we propose the Random Fourier Softmax (RF-softmax) method that utilizes the powerful Random Fourier Features to enable more efficient and accurate sampling from an approximate softmax distribution. We show that RF-softmax leads to low bias in estimation in terms of both the full softmax distribution and the full softmax gradient. Furthermore, the cost of RF-softmax scales only logarithmically with the number of classes.

Subfields: Artificial Intelligence, Representation & Self-Supervised Learning

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Representation & Self-Supervised Learning

Citations in the corpus

Socratic Models: Composing Zero-Shot Multimodal Reasoning with Language (2022, methodology): “The other approach, that combined the ideas from both MIP-search and linear attention systems, leverages the so-called random feature tree (RFT) data structure (Rawat et al. 2019).”
Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (2020, background): “Recent works (Blanc & Rendle, 2017; Rawat et al., 2019), have approximated softmax with a linear dot product of feature maps to speed up the training through sampling.”
Nyströmformer: A Nystöm-based Algorithm for Approximating Self-Attention (2021, methodology): “In (Rawat et al. 2019), the Random Fourier Softmax (RF-softmax) idea uses random Fourier features to perform efficient sampling from an approximate softmax distribution for normalized embedding.”
Random Feature Attention (2021, methodology): “R FA builds on a kernel perspective of softmax (Rawat et al., 2019).”
Hybrid Random Features (2021, methodology): “In this subsection, we will describe how one can use HRF to train a LSTM on the language modeling task on the PennTree Bank dataset, similar to the experiments carried out in (Rawat et al., 2019).”
FEDformer: Frequency Enhanced Decomposed Transformer for Long-term Series Forecasting (2022, methodology): “(Rawat et al., 2019) proposes the Random Fourier softmax (RF-softmax) method that utilizes the powerful Random Fourier Features to enable more efﬁcient and accurate sampling from an approximate softmax distribution.”
Choose a Transformer: Fourier or Galerkin (2021, methodology): “Another type of linearization exploits the low-rank nature of the matrix product using various methods such as sampling or projection [73, 11, 79, 92], or fast multipole decomposition [65].”
Visual Attention Methods in Deep Learning: An In-Depth Survey (2022, methodology): “RFA can be integrated into backbones easily to replace the normal softmax without increasing the number of parameters, only 0 .”