低秩词表矩阵实验：输入嵌入端反而降低测试损失？

写下我希望在开始学习某样东西时就能找到的那篇文章……

归档

分类

友情链接

我被勾起了好奇心！在我为上一篇文章创建的某个模型的 HF 讨论页上， AndrewThompson1233 问我是否考虑过尝试因式分解嵌入 ——这是他用于自己的模型 Maba 的方法，此前也被用于一些其他模型，比如 ALBERT。

这是个非常巧妙的想法；你使用与 LoRA 类似的技巧，来减少用于嵌入和输出头的参数量。而对于小型模型来说，这些参数可能占到总数的很大一部分。

例如，在我那 1.63 亿参数的 GPT-2 small 风格模型中，嵌入大约占3900 万；输出头大小相同，所以把这两者加在一起，模型的一半都只是用来把东西送进去和取出来，而不是真正进行思考。即使你使用权重绑定，就像原始 GPT-2 模型那样，你最终也会把 30% 的“预算”花在嵌入和输出头共享的单个矩阵上。

同样，虽然更大的模型在这方面花费的比例更小，但即使是中等规模的 MoE 模型，也可能把大量活跃参数花在嵌入上；我计算过，对于 Qwen 3.6 35B MoE，其活跃参数为 3B，其中大约有 1B 用在了输入嵌入和输出头上。

所以，任何能减少这一开销的方法——只要不以模型能力的高昂代价为代价——都值得考虑。你节省了参数量，这要么意味着模型更小、训练可能更快，要么允许你把节省下来的部分“投资”到更多思考参数上——也就是更宽、嵌入维度更高的网络，或者更深、层数更多的网络。

Andrew 报告说，使用这个技巧的影响微乎其微：

在 5 万词表、秩为 128 时，验证交叉熵的损失可以忽略不计（通常在 +0.02 到 +0.04 损失之间，或 <0.5 困惑度变化）。

现在，这表明他的损失大约在 2.5 左右（因为损失增加 0.04 意味着困惑度变化 0.5），这比我用 GPT-2 架构通常得到的要低得多（对我来说更像是 3.5），但尽管如此，也许我也能得到不错的结果？我决定训练几个模型看看会发生什么。

结果很有趣！我发现，在嵌入和输出头上都使用这个技巧，会导致损失增加比 Andrew 描述的略大——在我的第一次测试中约为 0.07，第二次为 0.10，第三次为 0.14。这些相对较大，尽管可能通过进一步训练或更大的模型来恢复。

但更令人惊讶的是，我发现只在输入嵌入上使用这个技巧似乎会降低测试损失。在开始之前，我原本以为它在输入侧的有害程度会低于输出侧，但看到改进确实出乎意料。通过减少参数量来降低损失，这可不是通常会发生的事！

所以，这绝对值得深入探讨一下。让我们从理论开始：这些分解嵌入是什么，它们是如何工作的？像LoRA一样，它们依赖于低秩因子，所以我首先定义这些。

低秩因子

我们的模型由大量矩阵组成——嵌入、输出头、注意力权重、FFN线性层等等。使它们变小显然会减少模型的参数数量，尽管你可能会预期这会有代价。低秩因子的思想是，你可以用一个较大的矩阵替换为两个较小的矩阵，这两个较小的矩阵将包含大部分重要信息。

假设你有一个大小为m×n的矩阵。现在，从矩阵乘法中，我们知道如果你将一个m×r矩阵乘以一个r×n矩阵，你会得到一个m×n的结果。如果我们这样做，那么我们就分解了矩阵（就像我们可以将12分解为3和4，因为3×4=12一样），值r被称为（逻辑上足够）分解的秩。如果r足够小于m和n，那么这两个矩阵——低秩因子——组合起来将包含比完整矩阵更少的参数。

让我们具体说明，使用一个没有权重绑定的GPT-2风格模型的输出头。它接收来自Transformer层（归一化后）的嵌入，并将它们转换为词汇表上的logits。对于GPT-2小尺寸，我们的嵌入是768维的，词汇表大小为50,257个标记。所以当我们使用线性层进行这种映射时，它有768×50257=38597376个参数。

如果我们用两个矩阵替换它——比如，一个768×128和一个128×50257，那么它们组合起来将占用768×128+50257×128=98304+6432896=6531200个参数。这几乎小了六倍！

所以：想法是，我们不是用m×n矩阵创建模型，而是用一对m×r和r×n创建它，并训练这些。如果一切顺利，这对矩阵将几乎能够学习我们想要的内容，就像完整矩阵一样，所以我们将得到足够接近的结果。

数学和实现也很简单。对于没有偏置的线性层，我们可以写出将输入X和权重矩阵W相乘并产生输出Z的矩阵乘法，像这样：¹

Z=XW

现在，如果我们使用一对低秩因子矩阵A和B而不是完整的W，我们可以说我们想要用于计算的“虚拟”权重是W′，使得：

W′=AB

所以这意味着我们对此神经网络的计算是这样的：

Z=X(AB)

矩阵乘法是可结合的，这意味着你可以将其重写为：

Z=(XA)B

……希望你能看出这等同于将X通过一个使用A作为权重的层，然后将结果通过第二个使用B的层。在代码中，你取了这样的东西：

        out_head = nn.Linear(d_emb, vocab_size, bias=False)
        ...
        logits = out_head(embeddings)

……并替换为这个：

        out_head_A = nn.Linear(d_emb, r, bias=False)
        out_head_B = nn.Linear(r, vocab_size, bias=False)
        ...
        intermediate = out_head_A(embeddings)
        logits = out_head_B(intermediate)

不过，从直觉上大概很明显，这是有代价的。如果你的参数更少，那么你能存储的信息就更少，所以模型的这一部分就“更笨”。如果还不清楚，想象一下上面代码中的 r 是 1。你会把（对于 GPT-2 small）768 维的嵌入，转换成一个单一数字，然后再把这个单一数字扩展成 50,257 维的 logits 集合。信息会丢失——低秩因子背后的想法是，只要你为 r 选择一个合适的值（r 通常被称为低秩瓶颈的宽度），你就不会丢失重要的东西。

