[論文レビュー] Simple Hardware-Efficient Long Convolutions for Sequence Modeling

Daniel Y. Fu, Elliot L. Epstein+6 more|arXiv (Cornell University)|Feb 13, 2023

Topic Modeling被引用数 6

ひとこと要約

本論文は、系列モデルにおける長距離畳み込みカーネルの直接的パrameterizationを提案し、単純な正則化（例：カーネル重みの圧縮）によって滑らかなカーネルが得られ、状態空間モデル（SSM）と同等の性能を達成できることを示している。FlashButterflyと呼ばれるI/Oに配慮したアルゴリズムを用いることで、長距離畳み込みはcuFFTよりも2.2倍高速化され、64K長の系列では7.2倍高速に学習が可能となり、Path256や脳fMRIデータを含む多数のベンチマークでSSMやTransformerを上回った性能を発揮した。

ABSTRACT

State space models (SSMs) have high performance on long sequence modeling but require sophisticated initialization techniques and specialized implementations for high quality and runtime performance. We study whether a simple alternative can match SSMs in performance and efficiency: directly learning long convolutions over the sequence. We find that a key requirement to achieving high performance is keeping the convolution kernels smooth. We find that simple interventions--such as squashing the kernel weights--result in smooth kernels and recover SSM performance on a range of tasks including the long range arena, image classification, language modeling, and brain data modeling. Next, we develop FlashButterfly, an IO-aware algorithm to improve the runtime performance of long convolutions. FlashButterfly appeals to classic Butterfly decompositions of the convolution to reduce GPU memory IO and increase FLOP utilization. FlashButterfly speeds up convolutions by 2.2$ imes$, and allows us to train on Path256, a challenging task with sequence length 64K, where we set state-of-the-art by 29.1 points while training 7.2$ imes$ faster than prior work. Lastly, we introduce an extension to FlashButterfly that learns the coefficients of the Butterfly decomposition, increasing expressivity without increasing runtime. Using this extension, we outperform a Transformer on WikiText103 by 0.2 PPL with 30% fewer parameters.

研究の動機と目的

• 直接的パrameterizationがなされた長距離畳み込みが、複雑な初期化や特化した実装を必要とせずに、状態空間モデル（SSM）と同等の性能を達成できるかどうかを調査すること。
• 現代のハードウェア、特にGPUのメモリI/Oボトルネックや未利用のFLOP容量の問題に起因する長距離畳み込みの実行時間の非効率性を是正すること。
• バタフライ分解のような構造的行列分解が、長距離畳み込みにおける効率性と表現力の両面を向上させられるかどうかを検討すること。
• 提案手法の一般化能力を、言語、視覚、神経画像解析といった多様なモダリティにわたって評価すること。
• 単純な正則化と効率的なアルゴリズム設計によって、標準畳み込みとSSMの性能格差を埋められ、長系列での高速学習を可能にできるかどうかを実証すること。

提案手法

• 周波数ドメインにおける滑らかさとスパarsityを促進するために、カーネル重みにSquash演算子を適用し、モデル品質を向上させる。
• 時間ドメインでSmooth演算子を用いることで、カーネル重みをさらに正則化し、周波数ドメインにおける滑らかさを強化する。
• FFTに基づく畳み込みをブロックスパースなバタフライ行列分解を用いて再構成することで、メモリI/Oを削減し、FLOPの利用効率を向上させるI/Oに配慮したアルゴリズム「FlashButterfly」を開発する。
• GPUの行列乗算ユニットを活用するため、バタフライ分解を介して畳み込みをスパース行列乗算の連鎖に再構成する。
• 分解係数を学習可能とする拡張を導入することで、推論コストを増加させずに表現力を向上させる。
• 得られたアーキテクチャをS4やH3などのモデルにおけるSSMレイヤーの即時置換として利用し、LRA、CIFAR、WikiText103、fMRIを含む多様なベンチマークで評価する。

実験結果

リサーチクエスチョン

• RQ1 複雑な初期化や特化したカーネルを必要とせずに、直接的パrameterizationがなされた長距離畳み込みが、長系列モデリングタスクでSSMと同等の性能を達成できるか。
• RQ2 長距離畳み込みをSSMと同等に競合可能にするために、どのような単純な正則化技術がカーネルの滑らかさを促進するか。
• RQ3 I/Oに配慮したアルゴリズム（例：FlashButterfly）が、現代のGPU上での長距離畳み込みの実行効率を顕著に向上させられるか。
• RQ4 構造的バタフライ分解を用いることで、長距離畳み込みに速度向上と高い表現力の両方をもたらせるか。
• RQ5 fMRI系列モデリングのような困難で現実的な応用において、長距離畳み込みがTransformerやSSMを上回れるか。

主な発見

• 単純な正則化（特にカーネル重みの圧縮と平滑化）により、周波数ドメインにおける滑らかなカーネルが得られ、Long Range Arena（LRA）ではSSMレベルの性能を回復し、ランダム初期化でもS4と同等の性能を達成した。
• 画像分類タスクでは、順方向CIFARではS4より0.3%高い正確性を達成し、2次元CIFARではS4ND-ISOと0.8%以内の差に収まった。
• 言語モデリングでは、PILEではH3と同等の性能（0.3 PPL以内）を達成し、OpenWebTextでも0.3 PPL以内の差に収まった。
• 脳fMRIモデリングでは、TransformerとSSMをそれぞれ0.14点および0.16点のMAE向上で上回った。
• FlashButterflyは、cuFFTよりも長距離畳み込みを2.2倍高速化し、Path256（64K長の系列）の学習を可能にした。29.1点の正確性向上と、先行研究比7.2倍の高速学習を達成し、最先端の性能を発揮した。
• 学習可能なバタフライ拡張により、実行コストを増加させずに表現力を向上させ、WikiText103では30%少ないパラメータ数でTransformerを0.2 PPLの差で上回った。