“兄弟包会的”--LLAMA从0到1原理＋代码详解（一）[Embedding、RMSNorm、旋转位置编码]

概述

长时间没更新，忙于学业和实习，接下来将从0开始讲解LLAMA代码，参考了大量文献和菜菜老师的课程，希望大家能多多交流，提出宝贵的意见。

1 Embedding层

解释：

Embedding层是神经网络中的一个重要组成部分，特别是在处理自然语言处理（NLP）任务时。它的主要作用是将离散的输入数据（如单词、用户ID等）转换成固定长度的连续向量表示。这些向量能够捕捉输入数据之间的语义关系，使得相似的输入在向量空间中的距离更近。简单来说，Embedding层就像是一个翻译器，将难以直接处理的离散数据转换成神经网络容易理解的向量形式。

具体来说，Embedding层通过训练学习到一个映射关系，将输入数据（通常是一个整数索引，代表某个词或实体的ID）转换成一个低维的实数向量。这个向量的维度是事先设定的，它决定了向量能够捕捉到的信息丰富程度。在训练过程中，Embedding层的参数（即映射关系）会逐渐优化，使得转换后的向量能够更好地表示输入数据的语义信息。

例如，在NLP任务中，Embedding层可以将每个单词转换成一个词向量。这些词向量能够捕捉单词之间的语义关系，使得相似的单词在向量空间中的距离更近。这样，神经网络就可以利用这些词向量来进行后续的文本分类、情感分析、机器翻译等任务。

注：

为了让计算机能够统一处理长短不一的句子，我们需要对句子中的token数量进行标准化。具体来说，对于较长的句子，我们会进行截断（truncation），即只保留前N个token，以确保句子长度不会超出预设的限制；而对于较短的句子，我们则会进行填充（padding），即在句子末尾添加特殊的填充token，使其长度达到预设的标准。

词嵌入（Embedding）过程是一个带有参数的学习过程。这些参数会根据神经网络的损失函数进行优化，以便生成能够更好地反映单词位置和语义信息的词向量。通过训练，词嵌入层能够学习到单词之间的关联和相似性，从而将相似的单词映射到向量空间中相近的位置。这样，即使句子中的单词顺序或具体表述有所不同，词嵌入层也能够捕捉到它们之间的相似性和关联，为后续的文本处理任务提供有力的支持。

代码：

import torch
import torch.nn as nn

# 定义Embedding层
#num_embeddings: 嵌入表的大小，即词汇表的大小或类别数。它定义了有多少个不同的“离散输入”可以映射到嵌入向量。
#embedding_dim: 每个离散输入（类别、单词等）将被映射到的连续向量的维度大小。
embedding = nn.Embedding(num_embeddings, embedding_dim) 
# 查看权重
print(embedding.weight)

注：

nn.Embedding和nn.Linear的区别，可以参考一下博客：nn.Linear与nn.Embedding区别及源码 – 知乎

2 RMSNorm

**解释：**RMSNorm是一种特殊的归一化方法，它在处理输入特征时，不计算特征的均值，而是专注于计算每个特征的均方根（RMS，Root Mean Square）。简而言之，RMSNorm跳过了均值计算这一步骤，直接利用每个特征的均方根值来对其进行归一化处理。这样，每个特征都会根据其自身的RMS值被缩放到一个适当的范围内，从而有助于提升后续模型处理的稳定性和性能。

**优点：**减少计算量，因为好多数最终变成一个数;保留了归一化效果，有更好的收敛性。

举例：

假设我们有一个形状为 (2, 3, 3) 的张量，即2个batch，3个token，3个d_mode

[
  [
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
  ],
  [
    [10, 11, 12],
    [13, 14, 15],
    [16, 17, 18]
  ]
]

步骤 1: 计算每个特征的平方

首先，我们需要计算每个元素的平方：

[
  [
    [1, 4, 9],
    [16, 25, 36],
    [49, 64, 81]
  ],
  [
    [100, 121, 144],
    [169, 196, 225],
    [256, 289, 324]
  ]
]

步骤 2: 计算每个特征的均值

然后，我们需要计算每个特征的均值。由于我们的张量形状是 (2, 3, 3)，我们可以将每个 3×3 子矩阵视为一个样本，共有 2 个样本，每个样本有 9 个特征。