这在许多情况下效果出奇地好——特别是 LoRA，它让你能够微调那些在你的硬件上难以完全训练的模型，或者更快地微调它们，它利用这一点来轻松地训练类似低秩因子“差分”的东西，加到你的权重矩阵上 ²。

但在这篇文章中，我们将尝试最简单的版本：如果我们把输入嵌入和输出头——也就是词汇矩阵——整体都替换为低秩瓶颈，会怎样？

介绍 LoRE

词汇矩阵的低秩分解有点拗口，所以基于“LoRA”这个名字，我决定把这个技巧称为 LoRE，即低秩嵌入。严格来说这并不准确（毕竟，我们既对 LLM 开头的嵌入矩阵做了这件事，也对末尾的输出头做了这件事），而且我不指望它会流行起来，但我挺喜欢它，并将在本文中使用它 :-)

其背后的论点是，嵌入只是一个简单的查找表，所以正是那种你期望能够通过只考虑巨大矩阵中重要部分来节省的地方；同样的道理可能也适用于末尾的输出头，不过我个人对这部分不太信服。

这种不对称性值得稍微研究一下。我的直觉是，虽然嵌入确实是一个查找表，但输出头在做一些更微妙的事情。它把从 Transformer 层出来的连续嵌入投影到 logits，我们（通过 softmax）将其解释为可能的下一个 token 的概率分布。这比从“猫”映射到适当的嵌入要复杂得多：作为一个（简化的）例子，如果你有一个嵌入，其含义接近“猫”、“狗”、“沙鼠”、“家养宠物”等等，那么把它投影到下一个 token 的适当值就不是一件简单的事。³

不过，在实现方面，这真的很简单。我决定扩展我从 Sebastian Raschka 的书“Build a Large Language Model (from Scratch)”中得到的 PyTorch 代码，在模型配置 JSON 中添加一个可选的 lore 部分。这将允许我们独立地为输入嵌入和输出头开启 LoRE 模式，并指定秩——也就是上面部分中的 r，它表示矩阵 A 和 B 之间的低秩瓶颈有多宽。

所以，在代码中，我的第一版改动把这个：

        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])

……变成了这个：

        if "lore" in cfg and cfg["lore"].get("input_embeddings", False):
            self.tok_emb = nn.Sequential(
                nn.Embedding(cfg["vocab_size"], cfg["lore"]["rank"]),
                nn.Linear(cfg["lore"]["rank"], cfg["emb_dim"], bias=False)
            )
        else:
            self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])

……以及这个：

        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

……变成了这个：

        if "lore" in cfg and cfg["lore"].get("output_head", False):
            self.out_head = nn.Sequential(
                nn.Linear(
                    cfg["emb_dim"], cfg["lore"]["rank"], bias=False
                ),
                nn.Linear(
                    cfg["lore"]["rank"], cfg["vocab_size"], bias=False
                ),
            )
        else:
            self.out_head = nn.Linear(
                cfg["emb_dim"], cfg["vocab_size"], bias=False
            )

这本身是对底层概念的一个完全合理的实现。 但它有一个问题。我想比较使用和不使用 LoRE 时的损失，还想 弄清楚仅输入嵌入和仅输出头使用 LoRE 会有什么效果。 但我在其他实验中发现，我的模型开始时所用的初始随机权重 会对最终模型的损失产生显著影响。上面的代码 在不同的配置下会给出相当不同的随机权重。不过修复方法 很简单：

        if "lore" in cfg:
            lore_tok_emb = nn.Sequential(
                nn.Embedding(cfg["vocab_size"], cfg["lore"]["rank"]),
                nn.Linear(cfg["lore"]["rank"], cfg["emb_dim"], bias=False)
            )
        normal_tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        if "lore" in cfg and cfg["lore"].get("input_embeddings", False):
            self.tok_emb = lore_tok_emb
        else:
            self.tok_emb = normal_tok_emb

        ...

        if "lore" in cfg:
            lore_out_head = nn.Sequential(
                nn.Linear(
                    cfg["emb_dim"], cfg["lore"]["rank"], bias=False
                ),
                nn.Linear(
                    cfg["lore"]["rank"], cfg["vocab_size"], bias=False
                ),
            )
        normal_out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )
        if "lore" in cfg and cfg["lore"].get("output_head", False):
            self.out_head = lore_out_head
        else:
            self.out_head = normal_out_head

有了那段代码，只要存在任何 LoRE 配置，我们就会创建相同的权重， 无论它是用于输入嵌入还是输出头。我 需要训练一个基线，它带有 LoRE 的配置，但只是把 input_embeddings 和 output_head 都设为 False——那实际上会是一个正常的、非 LoRE 的模型， 但在它与使用 LoRE 的等效模型共享的任何部分中，都会拥有相同的初始随机权重。

此外，Andrew 曾说过我们需要注意初始化：

训练期间要留意的主要失败模式是早期的梯度尖峰： 将瓶颈投影的初始化缩放为 1 / sqrt(r) 可以保持方差 均匀，并确保训练动态从第 1 步起就与全秩基线一致。

……以及后来：

在嵌入表上保持标准的 GPT-2 初始化，但将内部的线性投影矩阵 按 std = 1.0 / math.sqrt(128) 进行缩放。

（其中的 128 是因为我们当时在讨论 r=128 的设置， 而这实际上正是我最终使用的。）

