ロングコンテキスト：トランスフォーマーコンテキストウィンドウの拡張

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

以下が必要な場合はロングコンテキスト技術を使用してください：

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

主な技術：RoPE（Rotary Position Embeddings）、YaRN、ALiBi（Attention with Linear Biases）、位置補間

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

インストール

# HuggingFace Transformers (RoPE、YaRN対応を含む)
pip install transformers torch

# カスタム実装用
pip install einops  # テンソル操作
pip install rotary-embedding-torch  # スタンドアロンRoPE

# オプション：効率性のためのFlashAttention
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__()
        # 逆周波数を計算
        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):
        # 位置インデックス
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)

        # 周波数を計算
        freqs = torch.outer(t, self.inv_freq)  # (seq_len, dim/2)

        # sinとcosを計算
        emb = torch.cat((freqs, freqs), dim=-1)  # (seq_len, dim)
        return emb.cos(), emb.sin()

def rotate_half(x):
    """隠れ次元の半分を回転する。"""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin):
    """クエリとキーに回転埋め込みを適用する。"""
    # q, k形状: (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

# 使用方法
rope = RotaryEmbedding(dim=64, max_seq_len=8192)
cos, sin = rope(seq_len=2048, device='cuda')

# 注意層内で
q_rotated, k_rotated = apply_rotary_pos_emb(query, key, cos, sin)

ALiBi（線形バイアス付き注意）

def get_alibi_slopes(num_heads):
    """各注意ヘッドのALiBiスロープ値を取得する。"""
    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:
        # 最も近い2の累乗
        closest_power = 2 ** math.floor(math.log2(num_heads))
        slopes = get_slopes_power_of_2(closest_power)
        # 追加のスロープ
        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):
    """ALiBi注意バイアスを作成する。"""
    # 距離行列
    context_position = torch.arange(seq_len)
    memory_position = torch.arange(seq_len)
    relative_position = memory_position[None, :] - context_position[:, None]

    # スロープを取得
    slopes = torch.tensor(get_alibi_slopes(num_heads))

    # 距離にスロープを適用
    alibi = slopes[:, None, None] * relative_position[None, :, :]
    return alibi  # (num_heads, seq_len, seq_len)

# 注意内で使用
num_heads = 8
seq_len = 2048
alibi_bias = create_alibi_bias(seq_len, num_heads).to('cuda')

# 注意スコアにバイアスを追加
# attn_scores形状: (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

# 元のコンテキスト：2048トークン
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# 位置補間で32kに拡張
# RoPEベース周波数を変更
model.config.rope_scaling = {
    "type": "linear",
    "factor": 16.0  # 2048 * 16 = 32768
}

# またはダイナミックスケーリングを使用
model.config.rope_scaling = {
    "type": "dynamic",
    "factor": 16.0
}

# 長いドキュメントでファインチューニング（最小限のステップが必要）
# この設定変更後、位置補間は標準で機能します

コア概念

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（Yet another RoPE extensioN）

主な革新：

• NTK対応補間（Neural Tangent Kernel）
• 注意温度スケーリング
• 効率的なコンテキスト拡張（ベースラインより10倍少ないトークン）

パラメータ：

# YaRN設定
yarn_config = {
    "scale": 16,                    # 拡張係数
    "original_max_position": 2048,  # ベースコンテキスト
    "extrapolation_factor": 1.0,    # NTKパラメータ
    "attn_factor": 1.0,             # 注意スケーリング
    "beta_fast": 32,                # 高周波スケール
    "beta_slow": 1,                 # 低周波スケール
}

パフォーマンス：

• LLaMAを128kトークンに拡張
• ベースラインより2.5倍少ない学習ステップ
• 最先端のコンテキストウィンドウ拡張

3. ALiBi（線形バイアス付き注意）

コアアイデア：

• トークンに位置埋め込みを追加しない
• キー・クエリ距離の注意スコアに距離ペナルティを直接適用
• バイアスは距離に比例

式：

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

where m = 各注意ヘッドのスロープ

利点：

• 正弦埋め込みと比べて11%高速な学習
• 11%少ないメモリ使用量
• 強い長さ外挿（1kで学習、2k以上でテスト）
• 最新性へのインダクティブバイアス

4. 位置補間

技術：

• 位置インデックスを線形にダウンスケール
• 学習範囲内で補間（範囲外の外挿ではなく）
• 最小限のファインチューニングで対応

式：

# 元：位置インデックス [0, 1, 2, ..., L]
# 拡張：位置インデックス [0, 0.5, 1.0, ..., L/2]
# (2倍拡張の場合)

scaled_position[i] = i / extension_factor

結果：

• LLaMA 7B-65Bを32kトークンに拡張
• 1000ファインチューニングステップで十分
• 外挿より600倍安定性が高い

メソッド比較

メソッド	最大コンテキスト	必要な学習	メモリ	外挿	最適用途
RoPE	8k-32k	完全な事前学習	中程度	良好	新しいモデル
YaRN	32k-128k	最小限（10倍効率的）	中程度	優秀	既存モデルの拡張
ALiBi	無制限	完全な事前学習	低（-11%）	優秀	スクラッチからの学習
位置補間	32k以上	最小限（1kステップ）	中程度	低（設計上）	迅速な拡張

