Better Embeddings with Coupled Adam
Authors: name: Felix Stollenwerk; name: Tobias Stollenwerk
Published: 2025-02-12 14:32:17

Abstract: Despite their remarkable capabilities, LLMs learn word representations that
exhibit the undesirable yet poorly understood feature of anisotropy. In this
paper, we argue that the second moment in Adam is a cause of anisotropic
embeddings, and suggest a modified optimizer called Coupled Adam to mitigate
the problem. Our experiments demonstrate that Coupled Adam significantly
improves the quality of embeddings, while also leading to better upstream and
downstream performance on large enough datasets.