现在，我用的是 PyTorch 的默认初始化，而不是经典的 GPT-2 设置，但我还是添加了一个 smart_initialize 标志，按照 Andrew 描述的方式处理缩放：如果它被设置，那么在 模型创建的最后，我会将初始化 应用到模型中位于最内层的低秩因子矩阵上（也就是最靠近 Transformer 层的那些）：

        if "lore" in cfg:
            if cfg["lore"].get("smart_initialize", False):
                nn.init.normal_(
                    lore_tok_emb[1].weight,
                    mean=0.0,
                    std=1.0 / math.sqrt(cfg["lore"]["rank"])
                )
                nn.init.normal_(
                    lore_out_head[0].weight,
                    mean=0.0,
                    std=1.0 / math.sqrt(cfg["lore"]["rank"])
                )

尽管这并非对每个配置都会运行，但由于它发生在模型创建的 最后，它不会影响其他权重，所以我认为 这样是安全的。

你可以在这里看到与普通 GPT-2 代码的完整差异。

这些改动在我看来是正确的，我让 GPT-6 Astra（在“超高”思考模式下） 确认在讨论的语境下它看起来是否正确，它确认了。 做完这些后，我又与 Andrew 再次确认， 因为我真的很想让这次成为对他结果的一次干净的复现尝试。他确认了， 那么是时候训练一些模型了！

最初的几次训练运行

我决定要进行五次训练运行来彻底检验这一点：

1. 一个基线。正如我之前所说，LoRE 相关的东西会影响 初始权重所用的随机性，所以我想训练一个“启用”了 LoRE 但实际上并未使用它的模型。配置：model.json， train.json。
2. 一个在输入嵌入和输出头上都启用了 LoRE，但没有对低秩矩阵进行 智能初始化的模型。这只是为了向自己确认它是必需的。配置：model.json， train.json。
3. 一个两端都使用 LoRE、并带有智能初始化的模型——本质上， 就是 Andrew 所描述的那样。配置：model.json， train.json。
4. 仅在输入嵌入上使用 LoRE 的模型，如果事实证明有必要就使用智能初始化，如果不需要就不使用。我故意将配置设为无效，在 smart_initialize 标志中使用 TODO 而不是 true 或 false，这样我就会记得稍后填写它们。 配置：model.json， train.json。
5. 仅在输出头上使用 LoRE 的模型，智能初始化与 (4) 相同。 配置：model.json， train.json。

为了节省时间，我决定在云端训练它们：前三个，然后最后两个，分两批并行运行。我使用了 Lambda 上的 8x A100 机器，每块 GPU 有 40 GiB 显存。我将打破通常对每次训练运行进行详细说明的传统，因为（剧透）这篇文章中有相当多这样的运行！不过，为了可复现性，我为每次运行执行的命令看起来像这样：

git clone https://github.com/gpjt/ddp-base-model-from-scratch.git
cd ddp-base-model-from-scratch/
git checkout lore-experiment
./setup_lambda.sh
uv sync
uv run torchrun --nproc_per_node=8 ddp_train.py 8xa100m40-lore-1-baseline-disabled datasets

...将运行 ID 8xa100m40-lore-1-baseline-disabled 替换为给定 运行设置 对应的那个。

以下是前三次运行的结果：基线，以及两端都使用 LoRE 的模型，分别使用和不使用智能初始化。模型名称链接指向 Hugging Face 上上传的模型。

你可以看到我在最后一次运行中遇到了一些问题；可用实例短缺，我最终在日本启动了一个。从我在葡萄牙的家中连接到它的速度非常慢且不可靠；当我在训练运行后尝试下载模型时，scp 预测需要 10 小时（按照该机器每小时 15.92 美元的成本，那将是疯狂的）。然而，我能够快速将其复制到我的 PythonAnywhere 账户，也许是因为那是在美国的亚马逊数据中心，与日本的连接更好，然后我可以从那里将其复制到我的机器，同样相当快。不幸的是，在忙乱处理这些时，我没有捕获所有输出。

尽管如此，在这个阶段重要的数据是测试损失列。正如预期的那样，使用 LoRE 使测试损失变差。而智能初始化似乎确实有助于减少这种惩罚。

损失增量为 +0.07，比 Andrew 在他的消息中描述的 +0.02 到 +0.04 更差，但并非天差地别——当然足够小，以至于似乎可以通过将 64M 参数差异投资于更深或更宽的模型来恢复它，甚至可能更多。有希望！

成本和时间节省也是真实的：日本的运行因我复制模型时遇到的所有麻烦而成本增加，所以 8xa100m40-lore-2-both-ordinary-init 那个是最好的指示：38.55 美元而不是 43.96 美元，节省了超过 10%。虽然从略多于 2 小时 30 分钟降到 2 小时 14 分钟可能看起来不多，但对于更大的模型或更长的训练运行，近 14% 的减少本身可能就很有价值。（这还没有考虑将节省的参数“再投资”到更大模型的替代方案。）

所以那是相当好的消息！是时候进行消融运行，看看仅在 token 嵌入上或仅在输出头上使用 LoRE 是否会有所不同。我将剩余两个模型的 smart_initialize 设置切换为 true，然后启动了它们。

现在，这里面有很多有趣的东西，但让我们聚焦于测试损失。我关于这一部分的实验笔记写道：“嗯，这有点令人震惊！”相当令人惊讶——至少对我来说——我们使用仅输入的 LoRE 得到了更低的测试损失，为 3.526610，而使用无 LoRE 基线得到的测试损失为 3.599611。

当你比较不同版本（排除非智能初始化）之间的差异时，事情变得更加有趣，这些差异是相对于无 LoRE 基线的（为简洁起见，我将在下文中称之为 delta）。四舍五入到小数点后 3 位：

如果你将两个单 LoRE 损失 delta 相加——0.138 和 -0.073——你会得到 0.065，这令人惊讶地接近双 LoRE 选项的 delta 0.072，仅相差 0.007。

