论文重点总结

论文全文翻译

Abstract

主流的序列转导模型都是基于复杂的循环或卷积神经网络，包含一个编码器和一个解码器。表现最好的模型还通过一个注意力机制来连接编码器和解码器。我们提出了一种新的、简单的网络架构——Transformer，它完全基于注意力机制，完全摒弃了循环和卷积。在两个机器翻译任务上的实验表明，这些模型在质量上更优，同时具有更好的并行性，并且训练所需的时间显著减少。在 WMT 2014 英德翻译任务中，我们的模型取得了 28.4 的 BLEU 值，比以往包括集成模型在内的最佳结果还要高出 2.0 BLEU。在 WMT 2014 英法翻译任务中，我们的模型在八个 GPU 上训练了 3.5 天后，便创下了单一模型下 41.8 的 BLEU 值新纪录，而这只是文献中最佳模型训练成本的一小部分。我们还证明了 Transformer 能够很好地泛化到其他任务，并成功地将其应用于拥有大量和有限训练数据的英语成分句法分析任务中。

1. Introduction

循环神经网络（RNN），特别是长短期记忆（LSTM） 和门控循环（gated recurrent） 神经网络，已经在序列建模和如语言建模、机器翻译等转导问题中，被牢固地确立为最先进的方法。此后，大量的研究工作继续推动着循环语言模型和编码器-解码器架构的边界。

循环模型通常沿着输入和输出序列的符号位置进行计算。通过将位置与计算时间中的步骤对齐，它们生成一系列隐藏状态 \(h_t\)，作为前一个隐藏状态 \(h_{t-1}\) 和当前位置 t 输入的函数。这种固有的顺序性阻碍了训练样本内的并行化，这在序列长度较长时变得至关重要，因为内存限制了跨样本的批处理。最近的工作通过因式分解技巧和条件计算 在计算效率上取得了显著的改进，同时后者还提升了模型性能。然而，顺序计算这个根本性的制约依然存在。

注意力机制已经成为各种任务中一个不可或缺的部分，它使得模型能够对依赖关系进行建模，而不考虑它们在输入或输出序列中的距离。然而，在除了少数几个案例外，这种注意力机制都是与循环网络结合使用的。

在这项工作中，我们提出了 Transformer，一种摒弃了循环结构的模型架构，并完全依赖于注意力机制来绘制输入和输出之间的全局依赖关系。Transformer 允许进行显著更多的并行化，并且在八个 P100 GPU 上仅训练十二小时后，就可以在翻译质量上达到新的技术水平。

2. Background

减少顺序计算的目标也构成了 Extended Neural GPU、ByteNet 和 ConvS2S 的基础，所有这些模型都使用卷积神经网络作为基本构建块，为所有输入和输出位置并行计算隐藏表示。在这些模型中，关联来自两个任意输入或输出位置的信号所需的操作数量，会随着位置间距离的增加而增长，对于 ConvS2S 是线性增长，对于 ByteNet 是对数增长。这使得学习远距离位置之间的依赖关系变得更加困难。在 Transformer 中，这个操作数量被减少到了常数级别，尽管这是以牺牲因注意力加权位置平均化而导致的有效分辨率降低为代价的，但我们通过第 3.2 节中描述的多头注意力机制来弥补了这一影响。

自注意力（Self-attention），有时也称为内部注意力（intra-attention），是一种关联单个序列不同位置的注意力机制，目的是为了计算序列的表示。自注意力已在多种任务中成功应用，包括阅读理解、摘要概括、文本蕴含以及学习与任务无关的句子表示。

端到端的记忆网络是基于循环注意力机制而非序列对齐的循环，并已在简单语言问答和语言建模任务上表现良好。

然而，据我们所知，Transformer 是第一个完全依赖自注意力来计算其输入和输出表示的转导模型，而不使用序列对齐的 RNN 或卷积。在接下来的章节中，我们将描述 Transformer，阐述使用自注意力的动机，并讨论其相对于 和 等模型的优势。

3. Model Architecture

