前言
本文基于 llama 模型的源码,学习相对位置编码的实现方法,本文不细究绝对位置编码和相对位置编码的数学原理。
大模型新人在学习中容易困惑的几个问题:
-
为什么一定要在 transformer 中使用位置编码?
-
相对位置编码在 llama 中是怎么实现的?
-
大模型的超长文本预测和位置编码有什么关系?
1、为什么需要位置编码
很多初学者都会读到这样一句话:transformer 使用位置编码的原因是它不具备位置信息。大家都只把这句话当作公理,却很少思考这句话到底是什么意思?
这句话的意思是,如果没有位置编码,那么 “床前明月”、“前床明月”、“前明床月” 这几个输入,会预测出完全一样的文本。
也就是说,不管你输入的 prompt 顺序是什么,只要 prompt 的文本是相同的,那么模型 decode 的文本就只取决于 prompt 的最后一个 token。
import torch
from torch import nn
import math
batch = 1
dim = 10
num_head = 2
embedding = nn.Embedding(5, dim)
q_matrix = nn.Linear(dim, dim, bias=False)
k_matrix = nn.Linear(dim, dim, bias=False)
v_matrix = nn.Linear(dim, dim, bias=False)
x = embedding(torch.tensor([1,2,3])).unsqueeze(0)
y = embedding(torch.tensor([2,1,3])).unsqueeze(0)
def attention(input):
q = q_matrix(input).view(batch, -1, num_head, dim // num_head).transpose(1, 2)
k = k_matrix(input).view(batch, -1, num_head, dim // num_head).transpose(1, 2)
v = v_matrix(input).view(batch, -1, num_head, dim // num_head).transpose(1, 2)
attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(dim // num_head)
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
outputs = torch.matmul(attn_weights, v).transpose(1, 2).reshape(1, len([1,2,3]), dim)
print(outputs)
attention(x)
attention(y)
执行上面的代码会发现,虽然 x 和 y 交换了第一个 token 和第二个 token 的输入顺序,但是第三个 token 的计算结果完全没有发生改变,那么模型预测第四个 token 时,便会得到相同的结果。
如果有读者对矩阵运算感到混淆的话,可以看看下面的简单推导:

可以看出,当第一个 token 与第二个 token 交换顺序后,模型输出矩阵的第一维和第二维也交换了顺序,但输出的值完全没有变化。
第三个 token 的输出结果也是完全没有受到影响,这也就是前面说的:如果没有位置编码,模型 decode 的文本就只取决于 prompt 的最后一个 token。
不过需要注意的是,由于 attention_mask 的存在(前置位 token 看不到后置位 token),所以即使不加位置编码,transformer 的输出还是会受到 token 的位置影响。
2、相对位置编码的实现
我们以 modeling_llama.py 的源码为例,来学习相对位置编码的实现方法。
class LlamaRotaryEmbedding(torch.nn.Module):
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
super().__init__()
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
self.register_buffer("inv_freq", inv_freq)
# Build here to make `torch.jit.trace` work.
self.max_seq_len_cached = max_position_embeddings
t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
freqs = torch.einsum("i,j->ij", t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1)
self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
def forward(self, x, seq_len=None):
# x: [bs, num_attention_heads, seq_len, head_size]
# This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
if seq_len > self.max_seq_len_cached:
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
freqs = torch.einsum("i,j->ij", t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
return (
self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
)
def rotate_half(x):
"""Rotates half the hidden dims of the input."""
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
# The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
相对位置编码在 attention 中的应用方法如下:
self.rotary_emb = LlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
if past_key_value is not None:
# reuse k, v, self_attention
key_states = torch.cat([past_key_value[0], key_states], dim=1)
value_states = torch.cat([past_key_value[1], value_states], dim=1)
根据 value_states 矩阵的形状去调取 cos 和 sin 两个 tensor, cos 与 sin 的维度均是 batch_size * head_num * seq_len * head_dim;
利用 apply_rotary_pos_emb 去修改 query_states 和 key_states 两个 tensor,得到新的 q,k 矩阵
需要注意的是,在解码时,position_ids 的长度是和输入 token 的长度保持一致的,prompt 是 4 个 token 的话。
第一次解码时,position_ids: tensor([[0, 1, 2, 3]], device=‘cuda:0’),q 矩阵与 k 矩阵的相对位置编码信息通过 apply_rotary_pos_emb() 获得;
第二次解码时,position_ids: tensor([[4]], device=‘cuda:0’),当前 token 的相对位置编码信息通过 apply_rotary_pos_emb() 获得。
前 4 个 token 的相对位置编码信息则是通过 key_states = torch.cat([past_key_value[0], key_states], dim=1) 集成到 k 矩阵中;
……
……
以上代码的公式,均可以从苏神原文中找到。
这些代码可以从 llama 模型中剥离出来直接执行,如果感到困惑,可以像下面一样,将 apply_rotary_pos_emb() 的整个过程给 print 出来观察一下:
head_num, head_dim, kv_seq_len = 8, 20, 5
position_ids = torch.tensor([[0, 1, 2, 3, 4]])
query_states = torch.randn(1, head_dim, kv_seq_len, head_dim)
key_states = torch.randn(1, head_dim, kv_seq_len, head_dim)
value_states = torch.randn(1, head_dim, kv_seq_len, head_dim)
rotary_emb = LlamaRotaryEmbedding(head_dim)
cos, sin = rotary_emb(value_states, seq_len=kv_seq_len)
print(cos, sin)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
3、位置编码与长度外推
长度外推指的是,大模型在训练的只见过长度为 X 的文本,但在实际应用时却有如下情况:

我们假设 X 的取值为 4096,那么也就意味着,模型自始至终没有见到过 pos_id >= 4096 的位置编码,进而导致模型的预测结果完全不可控。
因此,解决长度外推问题的关键便是如何让模型见到比训练文本更长的位置编码。

以上关于文本外推的介绍均是比较大白话的理解,只是为了强调位置编码很重要这一观点。
最后
为了助力朋友们跳槽面试、升职加薪、职业困境,提高自己的技术,本文给大家整了一套涵盖AI大模型所有技术栈的快速学习方法和笔记。目前已经收到了七八个网友的反馈,说是面试问到了很多这里面的知识点。
由于文章篇幅有限,不能将全部的面试题+答案解析展示出来,有需要完整面试题资料的朋友,可以扫描下方二维码免费领取哦!!! 👇👇👇👇

面试题展示
1、请解释一下BERT模型的原理和应用场景。
答案:BERT(Bidirectional Encoder Representations from Transformers)是一种预训练的语言模型,通过双向Transformer编码器来学习文本的表示。它在自然语言处理任务中取得了很好的效果,如文本分类、命名实体识别等。
2、什么是序列到序列模型(Seq2Seq),并举例说明其在自然语言处理中的应用。
答案:Seq2Seq模型是一种将一个序列映射到另一个序列的模型,常用于机器翻译、对话生成等任务。例如,将英文句子翻译成法文句子。
3、请解释一下Transformer模型的原理和优势。
答案:Transformer是一种基于自注意力机制的模型,用于处理序列数据。它的优势在于能够并行计算,减少了训练时间,并且在很多自然语言处理任务中表现出色。
4、什么是注意力机制(Attention Mechanism),并举例说明其在深度学习中的应用。
答案:注意力机制是一种机制,用于给予模型对不同部分输入的不同权重。在深度学习中,注意力机制常用于提升模型在处理长序列数据时的性能,如机器翻译、文本摘要等任务。
5、请解释一下卷积神经网络(CNN)在计算机视觉中的应用,并说明其优势。
答案:CNN是一种专门用于处理图像数据的神经网络结构,通过卷积层和池化层提取图像特征。它在计算机视觉任务中广泛应用,如图像分类、目标检测等,并且具有参数共享和平移不变性等优势。
6、请解释一下生成对抗网络(GAN)的原理和应用。
答案:GAN是一种由生成器和判别器组成的对抗性网络结构,用于生成逼真的数据样本。它在图像生成、图像修复等任务中取得了很好的效果。
7、请解释一下强化学习(Reinforcement Learning)的原理和应用。
答案:强化学习是一种通过与环境交互学习最优策略的机器学习方法。它在游戏领域、机器人控制等领域有广泛的应用。
8、请解释一下自监督学习(Self-Supervised Learning)的原理和优势。
答案:自监督学习是一种无需人工标注标签的学习方法,通过模型自动生成标签进行训练。它在数据标注困难的情况下有很大的优势。
9、解释一下迁移学习(Transfer Learning)的原理和应用。
答案:迁移学习是一种将在一个任务上学到的知识迁移到另一个任务上的学习方法。它在数据稀缺或新任务数据量较小时有很好的效果。
10、请解释一下模型蒸馏(Model Distillation)的原理和应用。
答案:模型蒸馏是一种通过训练一个小模型来近似一个大模型的方法。它可以减少模型的计算和存储开销,并在移动端部署时有很大的优势。
11、请解释一下LSTM(Long Short-Term Memory)模型的原理和应用场景。
答案:LSTM是一种特殊的循环神经网络结构,用于处理序列数据。它通过门控单元来学习长期依赖关系,常用于语言建模、时间序列预测等任务。
12、请解释一下BERT模型的原理和应用场景。
答案:BERT(Bidirectional Encoder Representations from Transformers)是一种预训练的语言模型,通过双向Transformer编码器来学习文本的表示。它在自然语言处理任务中取得了很好的效果,如文本分类、命名实体识别等。
13、什么是注意力机制(Attention Mechanism),并举例说明其在深度学习中的应用。
答案:注意力机制是一种机制,用于给予模型对不同部分输入的不同权重。在深度学习中,注意力机制常用于提升模型在处理长序列数据时的性能,如机器翻译、文本摘要等任务。
14、请解释一下生成对抗网络(GAN)的原理和应用。
答案:GAN是一种由生成器和判别器组成的对抗性网络结构,用于生成逼真的数据样本。它在图像生成、图像修复等任务中取得了很好的效果。
15、请解释一下卷积神经网络(CNN)在计算机视觉中的应用,并说明其优势。
答案:CNN是一种专门用于处理图像数据的神经网络结构,通过卷积层和池化层提取图像特征。它在计算机视觉任务中广泛应用,如图像分类、目标检测等,并且具有参数共享和平移不变性等优势。
16、请解释一下强化学习(Reinforcement Learning)的原理和应用。
答案:强化学习是一种通过与环境交互学习最优策略的机器学习方法。它在游戏领域、机器人控制等领域有广泛的应用。
17、请解释一下自监督学习(Self-Supervised Learning)的原理和优势。
答案:自监督学习是一种无需人工标注标签的学习方法,通过模型自动生成标签进行训练。它在数据标注困难的情况下有很大的优势。
18、请解释一下迁移学习(Transfer Learning)的原理和应用。
答案:迁移学习是一种将在一个任务上学到的知识迁移到另一个任务上的学习方法。它在数据稀缺或新任务数据量较小时有很好的效果。
19、请解释一下模型蒸馏(Model Distillation)的原理和应用。
答案:模型蒸馏是一种通过训练一个小模型来近似一个大模型的方法。它可以减少模型的计算和存储开销,并在移动端部署时有很大的优势。
20、请解释一下BERT中的Masked Language Model(MLM)任务及其作用。
答案:MLM是BERT预训练任务之一,通过在输入文本中随机mask掉一部分词汇,让模型预测这些被mask掉的词汇。
由于文章篇幅有限,不能将全部的面试题+答案解析展示出来,有需要完整面试题资料的朋友,可以扫描下方二维码免费领取哦!!! 👇👇👇👇



3878

被折叠的 条评论
为什么被折叠?