计算均值：

第一个样本均值：

[ (1+4+9+16+25+36+49+64+81) / 9 ] = [ 30.6667 ]

第二个样本均值：

[ (100+121+144+169+196+225+256+289+324) / 9 ] = [ 202.6667 ]

步骤 3: 计算均值的平方根（RMS）

接下来，计算每个均值的平方根：

第一个样本的 RMS：

sqrt(30.6667) ≈ 5.5385

第二个样本的 RMS：

sqrt(202.6667) ≈ 14.2361

步骤 4: 标准化

最后，我们将每个元素除以其对应的 RMS 值：

第一个样本标准化后：

[
  [1/5.5385, 2/5.5385, 3/5.5385],
  [4/5.5385, 5/5.5385, 6/5.5385],
  [7/5.5385, 8/5.5385, 9/5.5385]
]
≈
[
  [0.1806, 0.3612, 0.5417],
  [0.7225, 0.9031, 1.0837],
  [1.2643, 1.4449, 1.6255]
]

第二个样本标准化后：

[
  [10/14.2361, 11/14.2361, 12/14.2361],
  [13/14.2361, 14/14.2361, 15/14.2361],
  [16/14.2361, 17/14.2361, 18/14.2361]
]
≈
[
  [0.7024, 0.7726, 0.8428],
  [0.9131, 0.9834, 1.0536],
  [1.1239, 1.1941, 1.2643]
]

因此，最终的 RMS 标准化后的张量为：

[
  [
    [0.1806, 0.3612, 0.5417],
    [0.7225, 0.9031, 1.0837],
    [1.2643, 1.4449, 1.6255]
  ],
  [
    [0.7024, 0.7726, 0.8428],
    [0.9131, 0.9834, 1.0536],
    [1.1239, 1.1941, 1.2643]
  ]
]

这就是对一个形状为 (2, 3, 3) 的张量进行 RMSNorm 标准化的具体过程。

class RMSNorm(torch.nn.Module):
    #eps为防止除0
    def __init__(self,dim:int,eps:float):
        super().__init__():
            self.eps = eps
            self.weight = nn.Parameter(torch.ones(dim))
    #_外部不能调用
    def _norm(self,x):
        #-1为在最后一个维度求均值，及为d_model；keepdim为维持住三维的结构
        return x*torch.rsqrt(x.pow(2).mean(-1,keepdim = True)+self.eps)
    def forward(self,x):
        #type_as为还原为原先的数据结构类型
        output = self._norm(x.float()).type_as(x)
        return output*self.weight
        

3 旋转位置编码Rotary Positional Embedding

我们不想了解到位置非常远的信息，但是想要了解到上下文位置的信息，即相对位置。

上图原理流程：

（1）x1 和 x2 是 token 的原始编码值。

（2）θ1（theta1） 是一个常数，为每两维度的编码设置。我们将

θ

1

theta_1

θ1​、

θ

2

theta_2

θ2​…

θ

d

/

2

theta_{d/2}

θd/2​这个序列总称为“频率”。

（3）m 是 position（位置），表示当前 token 在序列中的位置。

（4）通过m * θ 计算角度，并将 x1 和 x2 按照这个角度进行旋转，得到新的编码

x

1

∗

{x1}^*

 x1∗ 和

x

2

∗

{x2}^*

 x2∗。

注：

旋转位置编码与向量夹角

在旋转位置编码中，每个词（或token）都会根据其在句子中的位置被赋予一个特定的旋转角度。这种旋转是在向量空间中进行的，可以想象成每个词向量都像是一个小指针，在空间中按照其位置旋转到不同的方向。

• 当两个词在句子中紧密相连时，它们对应的向量在旋转后的夹角会很小，通常是一个锐角。这意味着这两个向量在方向上非常接近，因此它们在注意力机制中的相关性会很高。
• 相反，当两个词距离很远时，它们对应的向量在旋转后的夹角会更大。这反映了在自然语言中，远距离的词语通常关联度较低的现象。

注意力机制与点积计算

在注意力机制中，我们通常通过计算两个词向量之间的点积来评估它们的相关性。点积的大小与两个向量的夹角和长度有关。

