llama源码学习·model.py[3]ROPE旋转位置编码(2)旋转角度生成代码

一、源码注释

def precompute_freqs_cis(dim: int, end: int, theta: float = 1000.0):
    '''预先计算频率和复数的cosine和sine值，用于后续的Positional Encoding
    dim: 维度
    end: 一个序列的最大长度或位置的最大值
    theta: 用于计算频率的超参数，默认值为1000.0
    '''
    # 生成一个等比数列，即频率（frequencies），这种方法是基于 "Attention is All You Need" 论文中描述的位置编码
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    
    # 生成了一个从0到end的序列
    t = torch.arange(end, device=freqs.device)
    
    # 计算两个向量的外积
    # 结果矩阵的形状是（end, dim//2）
    # 这里的freqs 其实是旋转角度 theta
    freqs = torch.outer(t, freqs).float()
    
    # 将极坐标转换为复数形式
    # torch.polar(r, theta): 是一个函数，它接受两个参数：模 r 和相位 theta，然后返回一个复数，
    #                       该复数的实部为 r * cos(theta)，虚部为 r * sin(theta)。
    # torch.ones_like(freqs): 生成一个与 freqs 形状相同的张量，但所有元素都是1，这意味着模r为1。
    # freqs: 它表示每个位置的相位或角度。
    # freqs_cis: 是一个形状为（end, dim//2）的复数矩阵，每个元素都是一个复数，用于后续的位置编码。
    
    # 这行代码实际上为每个位置和每个频率生成了一个复数，其模为1，而相位为我们之前计算的频率。
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

二、源码与公式的对应

第一步：旋转嵌入生成

需要给定一个位置索引 $pos$ 和频率向量 $freq$， 来计算旋转角度 $\theta =pos\times freq$

freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

生成的这个等比数列就是频率向量，这是基于 “Attention is All You Need” 论文中描述的位置编码来实现的

 t = torch.arange(end, device=freqs.device)

这个长度为 $end$ 的数列是位置索引 $pos$

freqs = torch.outer(t, freqs).float()

这一行是在计算两个位置索引 $pos$ 和 频率向量 $freq$ 的外积生成旋转角度 $\theta$ ,不过旋转角度的信息在代码中依旧存储在 $freqs$ 这个变量中

freqs_cis = torch.polar(torch.ones_like(freqs), freqs)

freqs 是旋转角度向量，freqs_cis 使用复数表示的旋转矩阵

三、举例说明

1.假设函数参数

dim = 4 位置编码的维度是 4

end = 3 序列的最大长度是 3

2.生成频率向量 freq

freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) 

假设 x = torch.arange(0, dim, 2) 也就是从 0 ~ dim 步长为二的等比数列

$freqs=\frac{1}{{\theta }^{\frac{x}{dim//2}}}=\frac{1}{{\theta }^{\frac{[0,2]}{4//2}}}=\frac{1}{{\theta }^{[0,0.5]}}=\frac{1}{[1,\sqrt{\theta }]}=[1,\frac{1}{\sqrt{\theta }}]$

3.生成从 0 到 end 的位置索引

t = torch.arange(end, device=freqs.device) 

$t=[0,1,2]$

4.计算两个向量的外积得到旋转角度 theta

freqs = torch.outer(t, freqs).float()

5.将极坐标转换为复数形式

freqs_cis = torch.polar(torch.ones_like(freqs), freqs)

这里返回的 freqs_cis 是一个用复数表示的旋转矩阵