Long Context

orchestra-research/ai-research-skills/19-emerging-techniques/long-context

作者 orchestra-research773a52944ba4MIT13K 个星标收录于 2026年10月8日更新于 2026年10月8日仓库3个月前更新

Extend context windows of transformer models using RoPE, YaRN, ALiBi, and position interpolation techniques. Use when processing long documents (32k-128k+ tokens), extending pre-trained models beyond original context limits, or implementing efficient positional encodings. Covers rotary embeddings, attention biases, interpolation methods, and extrapolation strategies for LLMs.

仅含说明AI & Agents
AI 生成的概览

指导使用 RoPE、YaRN、ALiBi 和位置插值扩展 Transformer 上下文窗口，涵盖微调与部署。

功能
该技能提供扩展 Transformer 语言模型上下文窗口的参考指导与代码示例。内容涵盖旋转位置编码（RoPE）、YaRN、ALiBi 和位置插值，并给出 HuggingFace Transformers 配置片段、自定义注意力实现、微调流程以及推理与显存优化说明。此外还包含方法对比表、缩放系数建议和常见误区。
适用场景
适用于处理约 32k 至 128k 以上 token 的长文档、将预训练模型扩展到超出原始上下文限制，或实现高效的位置编码。也面向需要微调或部署可处理变长长输入模型的场景。
运行要求
需要 Python 以及 transformers 和 torch，可选安装 einops、rotary-embedding-torch、flash-attn 等包。部分示例涉及从 HuggingFace 下载模型并克隆外部 YaRN 仓库，因此可能需要网络访问和模型访问权限。该技能不附带脚本，仅为说明文档与参考资料。

Long Context: Extending Transformer Context Windows

When to Use This Skill

Use Long Context techniques when you need to:

  • Process long documents (32k, 64k, 128k+ tokens) with transformer models
  • Extend context windows of pre-trained models (LLaMA, Mistral, etc.)
  • Implement efficient positional encodings (RoPE, ALiBi)
  • Train models with length extrapolation capabilities
  • Deploy models that handle variable-length inputs efficiently
  • Fine-tune existing models for longer contexts with minimal compute

Key Techniques: RoPE (Rotary Position Embeddings), YaRN, ALiBi (Attention with Linear Biases), Position Interpolation

Papers: RoFormer (arXiv 2104.09864), YaRN (arXiv 2309.00071), ALiBi (arXiv 2108.12409), Position Interpolation (arXiv 2306.15595)

Installation

bash
# HuggingFace Transformers (includes RoPE, YaRN support)pip install transformers torch
# For custom implementationspip install einops  # Tensor operationspip install rotary-embedding-torch  # Standalone RoPE
# Optional: FlashAttention for efficiencypip install flash-attn --no-build-isolation

Quick Start

RoPE (Rotary Position Embeddings)

python
import torchimport torch.nn as nn
class RotaryEmbedding(nn.Module):    """Rotary Position Embeddings (RoPE)."""
    def __init__(self, dim, max_seq_len=8192, base=10000):        super().__init__()        # Compute inverse frequencies        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))        self.register_buffer("inv_freq", inv_freq)        self.max_seq_len = max_seq_len
    def forward(self, seq_len, device):        # Position indices        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
        # Compute frequencies        freqs = torch.outer(t, self.inv_freq)  # (seq_len, dim/2)
        # Compute sin and cos        emb = torch.cat((freqs, freqs), dim=-1)  # (seq_len, dim)        return emb.cos(), emb.sin()
def rotate_half(x):    """Rotate half the hidden dimensions."""    x1, x2 = x.chunk(2, dim=-1)    return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q, k, cos, sin):    """Apply rotary embeddings to queries and keys."""    # q, k shape: (batch, heads, seq_len, dim)    q_embed = (q * cos) + (rotate_half(q) * sin)    k_embed = (k * cos) + (rotate_half(k) * sin)    return q_embed, k_embed
# Usagerope = RotaryEmbedding(dim=64, max_seq_len=8192)cos, sin = rope(seq_len=2048, device='cuda')
# In attention layerq_rotated, k_rotated = apply_rotary_pos_emb(query, key, cos, sin)

ALiBi (Attention with Linear Biases)

python
def get_alibi_slopes(num_heads):    """Get ALiBi slope values for each attention head."""    def get_slopes_power_of_2(n):        start = 2 ** (-(2 ** -(math.log2(n) - 3)))        ratio = start        return [start * (ratio ** i) for i in range(n)]
    if math.log2(num_heads).is_integer():        return get_slopes_power_of_2(num_heads)    else:        # Closest power of 2        closest_power = 2 ** math.floor(math.log2(num_heads))        slopes = get_slopes_power_of_2(closest_power)        # Add extra slopes        extra = get_slopes_power_of_2(2 * closest_power)        slopes.extend(extra[0::2][:num_heads - closest_power])        return slopes
