arxiv論文を読む：トランスフォーマーの入力トークン数をRMTで百万トークン以上にスケーリングする

Scaling Transformer to 1M tokens and beyond with RMT

• Transformer技術の新たなブレイクスルーを予感する論文でしたので、再帰メモリトランスフォーマー拡張(RMT)による、Transformerの入力トークン数の拡張方法の論文の前半部分をまとめてみました。

• なお、論文後半の記憶タスクでの使用例や実験結果については、論文本文を参照ください。

• この技術は元の入力トークンをセグメント化して、別途保持するメモリーセグメントと合わせてトランスフォーマーの入力とすることで、入力トークンのサイズの制約を無くす画期的なアイデアのようです。

• 個人的には、単純な長期記憶的な知識を問うようなタスクには、やみくもに入力トークン数を増やすアプローチよりも、外部の埋め込みデータベースなどを確実に参照させる方法がシンプルで良いような気もしますが、例えば複雑な伏線を張った、長編小説の読解や生成、サイズが大きな動画などの用途には有効かもしれません。

• 記載漏れや誤り、理解不足などありましたら、ぜひコメント、ご助言いただけると幸いです。

1. 概要

• この論文では、事前トレーニング済みのBERTモデル（RMT: Recurrent Memory Transformer）を使用して、計算量の増加をリニアに抑えたままで、入力トークンの増加できる方法を提案。

• 本論文では、入力トークン数を 512 トークン x 最大 4096 セグメント= 2048k トークンまで増やしても高い検索精度を維持できることが確認されました。（注釈：GPT-4が通常8k、最大32kトークン）

• 入力トークン数が増えた場合でも、GPUの必要メモリサイズは 増加させずに1セグメント分のメモリーで実行できることが確認されました。

figure 1

2. RMT：再帰メモリートランスフォーマー

• 入力トークンをセグメントに分割し、分割されたセグメント(segment 1)は単一のバックボーンのTransformer Layersで処理します。出力は通常のTransformer Layersの出力とメモリートークン(mem)です。

• ここで更新されたメモリトークン(mem)を次のセグメント(segment 2)と合わせて、先ほどと同じグローバルTransformer Layer入力に渡すことを再帰的に繰り返します。(segment 3 ,4…)

• なお、トレーニング時は、現在のセグメントからメモリセグメント(mem)を介して、ひとつ前のセグメントにmemの勾配を流してバックボーンのTransformerをチューニングします。

• RMTメモリ拡張は、どのTransformerファミリーとも互換性あります。

計算効率

• 通常のTransformerの計算量は入力トークン長の二乗に比例して増大しまうので、あまりにトークン長を大きくすることは、計算量、必要メモリーの観点から現実的ではありませんが、RMT推論では入力トークンに対して計算量を直線的なスケーリングに抑えることができます。

• また、必要メモリーは1セグメントあたりのサイズに依存するのでセグメント数が増えても必要メモリが増えることはありません。