Transformer-XL

Transformer-XLは、前のセグメントの隠れ状態をキャッシュし、次のセグメントがそれらに注意を向けられるようにすることで、固定長のテキストセグメントをまたいでコンテキストを保持するTransformer言語モデルです。Carnegie Mellon UniversityとGoogle BrainのZihang Dai、Zhilin Yangらが、ACL 2019でTransformer-XL: Attentive Language Models Beyond a Fixed-Length Contextを発表しました。固定長のセグメントごとに独立して学習した通常のTransformerでは、各セグメントが別々のコンテキストウィンドウとなり、現在のセグメントの開始位置より前を参照できません。著者らはこの問題をコンテキストの断片化と呼びました。Transformer-XLは、リカレントニューラルネットワーク（RNN）より80%、通常のTransformerより450%長い依存関係を学習しました。

セグメント単位の再帰#

モデルがセグメントを処理するとき、各層は前のセグメントのキャッシュ済み隠れ状態と現在のセグメントの隠れ状態を連結します。クエリは現在のセグメントだけから取得し、キーと値は両方のセグメントにわたります。また、勾配はキャッシュの境界で停止するため、セグメントごとのトレーニングコストは一定です。PyTorchで単一のアテンション層を示すと、次のようになります。

import torch
from torch import nn

torch.manual_seed(0)
seg_len, dim = 4, 16
attn = nn.MultiheadAttention(dim, num_heads=2, batch_first=True)
hidden = torch.randn(1, 3 * seg_len, dim)  # hidden states for three consecutive segments

memory = torch.empty(1, 0, dim)  # nothing is cached before the first segment
for i, segment in enumerate(hidden.split(seg_len, dim=1), start=1):
    # Queries come from the current segment; keys and values also cover the cached one
    context = torch.cat([memory, segment], dim=1)
    out, weights = attn(segment, context, context)
    print(f"Segment {i}: {weights.shape[1]} queries attend over {weights.shape[2]} positions")
    # Cache this segment for the next one; detach() stops gradients as in Transformer-XL
    memory = segment.detach()

最初のセグメントでは、4つのクエリはそのセグメントの4つのトークンだけを参照します。それ以降の各セグメントでは、クエリは自分の4つとキャッシュ済みの4つ、合計8つを参照します。セグメント内の層nが、1つ前のセグメントの層n − 1の状態を読み込むため、各層で参照範囲がさらに1セグメント分広がり、セグメント長をN、層数をLとすると、約N × L個のトークンに及びます。

キャッシュは評価も高速化します。トークンを1つずつ予測する通常のTransformerでは、新しいトークンを生成するたびにウィンドウ全体を再計算する必要がありますが、Transformer-XLはキャッシュ済みの状態を再利用するため、評価時の速度が最大1,800倍以上になりました。Ultralytics YOLO26のトラッキングでも、スクリプトが動画を1フレームずつmodel.track()に渡すと、同じ断片化の問題が発生します。persist=Trueは呼び出し間でトラッカーの状態を保持するため、各物体はフレームをまたいで同じIDを保ちます。

相対位置エンコーディング#

キャッシュ済みの状態を再利用すると、絶対位置が機能しなくなります。キャッシュされたセグメントの最初のトークンと現在のセグメントの最初のトークンが、同じ位置埋め込みを持つことになるためです。Transformer-XLは代わりに、クエリとキーの内容に、両者の距離を示す正弦波エンコーディングを加え、内容用と位置用の2つの学習可能なバイアスベクトルをすべての位置で共有して、各アテンションスコアを計算します。スコアは距離のみに依存するため、評価時にはトレーニング時より長いキャッシュに注意を向けられます。現在のほとんどの大規模言語モデル（LLM）で使われる回転位置埋め込み（RoPE）も、アテンションスコアを相対位置に依存させます。

成果と影響#

Transformer-XLは5つのベンチマークで新たな最先端の結果を達成しました。enwik8では1文字あたり0.99ビット、text8では1.08ビット、WikiText-103ではパープレキシティ18.3、One Billion Wordでは21.8、Penn Treebankではファインチューニングなしで54.5でした。WikiText-103だけで学習し、数千トークンに及ぶ、かなり一貫性のある新しい記事を生成しました。同じ著者の数名が開発したXLNetは、Transformer-XLのアーキテクチャを自然言語理解の事前学習に採用しました。また、DeepMindのCompressive Transformerは、古いメモリを破棄せずに圧縮することでキャッシュを拡張しました。

Transformer-XLと関連モデル#

• 長短期記憶（LSTM）: 隠れ状態をトークンごとに更新するため、遠い過去の情報は、その間にあるすべてのステップを経由する必要があります。
• Longformer: ウィンドウ型アテンションとグローバルアテンションを使い、再帰なしで4,096トークンの文書全体を一度に読み込みます。
• KVキャッシュ: 生成中に以前のトークンのキーと値を保存して、再計算を防ぎます。Transformer-XLのメモリはトレーニング時にも使用され、セグメント境界で勾配が停止します。
• 状態空間モデル（SSM）: 履歴全体を固定サイズの状態に圧縮します。