Context Windows: управление контекстом в LLM
opensourceaillmcontext-windowtransformerit
Введение: Context Window
Что такое Context Window?
Context Window = максимальное количество токенов, которые модель может обработать
Традиционные LLM:
GPT-3: 2,048 токенов (~4,000 слов)
GPT-3.5: 4,096 токенов
Современные LLM:
GPT-4 Turbo: 128,000 токенов
Claude 3.5: 200,000 токенов
Gemini 1.5: 1,000,000+ токенов
1 токен ≈ 0.75 слова (английский)
1 токен ≈ 2 символа (русский)
Факт: Context window определяет, сколько текста модель может "помнить" за один раз.
Внутреннее устройство
Attention и O(n²) сложность
Self-Attention вычисляет:
Attention(Q, K, V) = softmax(QK^T / √d)V
Матрица attention: n × n
n = длина контекста
O(n²) сложность:
n=4K: 16M операций
n=32K: 1,024M операций
n=128K: 16,384M операций
Проблема: квадратичный рост с длиной контекста
class StandardAttention(torch.nn.Module):
"""Standard O(n²) self-attention"""
def forward(self, x):
B, T, C = x.size() # batch, time, channels
# Q, K, V projections
q = self.query(x) # (B, T, d)
k = self.key(x) # (B, T, d)
v = self.value(x) # (B, T, d)
# Attention scores
scale = 1.0 / (C ** 0.5)
scores = torch.matmul(q, k.transpose(-2, -1)) * scale # (B, T, T)
# Softmax
weights = torch.softmax(scores, dim=-1)
weights = torch.dropout(weights, self.dropout, self.training)
# Weighted sum
out = torch.matmul(weights, v) # (B, T, d)
return self.output(out)
Техники расширения контекста
1. RoPE Scaling
RoPE (Rotary Positional Embeddings):
Применяет вращение к векторам по позиции
pos = 10000^(2i/d) для i = 0, ..., d/2
Scaling methods:
- Linear: прямое масштабирование
- NTK: NTK-aware scaling
- Dynamic: динамическое масштабирование
class RoPEScaled(torch.nn.Module):
"""RoPE with scaling for longer contexts"""
def __init__(self, dim, base=10000, max_seq_len=8192, factor=1.0):
super().__init__()
self.dim = dim
self.base = base
self.max_seq_len = max_seq_len
self.factor = factor
# Compute frequencies
freqs = 1.0 / (base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
self.register_buffer('freqs', freqs)
def get_scaled_freqs(self, seq_len):
"""Compute scaled frequencies for longer sequences"""
if seq_len <= self.max_seq_len:
return self.freqs
# NTK-aware scaling
dim = self.dim
base = self.base
max_seq_len = self.max_seq_len
factor = seq_len / max_seq_len
ratio = base / (base * (factor ** (dim / (dim - 2))) - base)
new_base = base * (1 + ratio * (factor - 1))
new_freqs = 1.0 / (new_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
return new_freqs
def apply_rope(self, x, seq_len):
"""Apply RoPE to input embeddings"""
freqs = self.get_scaled_freqs(seq_len).to(x.device)
# Split into even/odd
x1 = x[..., 0::2] # even
x2 = x[..., 1::2] # odd
# Apply rotation
cos = torch.cos(freqs)
sin = torch.sin(freqs)
out1 = x1 * cos - x2 * sin
out2 = x1 * sin + x2 * cos
out = torch.stack([out1, out2], dim=-1).flatten(-2)
return out
2. Linear Attention
Linear Attention: O(n) вместо O(n²)
Обычный attention:
out = softmax(QK^T)V → O(n²d)
Linear attention:
out = softmax(Q) · softmax(K)^T · V
= (softmax(Q) · (softmax(K)^T · V))
Сначала: K^T · V → O(nd²)
Затем: Q · result → O(n²d)
Но можно аппроксимировать:
out ≈ φ(Q) · φ(K)^T · φ(V)
где φ — positive kernel (ReLU, exp, etc.)
class LinearAttention(torch.nn.Module):
"""O(n) linear attention"""
def __init__(self, dim, num_heads=8):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim ** -0.5
self.query = torch.nn.Linear(dim, dim)
self.key = torch.nn.Linear(dim, dim)
self.value = torch.nn.Linear(dim, dim)
self.output = torch.nn.Linear(dim, dim)
# Small epsilon for numerical stability
self.eps = 1e-6
def kernel_fn(self, x):
"""Positive kernel function (ReLU approximation)"""
return torch.nn.functional.relu(x)
def forward(self, x):
B, T, C = x.size()
q = self.query(x).view(B, T, self.num_heads, self.head_dim)
k = self.key(x).view(B, T, self.num_heads, self.head_dim)
v = self.value(x).view(B, T, self.num_heads, self.head_dim)
q = self.kernel_fn(q) + self.eps
k = self.kernel_fn(k) + self.eps
# Linear attention: O(n) complexity
# KV = sum(k_i ⊗ v_i) over sequence
KV = torch.einsum('bnd,bne->bde', k, v) # (B, d, d)
# Q · KV
q = q * self.scale
out = torch.einsum('btd,bde->bte', q, KV) # (B, T, d)
# Normalize by (Q · sum(K))
sum_k = k.sum(dim=-2) # (B, d)
norm = torch.einsum('btd,bd->bt', q, sum_k) + self.eps
out = out / norm.unsqueeze(-1)
out = out.reshape(B, T, C)
return self.output(out)
3. Sparse Attention
Sparse Attention: только важные паттерны
Patterns:
1. Local/Window: attention только в окне
2. Strided: каждые k шагов
3. Global: несколько глобальных токенов
4. Diagonal: attention по диагонали
BigBird (Google):
- Random attention (случайные пары)
- Window attention (соседние токены)
- Global attention (первые и центральные)
class SparseAttention(torch.nn.Module):
"""Sparse attention with multiple patterns"""
def __init__(self, dim, num_heads=8, window_size=64, num_global=4):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.window_size = window_size
self.num_global = num_global
self.query = torch.nn.Linear(dim, dim)
self.key = torch.nn.Linear(dim, dim)
self.value = torch.nn.Linear(dim, dim)
self.output = torch.nn.Linear(dim, dim)
def forward(self, x):
B, T, C = x.size()
q = self.query(x).view(B, T, self.num_heads, self.head_dim)
k = self.key(x).view(B, T, self.num_heads, self.head_dim)
v = self.value(x).view(B, T, self.num_heads, self.head_dim)
q = q.permute(0, 2, 1, 3) # (B, H, T, D)
k = k.permute(0, 2, 3, 1) # (B, H, D, T)
v = v.permute(0, 2, 1, 3) # (B, H, T, D)
# 1. Local window attention
local_out = self._window_attention(q, k, v)
# 2. Global attention (first + center tokens)
global_out = self._global_attention(q, k, v)
# 3. Random attention
random_out = self._random_attention(q, k, v)
# Combine
out = local_out + global_out + random_out
out = out.permute(0, 2, 1, 3).reshape(B, T, C)
return self.output(out)
def _window_attention(self, q, k, v):
"""Attention within local window"""
B, H, T, D = q.size()
out = torch.zeros_like(q)
for i in range(T):
start = max(0, i - self.window_size // 2)
end = min(T, i + self.window_size // 2 + 1)
scores = torch.matmul(q[:, :, i:i+1, :], k[:, :, start:end]) # (B, H, 1, window)
weights = torch.softmax(scores / (D ** 0.5), dim=-1)
out[:, :, i:i+1, :] = torch.matmul(weights, v[:, :, start:end])
return out
def _global_attention(self, q, k, v):
"""Attention for global tokens only"""
global_idx = list(range(self.num_global)) + [T // 2]
q_global = q[:, :, global_idx, :]
k_global = k[:, :, global_idx, :]
v_global = v[:, :, global_idx, :]
scores = torch.matmul(q_global, k_global)
weights = torch.softmax(scores / (D ** 0.5), dim=-1)
return torch.matmul(weights, v_global)
def _random_attention(self, q, k, v):
"""Random sparse attention"""
num_random = 8
random_idx = torch.randint(0, T, (num_random,))
k_random = k[:, :, random_idx, :]
v_random = v[:, :, random_idx, :]
scores = torch.matmul(q, k_random)
weights = torch.softmax(scores / (D ** 0.5), dim=-1)
return torch.matmul(weights, v_random)
4. KV Cache Compression
KV Cache Compression: сжимаем историю
Methods:
1. Quantization: 16bit → 8bit → 4bit
2. Pruning: удаляем неважные пары
3. Summarization: заменяем на summary tokens
4. Eviction: удаляем старые токены
class KVCacheCompressor:
"""Compress KV cache for long contexts"""
def __init__(self, method='quantization', **kwargs):
self.method = method
self.kwargs = kwargs
def compress(self, kv_cache, threshold=0.01):
"""Compress KV cache"""
if self.method == 'quantization':
return self._quantize(kv_cache)
elif self.method == 'pruning':
return self._prune(kv_cache, threshold)
elif self.method == 'summarization':
return self._summarize(kv_cache)
elif self.method == 'eviction':
return self._evict(kv_cache)
def _quantize(self, kv_cache, bits=8):
"""Quantize KV cache to fewer bits"""
def quantize_tensor(t, bits):
scale = t.abs().max() / (2**(bits-1) - 1)
q = torch.round(t / scale)
return q.clamp(-2**(bits-1), 2**(bits-1)-1) * scale
k_q = quantize_tensor(kv_cache['k'], bits)
v_q = quantize_tensor(kv_cache['v'], bits)
return {'k': k_q, 'v': v_q}
def _prune(self, kv_cache, threshold):
"""Remove low-importance KV pairs"""
# Compute importance scores
importance = self._compute_importance(kv_cache)
# Keep top-k
k = kv_cache['k']
v = kv_cache['v']
_, keep_idx = importance.topk(int(threshold * importance.size(-1)))
k_pruned = k.gather(-1, keep_idx.unsqueeze(-1).expand(-1, -1, -1, k.size(-1)))
v_pruned = v.gather(-1, keep_idx.unsqueeze(-1).expand(-1, -1, -1, v.size(-1)))
return {'k': k_pruned, 'v': v_pruned}
def _compute_importance(self, kv_cache):
"""Compute importance of each KV pair"""
k = kv_cache['k'] # (B, H, T, D)
# Importance = mean absolute value across heads/dims
importance = k.abs().mean(dim=-1).mean(dim=1) # (B, T)
return importance
def _summarize(self, kv_cache, summary_len=32):
"""Replace KV pairs with summary tokens"""
k = kv_cache['k']
v = kv_cache['v']
B, H, T, D = k.size()
# Group tokens into windows
window_size = T // summary_len
k_sum = k.view(B, H, summary_len, window_size, D).mean(dim=2)
v_sum = v.view(B, H, summary_len, window_size, D).mean(dim=2)
return {'k': k_sum, 'v': v_sum}
def _evict(self, kv_cache, keep_ratio=0.5):
"""Remove oldest tokens"""
k = kv_cache['k']
v = kv_cache['v']
B, H, T, D = k.size()
keep = int(T * keep_ratio)
# Keep most recent tokens
k_recent = k[:, :, -keep:, :]
v_recent = v[:, :, -keep:, :]
return {'k': k_recent, 'v': v_recent}
Практические советы
Управление контекстом
1. Chunking: разбиваем текст на части
- По абзацам
- По предложениям
- По sliding window
2. Summarization: суммируем старые части
- Summarize → Keep summary + new text
- Recursive summarization
3. Retrieval: ищем релевантные части
- Embed chunks
- Retrieve top-k
- Concatenate for context
4. Streaming: обрабатываем по частям
- Process chunk by chunk
- Maintain hidden state
class ContextManager:
"""Manage context window efficiently"""
def __init__(self, max_tokens=4096, model):
self.max_tokens = max_tokens
self.model = model
self.history = []
self.summary = ""
def add_message(self, message):
"""Add message with context management"""
self.history.append(message)
# Check if context exceeds limit
total_tokens = self._count_tokens(self.history)
if total_tokens > self.max_tokens:
# Compress context
self._compress_context()
def _compress_context(self):
"""Compress context when full"""
# Summarize older messages
half = len(self.history) // 2
old_messages = self.history[:half]
new_messages = self.history[half:]
# Generate summary
summary_prompt = self._build_summary_prompt(old_messages)
summary = self.model.generate(summary_prompt)
# Replace old messages with summary
self.history = [f"[Summary: {summary}]"] + new_messages
def _count_tokens(self, messages):
"""Count total tokens in messages"""
return sum(len(msg.split()) * 1.3 for msg in messages) # rough estimate
Заключение
Context Window — критический параметр LLM:
- Ограничения: O(n²) attention, память, стоимость
- RoPE Scaling: расширяем контекст без переобучения
- Linear Attention: O(n) вместо O(n²)
- Sparse Attention: только важные паттерны
- KV Cache Compression: сжимаем историю
Для LLM:
- Используйте chunking для длинных документов
- Применяйте retrieval для больших баз знаний
- Мониторьте использование токенов