Reformer

Reformerは、UC BerkeleyとGoogle ResearchのNikita Kitaev氏、Łukasz Kaiser氏、Anselm Levskaya氏がICLR 2020で発表したReformer: The Efficient Transformerに登場する、メモリ効率に優れたTransformerです。このモデルは、シーケンス長を制限する2つのコストに対処します。局所性鋭敏型ハッシュ（LSH）アテンションは、完全なn × nのアテンション行列を置き換え、長さLのシーケンスに対する計算コストをおよそL log Lに抑えます。また、可逆残差レイヤーでは、バックプロパゲーションに必要なアクティベーションをレイヤーごとではなくネットワーク全体で一度だけ保存します。Googleは、Reformerがメモリ16 GBの単一アクセラレーター上で最大100万語のコンテキストウィンドウを扱えると報告し、小説『罪と罰』全体を単一の学習例として処理しました。

LSHアテンション#

Softmaxの後、各クエリのアテンションは最も類似する少数のキーに支配されるため、Reformerはそれらだけを検索します。クエリとキーには1つの共有射影を使用し、すべてのベクトルを角度ベースの局所性鋭敏型ハッシュでハッシュします。ランダム行列Rを乗算し、[xR, −xR]の最大値の位置をバケットとして使用するため、似た方向を向くベクトルは通常同じバケットに入ります。次に、トークンをバケット順に並べ、ソートしたシーケンスを同じ大きさのチャンクに分割します。各トークンは、自身のチャンクと直前のチャンク内だけでアテンションを計算します。単一のラウンドでは分離される類似ペアを捉えるため、複数のハッシュラウンドを並列実行します。

LSHは近似最近傍のベクトル検索に用いられる古典的な手法です。厳密な検索より大幅に低いコストで、クエリのk近傍の大部分を見つけます。バケット化の処理は単独で実行できるほど短時間です。

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

クエリとのコサイン類似度が0.96の近傍ベクトルは、1,000ラウンドのうち72%でクエリと同じバケットに入り、類似度が0.04の無関係なベクトルは9%でした。したがって、アテンションをバケット内に限定しても、重要なペアの大部分を保持できます。

可逆レイヤーとチャンク化#

標準的なニューラルネットワークでは、逆方向のパスに備えて各レイヤーのアクティベーションをすべて保持するため、メモリ使用量は深さとシーケンス長の両方に応じて増えます。Reformerは可逆残差ネットワーク（RevNet）を応用し、アクティベーションを2つに分け、一方をもう一方の関数で更新します。

逆方向のパスでは、各レイヤーが出力から入力を復元します。x2 = y2 − FeedForward(y1)、続いてx1 = y1 − Attention(x2)となるため、保存するのは最終アクティベーションだけで、各レイヤーは1回だけ再計算されます。勾配チェックポイントも同様のトレードオフを行い、いくつかのチェックポイントでアクティベーションを保存し、その間のレイヤーを再計算します。Reformerは、隠れレイヤーがモデルより数倍広いため、フィードフォワードレイヤーも位置ごとのチャンクに分けて実行します。アクティベーションメモリは、Ultralytics YOLO26の学習時のバッチサイズの上限を決める要因です。batch=-1を使うと、AutoBatchは複数のバッチサイズでGPUメモリを測定し、直線を当てはめ、空きメモリの60%を使うと予測されるバッチサイズを選択します。

結果と実装#

Reformerは、enwik8の文字単位テキストとImageNet-64のピクセル単位画像生成において、メモリ使用量を大幅に抑え、長いシーケンスをより高速に処理しながら、標準的なTransformerと同等の性能を達成しました。Hugging Face TransformersではReformerModelとして実装されており、enwik8と『罪と罰』で学習した公開済みチェックポイントが用意されています。

Reformerと関連モデルの比較#

• Longformer： アテンションパターンを事前に固定します。各トークンの周囲にスライディングウィンドウを設け、少数のグローバルトークンを加えます。
• スパースアテンション： 一部のトークンペアのみをスコアリングする手法の総称です。LSHアテンションはそのうち、コンテンツベースの手法です。
• 線形アテンション： Softmaxをカーネル特徴マップに置き換え、アテンションを線形時間で計算できるようにします。
• Flash Attention： スコア行列を保存せず、GPUオンチップメモリのタイル上で厳密なアテンションを計算します。スコア数は二次のままですが、メモリ使用量は線形に増加します。