ABEJAでデータサイエンティストをしている藤原です。

今回も LLM のロングコンテキスト言語処理（Long-context language modeling; LCLM）に関連するブログになります。前回のブログはこちらです。

様々なコンテキスト長における LLM の Self-Attention の Query と Key の分析 - ABEJA Tech Blog

近年の LLM ではオープン・クローズド問わずより長大なコンテキストを正確に扱えるモデルが増えてきています。このようにロングコンテキストLLMが増えてくると、次はどのようにしてモデルのLCLM性能をどう評価すべきかが課題になります。しかしながら、たとえばモデルの生成結果を実際に自分で読んで確かめようと思っても、128kトークンもある入出力のテキストに目を通すのはかなり時間がかかってしまいます。

そこで今回は LCLM 評価をより手軽に行う方法がないか？ということで、モデルに入力する位置インデックス情報を操作してテキスト生成を行ってみようと思います。特に、位置インデックスを操作して短いテキストを LLM にだけ長文と錯覚させることで、現実的に人間が読めるテキストの長さでモデルの LCLM 性能を評価できるか？という観点で検証してみます。

1. はじめに
- 1.1 ロングコンテキストLLMの評価
- 1.2 位置インデックスを操作する
2. 実験内容
3. 実験結果
- 3.1 Qwen3 シリーズの内での比較
- 3.2 Qwen2.5 の派生モデル間での比較
4. まとめ
参考文献
We are hiring!

1. はじめに

1.1 ロングコンテキストLLMの評価

ロングコンテキストLLMの評価には、長大な入力を正確に処理する能力（Long context understanding などと呼ばれる）の評価と、長大な出力を一貫して生成する能力（Long form generation などと呼ばれる）の評価を行うものがあります。モデル開発の Technical Reports などで LCLM 性能の評価結果として記載されるのは、長大な入力を正確に処理する能力を評価するベンチマーク（RULER[1] や LongBench v2[2] 、MRCR[3] / OpenAI-MRCR [4] など）のスコアが多い印象です。本記事の対象も長大な入力を正確に処理する能力の評価になります。

1.2 位置インデックスを操作する

本記事では位置インデックスを操作して擬似的な長文入力を行った場合のモデルの生成結果を分析します。位置インデックスを操作するというと Position Interpolation や YaRN など手法を想像される方が多いかもしれません。これらは学習時の位置インデックスをスケーリングやリマッピングすることで、実際よりも長い入力をあたかも学習済みの範囲内に収めて処理するアプローチです。したがって「より長いコンテキストを扱えるようにする」という点が主眼になります。

一方で、今回検証するのは PoSE: Efficient Context Window Extension of LLMs via Positional Skip-wise Training で提案されたアプローチになります。この手法は入力テキスト系列そのものは短いままで、位置インデックスの値のリストにスキップを加えて「遠い位置関係を擬似的に体験させる」ことを目的としています。具体的には $indices = [0, 1, 2, 3, 4, 5$ ] という位置インデックスのリストを $indices = [0, 1, 2, 101, 102, 103$ ] のようにして、長距離の位置関係を擬似的に作り出します。 PoSE はモデルの学習時にこの方法を採用することで短いテキストでもロングコンテキストの学習ができるというものです。本記事では、学習に効果があるかを検証する前段として、そもそも既存のモデルでそのようなテキスト生成を実行した場合にどのような挙動をするのか検証してみようと思います。

2. 実験内容

ここからが本記事での検証内容です。

2.1 検証内容のイメージ

トークンIDの系列だけでは Attention ではトークン同士の位置関係を考慮できないため、入力トークンの並び順を数列（位置インデックス）で表現してそれを位置埋め込みで表現することで、位置関係を考慮した Attention 計算ができるようになっています。たとえば、分割位置は適当ですが、

トークン: ["こん", "にち", "は"]
位置    : [1, 2, 3]

という風に、3つのトークンを入力する場合、これらのトークンの ID だけでなく、その順番を表す [1, 2, 3] を一緒に LLM に入力することで、続きのテキストを生成することができます。 PoSE ではこの数値の配列を操作します。具体的には、この数字の列に手を加えて、たとえば [[”こん”, “にち”, “は”], [1, 10, 20]] のように数値の間隔を広げてみると、LLMからすれば「こん ……… にち ………. は」のように少し距離が空いたテキストのように見えることになります。

PoSE ではこれを学習時に適用することで、短いシーケンス長でモデルを学習しながらもロングコンテキスト対応の学習が可能である、という実験結果が得られています。今回の検証では、モデルの学習ではなく推論時にこれを適用することで、モデルの推論結果にどのような影響があるか分析します。

2.2 使用するモデル

使用するモデルは以下の 9 つです。