• 当两个向量的夹角是锐角时，点积值较大，表示它们之间有较强的相关性。
• 当夹角接近直角或更大时，点积值会减小，表示相关性减弱。

处理远距离词语的关联度问题

然而，有一个潜在的问题：如果两个远距离的词在旋转后的向量空间中形成了钝角或甚至反向的锐角，按照点积的计算方式，它们之间可能会表现出不应有的相关性。

为了解决这个问题，RoPE在设计时引入了一个巧妙的机制：随着位置值的增加，给词向量分配的频率逐渐减小。这意味着长距离的词在旋转时会受到一个额外的“衰减”效应。

• 即使两个远距离的词在旋转后的向量空间中形成了看似接近或反向的角度，由于它们对应的频率不同（特别是长距离词的频率较低），它们之间的点积值仍然会因为这种频率差异而减小。
• 这种衰减效应确保了即使两个远距离的词在向量空间中形成了不利的角度，它们也不会在注意力机制中获得过高的权重。

程序：

#定义频率计算,输入维度（1024，256，100000）
def precompute_pos_cis(dim: int, max_position: int, theta: float = 10000.0):
    #频率
    #[: (dim // 2)]为了防止维度是奇数，可将奇数除掉
    #/dim除以原始维度进行归一化，将所有的数据压缩到[0,1]之间；并且，允许列编号越大的特征有更大的值
    #torch.Size([512])
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    #位置编码m，device=freqs.device的意思为将m放到与fres相同的设备上计算，torch.Size([256])
    m = torch.arange(max_position, device=freqs.device)
    #频率乘以位置编码、外积，torch.Size([256, 512])
    freqs = torch.outer(m, freqs).float()
    #torch.Size([256, 512])
    pos_cis = torch.polar(torch.ones_like(freqs), freqs)
    return pos_cis

4 将旋转位置编码应用于q,k矩阵

解释：

在LLaMA模型中，旋转位置编码（RoPE）被特意设计来仅作用于**Query（查询）和Key（键）**矩阵，这是因为它的主要目的是在Transformer的自注意力机制中融入序列元素之间的相对位置信息。Query和Key矩阵在注意力权重的计算过程中起着关键作用，它们之间的点积决定了每个元素对其他元素的关注程度。通过将位置信息编码到Query和Key向量中，RoPE能够帮助模型更好地理解序列中元素之间的位置关系。

相比之下，Value（值）矩阵在自注意力机制中的作用是生成最终的输出表示，它并不直接参与注意力权重的计算。因此，将位置信息编码到Value向量中并不是必需的，也不会对注意力权重的计算产生直接影响。

这种设计选择不仅符合RoPE的初衷，即提供一种有效的方式来捕捉序列中的位置信息，而且也与Transformer模型的工作原理相契合，确保了模型能够高效地处理序列数据并捕捉到其中的关键信息。对于初学者来说，可以这样理解：RoPE就像是为Query和Key向量加上了一个“位置标签”，让它们在进行注意力计算时能够“知道”自己和其他元素之间的相对位置，而Value向量则专注于提供最终的输出信息，无需这样的“位置标签”。

代码：