大多数有竞争力的神经序列转导模型都具有编码器-解码器结构。在这里，编码器将一个由符号表示组成的输入序列 \((x_1, ..., x_n)\) 映射到一个连续表示序列 \(z = (z_1, ..., z_n)\)。给定 z，解码器随后逐个元素地生成一个输出序列 \((y_1, ..., y_m)\)。在每一步中，模型都是自回归的，即在生成下一个符号时，将先前生成的符号作为额外输入。

Transformer 遵循这种整体架构，使用堆叠的自注意力层和逐点全连接层来分别构建编码器和解码器，如图 1 的左半部分和右半部分所示。

3.1. Encoder and Decoder Stacks

编码器：编码器由 N = 6 个相同的层堆叠而成。每一层有两个子层。第一个是多头自注意力机制，第二个是一个简单的、按位置进行的全连接前馈网络。我们在每个子层周围都使用了一个残差连接，然后进行层归一化。也就是说，每个子层的输出是 LayerNorm(x + Sublayer(x))，其中 Sublayer(x) 是该子层本身实现的函数。为了方便这些残差连接，模型中所有的子层以及嵌入层，都产生维度为 dmodel = 512 的输出。

解码器：解码器同样由 N = 6 个相同的层堆叠而成。除了每个编码器层中的两个子层外，解码器还插入了第三个子层，该子层对编码器栈的输出执行多头注意力。与编码器类似，我们在每个子层周围使用残差连接，然后进行层归一化。我们还修改了解码器栈中的自注意力子层，以防止当前位置关注到后续位置。这种掩码机制，结合输出嵌入被偏移一个位置的事实，确保了对位置 i 的预测只能依赖于位置小于 i 的已知输出。

注意力函数可以被描述为将一个查询（query）和一组键值对（key-value pairs）映射到一个输出的过程，其中查询、键、值和输出都是向量。输出是根据值的加权和计算得出的，其中每个值的权重是通过查询与对应键的兼容性函数计算的。

3.2.1. Scaled Dot-Product Attention

我们将我们特定的注意力机制称为“缩放点积注意力”（图2）。输入由维度为 \(d_k\) 的查询和键，以及维度为 \(d_v\) 的值组成。我们计算查询与所有键的点积，然后将每个点积除以 \(\sqrt{d_k}\)，最后应用一个 softmax 函数来获得值的权重。

在实践中，我们同时对一组查询计算注意力函数，并将它们打包成一个矩阵 Q。键和值也分别打包成矩阵 K 和 V。我们计算输出矩阵如下：

$$ Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$$

两种最常用的注意力函数是加法注意力 和点积（乘法）注意力。除了缩放因子 \(\frac{1}{\sqrt{d_k}}\) 外，点积注意力与我们的算法是相同的。加法注意力使用一个带单个隐藏层的前馈网络来计算兼容性函数。虽然两者在理论上的复杂度相似，但点积注意力在实践中速度更快、空间效率更高，因为它可以使用高度优化的矩阵乘法代码来实现。

当 \(d_k\) 的值较小时，两种机制的表现相似，但对于较大的 \(d_k\) 值，加法注意力的表现优于没有缩放的点积注意力。我们怀疑，对于较大的 \(d_k\) 值，点积的量级会变得很大，从而将 softmax 函数推向梯度极小的区域。为了抵消这个影响，我们用 \(\frac{1}{\sqrt{d_k}}\) 来缩放点积。

我们发现，与其使用 \(d_{model}\) 维度的键、值和查询来执行单次注意力函数，不如将查询、键和值分别通过 h 次不同的、学习到的线性投影，将它们投影到 \(d_q\)、\(d_k\) 和 \(d_v\) 维度。然后，在每个投影版本的查询、键和值上并行执行注意力函数，产生 h 个 \(d_v\) 维的输出值。如图 2 所示，这些输出值被拼接起来，并再次进行投影，从而得到最终的值。

多头注意力允许模型在不同位置共同关注来自不同表示子空间的信息。而对于单一注意力头，平均化过程会抑制这一点。

$$MultiHead(Q, K, V)=Concat(head_1,...,head_h)W^O$$