Qwen3シリーズ
- Qwen3-0.6B
- Qwen3-1.7B
- Qwen3-4B
- Qwen3-8B
- Qwen3-14B
- Qwen3-32B
- Qwen3-235B-A22B
Qwen/Qwen2.5-7B-Instruct
Qwen/Qwen2.5-7B-Instruct-1M
abeja/ABEJA-Qwen2.5-7b-Japanese-v0.1
abeja/ABEJA-Qwen2.5-32b-Japanese-v1.0

以下に RULER ベンチマークでの評価結果を記載します。独自にベンチマーク評価した結果ではなく、Qwen2.5-1M[5] と Qwen3[6] のレポートに記載されているスコアです。Qwen3-0.6B と Qwen3-1.7B はスコアの記載がありませんでした。 ABEJA-Qwen2.5-7b-Japanese-v0.1 についてもベンチマークの全てのタスクでの評価を行っていないため正式なスコアは算出できていませんが、32k 以上のロングコンテキストでは Qwen2.5-7B-Instruct と比較して数% ~ 10% 程度劣る印象です。

モデル	コンテキスト長（モデルカード）	4K	8K	16K	32K	64K	128K	Avg.
Qwen3-235B-A22B	128K	97.7	97.2	96.4	95.1	93.3	90.6	95.0
Qwen3-30B-A3B	128K	96.5	97.0	95.3	92.4	89.1	79.2	91.6
Qwen3-32B	128K	98.4	96.0	96.2	94.4	91.8	85.6	93.7
Qwen3-14B	128K	98.0	97.8	96.4	96.1	94.0	85.1	94.6
Qwen3-8B	128K	96.3	96.0	91.8	91.2	82.1	77.4	89.1
Qwen3-4B	128K	95.1	93.6	91.0	87.8	77.8	66.0	85.2
Qwen2.5-7B-Instruct-1M	1M	96.8	95.3	93.0	91.1	90.4	84.4	91.8
Qwen2.5-7B-Instruct	1M	96.7	95.1	93.7	89.4	82.3	55.1	85.4

2.3 PoSE 的な位置インデックス操作の実装

各モデルに対して以下のパターンで生成結果を取得します。

通常の位置インデックス（HuggingFace Transformers が自動的に生成するもの）を使った生成
PoSE（人工的にギャップを挿入したposition_ids）を使った生成
- 質問ごとの位置を特定し、プロンプトを3つのチャンクに分割
  1. Q1の終わりまで
  2. Q2の初めからQ2の終わりまで
  3. Q3の始まり以降
- チャンク2, チャンク3 の位置インデックスの値を大きくする
  - 例えば、各チャンクの位置インデックスが [1, 2, 3] [4, 5, 6] [7, 8, 9] だとして、コンテキスト長を 1000 まで拡張させる場合は [1, 2, 3] [499, 500, 501] [997, 998, 999] といった操作を行う。
  - この 1000 の部分を変数として、0, 32, 1k, 4k, 16k, 32k, 64k, 128k, 256k, 512k, 1024k で変化させる。 0 の場合は拡張させないのでパターン1の生成結果と概ね同じになるはずである。（実装が正しいか確認する目的）

実装は以下のようになっています。この実装では自作のループ処理で回帰的なトークン生成を実行しています。 generate メソッドの実装にバグがあり、引数として position_ids を受け取る形式にはなっているものの、引数で与えても内部で自動的に連番の index を作成して使用してしまうため、今回は1トークンずつ生成するループを自前で動かしています。この辺りの issues （17283, #29149）で議論されていますが、優先度が高くないため今の所修正の予定はなさそうです。

▶︎ pose_inference.py

import gc
from pathlib import Path
import yaml

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ---------------- 0. モデル準備 ----------------
config_filepath = Path("./config.yml")
with config_filepath.open("r", encoding="utf-8") as f:
    config_dict = yaml.safe_load(f)

