This work focuses on the training dynamics of one associative memory module storing outer products of token embeddings. We reduce this problem to the study of a system of particles, which interact according to properties of the data distribution and correlations between embeddings. Through theory and experiments, we provide several insights. In overparameterized regimes, we obtain logarithmic growth of the ``classification margins.'' Yet, we show that imbalance in token frequencies and memory interferences due to correlated embeddings lead to oscillatory transitory regimes. The oscillations are more pronounced with large step sizes, which can create benign loss spikes, although these learning rates speed up the dynamics and accelerate the asymptotic convergence. In underparameterized regimes, we illustrate how the cross-entropy loss can lead to suboptimal memorization schemes. Finally, we assess the validity of our findings on small Transformer models.

翻译：本文聚焦于存储令牌嵌入外积的联想记忆模块的训练动态。我们将该问题简化为一个粒子系统的研究，这些粒子根据数据分布的特性及嵌入之间的相关性而相互作用。通过理论与实验，我们提供了若干见解。在过参数化场景中，我们获得了“分类边界”的对数增长。然而，我们表明令牌频率的不平衡以及由相关嵌入引起的记忆干扰会导致振荡过渡状态。当步长较大时，振荡更为显著，这可能会产生良性的损失尖峰，尽管这些学习率加速了动态过程并加快了渐近收敛。在欠参数化场景中，我们展示了交叉熵损失如何导致次优的记忆机制。最后，我们在小型Transformer模型上验证了我们的发现的有效性。