arXiv:2603.15923stat.MLcs.LG2026-03ICLR被引 1

揭示非正交嵌入下Transformer的记忆容量与数据量、维度、序列长度的乘积关系

Learning to Recall with Transformers Beyond Orthogonal Embeddings

  • 在随机嵌入下分析单层Transformer的早期梯度下降过程
  • 发现存储容量与样本数、维度、序列长度呈乘积关系
  • 适用于研究大模型知识记忆机制的研究者和工程师

现代大型语言模型在需要存储和检索知识的任务中表现优异,如事实回忆和问答。Transformer是这一能力的核心,因其能在训练中编码信息并在推理时检索。现有理论分析多基于理想假设,如无限数据或正交嵌入。但在现实场景中,模型在有限数据集上训练,且嵌入为非正交(随机)的。本文针对一个带有随机嵌入的单层Transformer,在简单的词元召回任务上进行分析,要求模型从长度为L的序列中识别出关键词元,并学习词元到标签的一一映射。通过追踪梯度下降的‘早期阶段’,我们推导出模型存储容量的显式公式,揭示了样本数N、嵌入维度d和序列长度L之间的乘积依赖关系。数值实验验证了这些缩放规律,并进一步给出了底层统计问题的下界,证明该乘积缩放在非正交嵌入下是固有的。

原文摘要 · Abstract (English)

Modern large language models (LLMs) excel at tasks that require storing and retrieving knowledge, such as factual recall and question answering. Transformers are central to this capability because they can encode information during training and retrieve it at inference. Existing theoretical analyses typically study transformers under idealized assumptions such as infinite data or orthogonal embeddings. In realistic settings, however, models are trained on finite datasets with non-orthogonal (random) embeddings. We address this gap by analyzing a single-layer transformer with random embeddings trained with (empirical) gradient descent on a simple token-retrieval task, where the model must identify an informative token within a length-$L$ sequence and learn a one-to-one mapping from tokens to labels. Our analysis tracks the ``early phase'' of gradient descent and yields explicit formulas for the model's storage capacity -- revealing a multiplicative dependence between sample size $N$, embedding dimension $d$, and sequence length $L$. We validate these scalings numerically and further complement them with a lower bound for the underlying statistical problem, demonstrating that this multiplicative scaling is intrinsic under non-orthogonal embeddings.

Transformer记忆容量梯度下降嵌入分析

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。