我们是否可能看到两个几乎独立的效果，分别对应 LoRE 应用位置的两个选项？仅输入 LoRE 的明显好处是真实的吗？还是由于偶然——也许我们在初始化 LoRE 权重时“运气好”，而在初始化非 LoRE 权重时“运气差”。如果 LoRE 能得到的最佳损失与无 LoRE 能得到的最差损失之间存在重叠，那么这个结果可能在噪声范围内。

我决定进行第二批训练运行，看看能否获得支持或反对的证据。但在我们深入之前，值得看看那两个表中的其他一些数字。

首先，让我们看看训练时间。训练机器在所有运行期间似乎都接近 100% 运行，所以暂时假设所用时间与计算量成正比。不出所料，基线运行的时间最长，为 9,339 秒。同样不出所料，我们有计时数据的完整 LoRE 运行要快得多，为 8,057 秒。用（小得多的）低秩因子对替换大矩阵，意味着前向和后向传播中的计算量都会减少，从而使训练运行更快，这当然是有道理的。

单 LoRE 运行中有些事最初让我感到惊讶。在输出头上使用 LoRE 大大加快了速度，但输入嵌入 LoRE 的影响要小得多。

然而，思考了一下之后，情况就清楚了。网络中的嵌入模块实际上在前向和后向传播中运行得非常快。虽然我喜欢从概念上将嵌入视为通过矩阵乘法将词汇空间中的独热向量投影到嵌入空间，并且我仍然认为这是解释所发生情况的一个极好模型，但在实践中，它更像是数组查找，几乎零成本。同样在后向传播中，我们只需要为选定的嵌入计算梯度，所以那也应该更快。

确实，考虑到所有这些，如果仅输入 LoRE 最终比基线稍慢，可能也不会令人惊讶，因为它至少需要通过实际矩阵进行前向和后向传播，即使这些矩阵相当小。

再做一次 delta 表，时间差以秒为单位（使用“普通初始化”运行作为双开数字）：

你可以看到，如果把单个 LoRE 选项的增量相加，得到 -1,251，这与双 LoRE 选项的增量非常接近，令人满意。这比损失结果要合理得多——有了训练计算量，你确实会期望这样的改进是可加的。

最后，初始训练损失中有一些有趣的东西。因为我的训练代码会打印出它在第一个全局步骤上得到的损失，所以我决定把它记下来。这些数据中有一些东西在之后会变得重要：在两次输出头上使用 LoRE 并且 使用智能初始化的训练运行中，损失都明显上升了。

通常，我的 GPT-2 风格模型的初始训练损失略低于 11。这与架构非常吻合；你会期望一个纯粹随机（严格来说，均匀）预测 token 的模型，其困惑度大致等于词表大小；GPT-2 分词器的词表大小是 50,257，而损失是困惑度的自然对数，所以这意味着损失大约为 10.82。

但两个在输出头上使用 LoRE 的智能初始化模型的初始损失都高了大约一个点，分别为 11.865 和 11.817，这意味着它们显著 差于 随机。这里可能发生了什么？我们稍后会揭晓 :-)

总之，目前来看，这是一组有趣的结果。我想进一步深入探究；特别是，如果我用不同的随机种子再试一次，仅输入 LoRE 模型那令人惊讶的测试损失改善是否会重现？

第二组模型

此时我有两个可行的假设；要么

1. 我们看到的是由于权重随机初始化的运气所致——LoRE 运气好，没有 LoRE 运气差。
2. 这里确实存在一种效应：仅输入 LoRE 实际上改善了所得模型的测试损失。

一个可靠且简单的测试来区分这两者，就是再用不同的随机种子训练另一组模型。如果我得到相同类型的结果，即输入 LoRE 击败非 LoRE，那将是支持 (2) 而非 (1) 的证据，而如果情况发生变化，则相反。

这次我决定不费心使用非智能初始化模型，所以我为四个新模型做了设置。我的训练代码 默认使用随机种子 42，但允许在 train.json 配置文件中覆盖它，所以我决定对这些模型使用 123，并创建了：

• 一个没有 LoRE 的新基线。配置：model.json，train.json。
• 一个仅输入 LoRE。我想把这个作为第一个非基线模型来运行，因为如果它的损失 更差（更高）于基线，那么看起来就很像我之前看到的效果是由于随机权重初始化造成的，所以我可能想考虑就此停止。配置：model.json，train.json。
• 如果我决定继续，下一个将在输入和输出两侧都使用 LoRE。配置：model.json，train.json。
• 最后，我会做一个仅输出头 LoRE 的模型。配置：model.json，train.json。

一旦配置就绪，就该训练模型了——但在 Lambda 上很难获得可用实例。我最终不得不使用我的 lambda-manager 脚本来启动，但几个小时后，我的手机收到警报，说有实例正在运行，可以开始了。

第二次仅输入 LoRE 的训练运行，其测试损失再次超过了基线，所以我做了全套运行，结果如下：

另一张相对于基线的差值矩阵：

所以，仅输入 LoRE 再次改善了测试损失（尽管 -0.013 的效应比上一组模型的 -0.073 要小）。

而且，仅输入和仅输出 LoRE 的差值之和再次接近组合 LoRE 模型的差值——实际上这次更接近，0.102 距离 0.101 仅差 0.001。

输出头上开启 LoRE 的模型训练时间结果和更高的初始训练损失也依然成立，这很令人欣慰。所以此时是时候收尾了——真的是这样吗？

smart_initialize 问题

在我进行第二批训练运行时，我开始撰写这个实验的总结。正如我通常所做的那样，我把草稿发给各种 AI 模型审阅，Claude Opus 5.5 在最大思考模式下发现了一些问题。

Andrew 曾写道：

在嵌入表上保持标准 GPT-2 初始化，但将内部线性投影矩阵按 std = 1.0 / math.sqrt(128) 缩放。

我把那个“内部”理解为低秩因子中“最内层”的矩阵——也就是最靠近 Transformer 层的那些——应该按他所说的方式初始化，即从标准差为 1/r、均值为零的正态分布中取值。

