Positional Encoding: как LLM знает порядок слов

opensourceaillmpositional-encodingtransformerit
← Back to Blog

Введение: проблема порядка слов

Transformer не рекуррентный — он не читает текст последовательно, как RNN. Значит, модель не знает порядок слов без явного добавления информации о позициях.

RNN: "Кот съел мышь"
  Шаг 1: "Кот" → hidden_1
  Шаг 2: "Кот съел" → hidden_2
  Шаг 3: "Кот съел мышь" → hidden_3
  → Порядок встроен в механизм

Transformer: ["Кот", "съел", "мышь"]
  Все токены обрабатываются параллельно
  → Без positional encoding порядок потерян!

Уровень 1: No Position (без позиционных кодирований)

Что будет без positional encoding?

Вход: ["Кот", "съел", "мышь"]
Embeddings: [e_cat, e_ate, e_mouse]

Attention:
  attention("Кот", "съел") = attention("Кот", "мышь")?
  → Нет! Attention вычисляется по значениям, не позициям
  
  Но: attention("Кот", "съел") = attention("мышь", "съел")
  → Модель не различит "Кот съел мышь" и "Мышь съел кот"

Результат: модель перестанет понимать грамматику и смысл.


Уровень 2: Absolute Positional Embeddings

Простая добавка позиций

class AbsolutePositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        # Позиционные эмбеддинги: [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_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)
        self.register_buffer('pe', pe.unsqueeze(0))  # [1, max_len, d_model]
    
    def forward(self, x):
        # x: [batch_size, seq_len, d_model]
        return x + self.pe[:, :x.size(1), :]

Как работают синусоидальные позиционные кодирования

d_model = 8, max_len = 10

Позиция 0: [sin(0), cos(0), sin(0), cos(0), sin(0), cos(0), sin(0), cos(0)]
          = [0, 1, 0, 1, 0, 1, 0, 1]

Позиция 1: [sin(1/10000^(2/8)), cos(...), ...]
          = [0.06, 0.99, 0.13, 0.99, 0.26, 0.96, 0.52, 0.85]

Позиция 2: [sin(2/10000^(2/8)), cos(...), ...]
          = [0.12, 0.99, 0.26, 0.96, 0.52, 0.85, 0.96, 0.28]

Ключевые свойства:

  1. Уникальность: каждая позиция имеет уникальный вектор
  2. Обобщение: позиции beyond training length работают (синусоида продолжается)
  3. Относительные позиции: sin и cos позволяют выразить относительные позиции через линейные преобразования

Уровень 3: Learned Positional Embeddings

Обучаемые позиционные эмбеддинги

class LearnedPositionalEmbedding(nn.Module):
    def __init__(self, max_len: int, d_model: int):
        super().__init__()
        # Просто таблица поиска
        self.weight = nn.Embedding(max_len, d_model)
    
    def forward(self, x):
        seq_len = x.size(1)
        positions = torch.arange(seq_len, device=x.device)
        return x + self.weight(positions)

Сравнение: синусоидальные vs обучаемые

Синусоидальные:
  + Не требуют обучения
  + Работают для позиций beyond training
  + Фиксированная функция
  - Могут не оптимальны для задачи

Обучаемые:
  + Оптимизируются под задачу
  - Требуют больше данных
  - Не работают для позиций beyond training
  - Нужно интерполировать для длинных контекстов

GPT-2 использует обучаемые, BERT — синусоидальные.


Уровень 4: Relative Positional Encoding

Проблема абсолютных позиций

Абсолютные позиции:
  "Кот сидит на ковре" (позиции 0, 1, 2, 3, 4)
  "Собака бегает по парку" (позиции 0, 1, 2, 3, 4)
  
  → Модель учитывает абсолютные позиции, а не отношения

Но в attention важнее относительные позиции:
  attention(i, j) зависит от (i - j), а не от i и j по отдельности

Relative Positional Encoding (Shaw et al., 2018)

class RelativePositionalBias(nn.Module):
    def __init__(self, num_buckets: int, max_distance: int, heads: int):
        super().__init__()
        # Таблица относительных позиций
        self.relative_attention_bias = nn.Embedding(num_buckets, heads)
    
    def _position_to_bucket(self, position, bias):
        # Абсолютная разница → индекс
        max_exact = position - max_distance // 2
        is_small = position < max_exact
        bucket = max_exact * 2 + (1 if is_small else 0)
        return bucket
    
    def forward(self, q_len, k_len):
        # q: [batch, heads, q_len, d], k: [batch, heads, k_len, d]
        positions = torch.arange(q_len).unsqueeze(0) - \
                   torch.arange(k_len).unsqueeze(1)
        buckets = self._position_to_bucket(positions, 0)
        bias = self.relative_attention_bias(buckets)
        return bias.permute(2, 0, 1)  # [heads, q_len, k_len]

