long-context

長文脈：Transformer のコンテキストウィンドウを拡張する

このスキルを使用する場面

次のような場合は、長文脈の手法を使用します。

• 長い文書（32k、64k、128k 以上のトークン）を Transformer モデルで処理する
• 事前学習済みモデル（LLaMA、Mistral など）のコンテキストウィンドウを拡張する
• 効率的な位置エンコーディング（RoPE、ALiBi）を実装する
• 長さの外挿能力を備えたモデルを学習する
• 可変長入力を効率的に処理するモデルをデプロイする
• 既存モデルを、最小限の計算量でより長いコンテキスト向けにファインチューニングする

主な手法：RoPE（回転位置埋め込み）、YaRN、ALiBi（線形バイアス付きアテンション）、位置補間

論文：RoFormer（arXiv 2104.09864）、YaRN（arXiv 2309.00071）、ALiBi（arXiv 2108.12409）、位置補間（arXiv 2306.15595）

インストール

# HuggingFace Transformers (includes RoPE, YaRN support)
pip install transformers torch

# For custom implementations
pip install einops  # Tensor operations
pip install rotary-embedding-torch  # Standalone RoPE

# Optional: FlashAttention for efficiency
pip install flash-attn --no-build-isolation

クイックスタート

RoPE（回転位置埋め込み）

import torch
import 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

# Usage
rope = RotaryEmbedding(dim=64, max_seq_len=8192)
cos, sin = rope(seq_len=2048, device='cuda')

# In attention layer
q_rotated, k_rotated = apply_rotary_pos_emb(query, key, cos, sin)

ALiBi（線形バイアス付きアテンション）

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 attention
num_heads = 8
seq_len = 2048
alibi_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_bias
attn_weights = torch.softmax(attn_scores, dim=-1)

LLaMA 向け位置補間

from transformers import LlamaForCausalLM, LlamaTokenizer

# Original context: 2048 tokens
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# Extend to 32k with position interpolation
# Modify RoPE base frequency
model.config.rope_scaling = {
    "type": "linear",
    "factor": 16.0  # 2048 * 16 = 32768
}

# Or use dynamic scaling
model.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

主な概念

1. RoPE（回転位置埋め込み）

仕組み：

• 回転行列を使って絶対位置をエンコードします
• アテンションに相対位置への依存性を持たせます
• 長さの外挿を可能にします

数式による定式化：

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)

利点：

• 距離に応じてトークン間の依存性が減衰します
• 線形アテンションと互換性があります
• 絶対位置エンコーディングより優れた外挿性能を発揮します

2. YaRN（別の RoPE 拡張手法）

主な革新点：

• NTK を考慮した補間（ニューラルタンジェントカーネル）
• アテンションの温度スケーリング
• 効率的なコンテキスト拡張（ベースラインと比べて必要なトークン数が 10 分の 1）

パラメーター：

# YaRN configuration
yarn_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
}

性能：

• LLaMA を 128k トークンまで拡張します
• ベースラインより学習ステップ数を 2.5 倍削減します
• 最先端のコンテキストウィンドウ拡張を実現します

3. ALiBi（線形バイアス付きアテンション）

基本的な考え方：

• トークンに位置埋め込みを追加しません
• 距離ペナルティをアテンションスコアに直接適用します
• バイアスはキーとクエリの距離に比例します

数式：

attention_bias[i, j] = -m * |i - j|

where m = slope for each attention head

利点：

• 正弦波埋め込みと比べて学習が 11% 高速です
• メモリ使用量を 11% 削減します
• 高い長さの外挿性能を備えます（1k で学習し、2k 以上でテスト）
• 直近の情報を重視する帰納バイアスを持ちます

4. 位置補間

手法：

• 位置インデックスを線形に縮小します
• 学習済みの範囲内で補間します（範囲外への外挿ではありません）
• 必要なファインチューニングは最小限です

数式：

# 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

結果：

• LLaMA 7B～65B を 32k トークンまで拡張します
• 1000 回のファインチューニングステップで十分です
• 外挿と比べて安定性が 600 倍向上します

手法の比較

手法	最大コンテキスト	必要な学習	メモリ	外挿	最適な用途
RoPE	8k-32k	完全な事前学習	中程度	良好	新しいモデル
YaRN	32k-128k	最小限（効率が 10 倍）	中程度	非常に優れている	既存モデルの拡張
ALiBi	無制限	完全な事前学習	低い（-11%）	非常に優れている	ゼロからの学習
位置補間	32k 以上	最小限（1k ステップ）	中程度	低い（仕様上）	手早い拡張

実装パターン

詳細、コード例、パラメーターテーブルについては、references/implementation-patterns.md を参照してください。この手順が該当する場合は、そちらをお読みください。

長文脈向けのファインチューニング

最小限のファインチューニング（位置補間）

