Reproduce “Frequency-aware SGD for Efficient Embedding Learning with Provable Benefits” on OpenResearch

About this paper

Yan Li, Dhruv Choudhary, Xiaohan Wei, Baichuan Yuan, Bhargav Bhushanam, Tuo Zhao, Guanghui Lan

Researchers from Georgia Institute of Technology and Facebook developed Frequency-aware SGD (FA-SGD) and Counter-based FA-SGD (CF-SGD), optimization algorithms for embedding learning that apply token-dependent learning rates based on token frequency. This work provides the first theoretical proof that such frequency-aware approaches offer provable convergence benefits over standard SGD in non-convex settings, while empirically demonstrating CF-SGD's competitive performance against adaptive methods like Adam and Adagrad on large-scale datasets, all while significantly reducing memory overhead.

Key insights

• Adaptive optimizers like Adam and Adagrad implicitly adjust learning rates based on token frequency, explaining their empirical success in embedding learning tasks.
• Infrequent tokens in embedding tables, often crucial for model quality, benefit significantly from higher, frequency-aware learning rates.
• Explicitly incorporating token frequency into learning rate schedules can overcome the limitations of general-purpose optimizers for highly sparse and imbalanced data.

Summary by alphaXiv.