Frequency-aware SGD for Efficient Embedding Learning with Provable Benefits

Abstract

Embedding learning has found widespread applications in recommendation systems and natural language modeling, among other domains. To learn quality embeddings efficiently, adaptive learning rate algorithms have demonstrated superior empirical performance over SGD, largely accredited to their token-dependent learning rate. However, the underlying mechanism for the efficiency of token-dependent learning rate remains underexplored. We show that incorporating frequency information of tokens in the embedding learning problems leads to provably efficient algorithms, and demonstrate that common adaptive algorithms implicitly exploit the frequency information to a large extent. Specifically, we propose (Counter-based) Frequency-aware Stochastic Gradient Descent, which applies a frequency-dependent learning rate for each token, and exhibits provable speed-up compared to SGD when the token distribution is imbalanced. Empirically, we show the proposed algorithms are able to improve or match adaptive algorithms on benchmark recommendation tasks and a large-scale industrial recommendation system, closing the performance gap between SGD and adaptive algorithms. Our results are the first to show token-dependent learning rate provably improves convergence for non-convex embedding learning problems.

arXiv commentAdditional experiments on Word2Vec embedding learning included

References in corpus (15)

• Adam: A Method for Stochastic Optimization
• Efficient Estimation of Word Representations in Vector Space
• Distributed Representations of Words and Phrases and their Compositionality
• The diameter of the world wide web
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• On the Convergence of Adam and Beyond
• DeepFM: A Factorization-Machine based Neural Network for CTR Prediction
• Pointer Sentinel Mixture Models
• Origins of power-law degree distribution in the heterogeneity of human activity in social networks
• On the Convergence of A Class of Adam-Type Algorithms for Non-Convex Optimization
• On the Convergence of Adaptive Gradient Methods for Nonconvex Optimization
• Why are Adaptive Methods Good for Attention Models?
• Deep Interest Network for Click-Through Rate Prediction
• A Simple Convergence Proof of Adam and Adagrad
• Understanding the Difficulty of Training Transformers