所以这意味着嵌入中的第二个矩阵，以及输出头中的第一个矩阵：

        if "lore" in cfg:
            if cfg["lore"].get("smart_initialize", False):
                nn.init.normal_(
                    lore_tok_emb[1].weight,
                    mean=0.0,
                    std=1.0 / math.sqrt(cfg["lore"]["rank"])
                )
                nn.init.normal_(
                    lore_out_head[0].weight,
                    mean=0.0,
                    std=1.0 / math.sqrt(cfg["lore"]["rank"])
                )

Claude 认为在两种情况下都应该是第二个低秩因子——也就是说，我们应该把

                    lore_out_head[0].weight,

……改为

                    lore_out_head[1].weight,

它为这种做法给出了非常有说服力的理由，而深入探究这一点很有用，因为它把这种“智能”初始化从盲目实现的“秘方”变成了真正有意义的东西。

让我们假设没有使用 LoRE，只考虑输出头。它是一个 PyTorch nn.Linear，这意味着：

……值从 𝒰(−k,k) 初始化，其中 k=1in_features

所以，我们有一个矩阵，可以称之为 W，它充满了从该概率分布中抽取的数字——一个在 −k 到 k 范围内的均匀（平坦）分布。

对于 LoRE，我们把 W 替换为两个独立的矩阵，称之为 A 和 B：

            self.out_head = nn.Sequential(
                nn.Linear(
                    cfg["emb_dim"], cfg["lore"]["rank"], bias=False
                ),
                nn.Linear(
                    cfg["lore"]["rank"], cfg["vocab_size"], bias=False
                ),
            )

假设 W 的形状是 m×n，这意味着 A 是 m×r，B 是 r×n。我们希望这两个矩阵的组合效果 W′=AB 一开始的随机数字——就其随机性而言——看起来像完整的矩阵 W。

但用上面的代码，我们“虚拟”矩阵 W′ 中的数字实际上是什么样子呢？表面上看，它听起来不太可能像 W 的概率分布。毕竟，W 是基于一个使用 in_features 的公式初始化的。你可以看到 A 会用相同的分布初始化，因为它有相同数量的输入特征，但 B 的输入特征数量是 r，所以它的随机数字会基于那个。它们组合成 W′ 时，大概会混合来自这两个不同概率分布的东西。

我们实际上可以算出那会是什么。数学来了：点击此处跳过。

根据 Wikipedia，两个独立随机标量变量（所以不是矩阵，但我们稍后会讨论这个）X 和 Y 的乘积的方差是这样的：

Var(XY)=(σX2+μX2)(σY2+μY2)−μX2μY2

...其中 σX2 是 X 的方差，μX 是它的均值，Y 同理。

现在，假设 X 和 Y 的均值都为零（就像 PyTorch 用于初始权重的分布 𝒰(−k,k) 那样）。这简化了上面的式子：

Var(XY)=σX2·σY2

这如何应用到我们的矩阵上呢？

好吧，让我们写出如何通过将 A 和 B 相乘来计算 W′，假设 m=3、n=4、r=2：

W′=AB =[a0,0a0,1a1,0a1,1a2,0a2,1][b0,0b0,1b0,2b0,3b1,0b1,1b1,2b1,3]

按照常规矩阵数学，结果中位置 0,0 处的值将是第一个矩阵的第零行（作为向量）与第二个矩阵的第零列的点积：

W0,0′=(a0,0a0,1)·(b0,0b1,0) =(a0,0·b0,0)+(a0,1·b1,0)

我们可以将其改写为适用于任意低秩瓶颈 r 的形式：

W0,0′=∑x=0r−1a0,x·bx,0

...更一般地，我们可以像这样定义 W′ 中的所有项：

Wi,j′=∑x=0r−1ai,x·bx,j

所以，每个元素都是 r 个乘积之和，其中每个乘积将一个来自用于初始化 A 元素的任何分布的数与一个来自用于 B 的任何分布的数相乘。

使用上面的符号，A 的元素具有方差 σA2，B 的元素具有方差 σB2。所以这意味着该和中每个分量的方差是 σA2·σB2。有 r 个这样的分量；如果你将独立变量相加，它们的方差会相加，所以这意味着我们得到的 W′ 的元素的方差是：

Var(W′)=r·σA2·σB2

现在，让我们记住，我们的矩阵 A 实际上已经按照我们希望非 LoRE 版本 W 所采用的方式进行了初始化。nn.Linear 的随机初始权重的分布完全取决于输入特征的数量——当然，A 的输入特征数量与 W 相同。

那么我们如何使 Var(W′) 等于 σA2 呢？我们需要让 σB2=1/r，然后上面的方程就像这样消去了：

Var(W′)=r·σA2·1r=σA2

现在，让我们看看代码（其中包含了 Claude 的修正）：

                nn.init.normal_(
                    lore_out_head[1].weight,
                    mean=0.0,
                    std=1.0 / math.sqrt(cfg["lore"]["rank"])
                )

我们将 B 矩阵的权重设置为从标准差为 1/r 的分布中抽取。标准差当然是方差的平方根（这就是为什么我们一直使用像 σB2 这样带有平方的符号来表示方差——单独的 σ 用于表示标准差）。所以在我们对其调用那个 normal_ 之后，矩阵的方差将如期望的那样为 1/r，而 W′ 将具有与直接创建 W 时相同的方差的随机初始权重。

这相当令人满意 :-)

所以这就是输出头。那么输入嵌入呢？它们是 nn.Embedding，初始化方式略有不同：

从 𝒩(0,1) 初始化

这是一个正态分布，均值为 0，方差（因此标准差）为 1。上面的数学适用于任何均值为零的分布——所以智能初始化也应该适用于这些。我们的 A 矩阵将具有正确的方差 1 和均值 0，因此 AB 也将如此，因为（对于嵌入）B 已经进行了智能初始化。

