LLM Attention and Rotary Position Embedding(旋转位置编码)

旋转位置编码(Rotary Position Embedding,RoPE)是一种能够将相对位置信息依赖集成Attention计算里的方法。就是在做词表映射的时候不是单一的进行一个embedding计算,还考虑位置信息。

一些资料

[1] https://arxiv.org/pdf/2104.09864

[2] https://arxiv.org/pdf/1706.03762

https://colab.research.google.com/drive/1rPk3ohrmVclqhH7uQ7qys4oznDdAhpzF

LLM Visualization

neural networks - What exactly are keys, queries, and values in attention mechanisms? - Cross Validated

Attention

从论文里: An attention function can be described as mapping a query and a set of key-value pairs to an output. Attention即把查询向量映射成一个输出的操作。在看之前先看下embedding和qkv等几个基本概念。

Embedding

embdding的作用是词表到特征向量的映射,即把一个int的index索引映射到一个向量空间表示。比如看llama的token: https://huggingface.co/meta-llama/Llama-2-7b-hf/blob/main/tokenizer.json

截图:

vocab里面是当前模型的所有词表,比如"<unk>"对应的索引是0。通过tokenizer.endocer的编码操

作就可以将一个字段编译成一组index vector。比如:

from transformers import LlamaForCausalLM, LlamaTokenizer

tokenizer = LlamaTokenizer.from_pretrained("/model/Llama-2-7b-chat-hf")


encode = tokenizer.encode("who")
print(encode)

# [1, 1058]
# "" -> 1
# "who" -> 1058

然后得到的[1, 1058]就是embedding的输入了。embeding的计算:

>>> import torch
>>> a = torch.tensor([1])   # 输入index是1就是前面encode后的数值

>>> import torch.nn as nn
>>> embedding = nn.Embedding(3, 3)
>>> embedding.weight
Parameter containing:
tensor([[ 0.0198, -0.5562,  0.8156],
        [-0.3192, -1.2203, -0.8307],
        [-0.1649, -0.2753, -0.9075]], requires_grad=True)
>>> o = embedding(a) # 从weight里取出来对应的第1行的数值。
>>> o
tensor([[-0.3192, -1.2203, -0.8307]], grad_fn=<EmbeddingBackward0>)

embeeding的操作是根据输入的index的数值(比如这里1)然后从对应的权重(embedding.weight, 3x3)里取出对应index行索引对应的权重vector(第一行,embedding.weight[1])。即一个简单的索引操作。所以这里就有一个限制,即输入的index的大小必须要小于embedding.weight的shape(0),超出了就会挂掉。另外embedding.weight的shape(1)的大小被称为hidden size,即feature的大小。后面的linear等操作都是根据这个feature来运算的。实际上,embedding将index变成了一个高维空间的表示,可以用于模型训练推理。正常下embedding的权重第一个维度大小是vocab_size,可以查看config.json配置文件。

query、keys and values

query: 希望取查询的文本对应的查找向量,即LLM的输入经过embedding映射后的向量就是一个query vector。

keys: 和query,是对输入的文本的一个向量映射。keys里存储了所有的之前LLM推理后的context的vector数值。The key vectors are used to compute how relevant each element in the input sequence is to the query.

values: values是和query一一对应的,一个key对应一个values, These values are weighted by the attention scores (computed from the query-key interaction) to determine how much each element contributes to the final output.

也就是根据key和value,来判断当前的query和哪个key的关联性最高,分数最高。那么是怎么计算的呢?参考qkv比较清楚计算方式了。

从图中可以看到对一个query([1, 0, 2])会分别和三组key/value计算,得到三组vactor。然后再把三组vector相加得到[2, 7, 1.5],这个vector在经过重复的mlp和attention和最后的logits处理得到llm模型的输出index, 比如:

第一次输入How to predict, 经过llm根据最后的softmax后vector选取最优的输出。每个有个score. (prefill的时候全部是query,没有key value的)decode的时候,query长度都是1,然后之前的context存储在所有的key value里。

Attention is all you need

An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. attention表示:

左图是最简单的attention操作,右边是multi head attention的图. 先看左图,输入QKV,经过计算Scaled Dot Product Attention,得到一组特征向量,然后每一组query和key/value得到一个vector(如上面qkv章节的图),然后将vector相加。计算公式 :

scaled dot product attention:

query和key的维度是d_k。如qkv章节的图,q和k点乘,然后经过softmax之后得到一个score分数。在将这个score和V做乘法,即得到一个基于当前组k/v的输出特征向量。

一个scaled dot product的pytorch实现:

class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):

        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))

        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)

        attn = self.dropout(F.softmax(attn, dim=-1))
        output = torch.matmul(attn, v)

        return output, attn

Multi-Head Attention(MHA)

多头注意力机制和单个scaled dot product attention的区别是先将输入的QKV分别进行多个linear操作,然后在最后的输出vector也经过多个linear映射。这样可以将qkv的向量映射到更高维度的空间上。每个linear的输出大小可以表示为head_dim, 有多少和linear可以用head_num来表示。多头(head)即指这里head_num有几个。一般在实现的时候由于linear是线性的,所以可以将多个head合并成一个linea

### 绝对位置编码的概念 绝对位置编码是一种用于表示序列中每个标记在其上下文中相对顺序的技术。在大型语言模型(LLM)中,这种技术对于捕获输入数据中的时间依赖性和结构至关重要[^1]。由于神经网络本身无法天然理解序列的位置信息,因此引入了位置编码来增强模型的学习能力。 #### 实现细节 在实际应用中,绝对位置编码可以通过多种方式实现。一种常见的方式是通过正弦和余弦函数生成固定的位置向量: ```python import torch import math def absolute_position_encoding(seq_len, d_model): pe = torch.zeros(seq_len, d_model) position = torch.arange(0, seq_len).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) # 偶数索引 pe[:, 1::2] = torch.cos(position * div_term) # 奇数索引 return pe.unsqueeze(0) seq_length = 512 embedding_dim = 768 position_encodings = absolute_position_encoding(seq_length, embedding_dim) print(f"Position Encoding Shape: {position_encodings.shape}") ``` 上述代码片段展示了如何基于正弦和余弦波形创建绝对位置编码矩阵。这种方法最初由 Vaswani 等人在 Transformer 模型中提出,并广泛应用于各种 LLM 的实现中。 #### 编码的作用 绝对位置编码的主要作用在于帮助模型区分不同位置上的相同单词或标记的意义差异。例如,在句子 “她喜欢苹果” 和 “苹果喜欢她” 中,“苹果”的含义因位置而异。通过加入位置编码,模型可以更好地学习到这些细微差别并提高预测准确性。 #### 替代方法 尽管绝对位置编码非常有效,但它也有一些局限性,比如难以处理超出训练期间最大长度的新序列。为此,研究人员提出了其他变体,如可学习的位置嵌入或者相对位置编码等方案作为补充或替代选择。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值