model_dirpath = Path(config_dict["model_dirpath"])
output_dirpath = Path(config_dict["output_dirpath"])
output_dirpath.mkdir(parents=True, exist_ok=True)
for model_name in config_dict["models"]:
    torch.cuda.empty_cache()
    gc.collect()

    print(model_name)
    model_path = model_dirpath / model_name
    tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(str(model_path), device_map="auto").eval()

    output_filepath = output_dirpath / f"{model_name}.txt"

    # ---------------- 1. 4 つの user 質問＋assistant 開始 ----------------
    enable_thinking = False
    max_new_tokens = 128
    messages = [
        {"role": "user", "content": p} for p in [
            "以下の質問に対して、それぞれ日本語で回答してください。\n",
            "  Q1. 日本の首都はどこですか？\n",
            "  Q2. 水は何でできていますか？\n",
            "  Q3. 富士山について教えてください。\n",
        ]
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_thinking,
        return_tensors="pt",
    )
    model_inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    input_ids = model_inputs["input_ids"]
    attention_mask = model_inputs["attention_mask"]

    # ==================== 3. 生成テスト ====================
    normal_out = model.generate(
        **model_inputs,  # input_ids + attention_mask
        position_ids=None,
        max_new_tokens=max_new_tokens,
    )
    with output_filepath.open("w", encoding="utf-8") as f:
        f.write("=== auto position ===\n")
        f.write(f"{tokenizer.decode(normal_out[0], skip_special_tokens=True)}\n\n")

    # ==================== 4. PoSE 検証 ====================
    # トークン列を復元
    decoded_text = tokenizer.decode(input_ids[0])

    # "Q2" の文字列位置（文字インデックス）
    char_idx_q2 = decoded_text.index("Q2.")
    char_idx_q3 = decoded_text.index("Q3.")

    # トークナイズされた各トークンの開始位置を取得
    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
    offsets = tokenizer(decoded_text, return_offsets_mapping=True)["offset_mapping"]

    # charインデックスに対応するトークンインデックスを見つける
    idx_q2 = next(i for i, (start, end) in enumerate(offsets) if start <= char_idx_q2 < end)
    idx_q3 = next(i for i, (start, end) in enumerate(offsets) if start <= char_idx_q3 < end)

    boundaries = [0, idx_q2, idx_q3, input_ids.size(1)]
    chunk_lengths = [boundaries[i + 1] - boundaries[i] for i in range(len(boundaries) - 1)]

    def make_pose_pos_ids_with_gaps(chunk_lengths: list[int], gaps: list[int]) -> torch.Tensor:
        """
        lengths: 各セクションの長さ（例: [20, 10, 10]）
        gaps: 各セクション間のギャップ（例: [500, 500]）→ len(gaps) = len(lengths) - 1
        """
        assert len(gaps) == len(chunk_lengths) - 1
        cursor = 0
        buf = []
        for i, length in enumerate(chunk_lengths):
            buf.append(torch.arange(length).to(model.device) + cursor)
            cursor += length
            if i < len(gaps):
                cursor += gaps[i]
        return torch.cat(buf).unsqueeze(0)

    # ---------------- PoSE 検証ループ ----------------
    eos_id = tokenizer.eos_token_id

    for extend_length in [0, 32, 1_000, 4_000, 16_000, 32_000, 64_000, 128_000, 256_000, 512_000, 1_024_000]:
        generated = input_ids.clone().to(model.device)
        attn_mask = attention_mask.clone().to(model.device)

        gap_total = max(extend_length - input_ids.size(1), 0)
        gap_each = gap_total // (len(chunk_lengths) - 1)
        gaps = [gap_each] * (len(chunk_lengths) - 1)
        current_pos_ids = make_pose_pos_ids_with_gaps(chunk_lengths=chunk_lengths, gaps=gaps).to(model.device)

        with torch.no_grad():
            for _ in range(max_new_tokens):
                outputs = model(input_ids=generated, attention_mask=attn_mask, position_ids=current_pos_ids)
                next_token = outputs.logits[:, -1].argmax(dim=-1, keepdim=True)

                if next_token.item() == eos_id:
                    break

                generated = torch.cat([generated, next_token], dim=-1)
                attn_mask = torch.cat([attn_mask, torch.ones_like(next_token)], dim=-1)

                # position_ids の最後に +1 を追加して伸ばす
                next_pos_id = current_pos_ids[0, -1] + 1
                current_pos_ids = torch.cat([current_pos_ids, next_pos_id.unsqueeze(0).unsqueeze(0).to(model.device)], dim=1)

        text = tokenizer.decode(generated[0], skip_special_tokens=True)
        with output_filepath.open("a", encoding="utf-8") as f:
            f.write(f"=== max length: {extend_length} ===\n")
            f.write(f"position_ids: {list(current_pos_ids)}\n")
            f.write(f"{text}\n")

    del model, tokenizer, model_inputs, input_ids, attention_mask
    torch.cuda.empty_cache()
    gc.collect()

2.4 実験条件

HuggingFace 上で公開されているモデルに対して PoSE 的な位置インデックス操作を施した入力を与え、その生成挙動を観察しました。入力プロンプトは事前に定義した以下の3つの質問（Q1〜Q3）を結合してチャットテンプレートを適用したものになります。

Q1. 日本の首都はどこですか？
Q2. 水は何でできていますか？
Q3. 富士山について教えてください。

生成トークン数：max 128 に限定し、デフォルトのコンテキスト長が32k以下のモデルは config.json で YaRN を設定し、コンテキスト長を 128k に拡張しています。Qwen3 の 0.6B, 1.7B はモデルカードで 32k まで対応となっているので、 YaRN を適用するのは本来非推奨だと思いますが、今回はこれらのモデルでもコンテキスト拡張を行いました。

3. 実験結果

3.1 Qwen3 シリーズの内での比較