def create_alibi_bias(seq_len, num_heads):    """Create ALiBi attention bias."""    # Distance matrix    context_position = torch.arange(seq_len)    memory_position = torch.arange(seq_len)    relative_position = memory_position[None, :] - context_position[:, None]
    # Get slopes    slopes = torch.tensor(get_alibi_slopes(num_heads))
    # Apply slopes to distances    alibi = slopes[:, None, None] * relative_position[None, :, :]    return alibi  # (num_heads, seq_len, seq_len)
# Usage in attentionnum_heads = 8seq_len = 2048alibi_bias = create_alibi_bias(seq_len, num_heads).to('cuda')
# Add bias to attention scores# attn_scores shape: (batch, num_heads, seq_len, seq_len)attn_scores = attn_scores + alibi_biasattn_weights = torch.softmax(attn_scores, dim=-1)

Position Interpolation for LLaMA

python
from transformers import LlamaForCausalLM, LlamaTokenizer
# Original context: 2048 tokensmodel = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
# Extend to 32k with position interpolation# Modify RoPE base frequencymodel.config.rope_scaling = {    "type": "linear",    "factor": 16.0  # 2048 * 16 = 32768}
# Or use dynamic scalingmodel.config.rope_scaling = {    "type": "dynamic",    "factor": 16.0}
# Fine-tune with long documents (minimal steps needed)# Position interpolation works out-of-the-box after this config change

Core Concepts

1. RoPE (Rotary Position Embeddings)

How it works:

  • Encodes absolute position via rotation matrix
  • Provides relative position dependency in attention
  • Enables length extrapolation

Mathematical formulation:

q_m = (W_q * x_m) * e^(imθ)k_n = (W_k * x_n) * e^(inθ)
where θ_j = base^(-2j/d) for j ∈ [0, d/2)

Advantages:

  • Decaying inter-token dependency with distance
  • Compatible with linear attention
  • Better extrapolation than absolute position encodings

2. YaRN (Yet another RoPE extensioN)

Key innovation:

  • NTK-aware interpolation (Neural Tangent Kernel)
  • Attention temperature scaling
  • Efficient context extension (10× less tokens vs baselines)

Parameters:

python
# YaRN configurationyarn_config = {    "scale": 16,                    # Extension factor    "original_max_position": 2048,  # Base context    "extrapolation_factor": 1.0,    # NTK parameter    "attn_factor": 1.0,             # Attention scaling    "beta_fast": 32,                # High-frequency scale    "beta_slow": 1,                 # Low-frequency scale}

Performance:

  • Extends LLaMA to 128k tokens
  • 2.5× less training steps than baselines
  • State-of-the-art context window extension

3. ALiBi (Attention with Linear Biases)

Core idea:

  • No positional embeddings added to tokens
  • Apply distance penalty directly to attention scores
  • Bias proportional to key-query distance

Formula:

attention_bias[i, j] = -m * |i - j|
where m = slope for each attention head

Advantages:

  • 11% faster training vs sinusoidal embeddings
  • 11% less memory usage
  • Strong length extrapolation (train 1k, test 2k+)
  • Inductive bias towards recency

4. Position Interpolation

Technique:

  • Linearly down-scale position indices
  • Interpolate within trained range (vs extrapolate beyond)
  • Minimal fine-tuning required

Formula:

# Original: position indices [0, 1, 2, ..., L]# Extended: position indices [0, 0.5, 1.0, ..., L/2]# (for 2× extension)
scaled_position[i] = i / extension_factor

Results:

  • LLaMA 7B-65B extended to 32k tokens
  • 1000 fine-tuning steps sufficient
  • 600× better stability than extrapolation

Method Comparison

MethodMax ContextTraining NeededMemoryExtrapolationBest For
RoPE8k-32kFull pre-trainingModerateGoodNew models
YaRN32k-128kMinimal (10× efficient)ModerateExcellentExtending existing models
ALiBiUnlimitedFull pre-trainingLow (-11%)ExcellentTraining from scratch
Position Interpolation32k+Minimal (1k steps)ModeratePoor (by design)Quick extension

Implementation Patterns

HuggingFace Transformers Integration

