Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch

あなたは、世界で最も美味しく複雑なケーキを焼こうとしていると想像してください。しかし、あなたのキッチンはとても狭く、窮屈です。レシピには、後でプロセスを完璧に逆転させて、ケーキがどうなったかを確認するために、混ぜ合わせたすべての材料、すべての温度変化、そしてすべての攪拌動作を記録しておく必要があります。問題は、あなたのキッチンカウンター（コンピュータのメモリ）が小さすぎて、すべてのメモを保持できないことです。もしすべてを書き留めようとすれば、カウンターが溢れ出し、作業を中断しなければなりません。これは、大規模な人工知能（AI）モデルを訓練している科学者たちが日々直面している苦闘です。彼らはAIを教えるために多くのステップを記憶する必要がありますが、コンピュータの容量が足りなくなってしまうのです。これを解決するために、彼らは「アクティベーション・チェックポインティング」と呼ばれる巧妙なトリックを使います。すべてのステップを書き留める代わりに、最も重要なステップを選んで保存し、重要度の低いものは後でやり直すことに合意するのです。これは、小さなフォトアルバムに残すべき写真を選び、忘れてしまった場合に何度でも撮り直せるものを選ぶようなものです。目標は、レシピの魔法を失うことなく、そのケーキ作りのプロセス全体を、あの小さなキッチンに収めることです。

長い間、多くのAI科学者がモデル構築に使用しているコンピュータプログラムであるPyTorchは、どのステップを保存するかを決定するための特定の方法を持っていました。それは、決定を下すことを「0/1ナップサック問題」という古典的なパズルとして扱っていました。想像してみてください。あなたは、一定の重さまでしか背負えないバックパックを持ったハイカーです。あなたの手元には、それぞれ重さと価値（それがどれほど役に立つか）を持つアイテムのリストがあります。あなたは、バックパックを壊さない範囲で、最大の価値をもたらすアイテムを選びたいと考えています。PyTorchのデフォルトの手法はこの問題を解くために、あらゆるアイテムの組み合わせを巨大な紙にすべて書き出そうとするようなものでした。この手法は完璧であり、最適な答えを見つけ出しますが、その紙があまりにも巨大になるため、コンピュータのメモリが爆発し、プログラムがクラッシュしてしまいます。研究者たちは、もし選ぶべきアイテムがわずか100個あったとしても、必要な紙のサイズがあまりに大きくなり、彼らのマシンの空き容量である64ギガバイトを遥かに超えてしまうことを発見しました。それは、部屋に収まりきらない完璧な解決策だったのです。

この論文で、著者はこのパズルを解くためのよりスマートな方法、すなわち dp_knapsack_sliding_hirschberg と呼ばれる手法を紹介しています。すべての巨大な紙を一気に書き出す代わりに、「スライディングウィンドウ」というトリックを使用します。長い本を読んでいるけれど、一度に2ページ分しか見ることができない小さな虫眼鏡を持っている状況を想像してください。その虫眼鏡を本に沿って滑らせ、2ページを見、次に次の2ページを見る、というように進めていきます。こうすることで、一度に意識に留めておくのは2ページ分だけで済み、膨大な精神的スペースを節約できます。しかし、単に2ページを見るだけでは物語全体を記憶するには不十分です。どのアイテムを選ぶべきかを知る必要があります。これを解決するために、彼らはスライディングウィンドウと、「ヒルシュバーグのアルゴリズム」と呼ばれる古く巧妙な戦略を組み合わせました。これは「分割統治（ディバイド・アンド・コンカー）」のゲームのようなものです。バックパックの問題を一度に解こうとするのではなく、アイテムのリストを半分に分割します。まず左半分を解き、次に右半分を解き、それから2つの最適な解決策をどのように組み合わせるかを判断します。これを再帰的に行い、問題を小さく、より小さく分解していき、簡単に解けるようになるまで繰り返します。これらすべてを、ごくわずかなメモリ量のみを使用して行います。

この新手法の結果は素晴らしいものです。著者は64ギガバイトのRAMを搭載したコンピュータでテストを行いました。旧来の手法は、わずか100個のアイテムを扱うだけでクラッシュしましたが、新手法は2,000個のアイテムの問題を、ピーク時58.4ギガバイトのメモリ使用量で正常に解決しました。これは、コンピュータが以前よりも20倍大きな問題を扱えるようになったことを意味します。さらに、この新手法は単なるメモリ節約術ではありません。より高速でもあります。著者がある特定のマシンで同じパズルを1,000回実行してテストしたところ、新しいソルバーは一貫して旧来のソルバーよりも速いことが分かりました。決定的なのは、答えを推測してわずかに誤りが生じる可能性がある他の「クイックフィックス（応急処置）」的な手法とは異なり、この新しい手法は毎回、正確で完璧な解を見つけるということです。これは旧来の手法と同じ精度を持ちながら、より効率的です。

この論文は、この新しいアプローチが単なる理論ではなく、PyTorchソフトウェアに正常に統合され、バージョン2.10で利用可能であることを裏付けています。著者は、スライディングウィンドウと分割統治を組み合わせることで、AIモデルの成長を阻んでいたメモリのボトルネックを解決できることを示しています。彼らはこれがこの種の問題を解決する唯一の方法であると主張しているわけではなく、また、あらゆる種類のコンピュータ・パズルに機能すると示唆しているわけでもありませんが、AIのステップを保存するという特定のタスクにおいては、証明された、正確で非常に効率的なアップグレードです。この論文は、旧来の手法が大規模なモデルに対して十分ではないことを明確に示すことで、アイテムの数が増えすぎると失敗することを証明しています。その代わりに、彼らは完璧な精度を維持しながらメモリ不足によるクラッシュを取り除き、科学者があの小さなキッチンで、より大きく複雑なAIケーキを焼けるようにする解決策を提示しているのです。