一、Transformer vs LLaMa

transformer主要分为解码器和编码器两部分。相较之下，LLaMA仅使用了Transformer的解码器部分你，采用了一个仅解码器的结构。在结构上，与transformer模型相比，llama2的主要变化是将其中的layerNorm替换为了均方根标准化（RMSNorm），多头注意力换成了分组查询注意力（GQA，在llama中则是多查询注意力MQA），并将位置编码替换为了旋转编码（RoPE）。

1、RMS Normalization均方根标准化

llama2为了提高训练的稳定性，对每个transformer层的输入进行归一化，而不是对输出进行归一化，同时使用RMS Norm归一化函数。RMSNorm的代码如下：

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, norm_size, limit, is_learn_p=False):
        super(RMSNorm, self).__init__()
        self.limit = limit
        self.norm_size = (norm_size,) if isinstance(norm_size, int) else tuple(norm_size)
        self.is_learn_p = is_learn_p
        if is_learn_p:
            self.gamma = nn.Parameter(torch.ones(size=self.norm_size))
        else:
            self.gamma = None
    def forward(self, x):
        dim = (-2, -1) if len(self.norm_size) == 1 else -1
        squared = torch.square(x)
        mean_squared = torch.mean(squared, dim=dim, keepdims=True)
        rms = torch.sqrt(mean_squared + self.limit)

        x_norm = x / rms
        if self.is_learn_p:
            x_norm = x_norm * self.gamma
        return x_norm

LayerNorm代码如下：

import torch
from torch import nn

class LayerNorm(nn.Module):
    def __init__(self, norm_size, limit, is_learn_p=False):
        super(LayerNorm, self).__init__()
        self.is_learn_p = is_learn_p
        self.norm_size = (norm_size,) if isinstance(norm_size,int) else tuple(norm_size)
        self.limit = limit

        if is_learn_p:
            self.gamma = nn.Parameter(torch.ones(size=self.norm_size))
            self.beta = nn.Parameter(torch.zeros(size=self.norm_size))
        else:
            self.gamma = None
            self.beta = None

    def forward(self,x):
        dim = (-2,-1) if len(self.norm_size) == 1 else -1

        mean = torch.mean(x, dim=dim, keep_dims=True)
        var = torch.var(x, dim=dim, keep_dims=True)

        x_norm = (x - mean) / torch.sqrt((var + self.limit))

        if self.is_learn_p:
            x_norm = x_norm * self.gamma + self.beta

        return x_norm

LayerNorm通过减均值、除标准差的方式，试图将每一层的输入分布固定下来。这个设计在早期非常有效，但随着模型变得更深、更复杂，尤其是当Transformer成为主流后，研究人员发现“减均值”这一步在某些情况下会引入新的不稳定性。
最核心的问题出现在与relu这类激活函数的配合上。relu会将所有负输入置零。LayerNorm的减均值操作，会强制将数据分布的中心拉到零点，这实际上人为地制造了更多负值，显著加剧了“神经元死亡”问题。这意味着更多的神经元输出为0，导致网络容量下降，并阻碍了梯度的有效回传，从而破坏了训练的稳定性。
而RMSNorm选择只进行缩放，即除以均方根值，不再减去均值。这样做的好处是，它完美地保留了输入数据各维度之间的相对比例关系——这些相对关系才是模型真正需要学习的模式。同时，它避免了因中心化而制造负值的问题，使得梯度流在网络中，尤其是在使用relu或类似激活函数的深层网络中，能够传递更加顺畅和稳定。
此外，从计算本质上看，缩放是归一化中控制数据幅度、稳定梯度的最关键环节。减去均值改变的是数据的位置，但对于稳定训练而言，控制幅度的优先级更高。RMSNorm抓住了这个主要矛盾，其更简洁的计算路径（只计算RMS，而非先算均值再算标准差）也带来了内在的数值稳定性优势。
因此，RMSNorm的稳定性优势并非否定LayerNorm的思想，而是对其进行的精准优化。它通过放弃在新架构下可能带来副作用的中心化步骤，以一种更简洁、更专注的方式，为深层模型提供了更稳定的训练基础。这也正是它在LLaMA等现代大模型中取得成功的原因。

2、分组查询注意力