Как relative positions добавляются в attention

Standard attention:
  attention(Q, K, V) = softmax(QK^T / sqrt(d))V

Relative attention:
  attention(Q, K, V) = softmax((Q + b_q)(K + b_k)^T / sqrt(d))V
  
  где b_q, b_k — относительные позиционные смещения

RoPE: Rotary Positional Embeddings

Идея: поворот векторов

RoPE (Su et al., 2021, Llama 3, Mistral) кодирует позиции через поворот векторов в комплексной плоскости:

import torch
import math

def rotate_half(x):
    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(q, k, positions, base=10000):
    """
    q, k: [batch, heads, seq_len, head_dim]
    positions: [seq_len]
    """
    dim = q.shape[-1]
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    
    t = positions.float() * inv_freq  # [seq_len, dim/2]
    
    # Синус и косинус для каждой позиции
    freqs = torch.cat((t, t), dim=-1)  # [seq_len, dim]
    cos = freqs.cos().unsqueeze(0).unsqueeze(0).unsqueeze(0)  # [1,1,1,1]
    sin = freqs.sin().unsqueeze(0).unsqueeze(0).unsqueeze(0)
    
    # Поворот
    q_emb = q * cos + rotate_half(q) * sin
    k_emb = k * cos + rotate_half(k) * sin
    
    return q_emb, k_emb

Визуализация RoPE

Вектор [x, y] на позиции 0:
  [x, y]

Вектор [x, y] на позиции 1:
  [x·cos(θ) - y·sin(θ), x·sin(θ) + y·cos(θ)]
  
  где θ = 1/base^(2/dim)

Вектор [x, y] на позиции 2:
  [x·cos(2θ) - y·sin(2θ), x·sin(2θ) + y·cos(2θ)]

Почему RoPE работает лучше?

1. Относительные позиции в attention:
   q(i)^T · k(j) = f(i - j)
   → Attention зависит только от разницы позиций!

2. Экстраполяция:
   RoPE хорошо работает для позиций beyond training length
   (хотя и не идеально)

3. Сохранение длины вектора:
   ||RoPE(q)|| = ||q||
   → Не меняет масштаб attention scores

NTK-aware Scale Positional Embedding

Проблема RoPE на длинных контекстах

RoPE обучалась на контексте 4K токенов
При контексте 32K токены на позиции 10K имеют:
  θ(10K) = 10K × (1/10000^(2/dim))
  → Очень быстрый оборот, потеря информации

NTK-aware решение (LongLoRA, 2023)

def get_ntk_scale(base_dim, seq_len, original_ctx=4096):
    """
    Вычисляет масштаб для NTK-aware positional interpolation
    """
    if seq_len <= original_ctx:
        return 1.0
    
    # NTK interpolation factor
    dim = base_dim
    scale = (base_dim / (dim * math.log(seq_len / original_ctx))) ** (dim / (dim - 2))
    return max(1.0, scale)

# Пример:
# original_ctx = 4096, new_ctx = 32768
# scale = 4.0 → используем base = 10000 * 4 = 40000
# → Более медленный оборот → лучше extrapolation

ALiBi: Attention with Linear Biases

Альтернатива: добавление bias в attention scores

class ALiBi(nn.Module):
    def __init__(self, num_heads: int, max_pos: int):
        super().__init__()
        # Коэффициенты для каждого head (отрицательные)
        slopes = torch.Tensor(self._get_slopes(num_heads))
        self.register_buffer('slopes', slopes)
        self.register_buffer('bias', 
            self._build_bias(max_pos))
    
    def _get_slopes(self, n):
        """Степенной закон: 2^(-8/h) для h > 8"""
        power = torch.Tensor(range(1, n + 1)) / n  # [1/n, 2/n, ..., 1]
        return 2 ** (-8 * power)  # [0.5, 0.25, ...]
    
    def _build_bias(self, max_pos):
        """Предвычисленные bias для каждой пары позиций"""
        bias = torch.arange(max_pos).unsqueeze(0) - \
               torch.arange(max_pos).unsqueeze(1)
        return bias  # [max_pos, max_pos], upper triangle = 0
    
    def forward(self, q, k):
        attention = torch.matmul(q, k.transpose(-2, -1))
        
        # Добавляем ALiBi bias
        bsz, num_heads, q_len, kv_len = attention.shape
        attention = attention + self.bias[:q_len, :kv_len].unsqueeze(0).unsqueeze(0)
        
        # Нормализуем по количеству head
        attention = attention / len(self.slopes)
        
        return torch.softmax(attention, dim=-1)