$$where\ head_i=Attention(QW^Q_i, QW^Q_i, QW^Q_i)$$

其中，投影是参数矩阵 \(Wᵢ^Q ∈ ℝ^{(dmodel×dk)}, Wᵢ^K ∈ ℝ^{(dmodel×dk)}, Wᵢ^V ∈ ℝ^{(dmodel×dv)}\) 和 \(W⁰ ∈ ℝ^{(h·dv×dmodel)}\)。

在这项工作中，我们采用 h = 8 个并行的注意力层，或称为“头”（heads）。对于每个头，我们使用 \(d_k = d_v = d_{model}/h = 64\)。由于每个头的维度减小，总计算成本与具有完整维度的单头注意力相似。

3.2.3. Applications of Attention in our Model

Transformer 以三种不同的方式使用多头注意力：

• 在“编码器-解码器注意力”层中，查询来自前一个解码器层，而记忆键和值来自编码器的输出。这使得解码器中的每个位置都能关注到输入序列中的所有位置。这模仿了序列到序列模型中典型的编码器-解码器注意力机制，如。
• 编码器包含自注意力层。在一个自注意力层中，所有的键、值和查询都来自同一个地方，即编码器中前一层的输出。编码器中的每个位置都可以关注到编码器前一层的所有位置。
• 类似地，解码器中的自注意力层允许解码器中的每个位置关注到解码器中直到并包括该位置在内的所有位置。我们需要防止解码器中的信息向左流动，以保持其自回归特性。我们在缩放点积注意力内部通过掩码（masking）实现这一点，即在 softmax 的输入中，将所有对应非法连接的值设置为-∞。参见图2。

除了注意力子层，我们编码器和解码器的每一层都包含一个全连接的前馈网络，该网络独立且相同地应用于每个位置。该网络由两个线性变换组成，中间有一个 ReLU 激活函数。

$$FFN(x)=max(0,xW_1+b_1)W_2+b_2$$

虽然线性变换在不同位置上是相同的，但它们在不同层之间使用不同的参数。描述这一点的另一种方式是，它是两个核大小为 1 的卷积。输入和输出的维度是 \(d_{model} = 512\)，内层的维度是 \(d_{ff} = 2048\)。

与其他序列转导模型类似，我们使用学习到的嵌入（embeddings）将输入词元（tokens）和输出词元转换为维度为 \(d_{model}\) 的向量。我们还使用常规的学习线性变换和 softmax 函数将解码器的输出转换为预测的下一个词元概率。在我们的模型中，我们共享两个嵌入层和 pre-softmax 线性变换之间的权重矩阵，类似于。在嵌入层中，我们将这些权重乘以 \(\sqrt {d_{model}}\)。

由于我们的模型不包含循环和卷积，为了让模型能够利用序列的顺序，我们必须注入一些关于序列中词元相对或绝对位置的信息。为此，我们将“位置编码”（positional encodings）添加到编码器和解码器栈底部的输入嵌入中。位置编码与嵌入具有相同的维度 \(d_{model}\)，因此两者可以相加。位置编码有多种选择，可以是学习的也可以是固定的。

在这项工作中，我们使用不同频率的正弦和余弦函数：

$$PE_{(pos,2i)}=sin(pos/10000^{2i/d_{model}})$$

$$PE_{(pos,2i+1)}=cos(pos/10000^{2i/d_{model}})$$

其中 pos 是位置，i 是维度。也就是说，位置编码的每个维度都对应一个正弦曲线。波长形成一个从 2π 到 10000·2π 的几何级数。我们选择这个函数是因为我们假设它能让模型轻松地学习通过相对位置来关注，因为对于任何固定的偏移量 \(k，PE_{pos+k}\) 都可以表示为 \(PE_{pos}\) 的线性函数。

我们也尝试了使用学习的位置嵌入，并发现两个版本产生了几乎相同的结果（见表3，行(E)）。我们选择了正弦版本，因为它可能允许模型外推到比训练期间遇到的序列更长的序列。

4. Why Self-Attention