在各种多头注意力机制比较重，原始的多头注意力（MHA）使得QKV三部分具有相等数量的“头”，并且它们之间是一一对应的。每一次计算注意力时，各个头部的QKV独立执行自己的计算，最后将所有头部的结果加在一起作为输出。
相对于MHA，多查询注意力（MQA）保持了原来的Query头数，但只为K和V各设置了一个头，即所有的Query头部都共享同一个K和V组合，因此得名“多查询”。据实验发现，这种机制通常可以提高30%-40%的吞吐量，对性能的影响较小。
分组查询注意力（GQA）综合了MHA和MQA，既避免了过多的性能损失，又能利用MQA的推理加速。在GQA中，把查询头分为G组，每个组内部的头部共享一个相同的K和V组合。当G设为1，即GQA-1，则所有Query共享同一组K和V，这时的GQA等效于MQA；而当G等于头的数量，即GQA-H，那么这时的GQA等效于MHA。

3、SwiGLU Activation Function

（1）激活函数的必要性：非线性能力，输出约束特性

观察下图具有单个隐藏层的MLP，并忽略激活函数列出$z^{(}2)$的表达式。

$$z^{(2)}=(x{W}^{(1)}+b_1)W^{(2)}+b_2=x{W}^{(1)}{W}^{(2)}+b_1{W}^{(2)}+b_2$$
在两层的神经网络中，如果不考虑激活函数，整个网络可以简化为一个线性的放射变换，即权重矩阵的乘积作用于输入x加上一个偏置项。因此，去除激活函数后，两层MLP就失去了非线性映射的能力，成为了一个线性模型。因此，是非线性激活函数的存在使得MLP能够表达更复杂的函数关系，去掉这些激活函数将使其无法解决非线性问题。给隐藏层的输出加上Sigmoid激活：
$$z^{(2)}=sigmoid(x{W}^{(1)}+b_1)W^{(2)}+b_2=h^{(1)}{W}^{(2)}+b_2$$
h^{(1)}作为一个经过非线性处理的隐藏状态，在BP学习过程中就能拟合一些复杂的非线性函数。除此之外，如果想用这个两层MLP完成一个二分类任务该怎么办？该模型的输出层有两个神经元，那么有监督的label可以用(1,0),(0,1)表示两个类别。那么目前来看输出的是输入x经过2次线性变换，1次非线性激活，加上偏置后的一个长度为2的状态向量，但不是一个概率分布向量，因此在最后加上softmax函数：
$$z^{(2)}=softmax(h^{(1)}{W}^{(2)}+b_2)$$
此时输出结果就是一个长度为2的概率分布向量，两个值分别代表预测为该类的可能性。
因此，激活函数在神经网络的核心作用：

• 非线性拟合能力：弥补线性神经元叠加还是线性的问题；
• 约束输出值的范围和特性，充当输出层分类器。

（2）常用激活函数

总体来说激活函数在以下指标中发展：

• 非线性连续函数：可以容忍某一点无导数，但不能容忍全部导数为0
• 正负对称性：输入是负的输出也得是负的，如果输出恒为正，那么后续的神经层就没有负输入了
• 导数不可为0：一旦归零就无法恢复，容易造成神经元死亡（ReLU）
• 计算开销：线性计算快（Leaky ReLU），指数计算慢（Tanh\ELU）
• 正负输出分别对待：
  • 神经元输出为正：几乎不处理，类似线性激活f(x)，均匀，不会出现接近某一区域梯度爆炸、消失（sigmoid）
  • 神经元输出为负：在负数区域输出近似为0，有助于实现神经网络的稀疏性。只有少数神经元对特定的输入产生响应。这种稀疏性有助于减少计算量，提高计算效率，并且在一定程度上能够防止过拟合。从生物学角度看，神经元在接收到足够的刺激时才会被激活并产生输出信号。这种激活过程与ReLU等函数在输入为正时产生输出、在输入为负时保持非激活状态的行为相吻合。

（2.1）阶跃Sgn

最符合生物神经元的0-1状态，但阶跃函数不连续、不平滑，不是一个好的函数。

（2.2）Sigmoid

是阶跃的替代，也是0-1之间，不过由二院离散变为光滑连续，且中间陡峭，即大部分值都会被映射到更靠近0-1的位置。
$$Sigmoid(x)=\sigma (x)=\frac{1}{1+e^{-x}}$$
缺点：

• 在靠近0和1的部分曲线平缓，梯度趋近于0，在BP过程容易导致梯度消失权重不更新；
• 非中心对称，即上下不对称，x<0时值还是大于0，导致后一层接收的输入存在恒为正的偏置；
• exp()指数函数与其他非线性激活函数相比，计算成本高昂，计算机运行起来速度较慢。
  

