原文：
towardsdatascience.com/einstein-notation-a-new-lens-on-transformers-761390a7960b
在这篇文章中，我们将踏上一次轻松愉快的旅程，探索 Transformer 的世界，使用爱因斯坦符号来揭示其架构的复杂性。
引言：
Transformer 模型已经彻底改变了自然语言处理（以及更广泛的领域），在各种任务上实现了最先进的结果。它们的表现令人印象深刻，但底层的数学运算可能很复杂，难以理解——尤其是在没有分解单个层的情况下。在这篇文章中，我提出使用爱因斯坦符号来表达 Transformer 模型内的数学运算。
注意，爱因斯坦符号通常用于物理学和数学，如广义相对论、电磁学、量子力学和流体力学，也用于线性代数，以更紧凑的形式表示矩阵运算。
目标是以简洁优雅的方式写出每一层的数学运算。通过利用重复索引上的隐式求和，爱因斯坦符号可以简化张量运算的表示，使其（可能）更容易理解，因此也更容易实现 Transformer 模型的各个层。更详细地说，我将展示如何将爱因斯坦符号应用于 Transformer 模型的各个组成部分，包括自注意力、前馈神经网络和层归一化。
为了简单起见，我将仅关注 Transformer 的解码器部分，这在目前是生成大型语言模型（LLMs）的常见最佳实践。
动机
到目前为止，现代 Transformer 模型依赖于计算密集型操作，特别是在自注意力机制中。换句话说，在研究和开发中，我们体验到随着序列/标记长度的增加，计算成本的指数增长成为了一个主要的瓶颈。因此，依赖于当前数学解决方案进行推理和训练，模型扩展变得困难。
回到过去，我在物理学中学习了爱因斯坦符号。爱因斯坦符号以其优雅和高效而闻名。最近，我试图探索其在简化并优化 Transformer 架构数学表示方面的潜力——从而将 Transformer 模型的数学转化为更广泛的受众，例如非机器学习研究人员——创造一个新的视角，这可能导致新的见解和优化。
自注意力机制的核心概念相对简单。然而，显式的矩阵运算和求和可能会掩盖底层结构，使得一些数学步骤难以追踪。
最终，本文的目标是通过采用其数学基础的新视角来增进对 Transformer 模型的理解。
让我们快速回顾一下爱因斯坦符号的基本知识：
内积的爱因斯坦符号（作者创建的方程图）
两个向量的内积是通过将每个向量中的对应元素配对，乘以这些配对，然后将所有结果相加得到的。
叉积的爱因斯坦符号（作者创建的方程图）
Levi-Civita 符号（ε）用于简洁地表达叉积。重复的索引 j 和 k 是隐式求和的。换句话说，爱因斯坦求和约定意味着我们对重复的索引（在这种情况下是 j 和 k）进行求和。在使用 Levi-Civita 符号和爱因斯坦符号时，叉积可以以紧凑和优雅的方式表达，而不需要明确写出求和。
两个向量 a 和 b 的叉积产生一个新向量，该向量垂直于 a 和 b。结果向量的长度等于 a 和 b 的长度的乘积乘以它们之间的角度的正弦值。
矩阵乘法的爱因斯坦符号（作者创建的方程图）
在矩阵乘法中，重复的索引 k 也是隐式求和的。这意味着要计算乘积矩阵 AB 中位置 (i, j) 的元素，我们需要将 A 的第 i 行的对应元素与 B 的第 j 列相乘，并将这些乘积相加。注意——矩阵乘法是 Transformer 模型中的基本操作——也是本文的灵感来源。
基于上述示例——可以清楚地看到爱因斯坦符号有一些优点。它更简洁，因为它减少了所需的符号和运算的数量。它更清晰，因为它突出了“张量”操作的基本结构，也使它们更直观。高效，因为它可能导致算法的实现更加容易，尤其是在处理高维矩阵时。最后，它更通用，因为它可以应用于广泛的“张量”操作，使其成为表达复杂数学关系的多功能工具。简而言之，利用爱因斯坦符号，研究人员和实践者可以在数学意义上利用这些优势来处理基于“张量”的模型和其他深度学习架构。
方法论
在本节中，我将以标准方式介绍 Transformer 模型（解码器）背后的数学。此外，我将演示如何使用爱因斯坦符号从不同角度表示数学运算。
爱因斯坦符号允许那些不熟悉最先进的机器学习研究和相应的数学符号的读者，以更标准的方式消费 Transformer 的数学基础。
令牌嵌入将输入令牌（单词或子单词，即“令牌”）转换为密集的数值表示，使模型能够处理和理解文本中的语义和句法信息。
令牌嵌入（由作者使用 FLUX1-schnell 创建）
令牌嵌入的爱因斯坦符号（方程图由作者创建）
e: 输入令牌 x 在索引 i 的嵌入向量 E: 嵌入矩阵，其中 i 是令牌的索引，j 是维度 x: 输入令牌 x 在索引 j 的一热编码表示 i: 令牌的索引 j: 嵌入维度中的维度索引
位置编码增加了对单词（或令牌）嵌入，并为模型提供了关于序列中每个单词的相对和绝对位置的信息。在这里，爱因斯坦符号对原始公式没有做任何改变。
位置编码（由作者使用 FLUX1-schnell 创建）
位置编码的爱因斯坦符号（方程图由作者创建）
PE(pos, i): 位置 pos 和维度 i 的位置编码 pos: 令牌在序列中的位置 i: 位置编码的维度 d: 模型维度或嵌入维度
注意力（由作者使用 FLUX1-schnell 创建）
注意力机制通过计算输入嵌入的加权求和来计算每个输入令牌对当前输出令牌的相关性——其中权重由来自每个头（i）的输入和输出令牌的查询（Q）、键（K）和值（V）向量的注意力分数确定。
爱因斯坦符号的注意力机制（方程图由作者创建）
带有 Softmax 函数的注意力机制的 Einstein 符号（由作者创建的方程图）
Q：查询矩阵 K：键矩阵 V：值矩阵 i，j，k，l：用于访问矩阵特定元素的索引 d_k：键向量的维度
前馈网络（FFN）（由作者使用 FLUX1-schnell 创建）
前馈网络（FFN）的重要性有两方面。首先，它通过激活函数引入非线性。在原始的“注意力即所有”论文中——使用了 ReLU 激活函数。如今，我们在当前的仅解码器大型语言模型中看到了更多先进的激活函数。简要回顾一下——非线性允许网络通过反向传播学习输入和输出之间的复杂映射。其次，FFN 作用于注意力层的输出，该层捕捉了长距离依赖关系——因此，FFN 有助于提取有意义的特征。最后，文献还指出，FFN 通过引入额外的层和参数增加了网络的能力。
前馈网络（FFN）使用 Einstein 符号（由作者创建的方程图）
xj：输入向量。W：第一层和第二层的权重矩阵，分别。b：第一层和第二层的偏置向量，分别。i，j，k：用于访问矩阵和向量特定元素的索引。
层归一化（由作者使用 FLUX1-schnell 创建）
层归一化是 LLMs 中注意力机制的重要组件，因此通过（1）稳定训练和（2）增强，在它们的有效性和稳定性中发挥着重要作用。层归一化的主要优势在于梯度裁剪的概念。换句话说，归一化有助于防止训练过程中梯度消失或爆炸问题——因为它将梯度保持在合理的范围内，使训练更加稳定。其次，层归一化将输入向量投影到一个空间，其中注意力查询可以平等地关注所有键。这减少了从注意力机制中学习这种行为的负担。进一步来说，通过将所有向量缩放到相同的范数，层归一化确保没有单个键可以主导注意力过程——从而避免对某些输入产生偏差。
x：形状为[批大小，序列长度，隐藏大小]的输入张量 μ：x 在隐藏大小维度上的均值 σ：x 在隐藏大小维度上的标准差 α，β：可学习参数（缩放和偏移）
结论（由作者使用 FLUX1-schnell 创建）
结论，局限性和未来研究方向
在本文中，我们探讨了爱因斯坦符号在变压器模型内部数学运算中的应用。为此，我们利用了重复索引上的隐式求和。因此，我们展示了涉及注意力、前馈神经网络和层归一化的复杂张量运算的更简洁、更优雅的表示。
虽然爱因斯坦符号为变压器模型提供了一个有价值的视角，但承认其局限性和未来研究的潜在领域同样重要。首先，学习爱因斯坦符号存在一定的曲线。尽管使用它可以简化复杂的表达式，但完全掌握其细微差别需要一定的数学成熟度。其次，从研究和分析的角度来看，爱因斯坦符号传达新视角是有意义的，但将爱因斯坦符号直接转换为高效代码可能具有挑战性，尤其是在大规模模型中。
本研究的未来方向可能包括探索编译器优化和硬件加速技术，以利用爱因斯坦符号的潜在性能优势。此外，结合爱因斯坦符号和传统矩阵符号的混合方法可能是有用的，以在简洁性和可读性之间取得平衡。
最重要的是，理论洞察力的生成可能是一个非常吸引人的未来研究方向，因为它可能导致对变压器模型潜在原理的更深入理解，并可能激发新的架构和优化技术。
显然，通过解决这些局限性和探索未来的研究方向，我们可以充分发挥爱因斯坦符号在推进我们对变压器模型的理解和开发方面的潜力。
参考文献
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin, “Attention Is All You Need,” in Advances in Neural Information Processing Systems 30 (2017)
扫一扫
170
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
