【ITニュース解説】Accelerating Spatio-Temporal Attention for Video Diffusion on TPUs

近年、AI技術は目覚ましい進歩を遂げ、テキストから高品質な画像を生成したり、動画まで作り出せるようになった。特に動画生成AIは、リアルな映像表現を可能にし、エンターテインメントやコンテンツ制作の現場で大きな注目を集めている。しかし、これらのAIモデルが高解像度の動画、例えばスマートフォンやPCの画面で美しく見える1440pのような動画を生成しようとすると、処理速度の面で大きな課題に直面していた。

この課題は、AIモデルが動画の各部分にどれくらい「注意」を払うべきかを判断する「自己注意機構」という仕組みに起因する。この機構は、動画の一貫性や自然な動きを表現するために不可欠だが、動画の解像度や長さが増えるにつれて、必要な計算量が急速に増加する「二次曲線的なボトルネック」という問題を引き起こす。例えば、動画サイズが2倍になると計算量が4倍になるような増加で、高解像度動画の実用的な生成を妨げていた。

この深刻なボトルネックを解消するため、開発者たちは「Sparse VideoGen (SVG)」という新しいアプローチを考案した。これは、自己注意機構の計算を「疎（Sparse）」に、つまり本当に必要な部分だけに集中させることで、無駄な計算を大幅に削減する技術だ。AIの「注意ヘッド」と呼ばれる複数の処理単位を、動的に空間的または時間的に限定された「疎なマスク」に割り当てる。具体的には、動画の同じフレーム内の特定の領域に注目する空間的マスクや、時間的に近いフレーム間の関係性に注目する時間的マスクを、状況に応じて柔軟に切り替える。これにより、動画全体で膨大な関係性を計算するのではなく、効率的に注意の焦点を絞り、計算量を大きく削減する。

アルゴリズムレベルでの計算量削減は重要だが、それをAI専用の高速処理チップであるTPU（Tensor Processing Unit）上で最大限に活かすには、ハードウェアの特性に合わせた最適化が不可欠となる。TPUは、行列計算を高速に実行することに特化したプロセッサで、効率的なデータ処理のためにメモリを「タイル」と呼ばれる小さなブロックに分割して扱う。Sparse VideoGenのアルゴリズム的な効率性を、TPUの物理的な処理速度向上に結びつけるため、「Splash Attentionカーネル」というTPU上で注意計算を行うための低レベルなソフトウェアが最適化された。この最適化では、主に三つの工夫が凝らされた。一つ目は、計算が不要な「空のメモリタイル」をTPUがスキップするようにし、無駄な処理時間をなくすこと。二つ目は、精密な座標指定によるマスキング処理を、本当に必要な「境界タイル」だけに厳密に限定すること。これにより、不必要なマスキング処理によるオーバーヘッドを防いだ。そして三つ目は、動画のデータをメモリ上で「時間軸優先（temporal-major）」の順序に並べ替えること。これにより、TPUが連続したデータに効率的にアクセスできるようになり、データの読み込み速度が向上した。

これらのソフトウェアとハードウェア双方にわたる綿密な最適化は、疎なマスクの適用とTPUのタイル処理が実際の計算と密接に連携するように設計された。その結果、無駄な行列演算が大幅に削減され、1440pという高解像度の動画生成において、エンドツーエンド（最初から最後まで）で最大1.69倍もの推論速度向上を達成した。この技術革新は、高解像度かつ高品質な動画をより速く、より効率的に生成できることを意味する。AIモデルの設計におけるアルゴリズムの工夫と、それを支えるハードウェア特性を最大限に引き出す最適化の両輪が揃うことで、これまで難しかった大規模なAIタスクの実現に大きく貢献し、今後の動画生成AIの可能性をさらに広げる重要な一歩となった。