（2.3）Tanh

双曲正切，解决了sigmoid中当x<0时y依旧为正数导致的其后一层的神经元的输入发生偏置偏移（bias shift），但没解决梯度消失的问题。
$$Tanh(x)=tanh(x)=\frac{e^x-e^{-x}}{e^x+e^{-x}}$$

（2.4）softmax

主要用于归一化，但也能用于激活，因为都是映射到0~1。

（2.5）ReLU

整流线性单元，解决了sigmoid的梯度消失和输出偏置，线性函数计算快，但由于x<0时y为0，会造成神经元死亡。
$$ReLU(x)=max(0,x)$$

（2.6）Leaky ReLU

解决了ReLU的负数部分神经元死亡问题，保留了线性计算快的特点。
$$LeakyReLU(x)=max(0,x)+negativ{e}_{slope}\ast min(0,x)$$

（2.7）ELU

指数线性单元，解决了ReLU负数部分的问题，继承其所有优点，但负数部分计算量过大。

（2.8）GELU

高斯误差线性单元，下式的$\Phi (x)$是高斯分布函数
$$GELU(x)=x\ast \Phi (x)$$

（2.9）GLU

门控线性单元，两个线性变换的分量积，其中一个线性变换由sigmoid激活。不是传统的激活，所以无法画图。矩阵克罗内克积，两个矩阵shape任意，mxn矩阵和pxq矩阵积是一个mpxnq的分块矩阵。
$$GLU(x)=Sigmoid(W_1+b)\otimes (Vx+c)$$

（2.10）Swish/SiLU

Sigmoid Gated Linear Unit，sigmoid和ReLU的改进版。比ReLU在0附近提供更平滑转换，$\beta$为可学习参数。
$$Swish(x)=x\ast sigmoid(\beta x)$$

（2.11）SwiGLU

结合Swish和GLU两者特点，是一个GLU，但不是将sigmoid作为激活函数，而是使用$\beta =1$的Swish，也无法画图。
$$SwiGLU(x)=Swish(W_{1}x+b)\otimes (Vx+c)$$

二、大模型数据清洗过程

1、文本清洗主要目标

（1）数据质量改进

文本数据通常包含错误、不一致和不相关的内容。清理有助于确保数据准确、可靠和一致。

（2）降噪

文本数据中的噪声可能包括特殊字符、HTML标签、标点符号和其他对分析或建模目标无益的元素。清洁可以消除或减少这种噪音。

（3）标准化

文本清洗通常包括标准化文本，例如将所有文本转换为小写，以确保一致性并防止与案例相关的问题影响分析或建模。

（4）停用词删除

停用词是注入“the”、“and”或“in”之类的常见单词，在文本清洗过程中经常被删除，因为它们对许多任务来说没有重要意义。

（5）词干提取和词形还原

将单词简化为词根形式，有助于对相似的单词进行分组。词干提取和词形还原对于文本分析任务特别有用，其中单词变体应被视为一个单词。

（6）处理缺失数据

文本数据可能包含缺失值或不完整的句子。文本清洗可能涉及填充缺失数据或解决不完整文本的策略。

（7）重复数据删除

删除重复或接近重复的文本条目对于确保数据完整性并防止分析或建模中的偏差至关重要。

（8）处理嘈杂的文本

嘈杂的文本数据可能包括拼写错误、缩写或非标准语言用法。文本清洗策略有助于减轻此类噪音的影响。

三、语言模型的困惑度PPL