# 将频率（位置编码）应用于查询矩阵 q 和键矩阵 k
def apply_rotary_emb(xq, xk, pos_cis):
    # 内部函数：用于调整 pos_cis（位置编码）的形状，使其与输入张量 x 的形状匹配
    def unite_shape(pos_cis, x):
        # 注意这里输入的x是已经转变为复数的Q和K矩阵
        # 复数Q、K矩阵的维度与实数Q、K矩阵的维度有区别
        # 例如，当实数Q矩阵的结构为 (10,128,512) 时
        # 复数Q矩阵的结构为(10,128,256,2)，其中后两位代表复数的实部和虚部
        # 此时如果对Q矩阵取最后一维索引，会得到最后一个实部，也就是256
        # 获取输入张量的维度数量，x->(batch,seq_len,d_model)
        #ndim是有几个维度就是几，例如x.shape->(2,3,4),ndim=3
        ndim = x.ndim
        # 确保输入张量的维度数是有效的
        assert 0  1  ndim
        # 确保 pos_cis 的形状与输入 x 的形状中的seq_len, d_model维度匹配
        assert pos_cis.shape == (x.shape[1], x.shape[-1])
        # 构造新的形状，除了第二维度和最后一维度之外，其他维度都设置为 1
        # 这是为了广播 pos_cis 以匹配输入 x 的形状
        shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
        # 调整 pos_cis 的形状为新的 shape
        return pos_cis.view(*shape)
    # 将查询张量 xq 的最后一个维度视为复数的一部分，形状变为 (*xq.shape[:-1], -1, 2)
    # 这意味着将最后一维度按 2 拆分，转换为复数表示（因为一个复数由实部和虚部组成）
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    # 对键张量 xk 做同样的处理，将其转换为复数形式
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    # 调整 pos_cis 的形状，使其与查询矩阵 xq_ 的形状匹配
    pos_cis = unite_shape(pos_cis, xq_)
    # 将旋转嵌入应用到查询矩阵，复数乘法会影响幅度和相位
    # 然后将复数结果转换回实数形式并将其展平（恢复到原来的维度结构，即第三个维度和第四的维度合并）
    xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
    # 对键矩阵做同样的操作，应用旋转嵌入
    xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
    # 返回处理后的查询矩阵和键矩阵，且类型与输入张量相同
    return xq_out.type_as(xq), xk_out.type_as(xk)

注：

(1)为什么不直接使用正余弦位置编码？旋转位置编码有什么好处？

在LLaMA中，不直接使用传统的正余弦位置编码（如BERT模型所采用的位置编码方式），而是选择了旋转位置编码（RoPE, Rotary Position Embedding），这背后有着多方面的考虑。旋转位置编码相比正余弦位置编码具有一些显著的优势，这些优势使得RoPE更适合LLaMA这类大型语言模型。

正余弦位置编码的局限性

正余弦位置编码虽然简单且易于实现，但它存在一些局限性，尤其是在处理长序列时：

• 固定长度限制：正余弦位置编码通常是在训练时预先计算好的，其长度受限于训练时序列的最大长度。这限制了模型在处理比训练时更长的序列时的泛化能力。
• 线性衰减问题：正余弦位置编码在表示位置信息时，随着位置的增加，其编码的区分度会逐渐降低，这可能导致模型在处理长序列时无法准确捕捉到远距离元素之间的依赖关系。

旋转位置编码的优势

旋转位置编码针对正余弦位置编码的局限性进行了改进，其优势主要体现在以下几个方面：

• 外推性：RoPE能够处理任意长度的序列，而不需要重新计算位置编码。这使得模型在处理比训练时更长的序列时仍然能够保持较好的性能。
• 相对位置编码：RoPE通过旋转变换将位置信息编码到Query和Key向量中，使得模型能够捕捉到序列中元素的相对位置关系。这对于语言建模等任务来说是非常重要的，因为语言的很多特性都依赖于元素之间的相对位置。
• 计算效率高：RoPE的计算过程相对简单且高效，它不需要像正余弦位置编码那样为每个位置单独计算编码向量。这使得模型在处理长序列时能够更快地计算出注意力权重。

LLaMA选择RoPE的原因

LLaMA模型选择RoPE作为其位置编码方式，主要是出于以下几个方面的考虑：

• 提高模型性能：RoPE能够帮助模型更准确地捕捉到序列中元素的相对位置关系，从而提高模型在处理语言任务时的性能。
• 支持长序列处理：RoPE的外推性使得LLaMA能够处理任意长度的序列，而不需要担心位置编码的长度限制问题。
• 提升计算效率：RoPE的高效计算过程有助于加快LLaMA的训练和推理速度，使其能够更好地应对大规模语言处理任务。

(2)复数维度变化。

复数的维度变化跟实数的维度变化不太相同，大家需要注意。

文章来源于互联网:“兄弟包会的”–LLAMA从0到1原理＋代码详解（一）[Embedding、RMSNorm、旋转位置编码]

相关推荐: 【AI绘画】Stable Diffusion写真完整教程

前言 最近自己对AI非常痴迷，并且今后也会一直在这个领域深耕，所以就想着先入门，因此花时间研究了一番，还好，出了点小成果，接下来给大家汇报一下。 AI绘画 提到AI绘画，大家可能立马会想到midjourney，它的威力我就不多说了，确实很强。但是使用门槛略高（…