from transformers import Trainer, TrainingArguments

# Extend model config
model.config.max_position_embeddings = 32768
model.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 documents
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=long_document_dataset,  # 32k token sequences
)

trainer.train()

YaRN のファインチューニング

# Clone YaRN implementation
git clone https://github.com/jquesnelle/yarn
cd yarn

# Fine-tune LLaMA with YaRN
python 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

ベストプラクティス

1. 適切な手法を選ぶ

# For NEW models (training from scratch)
use_method = "ALiBi"  # Best extrapolation, lowest memory

# For EXTENDING existing RoPE models
use_method = "YaRN"  # Most efficient extension (10× less data)

# For QUICK extension with minimal compute
use_method = "Position Interpolation"  # 1000 steps

# For MODERATE extension with good efficiency
use_method = "Linear RoPE Scaling"  # Built-in, simple

2. スケーリング係数の選択

# 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 → 64k
scaling_factor = 16.0  # 8k → 128k

# Rule: Larger factors need more fine-tuning steps
steps_needed = 100 * scaling_factor  # Rough estimate

3. ファインチューニングデータ

# ✅ Good: Long documents matching target length
train_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. よくある落とし穴を避ける

# ❌ Bad: Applying position interpolation without fine-tuning
model.config.rope_scaling = {"type": "linear", "factor": 16.0}
# Model will perform poorly without fine-tuning!

# ✅ Good: Fine-tune after scaling
model.config.rope_scaling = {"type": "linear", "factor": 16.0}
fine_tune(model, long_documents, steps=1000)

# ❌ Bad: Too aggressive scaling without data
scale_to_1M_tokens()  # Won't work without massive fine-tuning

# ✅ Good: Incremental scaling
# 8k → 16k → 32k → 64k (fine-tune at each step)

本番環境へのデプロイ

長文脈での推論

from transformers import AutoModelForCausalLM, AutoTokenizer

# Load long-context model
model = 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 document
long_text = "..." * 30000  # 30k tokens
inputs = tokenizer(long_text, return_tensors="pt", truncation=False).to('cuda')

# Generate
outputs = model.generate(
    **inputs,
    max_new_tokens=512,
    temperature=0.7,
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)

メモリの最適化

# Use gradient checkpointing for fine-tuning
model.gradient_checkpointing_enable()

# Use Flash Attention 2
model = 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
)

リソース

• RoPE 論文: https://arxiv.org/abs/2104.09864 (RoFormer)
• YaRN 論文: https://arxiv.org/abs/2309.00071
• ALiBi 論文: https://arxiv.org/abs/2108.12409 (短く学習し、長くテストする)
• Position Interpolation: https://arxiv.org/abs/2306.15595
• HuggingFace RoPE Utils: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py
• YaRN の実装: https://github.com/jquesnelle/yarn
• Together AI のブログ: https://www.together.ai/blog/llama-2-7b-32k

関連項目

• references/rope.md - RoPE の実装と理論の詳細
• references/extension_methods.md - YaRN、ALiBi、Position Interpolation の比較
• references/fine_tuning.md - コンテキスト拡張のためのファインチューニング完全ガイド

エージェントの運用手順

1. 環境を確認します。 GPU メモリ、モデルのチェックポイントとライセンス、フレームワークのバージョン、入力形式を確認します。
2. 入力を明確にします。 データまたはユーザーから、形式、識別子、パラメーターを確認します。結果が変わる値は推測せず、質問してください。
3. まず小規模で実行します。 最初に、低解像度または短い長さで単一サンプルの推論を実行します。
4. 上記の手順と参考資料を使って、タスク全体を実行します。
5. 結果を検証します。 出力を定性的に確認し、タスクの指標で評価します。生成コンテンツに対するライセンス上の制限も確認します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の事項を伝えます。

この状況が発生した場合	対応
大きな入力でメモリ不足になる	解像度またはシーケンス長を下げる、半精度を使用する、または入力をタイルに分割します。
インストール済みバージョンに、これらの手順に記載されている関数、フラグ、またはエンドポイントがない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認し、対応を調整してユーザーに伝えます。API を決して捏造しないでください。
必須の入力、識別子、またはパラメーターが曖昧である	ユーザーに質問するか、実行前に仮定を明確に伝えます。

インテグリティに関するルール

• 結果、パラメーター、識別子、引用、統計情報を決して捏造しないでください。実行または検証できない場合は、その旨を明確に伝えてください。
• 商用利用または公開利用の前に、モデルとデータセットのライセンスを確認してください。
• ここに記載されているバージョン固有の詳細は古くなっている可能性があるため、インストール済みバージョンの公式ドキュメントで確認してください。
• 費用が発生する操作、共有 GPU やクラウドのクォータを消費する操作、個人データや患者データに触れる操作、または元に戻せない操作を行う前に、確認を求めてください。