在本节中，我们将自注意力层的各个方面与循环层和卷积层进行比较，这些层通常用于将一个可变长度的符号表示序列 \((x_1, ..., x_n)\) 映射到另一个等长的序列 \((z_1, ..., z_n)\)，其中 \(x_i, z_i ∈ ℝᵈ\)，例如在一个典型的序列转导编码器或解码器的隐藏层中。我们使用自注意力的动机主要考虑三个方面。

一是每层的总计算复杂度。另一个是可以并行化的计算量，以所需的最小顺序操作数量来衡量。

第三是网络中远距离依赖关系之间的路径长度。学习远距离依赖关系是许多序列转导任务中的一个关键挑战。影响学习这种依赖关系能力的一个关键因素是前向和后向信号在网络中必须传播的路径长度。输入和输出序列中任意位置组合之间的路径越短，学习远距离依赖关系就越容易。因此，我们还比较了由不同层类型组成的网络中任意两个输入和输出位置之间的最大路径长度。

如表1所示，自注意力层用常数数量的顺序执行操作连接所有位置，而循环层需要 \(O(n)\) 的顺序操作。在计算复杂度方面，当序列长度 n 小于表示维度 d 时，自注意力层比循环层快，这在机器翻译中最先进的模型所使用的句子表示中通常是这种情况，例如词片（word-piece） 和字节对（byte-pair） 表示。为了提高处理非常长序列任务的计算性能，可以将自注意力限制为只考虑输入序列中以相应输出位置为中心的大小为 r 的邻域。这将使最大路径长度增加到 \(O(n/r)\)。我们计划在未来的工作中进一步研究这种方法。

一个核宽 k < n 的单一卷积层并不能连接所有输入和输出位置对。要做到这一点，需要一个由 \(O(n/k)\) 个卷积层组成的栈（对于连续核），或者 \(O(log_k(n))\) 个（对于扩张卷积），这增加了网络中任意两个位置之间最长路径的长度。卷积层通常比循环层计算成本更高，大约是 k 倍。然而，可分离卷积 将复杂度显著降低到 \(O(k·n·d + n·d²)\)。然而，即使 k=n，可分离卷积的复杂度也等于我们模型中采用的方法，即自注意力层和逐点前馈层的组合。

作为一个附带的好处，自注意力可以产生更具可解释性的模型。我们从我们的模型中检查注意力分布，并在附录中展示和讨论了例子。不仅单个注意力头清晰地学会了执行不同的任务，许多头似乎还表现出与句子的句法和语义结构相关的行为。

附录 注意力可视化

5. Training

5.1. Training Data and Batching

我们在标准的 WMT 2014 英德翻译数据集上进行了训练，该数据集包含约 450 万个句子对。句子使用字节对编码（byte-pair encoding） 进行编码，该编码有一个约 37000 个词元的共享源-目标词汇表。对于英法翻译，我们使用了更大的 WMT 2014 英法数据集，该数据集包含 3600 万个句子，并将词元拆分为一个 32000 个词片的词汇表。句子对根据近似的序列长度进行批处理。每个训练批次包含一组句子对，其中大约有 25000 个源词元和 25000 个目标词元。

5.2. Hardware and Schedule

我们在一台配备 8 个 NVIDIA P100 GPU 的机器上训练我们的模型。对于使用本文描述的超参数的基础模型，每个训练步骤大约需要 0.4 秒。我们总共训练了 100,000 步，即 12 个小时。对于我们的大模型（在表3的最后一行描述），步长为 1.0 秒。大模型训练了 300,000 步（3.5天）。

5.3. Optimizer

我们使用了 Adam 优化器，其中 \(β₁ = 0.9，β₂ = 0.98，ε = 10⁻⁹\)。我们根据以下公式在训练过程中改变学习率：

$$lrate=d_{model}^{-0.5}\cdot min(step\_num^{-0.5},\ step\_num\cdot warmup\_steps^{-1.5})$$

这对应于在第一个 warmup_steps 训练步骤中线性增加学习率，然后按步数的平方根倒数比例减小它。我们使用了 warmup_steps = 4000。

5.4. Regularization

