🤖 machine learning

Kronecker Embeddings: Byte-Level Structured Token Representations for Parameter-Efficient Language Models

本文介绍了克朗内克嵌入，这是一种确定性的字节级分解方法，用固定编码器和可学习投影替代传统的大型嵌入表，从而在大型语言模型中消除了超过 90% 的输入侧可训练参数，同时提升了训练效率、拼写鲁棒性和运行时内存使用效率。

原作者： Rohan Shravan

发布于 2026-05-29

📖 1 分钟阅读☕ 轻松阅读

原作者： Rohan Shravan

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一座巨型图书馆，其中每本书都拥有一张独特的身份证。在标准的语言模型（AI）中，当计算机看到一个单词时，它所做的第一件事就是在一本庞大的、预先写好的电话簿中查找该单词的身份证。这本电话簿被称为嵌入表。

对于小型模型而言，这本电话簿尚可管理。但对于驱动当今最先进 AI 的庞大“前沿”模型来说，这本电话簿已变成一个臃肿的怪物。它占用了价值数亿美元的计算机内存，并要求 AI 在开始思考句子含义之前，先耗费巨大的算力去死记硬背哪个身份证对应哪个单词。

克罗内克嵌入（Kronecker Embeddings） 是一种构建这座图书馆的新方法，它彻底摒弃了那本巨大的电话簿。

以下是其工作原理，使用简单的类比来说明：

1. 旧方法：巨大的电话簿

将标准 AI 想象成一名学生，必须死记硬背一本字典，其中从"apple"到"zephyr"的每一个单词都被分配了一个特定的随机数字。

问题所在：如果字典有 130,000 个单词，学生就必须死记硬背 130,000 个随机数字。这会占用其大脑（内存）中的巨大空间，并拖慢其速度。
怪癖：该研究发现，这本“电话簿”实际上非常不擅长理解词族。如果你问 AI 关于"run"的问题，它的电话簿会显示最接近的单词是"Run"、"run"和".run"（仅仅是大小写或标点符号不同的同一个词）。它将"run"和"running"视为陌生人，尽管它们彼此相关。

2. 新方法：乐高蓝图

克罗内克嵌入扔掉了电话簿。相反，它给 AI 提供了一份乐高蓝图。

工作原理：每个单词只是一串微小的积木块（字节）。蓝图规定：“如果你在第一个位置看到'b'，就使用这块特定的乐高积木。如果你在第二个位置看到'a'，就使用那块积木。”
神奇之处：AI 不需要死记硬背单词"run"。它通过按特定位置将'r'、'u'和'n'的积木块拼接起来，即时构建出单词"run"。
结果：AI 不再需要一本巨大的电话簿。它只需要一本微小的说明书（投影矩阵），就能知道如何将那些乐高积木拼接成有意义的形状。这节省了模型输入端通常所需的**91–94%**的内存。

3. 论文实际发现的内容

研究人员将这种新方法与旧方法进行了测试，发现了一些令人惊讶的结果：

对拼写错误更智能：因为新方法是从单个字母（字节）构建单词的，所以它理解"netwrok"与"network"非常相似。如果你打错了字，AI 仍然能识别出单词的轮廓。而依赖巨大电话簿的旧方法，在出现拼写错误时，往往会将单词分解成令人困惑的碎片。
- 类比：如果你将"cat"误拼为"cst"，旧 AI 可能会认为你在谈论一个完全不同的、未知的物体。而新 AI 能看到"c"和"t"处于正确的位置，从而知道你可能指的是"cat"。
学习速度更快：在测试中，使用乐高蓝图（克罗内克）的 AI 在预测句子中下一个单词方面的表现，比使用电话簿的 AI 高出约2.5%。而且，它达到同等技能水平所使用的训练步数减少了 43%。
不会“忘记”新单词：如果你让 AI 谈论一个生造词，比如"kronekticus"，旧 AI 会感到困惑并编造一个虚假的定义。而新 AI 由于是从字母构建单词，可以完美地将这个生造词复述给你，保留你给出的确切拼写。

4. 权衡（潜在缺陷）

该论文诚实地指出了缺点。

混淆的“双胞胎”：因为新方法关注字母，所以它有时会将外观相似的单词视为相关，即使它们的含义不同。例如，"compute"和"commute"在字母层面上看起来非常相似。新 AI 可能会认为它们是近亲，尽管一个关乎数学，另一个关乎旅行。AI 必须利用其“大脑”（模型的其余部分）根据上下文来区分差异。
无法“系结”节点：在旧系统中，AI 可以使用同一本电话簿进行读写以节省空间。新系统的结构差异太大，无法这样做，因此它需要一本独立的小手册用于写作。不过，论文指出，最大的 AI 模型本身就已经在摒弃这种“系结”做法了。