那么，这是否意味着我们通过将 A 和 B 相乘创建的 W′“虚拟”矩阵在随机性方面与它所替代的矩阵 W 完全相同？嗯，并不完全如此。

记住，输出头的数值是从均匀分布中抽取的，而嵌入是从正态分布中抽取的。我们将作为“输入”低秩因子矩阵的 A 矩阵乘以从正态分布初始化的 B 矩阵。正如我们所展示的，A 矩阵的方差通过智能初始化得以保持，均值（为零）也会干净地传递过来。但分布形状会发生变化。

我稍微挖掘了一下，感觉像是掉进了一个兔子洞，但我能确定的一件事是，从均匀分布中抽取的变量乘以从正态分布中抽取的另一个变量，得到的结果既不是均匀分布也不是正态分布，而正态分布乘以正态分布会得到一种叫做正态乘积分布（或乘积正态分布）的东西，它同样既不是均匀分布也不是正态分布。

记住，我们 W′ 矩阵的每个元素由下式定义：

Wi,j′=∑x=0r−1ai,x·bx,j

所以每个元素都是 r 个数的和，每个数都来自乘积正态分布（对于嵌入）或来自均匀分布×正态分布（不管它叫什么）。那会是什么形状？再说一次：兔子洞，但中心极限定理指出（根据 Wolfram MathWorld）：

具有有限方差的独立随机变量的归一化和趋近于正态分布

因此，对于较大的 r，我们可以说结果将接近正态分布。⁴

智能初始化的作用是确保我们的“虚拟”权重 W′=AB 与我们试图近似的 W 具有相同的均值和方差。但它不会保持它们所来自的分布的相同形状——对于输出头来说肯定不是，它从均匀分布变成了近似正态分布。嵌入受到的影响较小（尽管并非完全不受影响），因为它们原本是正态分布，现在仍然几乎如此。

不过，这仍是一个巧妙的技巧！这真的是 Andrew 的本意吗？

修正智能初始化

在说服自己我一直在使用的智能初始化代码是错误的之后，我决定和 Andrew 核实一下。我首先想到的是，也许他使用了权重绑定——记住，使用权重绑定时，嵌入矩阵只是以转置形式重新用于输出头。那意味着嵌入将是 W′=AB，你会对 B 使用智能初始化，但输出头——作为那个转置——将对相同的 B 和 A 使用 W′T=BTAT，所以智能初始化会反过来。

但事实证明，这只是一个沟通误会。他说他的权重是不绑定的，并且他在“内部线性投影矩阵”中使用“内部”一词的含义与我理解的有所不同。他确认真正的智能初始化代码应该始终应用于第二个矩阵，正如数学（以及 Claude）所建议的那样。

所以那只是一个简单的改动。你可以看到这里的差异，但简而言之：

• 我修改了 smart_initialize 配置参数，使其不再只是 False 或 True，而可以是 False、"original" 表示原来错误的解释，或者 "corrected" 表示修正后的解释。
• 我将所有现有的、将 smart_initialize 设置为 True 的训练配置更新为 "original"。
• 我又添加了更多配置用于更多训练运行以测试这一点，将 smart_initialize 设置为新的 "corrected" 选项。

现在，好消息是（无论是对我自己的理智还是钱包而言），测试这个修复并不需要再进行四次训练。我的训练代码本质上是确定性的——在模型创建之后没有使用 dropout 或任何其他随机性⁵。因为智能初始化代码位于模型创建的末尾，而输出头的智能初始化就在最后，这意味着更改它不会影响任何其他使用随机性的代码。

所以，如果我进行两次训练——一次仅在输出头上启用 LoRE，另一次在嵌入和输出上都启用——那么这些将与之前完全没有 LoRE 以及仅在嵌入上启用 LoRE 的训练运行“兼容”。你可以在这里看到配置（没有显式随机种子，因此它们将使用我的默认值 42）：

• 两者：model.json，train.json。
• 仅输出：model.json，train.json。

以下是结果；我将它们与原始种子 42 训练运行的对应结果放在一个表格中：

在我们查看测试损失之前，请看看那个“初始训练损失”列。我之前注意到的奇怪现象，即该值比均匀分布差得多，已经消失了！这其实有一定道理。由于智能初始化应用在低秩输出头的错误部分，我们的方差完全错误，因此模型最初产生了特别差的结果。⁶

所以这很有希望！但测试损失结果更有趣。在我们最初的种子 42 结果中，使用错误初始化，仅输出头 LoRE 的模型测试损失为 3.737128，但这里我们得到了 3.785209。实际上，使用修正后的智能初始化后更差了！而且这也延续到了在输入和输出上都使用 LoRE 的模型结果；使用错误初始化的原始模型得到了 3.671585，而这里我们得到了 3.741209——恶化甚至更大。

将这两者结合起来，我喜欢用损失景观来描述的一个画面是：之前我们是从一座靠近深谷的山上开始——我们以高损失开始，但附近有一个低损失的地方。但是修复了智能初始化后，我们现在是从一座小山上开始，附近有一个不那么深的谷。

那么，如果我们做一个测试损失差值表，显示每个模型相对于基线的表现，会发生什么？