実装パターン

HuggingFace Transformers統合

from transformers import AutoModelForCausalLM, AutoConfig

# YaRNスケーリング付きRoPE
config = 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)

# 位置補間（シンプル）
config.rope_scaling = {
    "type": "linear",
    "factor": 4.0
}

# ダイナミックスケーリング（入力長に基づいて調整）
config.rope_scaling = {
    "type": "dynamic",
    "factor": 8.0
}

カスタムRoPE実装

class LongContextAttention(nn.Module):
    """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投影
        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

        # Q、K、Vに投影
        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)

        # マルチヘッド用に形状変更
        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)

        # RoPEを適用
        cos, sin = self.rotary_emb(seq_len, device=hidden_states.device)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # 標準注意
        attn_output = F.scaled_dot_product_attention(q, k, v)

        # 形状変更と投影
        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

ロングコンテキスト向けファインチューニング

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

from transformers import Trainer, TrainingArguments

# モデル設定を拡張
model.config.max_position_embeddings = 32768
model.config.rope_scaling = {"type": "linear", "factor": 16.0}

# 学習引数（最小限のステップが必要）
training_args = TrainingArguments(
    output_dir="./llama-32k",
    num_train_epochs=1,
    max_steps=1000,           # わずか1000ステップ！
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=2e-5,
    warmup_steps=100,
    logging_steps=10,
    save_steps=500,
)

# 長いドキュメントで学習
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=long_document_dataset,  # 32kトークンシーケンス
)

trainer.train()

YaRNファインチューニング

# YaRN実装をクローン
git clone https://github.com/jquesnelle/yarn
cd yarn

# YaRNでLLaMAをファインチューニング
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. 適切なメソッドを選択する

# 新しいモデルの場合（スクラッチから学習）
use_method = "ALiBi"  # 最高の外挿、最小メモリ

# 既存RoPEモデルの拡張の場合
use_method = "YaRN"  # 最も効率的な拡張（10倍少ないデータ）

# 最小限の計算で迅速に拡張する場合
use_method = "位置補間"  # 1000ステップ

# 中程度の拡張と良好な効率の場合
use_method = "線形RoPEスケーリング"  # 組み込み、シンプル

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

# 保守的（より安全、品質が良い）
scaling_factor = 2.0  # 8k → 16k

# 中程度（良好なバランス）
scaling_factor = 4.0  # 8k → 32k

# 積極的（より多くのファインチューニングが必要）
scaling_factor = 8.0  # 8k → 64k
scaling_factor = 16.0  # 8k → 128k

# ルール：より大きな係数はより多くのファインチューニングステップが必要
steps_needed = 100 * scaling_factor  # 大まかな推定値

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

# ✅ 良好：ターゲット長に合わせた長いドキュメント
train_data = [
    {"text": long_doc_32k_tokens},  # フル32k
    {"text": long_doc_24k_tokens},  # 様々な長さ
    {"text": long_doc_16k_tokens},
]

# ❌ 不良：短いドキュメント（長いコンテキストを学習しない）
train_data = [
    {"text": short_doc_2k_tokens},
]

# 次のようなデータセットを使用：
# - PG-19（本、長いテキスト）
# - arXiv論文
# - 長編の会話
# - GitHubリポジトリ（ファイルを連結）

4. 一般的な落とし穴を避ける

# ❌ 不良：ファインチューニング無しで位置補間を適用
model.config.rope_scaling = {"type": "linear", "factor": 16.0}
# ファインチューニング無しではモデルのパフォーマンスが低い！

# ✅ 良好：スケーリング後にファインチューニング
model.config.rope_scaling = {"type": "linear", "factor": 16.0}
fine_tune(model, long_documents, steps=1000)

# ❌ 不良：データなしで過度なスケーリング
scale_to_1M_tokens()  # 大規模なファインチューニングなしでは機能しない

# ✅ 良好：段階的なスケーリング
# 8k → 16k → 32k → 64k（各ステップでファインチューニング）

本番環境へのデプロイ

ロングコンテキストでの推論

from transformers import AutoModelForCausalLM, AutoTokenizer

# ロングコンテキストモデルをロード
model = AutoModelForCausalLM.from_pretrained(
    "togethercomputer/LLaMA-2-7B-32K",  # 32kコンテキスト
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("togethercomputer/LLaMA-2-7B-32K")

# 長いドキュメントを処理
long_text = "..." * 30000  # 30kトークン
inputs = tokenizer(long_text, return_tensors="pt", truncation=False).to('cuda')

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

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

メモリ最適化

# ファインチューニング用の勾配チェックポイントを使用
model.gradient_checkpointing_enable()

# Flash Attention 2を使用
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    attn_implementation="flash_attention_2",  # 2～3倍高速
    torch_dtype=torch.float16
)

# ページド注意を使用（vLLM）
from vllm import LLM

llm = LLM(
    model="togethercomputer/LLaMA-2-7B-32K",
    max_model_len=32768,  # 32kコンテキスト
    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（Train Short, Test Long）
• 位置補間：https://arxiv.org/abs/2306.15595
• HuggingFace RoPEユーティリティ：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、位置補間の比較
• references/fine_tuning.md - コンテキスト拡張のための完全なファインチューニングガイド