5. 这对未来的意义

该论文表明，通过移除巨大且沉重的电话簿，我们可以释放那部分庞大的计算机内存。

重新分配：不再将内存浪费在静态的单词列表上，而是可以将节省下来的空间用于扩大 AI 的“大脑”（Transformer 主体），使其更聪明，或能够阅读更长的文档。
边缘设备：对于在手机或小型设备上运行 AI 而言，这是一个游戏规则的改变者。新方法允许 AI 在发货时不带有多 GB 的单词列表。它可以从一个仅 4.5 MB 的小文件中重建单词，而不是一个 2 GB 的文件，这使得在小型硬件上运行强大的 AI 变得容易得多。

总结：
克罗内克嵌入用一套灵活的、逐字母的构建套件，取代了庞大且僵化的字典。它节省了海量内存，更好地处理拼写错误，学习速度更快，并允许 AI 处理它从未见过的单词，同时保持 AI 其余架构完全不变。

技术摘要：Kronecker 嵌入

1. 问题陈述

现代大型语言模型（LLM）依赖于一个形状为 $|V| \times d_{model}$ 的已学习输入嵌入表，用于将 token 标识符映射到连续向量。在前沿规模下，该表在发生任何上下文计算之前，就会消耗数亿至数十亿个可训练参数。例如，一个拥有 13 万词汇量和 $d_{model}=4096$ 的模型，仅输入嵌入就分配了约 5.37 亿个参数。这造成了显著的工程摩擦：该表必须在数据并行 rank 之间进行分片，在张量并行设置中产生通信开销，并且需要巨大的内存来存储优化器状态（在混合精度 Adam 中，每个参数约需 16 字节）。此外，该论文提出，从 token ID 到向量的这种固定映射是否必须通过学习获得，仍是一个未解之谜，特别是考虑到工程成本以及这些嵌入实际形成的结构性质。

2. 方法论：Kronecker 嵌入

作者提出了Kronecker 嵌入，这是一种确定性的、结构化的已学习输入嵌入表的替代方案。该方法保留了标准的 BPE 或 SentencePiece 分词，但用固定的编码器和单个已学习的投影层替换了查找表。

核心机制

对于具有 UTF-8 字节序列 $b = (b_1, \dots, b_L)$ 的 token，其嵌入计算如下：

字节 - 位置编码：将 token 分解为其组成字节及其位置。
Kronecker 积：构建一个确定性的编解码向量 $\kappa(b) \in \mathbb{R}^D$ $κ (b) \in R^{D}$ ，其中 $D = d_c \times d_p$ $D = d_{c} \times d_{p}$ 。
- $d_c = 256$ （完整字节字母表）。
- $d_p$ 是最大字节位置（例如 32）。
- 编解码器是 one-hot 字节向量和 one-hot 位置向量的 Kronecker 积之和： $\kappa(b) = \frac{1}{\sqrt{L}} \sum_{p=1}^L c_{b_p} \otimes p_p$ 。
归一化：编解码器输出经过每个 token 的 $z$ -归一化（均值为 0，标准差为 1），以稳定早期训练。
已学习投影：单个可训练线性层 $W_{proj} \in \mathbb{R}^{D \times d_{model}}$ 将编解码器向量投影到模型的隐藏维度。

操作变体

gpu_table：在启动时预计算完整的 $|V| \times D$ 编解码表。计算量可忽略（仅 gather 操作），但内存占用很高（数 GB）。
gpu_dynamic：在前向传播过程中使用紧凑的字节缓冲区（13.1 万词汇量下约 4.5 MB）即时重新计算编解码器。这以微小的计算开销（步骤时间的 0.01–0.24%）换取了巨大的内存节省。

归纳偏置

该方法引入了字节级局部性：两个字符串仅当它们在相同位置共享字节时，才会获得相似的嵌入。这保留了大小写区分（例如 swift 与 SWIFT）和对拼写错误的鲁棒性（例如 seperate 与 separate），而标准的 BPE 嵌入倾向于将排版变体聚类在一起。

3. 主要贡献

A. 跨模型探测发现

作者分析了六个现代语言模型（参数规模从 1.35 亿到 6710 亿，包括 Llama-3.2、Qwen3、DeepSeek-V3 等）的输入嵌入：