PPL(perplexity）：PPL是语言模型对测试样本中每个token预测概率的几何平均数的导数。困惑度的核心思想是：衡量一个语言模型在看到一个词序列后，对其“下一个词”的预测能力有多“困惑”
$$PPL=exp(-\frac{1}{N}\sum _{i=1}^{N}logP(w_i|w_1,w_2,...,w_{i-1}))$$

其中：

• N是测试样本中的总token数
• $P(w_i|w_1,w_2,...,w_{i-1})$是模型预测下一个词$w_i$的概率
  等价：
  $$PPL=({\prod _{i=1}^N\frac{1}{P(w_i)})}^{\frac{1}{N}}$$

1、PPL的性质

• 对于高质量数据集，若模型预测的PPL越小，则模型性能越好。对于低质量数据集，则未必。
• 若模型能够完美预测每个token，则PPL为1。
• 若模型预测样本中每个token的概率无限接近0，则PPL为无穷大。
• 若模型预测的为均匀分布，则PPL等于词表的大小。
• 同一模型和参数配置下，文本长度对PPL的影响。理论上，PPL是归一化的指标，通常以每个token为单位计算，因此文本长度本身不应直接影响PPL的大小。但在实践中，文本长度可能间接影响PPL，原因包括：
  • 文本内容的多样性：较长的文本可能包含更多复杂的句式和稀有词汇，增加模型预测的难度，从而提高PPL。
  • 上下文的利用：对于一些模型，较长的上下文可能提供更多信息，有助于模型更准确地预测下一个词，可能降低PPL。
• 受词汇表大小和分词方式影响：较大的词汇表可能增加模型的PPL，因为稀有词的预测更具挑战性。

2、PPL的作用

• 评估语言模型的能力
• 评估训练数据的质量
• 构造评估数据集的其他指标，比如IFD
• 监测PPL的变化趋势可以为模型的进一步改进提供线索。若PPL明显升高，可能需要检查训练数据或训练策略。
  

3、用base模型还是instruct模型计算

计算PPL时，通常使用base模型更为合适，原因如下：
（1）PPL的定义和用途：PPL主要用于评估语言模型在预测下一个词上的性能，是衡量模型对语言数据的概率分布拟合程度的指标。base模型经过大量无监督文本数据的训练，旨在学习语言结构和模式，因此更适合PPL的计算。
（2）instruct模型的优化目标不同：instruct模型是在base模型的基础上，通过指令微调得到的，目的是使模型更好地理解和执行人类指令。这种微调过程可能改变模型的概率分布和生成行为，使其在PPL指标上不再可比。
（3）评估的一致性和可比性：为了确保评估结果的公平和一致，使用base模型计算PPL可以更准确地反应模型在纯语言建模任务上的能力。由于instruct模型的特殊训练目标，其PPL可能受到微调数据和方式的影响，导致评估结果偏离实际语言建模性能。

4、困惑度和交叉熵的关系

在语言模型中，交叉熵用于衡量模型预测分布与真实数据分布之间的差异：
$$H=-\frac{1}{N}\sum _{i=1}^{N}logP(w_i|w_1,w_2,...,w_{i-1})$$

困惑度是交叉熵的指数形式：

• $PPL=e^H$
• $PPL=2^H$

交叉熵衡量的是模型预测分布和真实数据分布之间的不匹配程度。交叉熵值越小，表示模型预测越准确。
困惑度表示模型在预测下一个词时的平均不确定度。它是交叉熵的指数，因此也反映了模型的预测性能。困惑度越低，表示模型对下一个词的预测越有把握。

PPL、交叉熵、KL散度

（1）交叉熵

用于比较模型预测分布Q和真实分布P的差异，通过最小化交叉熵，使模型预测更接近真实标签。
公式：

代码：

import numpy as np
def cross_entropy(P, Q, eps=1e-12):
    Q = np.clip(Q, eps, 1-eps)
    return -np.sum(P * np.log(Q))

（2）困惑度

量化模型对测试数据的预测不确定性。
公式：

代码：

def perplexity(cross_entropy_loss):
    return np.exp(cross_entropy_loss)

（3）KL散度（相对熵）

衡量两个分布P（真实）和Q（预测）之间的非对称差异。
公式：

代码：

def kl_divergence(P, Q, eps=1e-12):
    P = np.clip(P, eps, 1-eps)
    Q = np.clip(Q, eps, 1-eps)
    return np.sum(P * log(P/Q))

四、位置编码

1、为什么需要位置编码

transformer使用位置编码的原因是它不具备位置信息。
若没有位置编码，那么“床前明月”、“前床明月”、“前明床月”这几个输入，会预测出完全一样的文本。
即不管输入的prompt顺序是什么，只要prompt的文本是相同的，那么模型decode的文本就只取决于prompt的最后一个token。

2、绝对位置编码

在每个输入序列的元素上添加一个位置向量，以表示该元素在序列中的具体位置。这个位置向量通常通过固定的函数生成，与输入数据无关。通常使用的是正弦和余弦函数，这样生成的编码具有很强的周期性，能够捕捉序列中的相对位置信息。
具体来说，对于序列中的第pos个位置，绝对位置编码向量的第i个维度的值定义如下：
$$PE(pos,2i)=sin(\frac{pos}{10000^{\frac{2i}{d}}})$$
$$PE(pos,2i+1)=cos(\frac{pos}{10000^{\frac{2i}{d}}})$$
优点：

• 正余弦函数的范围是[-1,1]，导出的位置编码与原词嵌入相加，不会使得结果偏离过远而破坏原有单词的语义信息。
• 依据三角函数的基本性质，可以得知第pos+k个位置的编码是第pos个位置的编码的线性组合，这意味着绝对位置编码中蕴含着单词之间的距离信息。

缺点：在处理变长序列或长距离依赖时，绝对位置编码可能无法充分表达复杂的位置信息。

代码：

class PositionalEncoding(nn.Module):
    def __init__(self,d_model,max_seq_len=80):
        super().__init__()
        self.d_model = d_model
        pe = torch.zeros(max_seq_len,d_model)
        for pos in range(max_seq_len):
            for i in range(0,d_model,2):
                pe[pos, i] = math.sin(pos / 10000 ** (2 * i / d_model))
                pe[pos, i + 1] = math.cos(pos / 10000 ** (2 * i / d_model))
        pe = pe.unsqueeze(0)
        self.register_buffer('pe',pe)
    def forward(self,x):
        x = x * math.sqrt(self.d_model)
        seq_len = x.size(1)
        x = x + Variable(self.pe[:,:seq_len],requires_grad=False).cuda()
        return x

3、相对位置编码

相对位置编码并不直接位每个位置分配一个唯一的编码，而是关注序列中各元素之间的相对位置。相对位置编码的核心思想是通过计算序列中元素之间的距离，来表示它们之间的相对关系。这种方法尤其适合处理需要捕捉长距离依赖关系的任务，因为它能够更加灵活地表示序列中的结构信息。
相对位置编码最常用的方法之一是将位置差值与注意力权重相结合，即在计算自注意力时，不仅考虑内容，还考虑位置差异。这样，模型能够根据元素之间的距离调整它们之间的交互强度。

对于一个序列长度为L的输入，我们定义位置i和j的相对位置为$r_{i,j}=j-i$。在多维空间中，我们需要为每个维度都添加相对位置编码。
假设序列的维度为d，位置编码的维度为$d_{model}$，可以将每个相对位置表示为$r_{i,j}\in {\mathbb{R}}^{d_{model}}$，并且每个维度都单独考虑相对位置。
在transformer模型中，Attention Scores是通过查询向量Q和键向量K的点积计算的。在使用相对位置编码时，将公式改写为：
$AttentionScor{e}_{i,j}=(Q_i+b_i)\cdot (K_j+a_j+r_{i,j})$
相对位置编码引入了额外的偏置项和位置信息，通过这种方式捕捉到序列中的长距离依赖关系。
代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class RelativePositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len):
        """
        相对位置编码初始化
        :param d_model: 模型向量维度，比如512,768
        :param max_len: 支持的最大句子长度
        """
        super().__init__()
        self.d_model = d_model
        self.max_len = max_len
        
        # =================== 第一步：生成相对位置矩阵 ===================
        position_ids = torch.arange(max_len)
        # 计算两两之间的相对距离：行 - 列 = 第i个词相对于第j个词的距离
        relative_position = position_ids[None, :] - position_ids[:, None]
        # 把距离全部变成非负数（方便查表，避免负数索引）
        relative_position += max_len - 1
        # 注册成模型的缓冲区（不参与训练，固定不变）
        self.register_buffer("relative_position", relative_position)
        
        # =================== 第二步：创建位置向量表 =================== 
        # 表的大小：[所有可能的距离数，向量维度]
        num_relative_positions = 2 * max_len - 1
        self.embedding = nn.Embedding(num_relative_positions, d_model)
    
    def forward(self, seq_len):
        """
        前向传播：给指定长度的句子生成相对位置编码
        :param seq_len: 当前句子长度
        :return: 相对位置编码，形状[seq_len, seq_len, d_model]
        """
        # 从最大矩阵中截取当前需要的长度的相对位置
        pos_ids = self.relative_position[:seq_len, seq_len]
        # 查表，把位置数字 -> 位置向量
        pos_emb = self.embedding(pos_ids)
        return pos_emb