我们在训练期间采用了三种类型的正则化：

残差丢弃 (Residual Dropout): 我们将 dropout 应用于每个子层的输出，在它被加到子层输入并进行归一化之前。此外，我们还将 dropout 应用于编码器和解码器栈中嵌入和位置编码的总和。对于基础模型，我们使用 \(P_{drop} = 0.1\) 的丢弃率。

标签平滑 (Label Smoothing): 在训练期间，我们采用了值为 \(ε_{ls} = 0.1\) 的标签平滑。这会损害困惑度，因为模型学会了变得更不确定，但它提高了准确性和 BLEU 分数。

6. Results

6.1. Machine Translation

在 WMT 2014 英德翻译任务中，大型 Transformer 模型（表2中的 Transformer (big)）的表现超过了之前报道过的所有最佳模型（包括集成模型），性能提升超过 2.0 BLEU，创造了 28.4 的新 BLEU 分数记录。该模型的配置在表3的最后一行中列出。训练在 8 个 P100 GPU 上耗时 3.5 天。即使是我们的基础模型，也超过了所有先前发布的模型和集成模型，而其训练成本只是任何一个竞争模型的一小部分。

在 WMT 2014 英法翻译任务中，我们的大型模型取得了 41.8 的 BLEU 分数，超过了所有先前发布的单一模型，而训练成本不到之前最先进模型的 1/4。为英法翻译任务训练的 Transformer (big) 模型使用了 \(P_{drop} = 0.1\) 的丢弃率，而不是 0.3。

对于基础模型，我们通过平均最后 5 个检查点（checkpoints）来获得一个单一模型，这些检查点是以 10 分钟为间隔保存的。对于大型模型，我们平均了最后 20 个检查点。我们使用了束搜索（beam search），束大小为 4，长度惩罚 \(α = 0.6\)。这些超参数是在开发集上进行实验后选择的。我们将推理时的最大输出长度设置为输入长度 + 50，但在可能的情况下会提前终止。

表2总结了我们的结果，并将我们的翻译质量和训练成本与文献中的其他模型架构进行了比较。我们通过将训练时间、使用的 GPU 数量以及每个 GPU 的持续单精度浮点运算能力的估计值相乘，来估算训练模型所用的浮点运算次数。

6.2. Model Variations

为了评估 Transformer 不同组件的重要性，我们通过不同方式改变我们的基础模型，测量其在英德翻译开发集 newstest2013 上的性能变化。我们使用前一节描述的束搜索，但没有进行检查点平均。我们在表3中展示了这些结果。

在表3的行(A)中，我们改变了注意力头的数量以及注意力的键和值维度，同时保持计算量不变，如第3.2.2节所述。虽然单头注意力的 BLEU 值比最佳设置差 0.9，但当头的数量过多时，质量也会下降。

在表3的行(B)中，我们观察到减小注意力键的大小 \(d_k\) 会损害模型质量。这表明确定兼容性并不容易，一个比点积更复杂的兼容性函数可能是有益的。我们在行(C)和(D)中进一步观察到，正如预期的那样，更大的模型更好，并且 dropout 对于避免过拟合非常有帮助。在行(E)中，我们用学习的位置嵌入替换了我们的正弦位置编码，并观察到与基础模型几乎相同的结果。

6.3. English Constituency Parsing

为了评估 Transformer 是否能泛化到其他任务，我们对英语成分句法分析进行了实验。这项任务提出了特定的挑战：输出受到强烈的结构性约束，并且明显长于输入。此外，RNN 序列到序列模型在小数据场景下未能达到最先进的结果。

我们训练了一个 4 层的 Transformer，其 \(d_{model} = 1024\)，在宾夕法尼亚树库（Penn Treebank） 的华尔街日报（WSJ）部分上进行训练，该部分约有 4 万个训练句子。我们还在半监督环境下进行了训练，使用了更大的高置信度和 BerkleyParser 语料库，其中包含约 1700 万个句子。我们对仅使用 WSJ 的设置使用了 16K 词元的词汇表，对半监督设置使用了 32K 词元的词汇表。

