機械翻訳などで現在主流となっているTransformerベースのモデルは、長い文章でも前後の内容をふまえた翻訳文を生成できる。

これは、最大数千ワードの範囲（コンテキストウインドウ）で、すべての単語同士を比較し、近似性に応じた重みづけ（Attention）が行われているためだ。

ただ長い文章といっても、複数のWebページの要約を作成することができる程度で、小説の要約まではむつかしい。そこでGoogleが開発したのがReformerというモデルだ。Transformerベースの同モデルは、最大100万ワードの長文が扱えるという。

・Attentionとメモリの問題を解決

同モデルでは、Transformerの2つの問題を解決している。Attentionおよびメモリの問題だ。

コンテキストウインドウ内のワードを比較する過程では、ワード数が増えるほど計算負荷が高くなる。このため、Reformerではあらかじめワードにハッシュ値を割り当て、これに応じてセグメントわけしておくようだ。同じ、または隣接するセグメント内のワードを比較することで、計算負荷は低くなる。

また、通常メモリ消費量はレイヤー移行につれて累積的に増加していく。キャッシュはトレーニング時にアクティベーションの復元で使用され、Reformerで扱うキャッシュはレイヤーごとに数GB単位になるという。

典型的なTransformerモデルには12個以上のレイヤーがあるため、当然メモリ不足に陥るわけだ。

・16GBのメモリで100万ワードに文章を処理

これを避けるためにReformerでは、リバーシブルネットワークを追加し、通常のアクティベーション更新のほかに差分をキャプチャするアクティベーションを設けている。

キャッシュを累積していく代わりにレイヤー移行ごとに入力の再計算が実施され、アクティベーションの復元では、ネットワークを逆方向に実行して差分を差し引くだけで済む。

こうした2つのアプローチを実装することで、Reformerではわずか16GBのメモリで、最大100万ワードの文章を扱えるようになった。

Google Colabのノートブックに、『罪と罰』をまるごと処理した例が示されているほか、同モデルは断片画像の復元といった画像処理なんかにも役立つとのこと。

参照元：Reformer: The Efficient Transformer/ Google AI Blog