4、旋转位置编码RoPE

旋转位置编码具有更好的外推性。外推性是指大模型在训练时和预测时的输入长度不一致，导致模型的泛化能力下降的问题。例如，如果一个模型在训练时只使用了512个token的文本，在预测时如果输入超过512个token，模型可能无法正确处理，这就限制了模型在处理长文本和多轮对话等任务的效果。
rope的设计思路可以这么理解，我们通常会通过向量q和k的内积来计算注意力分数，如果能够对q、k向量注入位置信息，然后用更新后的q、k向量做内积就会引入位置信息了。
rope就是通过复数乘法实现向量旋转，把相对位置信息编码进query和key之中。它把每个token的Q、K向量在最后一维两两分组，比如Q向量一般现状是[batch_size, seq_len, num_heads, head_dim]，假如现在Q向量的形状为[1, 4, 12, 64]，rope只动最后的维度head_dim，即64，那么两两分组后，形状就变成了[1,4, 12, 32, 2]，之所以是两两分组，是因为复数是实部和虚部两部分组成的，因此分组后，一个元素作为实部，一个元素作为虚部。
每组二维向量会被视为一个复数，然后和一个旋转因子相乘。这个旋转因子由复数表示，形式为$e^{im\theta }$，其中m是当前token的绝对位置，${\theta }_i$是随维度变化的角度，计算公式为${\theta }_i=1/10000^{2i/d}$。不同维度的旋转频率不同，低维度旋转快，高维度旋转极慢，保证了在长序列下位置编码不会重复。
从几何意义上，复数乘法等价于二维向量旋转。对位置m的Q、K分别旋转后，再计算它们的注意力分数，会得到一个关键数学性质：旋转后的QK内积，只与token之间的相对距离m-n有关，与绝对位置m、n无关。
正因如此，rope具备优秀的长度外推能力。模型学习的是相对距离关系，而不是固定的绝对位置下标，即便推理时序列长度超过训练时最大值，旋转角度依然可以按照公式平滑延续下去，不会出现信息断层。