python
from transformers import AutoModelForCausalLM, AutoConfig
# RoPE with YaRN scalingconfig = AutoConfig.from_pretrained("mistralai/Mistral-7B-v0.1")config.rope_scaling = {    "type": "yarn",    "factor": 8.0,    "original_max_position_embeddings": 8192,    "attention_factor": 1.0}
model = AutoModelForCausalLM.from_config(config)
# Position interpolation (simpler)config.rope_scaling = {    "type": "linear",    "factor": 4.0}
# Dynamic scaling (adjusts based on input length)config.rope_scaling = {    "type": "dynamic",    "factor": 8.0}

Custom RoPE Implementation

python
class LongContextAttention(nn.Module):    """Multi-head attention with RoPE."""
    def __init__(self, hidden_size, num_heads, max_seq_len=32768):        super().__init__()        self.num_heads = num_heads        self.head_dim = hidden_size // num_heads
        # Q, K, V projections        self.q_proj = nn.Linear(hidden_size, hidden_size)        self.k_proj = nn.Linear(hidden_size, hidden_size)        self.v_proj = nn.Linear(hidden_size, hidden_size)        self.o_proj = nn.Linear(hidden_size, hidden_size)
        # RoPE        self.rotary_emb = RotaryEmbedding(            dim=self.head_dim,            max_seq_len=max_seq_len        )
    def forward(self, hidden_states):        batch_size, seq_len, _ = hidden_states.shape
        # Project to Q, K, V        q = self.q_proj(hidden_states)        k = self.k_proj(hidden_states)        v = self.v_proj(hidden_states)
        # Reshape for multi-head        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        # Apply RoPE        cos, sin = self.rotary_emb(seq_len, device=hidden_states.device)        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        # Standard attention        attn_output = F.scaled_dot_product_attention(q, k, v)
        # Reshape and project        attn_output = attn_output.transpose(1, 2).contiguous()        attn_output = attn_output.view(batch_size, seq_len, -1)        output = self.o_proj(attn_output)
        return output

Fine-tuning for Long Context

Minimal Fine-tuning (Position Interpolation)

python
from transformers import Trainer, TrainingArguments
# Extend model configmodel.config.max_position_embeddings = 32768model.config.rope_scaling = {"type": "linear", "factor": 16.0}
# Training args (minimal steps needed)training_args = TrainingArguments(    output_dir="./llama-32k",    num_train_epochs=1,    max_steps=1000,           # Only 1000 steps!    per_device_train_batch_size=1,    gradient_accumulation_steps=16,    learning_rate=2e-5,    warmup_steps=100,    logging_steps=10,    save_steps=500,)
# Train on long documentstrainer = Trainer(    model=model,    args=training_args,    train_dataset=long_document_dataset,  # 32k token sequences)
trainer.train()

YaRN Fine-tuning

bash
# Clone YaRN implementationgit clone https://github.com/jquesnelle/yarncd yarn
# Fine-tune LLaMA with YaRNpython scripts/train.py \    --model meta-llama/Llama-2-7b-hf \    --scale 16 \    --rope_theta 10000 \    --max_length 32768 \    --batch_size 1 \    --gradient_accumulation 16 \    --steps 400 \    --learning_rate 2e-5

Best Practices

1. Choose the Right Method

python
# For NEW models (training from scratch)use_method = "ALiBi"  # Best extrapolation, lowest memory
# For EXTENDING existing RoPE modelsuse_method = "YaRN"  # Most efficient extension (10× less data)
# For QUICK extension with minimal computeuse_method = "Position Interpolation"  # 1000 steps
# For MODERATE extension with good efficiencyuse_method = "Linear RoPE Scaling"  # Built-in, simple

2. Scaling Factor Selection

python
# Conservative (safer, better quality)scaling_factor = 2.0  # 8k → 16k
# Moderate (good balance)scaling_factor = 4.0  # 8k → 32k
# Aggressive (requires more fine-tuning)scaling_factor = 8.0  # 8k → 64kscaling_factor = 16.0  # 8k → 128k
# Rule: Larger factors need more fine-tuning stepssteps_needed = 100 * scaling_factor  # Rough estimate

3. Fine-tuning Data

python
# ✅ Good: Long documents matching target lengthtrain_data = [    {"text": long_doc_32k_tokens},  # Full 32k    {"text": long_doc_24k_tokens},  # Varied lengths    {"text": long_doc_16k_tokens},]
# ❌ Bad: Short documents (won't learn long context)train_data = [    {"text": short_doc_2k_tokens},]
# Use datasets like:# - PG-19 (books, long texts)# - arXiv papers# - Long-form conversations# - GitHub repositories (concatenated files)