如果我们将仅输入测试损失差值 -0.073 和仅输出差值 0.186 相加，我们得到 0.113——与实际两者都使用时得到的差值 0.142 相当不同。我们曾经拥有的那个良好的“可加性”属性，即仅输入 LoRE 的差值加上仅输出 LoRE 的差值几乎正好等于两端都使用 LoRE 的差值，似乎已经消失了 :-(

尽管如此，这是一组有趣的结果！让我们把所有内容汇总起来。

结果

以下是所有训练运行汇总在一个表格中。为了保持可管理性，我删除了成本、最终训练损失和训练时间的列，并按测试损失排序。

有几件事很突出：

• 仅输入端的 LoRE 模型领先——尽管差距很小，而且 8xa100m40-lore-7-new-seed-input-only 与 8xa100m40-lore-6-new-seed-baseline-disabled 之间的差距在 0.013 的情况下很可能在噪声范围内。在那篇文章中，我用不同的权重初始化种子训练了三个模型，得到的模型差异最高可达约 0.017。
• 使用修正后的智能初始化的两个模型表现最差——甚至比没有额外初始化的模型 8xa100m40-lore-2-both-ordinary-init 还要差。

那么这一切告诉了我们什么？

总结

上面的数字表明了三件事——但下面有重要的注意事项：

1. LoRE 可以大幅减少这些模型中的参数数量，即使只应用于输入嵌入，而且这可能带来测试损失的改善（一次测试中为 0.073，第二次为 0.013）。
2. 当应用于输出头时，参数数量的减少幅度相同，但会导致测试损失的恶化（在三次测试中分别为 0.115、0.138 和 0.186）。
3. 两端都应用 LoRE 的结果似乎大致接近（在三种情况中的两种里非常接近）仅输入端和仅输出端增量之和，尽管这似乎并不是严格的加性关系。

更需谨慎地说，我们可以这样表述：

1. 我称之为“智能初始化”的数学上更干净的实现，实际上比我最初使用的错误版本效果更差，不过这仅用一个种子进行了测试，因此证据要弱得多。

仅输入端 LoRE 带来的两个明显改善中较大的那个，以及仅输出端带来的恶化，都不算小——当我在 GPT-2 模型中尝试各种干预措施时（比如移除 dropout、调度学习率、梯度裁剪等），其中许多措施的效果与之相当或更小。类似地，在双倍 Chinchilla 最优 token 数（即每个参数 40 个 token 而非 20 个）上训练一个可比的非 LoRE 模型，损失改善了 0.09，而将相同的计算资源投入到在更大模型上进行 Chinchilla 最优训练，相比基线改善了大约 0.13（原始数据见这篇文章）。

那么，仅输入端 LoRE 是改善测试损失的“一个奇怪技巧”吗？当然不是。仅输出端 LoRE 和两端 LoRE 是坏事吗？未必。

首先要记住的是，这是一组非常有限的实验。尽管我最终在训练运行上花费了大约 450 美元（我真应该投资把我的机器 poppy 升级为更好的训练机，以节省 Lambda 费用），但我只能测试：

• 针对我自己的 GPT-2 风格模型（与原始版本的重要差异：QKV 权重上没有权重绑定或偏置，没有 dropout，使用 PyTorch 默认权重初始化）。
• 用两个种子。
• 用两种智能初始化选项，其中一种只用一个种子。
• 在选定的 token 数量上，该数量对基线非 LoRE 模型是 Chinchilla 最优的。⁷

即使只是要对 LoRE 是否适用于这个特定架构做出任何确定性的结论，也需要在大量随机种子和不同训练运行长度（以 token 数量计）上进行测试：

• 也许这个效果仅限于我测试的那些随机种子？两个总比一个好，但三个会更好，四个还要更好，以此类推。
• 也许在更长的训练运行中，这种效果会消失？或许 LoRE 会导致某种容量损失，使模型在超过某个点后无法继续改进。直觉上这说得通——如果你把训练想象成在某种意义上将训练数据中的信息压缩到参数中，那么参数更少显然会在某种意义上限制你的容量。

要获得关于智能初始化是否是一件好事的可靠结果，你需要尝试多个随机种子，看看“正确的”智能初始化相对于“不正确的”初始化表现不佳是否是一个普遍结果。如果是，那又是为什么呢？

除此之外：仅输入 LoRE 带来的损失上的明显改进，是这些实验中出现的一个现象。也许它最终会被证明是真实的，也许不是，但 Andrew 在他的原始消息中提到的真正好处是关于完全不同的东西：LoRE 节省了大量参数，而（正如他后来所说）只付出了少量损失。如果你把这些节省下来的参数重新投入到具有更多层的更深网络中，你就有望弥补你，呃，损失的损失，并且还有额外收获——最终得到一个能力更强的模型。

因此，重要的进一步测试（我想我会做，但稍后）将是：

• Isoparameter tests:
  • 仅输入 LoRE 节省了 3200 万个参数，那么如果你把这些“花”在额外的层上会怎样？或者让模型变得更宽（就其嵌入维度而言）？或者两者兼而有之？
  • 仅输出 LoRE 同样如此。
  • 当然，还有输入和输出都使用的情况。
• 等计算量测试。值得注意的是，正如我之前提到的，仅输入 LoRE 并不能节省多少计算量——嵌入层在计算方面很便宜，尽管在参数方面不便宜。所以在这里，也许更有趣的是看看如果你使用仅输出 LoRE，但在更多数据上训练模型，使总训练计算预算相匹配，会发生什么。你能弥补损失的损失吗？

值得注意的是，在我的 Chinchilla 实验中，我将模型扩大了 2 倍。这大约增加了 41%，非常接近我们在两端使用 LoRE 所节省的 39%。所以如果损失改进能够延续，也许它会奏效？另一方面，那些模型也是在多 41% 的 token 上训练的，所以……

因此，需要大量的进一步实验来充实这些结果——而这甚至还没有超出这个特定架构。Andrew 报告说他的设置取得了可靠的结果，而且他的模型 Maba 看起来确实相当不同。但它会扩展到 Qwen 风格的模型吗？或者 Mistral？

接下来，还有缩放的问题。LoRE 在小模型中很有价值，因为嵌入和输出头在总参数数量中占很大比例。我的猜测是，随着模型变得更大，它会变得越来越无用。

举一个极端的例子，看看 Kimi K3 论文，它的词表大小为 160k，隐藏维度（我假设这就是嵌入层使用的维度）为 7,168。这意味着嵌入层大约使用了 1,146,880,000 个参数，输出头也是如此，总计约 23 亿参数（假设它们没有使用权重绑定，对于这样一个庞然大物来说，这似乎很可能）。 它每个 token 有 104.2B 活跃参数，所以如果使用 LoRE，你大约能节省活跃参数数量的 2%。但就“再投资”到更多层而言，真正的比较点是总参数量，而不是活跃参数，在约 2.8T 的总量中，用于嵌入的参数不到 0.1%。听起来帮助不大。

当然，除非仅输入 LoRE 带来的损失改进确实是真的，并且没有随着模型规模扩大而减弱……

但无论如何，这只是推测。我认为一组好的实验是尝试与之前相同的训练运行，但使用更大的模型。也许先试试 GPT-2 medium 和 large 规模会是个不错的开始？

工作量很大。如果我能使用 Google 的资源，我想我会启动一堆训练运行。另一方面，如果我能使用 Google 的资源，我大概就不会写这篇博客了……

我想我会从中得到一点收获。LoRE 是一种有趣的干预——我认为在任何词表矩阵占参数显著比例的新模型上都值得一试，看看是否有帮助。我将来会这样做。

如果其他人想尝试 LoRE 风格的训练运行，我很想听到你的消息。任何结果，无论正面还是负面，都会非常有趣。

感谢阅读！非常感谢 Andrew 提出这个想法来尝试——这很有趣 :-)

