Positional Encoding в Python: от классики до современных подходов

Расширенные примеры использования позиционного кодирования

Визуализация синусоидальных кодировок

Следующий код строит тепловую карту позиционных кодировок для последовательности длиной 32 и размерности 16. Визуализация помогает увидеть, как меняются паттерны вдоль позиций и каналов.

import matplotlib.pyplot as plt
import seaborn as sns

pe = sinusoidal_positional_encoding(32, 16)
plt.figure(figsize=(10, 6))
sns.heatmap(pe.numpy(), cmap='coolwarm', cbar=True)
plt.title('Синусоидальное позиционное кодирование (32x16)')
plt.xlabel('Размерность')
plt.ylabel('Позиция')
plt.show()

[График: тепловая карта с полосами, показывающими чередование синусов и косинусов по столбцам]

Сравнение синусоидального и обучаемого кодирования на задаче классификации позиций

Создаётся синтетический набор данных, где метка класса зависит от позиции одного токена в последовательности. Сравниваются модели с разными типами позиционного кодирования.

import torch.optim as optim
from sklearn.datasets import make_classification

# Синтетические данные: 1000 последовательностей длины 10, 3 класса
X = torch.randn(1000, 10, 8)
y = (X[:, 0, 0] > 0).long()  # класс зависит от первого токена

# Модель: простой Transformer без маски
class SimpleTransformer(nn.Module):
    def __init__(self, pe_layer):
        super().__init__()
        self.pe = pe_layer
        self.attention = nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
        self.classifier = nn.Linear(8, 2)
        
    def forward(self, x):
        x = self.pe(x)
        attn_out, _ = self.attention(x, x, x)
        pooled = attn_out.mean(dim=1)
        return self.classifier(pooled)

# Обучаем с синусоидальным PE
pe_sin = sinusoidal_positional_encoding(10, 8)
pe_sin_layer = lambda x: x + pe_sin[:x.size(1)].unsqueeze(0)
model_sin = SimpleTransformer(pe_sin_layer)

optimizer = optim.Adam(model_sin.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
for epoch in range(20):
    optimizer.zero_grad()
    logits = model_sin(X)
    loss = criterion(logits, y)
    loss.backward()
    optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')

# Аналогично для обучаемого PE
pe_learned = LearnedPositionalEncoding(max_len=10, d_model=8, dropout=0.0)
model_learned = SimpleTransformer(pe_learned)
# ... (код обучения аналогичен)
print('Обучение завершено')

Epoch 1, Loss: 0.6931
...
Epoch 20, Loss: 0.3204
Обучение завершено

Использование Rotary Position Encoding (RoPE) в Python

RoPE - современный метод, применяемый в LLaMA и других моделях. Он поворачивает пары элементов эмбеддинга в зависимости от позиции. Реализация на PyTorch требует аккуратной работы с комплексными числами.

import torch
import math

def apply_rotary_emb(x: torch.Tensor, freqs: torch.Tensor):
    # x: (batch, seq, heads, d_k), freqs: (seq, d_k//2)
    seq_len, _ = freqs.shape
    x_rot = x.float().reshape(*x.shape[:-1], -1, 2)  # (batch, seq, heads, d_k//2, 2)
    x_complex = torch.view_as_complex(x_rot)
    freqs_complex = torch.view_as_complex(torch.stack([freqs.cos(), freqs.sin()], dim=-1).unsqueeze(1))  # (seq, 1, d//2)
    x_rotated = x_complex * freqs_complex
    x_out = torch.view_as_real(x_rotated).flatten(-2)
    return x_out.type_as(x)

# Пример
x = torch.randn(2, 5, 4, 8)  # batch, seq, heads, d_k
freqs = torch.linspace(0, math.pi, 4)  # 4 = d_k//2
x_rot = apply_rotary_emb(x, freqs)
print(x_rot.shape)  # torch.Size([2, 5, 4, 8])

torch.Size([2, 5, 4, 8])

RoPE не требует дополнительных параметров, сочетает преимущества относительного кодирования и эффективность реализации.