当前位置：首页 > news >正文

LLMs-from-scratch ：embeddings 与 linear-layers 的对比

news 2025/10/17 10:13:28

代码链接：https://github.com/rasbt/LLMs-from-scratch/blob/main/ch02/03_bonus_embedding-vs-matmul/embeddings-and-linear-layers.ipynb

《从零开始构建大型语言模型》一书的补充代码，作者：Sebastian Raschka

代码仓库：https://github.com/rasbt/LLMs-from-scratch

理解嵌入层和线性层之间的区别

PyTorch 中的嵌入层与执行矩阵乘法的线性层实现相同的功能；我们使用嵌入层的原因是计算效率
我们将使用 PyTorch 中的代码示例逐步了解这种关系

import torchprint("PyTorch version:", torch.__version__)

PyTorch version: 2.5.1+cu124

使用 nn.Embedding

# 假设我们有以下 3 个训练示例，
# 它们可能代表 LLM 上下文中的标记 ID
idx = torch.tensor([2, 3, 1])# 嵌入矩阵中的行数可以通过
# 获取最大标记 ID + 1 来确定。
# 如果最高标记 ID 是 3，那么我们需要 4 行，用于可能的
# 标记 ID 0, 1, 2, 3
num_idx = max(idx)+1# 所需的嵌入维度是一个超参数
out_dim = 5

idx

tensor([2, 3, 1])

num_idx

tensor(4)

让我们实现一个简单的嵌入层：

# 我们使用随机种子来保证可重现性，因为
# 嵌入层中的权重是用小的随机值初始化的
torch.manual_seed(123)embedding = torch.nn.Embedding(num_idx, out_dim)

我们可以选择性地查看嵌入权重：

embedding.weight

Parameter containing:
tensor([[ 0.3374, -0.1778, -0.3035, -0.5880,  1.5810],[ 1.3010,  1.2753, -0.2010, -0.1606, -0.4015],[ 0.6957, -1.8061, -1.1589,  0.3255, -0.6315],[-2.8400, -0.7849, -1.4096, -0.4076,  0.7953]], requires_grad=True)

然后我们可以使用嵌入层来获取 ID 为 1 的训练示例的向量表示：

embedding(torch.tensor([1]))

tensor([[ 1.3010,  1.2753, -0.2010, -0.1606, -0.4015]],grad_fn=<EmbeddingBackward0>)

下面是底层发生的事情的可视化：

类似地，我们可以使用嵌入层来获取 ID 为 2 的训练示例的向量表示：

embedding(torch.tensor([2]))

tensor([[ 0.6957, -1.8061, -1.1589,  0.3255, -0.6315]],grad_fn=<EmbeddingBackward0>)

现在，让我们转换之前定义的所有训练示例：

idx = torch.tensor([2, 3, 1])
embedding(idx)

tensor([[ 0.6957, -1.8061, -1.1589,  0.3255, -0.6315],[-2.8400, -0.7849, -1.4096, -0.4076,  0.7953],[ 1.3010,  1.2753, -0.2010, -0.1606, -0.4015]],grad_fn=<EmbeddingBackward0>)

在底层，它仍然是相同的查找概念：

使用 nn.Linear

现在，我们将演示上面的嵌入层与在 PyTorch 中对独热编码表示使用 nn.Linear 层完全相同
首先，让我们将标记 ID 转换为独热表示：

onehot = torch.nn.functional.one_hot(idx)
onehot

tensor([[0, 0, 1, 0],[0, 0, 0, 1],[0, 1, 0, 0]])

接下来，我们初始化一个 Linear 层，它执行矩阵乘法 $W^\top$ ：

torch.manual_seed(123)
linear = torch.nn.Linear(num_idx, out_dim, bias=False)
linear.weight

Parameter containing:
tensor([[-0.2039,  0.0166, -0.2483,  0.1886],[-0.4260,  0.3665, -0.3634, -0.3975],[-0.3159,  0.2264, -0.1847,  0.1871],[-0.4244, -0.3034, -0.1836, -0.0983],[-0.3814,  0.3274, -0.1179,  0.1605]], requires_grad=True)

请注意，PyTorch 中的线性层也是用小的随机权重初始化的；为了直接与上面的 Embedding 层进行比较，我们必须使用相同的小随机权重，这就是我们在这里重新分配它们的原因：

linear.weight = torch.nn.Parameter(embedding.weight.T)

linear.weight

Parameter containing:
tensor([[ 0.3374,  1.3010,  0.6957, -2.8400],[-0.1778,  1.2753, -1.8061, -0.7849],[-0.3035, -0.2010, -1.1589, -1.4096],[-0.5880, -0.1606,  0.3255, -0.4076],[ 1.5810, -0.4015, -0.6315,  0.7953]], requires_grad=True)

现在我们可以在输入的独热编码表示上使用线性层：

linear(onehot.float())

tensor([[ 0.6957, -1.8061, -1.1589,  0.3255, -0.6315],[-2.8400, -0.7849, -1.4096, -0.4076,  0.7953],[ 1.3010,  1.2753, -0.2010, -0.1606, -0.4015]], grad_fn=<MmBackward0>)

正如我们所看到的，这与我们使用嵌入层时得到的结果完全相同：

embedding(idx)

tensor([[ 0.6957, -1.8061, -1.1589,  0.3255, -0.6315],[-2.8400, -0.7849, -1.4096, -0.4076,  0.7953],[ 1.3010,  1.2753, -0.2010, -0.1606, -0.4015]],grad_fn=<EmbeddingBackward0>)