1. 这通常写作 Z=XWT，以反映权重通常以 dout×din 格式保存的事实，至少在 PyTorch 中是这样，但这实际上是实现细节。 ↩

2. 我将来会写更多关于 LoRA 的内容——我一直在玩它——但本质上，你要做的是：

   • 冻结构成你参数的所有矩阵（例如，在 PyTorch 中，将 requires_grad 设置为 False）。

   • 对于每个你想训练的矩阵（可能不是全部——对于微调，人们通常只训练一个子集），你添加一个低秩适配器，它是一对矩阵，就像上面的 A 和 B，它“坐在”现有矩阵旁边，像这样：

     

     梯度跟踪在低秩适配器中的矩阵上启用，其结果被添加到原始权重的结果中。

   这个想法是，适配器对学习为原始权重的结果创建一个“差异”，从而允许对模型的那部分进行一定量的训练。

   而且它效果出奇地好。再次强调：我将来会更深入地写这个。 ↩

3. 当我和一个 LLM 讨论这件事时——遗憾的是我忘了是哪个——这里提到了 softmax 瓶颈，后来与 Andrew 的通信中也提到了。那指的是这篇论文。我需要在以后某个时候仔细研究一下，但相关的要点似乎是，从 softmax 出来的值的秩（我们可以粗略地理解为预测可以从一个上下文到另一个上下文变化的独立方式有多少种）受到输入其中的值的秩的限制。如果我们强制让（在我的 GPT-2 small 情况下）768 维嵌入通过一个 128 维的瓶颈，我们就在降低它们固有的秩，而正是这个被降低的秩会驱动 softmax 的秩，所以结果会比用原来的 768 维得到的结果更差。看起来足够合乎逻辑。

   防链接失效的参考文献：Yang, Z., Dai, Z., Salakhutdinov, R., & Cohen, W. W. (2018). Breaking the softmax bottleneck: A high-rank RNN language model. In International Conference on Learning Representations (ICLR). ↩

4. 感谢 Claude Opus 5.5 向我指出中心极限定理。我确实需要温习一下我的统计学，而且我肯定在 90 年代的某个时候知道它…… ↩

5. 各种形式的 CUDA 抖动除外——不过即便如此，前一阵子我在 8 台 A100 机器上用相同配置做了两次训练，相隔数月，Safetensors 文件逐位完全相同。所以我认为在我的设置中这种影响微乎其微。 ↩

6. 想要深入探究其原因细节，以及那个错误的智能初始化到底对“虚拟”权重 W′ 的方差做了什么，这种诱惑很强烈。但我会忍住。 ↩

7. 当我把这篇文章的草稿拿给我的LLM 编辑委员会过目时，Grok 指出，仅在输入嵌入上使用 LoRE 所显现出的好处，也许是因为那个模型每个参数看到的 token 明显更多。这是个有趣的想法！如果我们算一下，仅输入和仅输出的 LoRE 模型是在每个参数约 25 个 token 上训练的，两端都有的模型是在约 33 个上训练的，而非 LoRE 模型是 20 个。

   然而，在我的 Chinchilla 实验中，用每个参数 40 个 token 而不是 20 个来训练，我只得到了 0.09 的损失改善。现在，我确信这个增益主要发生在训练早期——过度训练会遭遇收益递减——但仅输入 LoRE 的 0.073 结果意味着几乎全部增益都发生在每参数额外的前 5 个 token 中，这似乎很可疑。

   Those are different training runs with different seeds, and were using my JAX code, which has a better baseline loss than my PyTorch code, so we can't compare them directly. But it feels unlikely that the effect would be drastically different.

   Also, if we use reductio ad absurdum on Grok's logic, it does sound a bit weak -- if we get rid of parameters but train on the same number of tokens, the higher tokens per parameter will not in general counteract the damage from the smaller model. Imagine, for example, removing 32M parameters worth of Transformer layers and training on the Chinchilla-optimal tokens for the size of the model before you did that. I'm not going to run the experiment, but I'd be... somewhat more than surprised if it improved the model. ↩

Citing this post

This is a blog, and if you want to link to this post then please do :-) However, if you're writing something more academic and need to do a proper citation, then here's a BibTeX block to make things easier.

@misc{thomas2026oct-low-rank-vocab-matrices,
  author       = {Thomas, Giles},
  title        = {{Fun with low-rank vocab matrices (and a bonus test loss reduction?)}},
  year         = {2026},
  month        = oct,
  howpublished = {Blog post},
  url          = {https://www.gilesthomas.com/2026/10/low-rank-vocab-matrices},
}