4. Avoid Common Pitfalls

python
# ❌ Bad: Applying position interpolation without fine-tuningmodel.config.rope_scaling = {"type": "linear", "factor": 16.0}# Model will perform poorly without fine-tuning!
# ✅ Good: Fine-tune after scalingmodel.config.rope_scaling = {"type": "linear", "factor": 16.0}fine_tune(model, long_documents, steps=1000)
# ❌ Bad: Too aggressive scaling without datascale_to_1M_tokens()  # Won't work without massive fine-tuning
# ✅ Good: Incremental scaling# 8k → 16k → 32k → 64k (fine-tune at each step)

Production Deployment

Inference with Long Context

python
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load long-context modelmodel = AutoModelForCausalLM.from_pretrained(    "togethercomputer/LLaMA-2-7B-32K",  # 32k context    torch_dtype=torch.float16,    device_map="auto")tokenizer = AutoTokenizer.from_pretrained("togethercomputer/LLaMA-2-7B-32K")
# Process long documentlong_text = "..." * 30000  # 30k tokensinputs = tokenizer(long_text, return_tensors="pt", truncation=False).to('cuda')
# Generateoutputs = model.generate(    **inputs,    max_new_tokens=512,    temperature=0.7,)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)

Memory Optimization

python
# Use gradient checkpointing for fine-tuningmodel.gradient_checkpointing_enable()
# Use Flash Attention 2model = AutoModelForCausalLM.from_pretrained(    "meta-llama/Llama-2-7b-hf",    attn_implementation="flash_attention_2",  # 2-3× faster    torch_dtype=torch.float16)
# Use paged attention (vLLM)from vllm import LLM
llm = LLM(    model="togethercomputer/LLaMA-2-7B-32K",    max_model_len=32768,  # 32k context    gpu_memory_utilization=0.9)

Resources

See Also

  • references/rope.md - Detailed RoPE implementation and theory
  • references/extension_methods.md - YaRN, ALiBi, Position Interpolation comparisons
  • references/fine_tuning.md - Complete fine-tuning guide for context extension

来源与署名

来源：orchestra-research/ai-research-skills位于19-emerging-techniques/long-context提交773a529

许可证： MIT

内容归原作者所有。SourceWeft 从公开仓库中收录这些内容。

举报或申请下架

更多来自 orchestra-research/ai-research-skills 的技能

Speculative Decoding

orchestra-research

指导使用推测解码、Medusa 多头预测和前瞻解码来加速大语言模型推理。

AI & Agents13K3个月前更新

Systems Paper Writing

orchestra-research

为 OSDI、SOSP、ASPLOS、NSDI 等系统会议论文写作提供段落级结构与篇幅分配指导。

Writing & Content13K3个月前更新

Moe Training

orchestra-research

指导使用 DeepSpeed 或 HuggingFace 训练与推理混合专家模型，涵盖路由、负载均衡与专家并行。

AI & Agents13K3个月前更新

Presenting Conference Talks

orchestra-research

Generates conference presentation slides (Beamer LaTeX PDF and editable PPTX) from a compiled paper with speaker notes and talk script. Use when preparing oral talks, spotlight presentations, or invited talks for ML and systems conferences.

待分类13K3个月前更新

Model Pruning

orchestra-research

指导使用 Wanda、SparseGPT、幅度剪枝和 N:M 方法压缩 LLM，减小模型体积并加速推理。

AI & Agents13K3个月前更新

Outlines

orchestra-research

指导使用 Outlines 库，通过本地模型进行受结构约束的文本生成。

AI & Agents13K3个月前更新

更多AI & Agents技能

Ppt Template Creator

anthropics

精选

将用户的 PowerPoint 模板转化为可复用技能，用于生成品牌演示文稿。

AI & Agents39K2周前更新

Skill Development

anthropics

精选

指导创建 Claude Code 插件技能，涵盖结构、描述、渐进式披露与验证。

AI & Agents37K今天更新

Plugin Structure

anthropics

精选

指导 Claude Code 插件的结构、清单与组件布局。

AI & Agents37K今天更新

Command Development

anthropics

精选

指导创建 Claude Code 斜杠命令，涵盖结构、YAML frontmatter、参数与插件功能。

AI & Agents37K今天更新

Claude Md Improver

anthropics

精选

审查仓库中的 CLAUDE.md 文件，评估其质量，并在获得批准后应用有针对性的改进。

AI & Agents37K今天更新

Find Skills

vercel-labs

精选

指导通过 npx skills 命令行工具发现、评估并安装智能体技能。

AI & Agents33K今天更新