rotary-embedding-torch - 为Transformer模型提供旋转位置编码的PyTorch库

旋转嵌入 - Pytorch

这是一个独立的库，用于在Pytorch中为Transformer添加旋转嵌入，这种方法作为相对位置编码已取得成功。具体来说，它将使得在张量的任何轴上旋转信息变得简单高效，无论是固定位置还是学习得到的。这个库将以很小的代价为你提供最先进的位置嵌入结果。

我的直觉告诉我，在人工神经网络中还有更多关于旋转的特性可以被利用。

安装

$ pip install rotary-embedding-torch

使用方法

import torch
from rotary_embedding_torch import RotaryEmbedding

# 在你的Transformer中实例化位置嵌入，并传递给所有注意力层

rotary_emb = RotaryEmbedding(dim = 32)

# 模拟查询和键 - 维度应以（序列长度，特征维度）结尾，前面可以有任意数量的维度（批次、头部等）

q = torch.randn(1, 8, 1024, 64) # 查询 - （批次，头部，序列长度，头部维度）
k = torch.randn(1, 8, 1024, 64) # 键

# 在头部分离之后，但在点积和后续softmax（注意力）之前，对查询和键应用旋转

q = rotary_emb.rotate_queries_or_keys(q)
k = rotary_emb.rotate_queries_or_keys(k)

# 然后像往常一样用你的查询（q）和键（k）进行注意力计算

如果你正确地执行了上述所有步骤，在训练过程中应该会看到显著的改善。

推理时的键值缓存

在处理推理时的键/值缓存时，查询位置需要用key_value_seq_length - query_seq_length进行偏移。

为了简化这个过程，可以使用rotate_queries_with_cached_keys方法

q = torch.randn(1, 8, 1, 64)     # 每次只有一个查询
k = torch.randn(1, 8, 1024, 64)  # 连接了缓存的键/值

q, k = rotary_emb.rotate_queries_with_cached_keys(q, k)

你也可以手动这样做

q = rotary_emb.rotate_queries_or_keys(q, offset = k.shape[-2] - q.shape[-2])

轴向旋转嵌入

用于简化n维轴向相对位置嵌入的使用，例如视频Transformer

import torch

from rotary_embedding_torch import (
    RotaryEmbedding,
    apply_rotary_emb
)

pos_emb = RotaryEmbedding(
    dim = 16,
    freqs_for = 'pixel',
    max_freq = 256
)

# 用于旋转频率的查询和键
# 假设是一个有8帧的视频，图像为矩形（特征维度在最后）

q = torch.randn(1, 8, 64, 32, 64)
k = torch.randn(1, 8, 64, 32, 64)

# 获取轴向频率 - (8, 64, 32, 16 * 3 = 48)
# 将自动进行部分旋转

freqs = pos_emb.get_axial_freqs(8, 64, 32)

# 旋转频率

q = apply_rotary_emb(freqs, q)
k = apply_rotary_emb(freqs, k)

长度可外推的旋转嵌入

在这篇论文中，他们通过给旋转嵌入一个类似于ALiBi的衰减，成功解决了长度外推问题。他们将这种技术命名为XPos，你可以在初始化时设置use_xpos = True来使用它。

这只能用于自回归Transformer

import torch
from rotary_embedding_torch import RotaryEmbedding

# 在你的Transformer中实例化位置嵌入，并传递给所有注意力层
rotary_emb = RotaryEmbedding(
    dim = 32,
    use_xpos = True   # 将此设置为True可使旋转嵌入更好地外推到超过训练时使用的序列长度
)

# 模拟查询和键 - 维度应以(序列长度, 特征维度)结尾，并可有任意数量的前置维度(批次, 头部等)

q = torch.randn(1, 8, 1024, 64) # 查询 - (批次, 头部数, 序列长度, 头部维度)
k = torch.randn(1, 8, 1024, 64) # 键

# 在头部拆分后但点积和后续softmax(注意力)之前，对查询和键应用旋转

# 不使用`rotate_queries_or_keys`，而是使用`rotate_queries_and_keys`，其余部分会自动处理

q, k = rotary_emb.rotate_queries_and_keys(q, k)

## 序列位置插值

这篇MetaAI论文提出，只需对序列位置的插值进行微调，就可以将预训练模型扩展到更长的上下文长度。他们表明，这比仅仅在相同序列位置上进行微调但进一步扩展的效果要好得多。

你可以通过在初始化时将`interpolate_factor`设置为大于`1.`的值来使用此功能（例如，如果预训练模型在2048上训练，设置`interpolate_factor = 2.`将允许微调到`2048 x 2. = 4096`）

更新：社区中有人报告称效果不佳。如果您看到正面或负面结果，请发邮件告诉我

import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(
    dim = 32,
    interpolate_factor = 2.    # 在预训练模型中添加这行代码并微调约1000步，如论文所示
)

## 引用

（保留原文引用内容）