生成結果を全て表示すると長くなり過ぎてしまうため、先に結果と分析について記載し、最後に出力結果を折り畳んで添える形にします。まずは、各モデルの生成結果は以下のような傾向にありました。Q1, 2, 3 の順に回答を行っており、回答内容も質問に対して概ね適切など、生成結果に問題がなかった部分は「-」を記入しています。また、Q1,Q2を飛ばしてQ3から応答を開始するケースが頻出したため、そのような生成結果の部分は「Skip-to-Q3」と記載しています。

モデル名	\~4k	16k	32k	64k	128k	256k	512k	1024k (1M)
Qwen3-0.6B	-	-	-	-	-	-	Skip-to-Q3	意味不明な記号の羅列
Qwen3-1.7B	-	-	-	-	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	謎の数字の羅列
Qwen3-4B	-	-	-	-	-	Skip-to-Q3	Skip-to-Q3	意味不明な記号の羅列
Qwen3-8B	-	-	-	-	-	-	Q3質問の繰り返しや `<think>` のみ生成	Q3質問の繰り返しや `<think>` のみ生成
Qwen3-14B	-	-	-	-	-	Skip-to-Q3	Skip-to-Q3	「東」の後に謎の数字の羅列
Qwen3-32B	-	-	-	-	-	Skip-to-Q3	Skip-to-Q3	「です」を繰り返すのみ
Qwen3-235B-A22B	-	-	-	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	日本語は出るがプロンプトを忘れて定型文→途中で日本語崩壊

モデル間で比較すると、モデルサイズの大小と耐えられる擬似コンテキスト長の大きさには多少の相関があるが、小さいモデルでも大きな擬似コンテキスト長で正常に動くことがありました。また Qwen3-0.6B, Qwen3-1.7B は推奨設定では 32k までしかサポートされていないが、より大きなコンテキスト長まで正常に動くケースがありました。パラメータ数が最も多い Qwen3-235B-A22B は超ロングコンテキストでも自然な日本語の生成自体は維持できた一方で、小型のモデルに比べてもかなり早い段階でプロンプトの前半を忘れるような挙動が現れました。この結果はベンチマークのスコアの傾向と似ている部分もありますが、必ずしも一致しない結果といえます。

3.2 Qwen2.5 の派生モデル間での比較

以下が結果のまとめと分析です。

モデル名	\~4k	16k	32k	64k	128k	256k	512k	1024k (1M)
Qwen2.5-7B-Instruct	-	-	Skip-to-Q3	-	-	Skip-to-Q3	Skip-to-Q3	Q3「富士」に応答するが途中で日本語崩壊
Qwen2.5-7B-Instruct-1M	-	-	-	-	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	Q3に正しく回答するが同じテキストを3回繰り返す
abeja/ABEJA-Qwen2.5-7b-Japanese-v0.1	-	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	Q3「富士山」に触れるが回答としては不十分
abeja/ABEJA-Qwen2.5-32b-Japanese-v1.0	-	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	Skip-to-Q3	何も生成されない	何も生成されない	何も生成されない

Qwen2.5-7B-Instruct と Qwen2.5-7B-Instruct-1M を比較すると、Qwen2.5-7B-Instruct-1M の方が擬似コンテキスト長 1M でも生成の破綻は小さくなりました。しかしながら、完全に安定した動作をしているわけではありませんし、より短い擬似コンテキスト長の段階で前半の質問への回答をスキップしてしまいました。64k, 128k ではベンチマークスコアにはかなり差がありますが、この実験の結果としてはそれほど大きな差は生まれませんでした。

次に、 Qwen2.5-7B-Instruct と ABEJA-Qwen を比較すると、 ABEJA-Qwen は短い擬似コンテキスト長でもプロンプト前半を忘れやすい傾向が見られました。要因としては ABEJA-Qwen はシーケンス長 8k で継続事前学習を行っているため、長文処理の知識を少し忘却している可能性があります。一方で、 ABEJA-Qwen-7B では擬似コンテキスト長 1Mでの生成が Qwen2.5-7B-Instruct よりも自然な日本語を生成できており、日本語処理性能が向上している効果が見られました。

▶︎ Qwen2.5-7B-Instruct

▶︎ Qwen2.5-7B-Instruct-1M

▶︎ ABEJA-Qwen2.5-7b-Japanese-v0.1

▶︎ ABEJA-Qwen2.5-32b-Japanese-v1.0

4. まとめ

今回は LLM に入力する位置インデックスを操作して擬似的なロングコンテキスト下でのテキスト生成を検証しました。結果としては、ベンチマークのスコアほど細かく性能の差を見ることはできませんでしたが、大まかなモデルの挙動を実際に見ることはできたような気がします。とはいえ、この結果だけでモデルの LCLM 性能を判断することはできないため、あくまで参考程度に捉えていただきたいです。