我们只进行了少量实验来选择 dropout（包括注意力和残差，第5.4节）、学习率和束大小，这些实验是在第 22 节开发集上进行的，所有其他参数都与英德基础翻译模型保持不变。在推理过程中，我们将最大输出长度增加到输入长度 + 300。我们对仅 WSJ 和半监督设置都使用了 21 的束大小和 \(α = 0.3\)。

我们在表4中的结果显示，尽管缺乏针对特定任务的调优，我们的模型表现得惊人地好，除了循环神经网络语法（Recurrent Neural Network Grammar） 外，其结果优于所有先前报道的模型。

与 RNN 序列到序列模型 相比，即使只在 4 万个句子的 WSJ 训练集上训练，Transformer 的表现也优于 BerkeleyParser。

7. Conclusion

在这项工作中，我们提出了 Transformer，这是第一个完全基于注意力的序列转导模型，用多头自注意力取代了编码器-解码器架构中最常用的循环层。

对于翻译任务，Transformer 的训练速度比基于循环或卷积层的架构要快得多。在 WMT 2014 英德和 WMT 2014 英法翻译任务中，我们都取得了新的技术水平。在前一个任务中，我们最好的模型甚至超过了所有先前报道的集成模型。

我们对基于注意力的模型的未来感到兴奋，并计划将它们应用于其他任务。我们计划将 Transformer 扩展到涉及文本以外的输入和输出模态的问题，并研究局部的、受限的注意力机制，以有效处理如图像、音频和视频等大型输入和输出。使生成过程不那么顺序化是我们的另一个研究目标。

我们用于训练和评估模型的代码可在 tensor2tensor 获取。

论文重点实现解析

本文是对哈佛NLP团队实现的Pytorch版Transformer的源码解析。

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."

    def __init__(self, features, eps=1e-6):
        """
        features: d_model 512||1024
        """
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
                

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn
                

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h # 64
        self.h = h # 8
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = [
            lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        ]

        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout
        )

        # 3) "Concat" using a view and apply a final linear.
        x = (
            x.transpose(1, 2)
            .contiguous()
            .view(nbatches, -1, self.h * self.d_k)
        )
        del query
        del key
        del value
        return self.linears[-1](x)
                

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."

    def __init__(self, d_model, d_ff, dropout=0.1):
        """
        d_model = 512
        d_ff = 2048
        """
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))
                

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model # 512

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)
                

# 分词
word = 'Hi, 你好~'
tokens = tokenizer.tokenize(word)
print(f'{word} 分词: {tokens}')
# 输出：Hi, 你好~ 分词: ['Hi', ',', 'Ġ', 'ä½łå¥½', '~']

token_ids = tokenizer.convert_tokens_to_ids(tokens)
print(f'{word} Token IDs: {token_ids}')
# 输出：Hi, 你好~ Token IDs: [13048, 11, 220, 108386, 93]

# 获取该单词的嵌入向量
word_embedding = embeddings.weight[token_ids]
print(f'{word} 的嵌入形状：{word_embedding.shape}')
# 输出：Hi, 你好~ 的嵌入形状：torch.Size([5, 1536])

print(f'{word} 的嵌入内容: {word_embedding}')
# 输出：
# Hi, 你好~ 的嵌入内容: tensor([[ 0.0261,  0.0048, -0.0043,  ...,  0.0193, -0.0493, -0.0020],
#         [-0.0303, -0.0159, -0.0107,  ..., -0.0198, -0.0020, -0.0129],
#         [-0.0236, -0.0254,  0.0325,  ..., -0.0317, -0.0082,  0.0137],
#         [ 0.0270,  0.0042,  0.0014,  ...,  0.0425, -0.0195,  0.0011],
#         [-0.0205, -0.0408, -0.0013,  ...,  0.0272, -0.0060,  0.0032]],
#        dtype=torch.bfloat16, grad_fn=<\IndexBackward0>)
                

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term) # 偶数索引列填充
        pe[:, 1::2] = torch.cos(position * div_term) # 奇数索引列填充
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, x):
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)