Как ALiBi работает

Attention scores до bias:
  [[2.3, 1.8, 1.2, 0.9],
   [1.5, 2.1, 1.4, 1.0],
   [1.0, 1.3, 1.9, 1.5],
   [0.8, 1.0, 1.2, 1.8]]

ALiBi bias (slopes = [0.5, 0.25]):
  [[ 0, -0.5, -1.0, -1.5],
   [ 0,  0, -0.5, -1.0],
   [ 0,  0,  0, -0.5],
   [ 0,  0,  0,  0]]

Attention scores после bias:
  [[2.3, 1.3, 0.2, -0.6],
   [1.5, 2.1, 0.9, 0.0],
   [1.0, 1.3, 1.9, 1.0],
   [0.8, 1.0, 1.2, 1.8]]

→ Дальние токены получают больший штраф
→ Модель фокусируется на ближайших токенах

Сравнение методов позиционных кодирований

Метод              | Модели           | Плюсы                    | Минусы
-------------------|------------------|--------------------------|--------------------------
Absolute (sinusoid)| BERT             | Простой, фиксированный   | Не обучается
Absolute (learned) | GPT-2, GPT-3     | Оптимизируется           | Плохая экстраполяция
Relative           | T5, DeBERTa      | Учитывает отношения      | Больше параметров
RoPE               | Llama, Mistral   | Относительные, плавный   | Сложнее реализовать
ALiBi              | PaLM, OPT        | Нет параметров           | Штрафует дальние токены
NTK-aware RoPE     | LongLoRA         | Длинный контекст         | Требует tune scale

Практические аспекты

Как выбрать метод?

Короткий контекст (< 4K):
  → Absolute learned (просто и эффективно)

Средний контекст (4K-32K):
  → RoPE (хороший баланс)

Длинный контекст (> 32K):
  → RoPE + NTK-aware или YaRN

Без дополнительных параметров:
  → ALiBi

Для задач с важными отношениями:
  → Relative positional encoding

Экстраполяция позиций

# YaRN (Yet another RoPE) — улучшение экстраполяции
class YaRNScaledRoPE(nn.Module):
    def __init__(self, base, dim, max_pos=4096, scale=8):
        super().__init__()
        self.original_base = base
        self.dim = dim
        self.scale = scale
        
        # Корректируем base для нового контекста
        correction_factor = (
            0.1 * math.log(scale) + 1.0
        ) if scale > 1 else 1.0
        self.effective_base = base * correction_factor ** (dim / (dim - 2))
    
    def forward(self, x, positions):
        # Используем effective_base вместо original_base
        # → Более медленный оборот на больших позициях
        ...

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

Паттерны attention с разными позициями

Без positional encoding:
  attention matrix = случайная (нет информации о позициях)

С absolute positional:
  attention matrix = диагональная (близкие позиции привлекаются)

С RoPE:
  attention matrix = плавная decay (ближние > дальние)

С ALiBi:
  attention matrix = строгий decay (экспоненциальный штраф)

Heatmap позиционных весов

Позиция →
0    1    2    3    4    5    6    7
[1.0, 0.8, 0.5, 0.3, 0.1, 0.05, 0.02, 0.01]  ← RoPE
[1.0, 0.9, 0.8, 0.7, 0.6, 0.5,  0.4,  0.3]  ← Absolute
[1.0, 0.6, 0.3, 0.1, 0.05,0.02, 0.01, 0.005] ← ALiBi

Итоги

  • Transformer без positional encoding не знает порядок слов
  • Absolute positional encoding добавляет уникальные векторы для каждой позиции
  • Relative positional encoding учитывает отношения между позициями
  • RoPE (Llama) кодирует позиции через поворот векторов
  • ALiBi (PaLM) добавляет линейный bias в attention scores
  • NTK-aware и YaRN улучшают экстраполяцию на длинные контексты
  • Выбор метода зависит от длины контекста и задачи
  • Позиционные кодирования — критический компонент для понимания порядка