排版聚类：训练后的 BPE 嵌入将排版变体（例如 run、Run、.run）的聚类程度远高于形态学亲属词（例如 running、runner）。这种模式在不同分词器家族和规模中均成立。
Kronecker 逃逸：Kronecker 嵌入逃逸了这种排版聚类，转而检索字节相似的字符串。
形态结构：两种方法在嵌入层均未发展出严格的形态学聚类。在对 1.24 亿参数检查点的分层探测中，BPE 分支的第一个 Transformer 层用共现/上下文几何结构交换了形态学几何结构，而 Kronecker 在早期层中保留了字节级几何结构。

B. 受控训练比较

在 FineWeb-Edu 的 25 亿 token 上，对 nanoGPT GPT-2 124M 架构进行的三种子训练比较：

性能：Kronecker 的验证损失比 BPE 绑定基线低 2.5 ± 0.2%（差距为 0.083 nats，困惑度降低约 9%）。
收敛性：损失差距在训练过程中扩大，并在收敛时稳定，而非缩小。
效率：Kronecker 达到 BPE 基线收敛损失所需的优化器步数减少了 1.43 倍。
参数量：在 1.24 亿规模下，输入侧可训练参数减少了91%（从 3860 万降至 310 万），尽管由于需要未绑定的输出头（权重绑定在架构上与 Kronecker 的 $D \neq d_{model}$ 不兼容），总可训练数量略有增加（约 310 万）。在前沿规模下（例如 1200 亿 MoE），输入侧的减少幅度达到94%（节省约 5.03 亿参数）。

C. 鲁棒性与保真度探测

拼写鲁棒性：在 110 对干净/拼写错误提示中，Kronecker 在**55.5%**的情况下保留了 top-1 预测，而 BPE 为 47.3%（+8.2 个百分点）。它还显示出干净分布与拼写错误分布之间更低的 KL 散度（-7.6%）以及最终隐藏状态中更高的余弦相似度（+2.6%）。
生成保真度：定性探测显示，Kronecker 模型通过自回归生成“回声”字节新颖的字符串和拼写错误（例如 kronekticus、netwrok），而 BPE 模型则将其碎片化并遗忘。

D. 机制观察

稳定目标：BPE 嵌入范数在训练期间从 std=0.020 漂移到 0.026，而 Kronecker 投影范数稳定在 ~1.0。这表明 Kronecker 为 Transformer 主体提供了稳定的表示目标。
内存效率：即时变体将生产词汇量下的输入路径内存占用从 ~2.15 GB（预计算表）减少到 ~4.5 MB（字节缓冲区）。

4. 意义与主张

该论文声称，Kronecker 嵌入提供了一种参数高效、即插即用的标准输入嵌入替代方案，能够：

消除嵌入瓶颈：在前沿规模下将输入侧可训练参数减少 91–94%，释放容量用于 Transformer 主体或减少优化器状态内存。
改善训练动态：在受控设置下实现更低的验证损失和更快的收敛，这可能是由于稳定的表示目标和参数高效优化所致。
增强鲁棒性：提供对排版错误的固有鲁棒性，并在生成中保留字节级保真度，这是标准 BPE 嵌入所缺乏的特性。
重构嵌入角色：挑战了已学习嵌入必须捕捉形态学的假设，表明它们主要捕捉排版变体，而 Kronecker 提供了一种独特的字节级归纳偏置。

作者明确指出，该方法不是分词器的替代品，也不是字符级模型；它作用于 token 序列，但替换了嵌入查找。他们承认了局限性，包括无法使用权重绑定、语义遥远但字节相似的单词（例如 compute/commute）的聚类，以及需要在前沿规模（10 亿 + 参数）和多语言数据集上进行进一步验证。

5. 局限性与未来工作

该论文指出了几个开放性问题：

扩展性：验证损失优势和分层发现是否在 10 亿至 70 亿 + 规模下依然成立。
未绑定基线：受控比较使用了 BPE 绑定基线；未来的工作应与 BPE 未绑定基线进行比较。
机制：性能提升的确切机制（稳定目标 vs. 参数效率 vs. 归纳偏置）仍然是一个未解之谜。
输出侧：当前设计将输出头保留为标准的已学习矩阵。作者推测未来的工作可以将 Kronecker 扩展到输出侧以实现无词汇表解码，但这尚未得到测试。

该工作发布了代码、可复现脚本以及四个预训练的 LightningLM 模型（20 亿至 1200 亿 MoE），以展示其在规模上的稳定性。