有了这一形式后，具体实现有两种方式：

• 转到复数域，对两个向量进行旋转，再转回实数域；
• 由于上述矩阵 Rn 具有稀疏性，因此可以使用逐位相乘 ⊗ 操作进一步加快计算速度，直接在实数域通过向量和正余弦函数的乘法进行运算，也就是下面这个公式：
  
  代码：

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
	# 计算旋转角度基数(1 / θ ^ (2i / d))
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成位置索引[0,1,...,end-1]
    t = torch.arange(end, device=freqs.device) 
    # 外积计算所有位置和维度的角度矩阵[end, dim//2]
    freqs = torch.outer(t, freqs).float()
    # 转换为复数形式：cis(θ)=cos(θ) + i ·sin(θ)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis


def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    # 确保对齐序列位置维（1）和特征维（-1）
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    # 构建广播友好形状：[1, seq_len, 1, ..., 1, feat_dim]
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)


def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
	# 将输入重塑为复数形式（最后一维拆分为实部/虚部）
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    # 调整旋转因子形状用于广播
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    # 复数乘法实现旋转：(xq_ * freqs_cis) = |xq|·e^{i(φ + θ)} 
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

5、相对位置编码AliBi

收到T5 bias的启发，Press等人提出了ALiBi算法，是一种预定义的相对位置编码。与传统方法不同，ALiBi不向单词embedding中添加位置embedding，而是根据token之间的距离给attention score加上一个预设好的偏置矩阵，比如和相对位置差1的就加上一个-1的偏置，两个token距离越远这个负数就越大，代表他们的相互贡献越低。由于注意力机制一般会有多个head，这里针对每一个head会乘上一个预设好的斜率项（slope）。


ALiBi对最近性具有归纳偏差，它对远程查询-键对之间的注意力分数进行惩罚，随着键和查询之间的距离增加，惩罚增加。不同的注意力头以不同的速率增加其惩罚，这取决于斜率幅度。

代码：

import math
import torch
from torch import nn

def get_slopes(n_heads: int):
    n = 2 ** math.floor(math.log2(n_heads))
    m_0 = 2.0 ** (-8.0 / n)
    m = torch.pow(m_0, torch.arange(1, 1 + n))

    if n < n_heads:
        m_hat_0 = 2.0 ** (-4.0 / n)
        m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (n_heads - n), 2))
        m = torch.cat([m, m_hat])

    return m

@torch.no_grad()
def get_alibi_biases(n_heads: int, mask: torch.Tensor):
    m = get_slopes(n_heads).to(mask.device)
    seq_len = mask.size(0)
    distance = torch.tril(torch.arange(0, -seq_len, -1).view(-1, 1).expand(seq_len, seq_len))
    print(distance)

    return distance[:, :, None] * m[None, None, :]