핵심 개념
本稿では、ニューラルネットワークのトレーニングと推論におけるメモリフットプリントと計算コストを削減するため、テンソルタッカーフォーマットにおける因子分解層の形状認識トレーニングという新しいアプローチを提案する。
本稿では、ディープラーニングモデルのメモリフットプリントと計算コストを削減するための新しいアプローチである、テンソルタッカーフォーマットにおける因子分解層の形状認識トレーニングについて述べられています。
背景
最新のディープラーニングアーキテクチャは、その優れた性能の一方で、推論とトレーニングに膨大なメモリと計算コストを必要とします。この課題に対処するために、重みのスパース化や量子化など、モデルのサイズと計算の複雑さを軽減するための様々な技術が開発されてきました。しかし、スパース化による枝刈りは、密行列向けに設計されたGPUハードウェアを活用することが難しく、量子化を用いる場合はモデルの性能に関する誤差推定が困難です。さらに、これらの手法は推論に必要なリソースを削減することはできますが、性能に影響を与えることなくトレーニング中のメモリ削減を実現することは困難です。
因子分解層
スパース化や量子化と並んで、層の因子分解も広く用いられているモデル圧縮手法です。層の重みを様々な行列やテンソルの因子分解を用いて表現することで、モデルの性能と頑健性を維持しながら、メモリを大幅に削減することができます。最近の多くの研究で、過剰にパラメータ化されたネットワークの層の重みは低ランクになる傾向があり、小さな特異値を削除することで、モデルのサイズを大幅に削減しながら、性能を向上させることができる可能性があることが、理論的にも実験的にも示唆されています。低ランク因子分解の大きな利点の一つは、低パラメータの因子分解モデルをトレーニング全体およびファインチューニングの段階全体で使用できることです。
課題
層の因子を直接トレーニングする方法はディープラーニングで広く用いられていますが、この方法には2つの大きな欠点があります。
因子分解のランクrを事前に選択する必要があり、圧縮モデルの性能はランクに大きく依存する可能性があります。
トレーニングフローは初期化の選択に非常に敏感であり、その結果、損失の振動が大きくなり、収束が遅くなり、性能が最適ではなくなり、ランクを下げる前にフルモデルをトレーニングするウォームアップの段階が必要になる場合があります。
提案手法
本稿では、これらの課題を克服するために、基礎となるリーマン構造を利用してタッカーフォーマットで因子分解されたテンソル層をトレーニングする、ランク適応型の形状認識トレーニングアルゴリズムを導入しています。
動的低ランク近似
ランクρ = (r1, ..., rd)の場合、集合Mρ = {W : rank(Mati(W)) = ri, i = 1, ..., d}は、任意の点W = C ×di=1 Ui ∈ Mρにおいて接空間を持つ多様体です。低ランクのタッカー因子Cと{Ui}iのみを用いてMρ内の層の重みを計算する戦略を設計するために、接空間に射影された連続時間勾配フローとしてトレーニング問題を定式化します。
再パラメータ化
動的低ランク近似フレームワークは、タッカーフォーマットの層の勾配フローを提供します。しかし、これらの発展方程式をネットワークのトレーニングに直接使用すると、タッカーテンソルの多様体の曲率が高いため、学習率が非常に小さくなってしまいます。これを克服するために、再パラメータ化を導入し、再パラメータ化された因子に対してロバストなダイナミクスを定式化します。
ランク適応トレーニング戦略
得られた勾配系を数値的に積分し、基底の拡大ステップと組み合わせることで、タッカーフォーマットのネットワークに対して形状認識型のランク適応トレーニング戦略を提案します。
計算コストの削減
提案されたアプローチは、d次テンソルに対してネットワークと勾配テープをd+1回評価する必要があるため、計算コストが大幅に増加します。しかし、拡張基底の計算を大幅に簡略化できることにより、計算コストを大幅に削減することができます。
アルゴリズム
上記の戦略と計算コスト削減を組み合わせることで、タッカーフォーマットのテンソルに対する効率的なランク適応型の形状認識トレーニング手法であるアルゴリズム1が得られます。
理論的解析
提案手法の以下の3つの主要な理論的特性を示します。
損失の降下
期待値における定常点への収束
フルモデルへの近似を示すロバストな境界
実験
提案手法の性能を、フルモデル、標準的な層の因子分解、モデルの枝刈りのベースラインと比較して評価するため、一連の実験を行いました。
圧縮性能
CIFAR10とtiny-imagenetを用いて、提案手法の圧縮性能を評価しました。その結果、提案手法は、競合する因子分解ベースラインよりも高い精度で高い圧縮率を達成し、重みの初期化間のばらつきも小さいことがわかりました。
最適化のロバスト性
分解の因子に対する標準的な同時勾配降下と比較して、提案手法の利点をさらに明確にするために、提案手法とタッカー分解およびCP分解を用いた標準的なトレーニングを用いた場合の、MNISTにおけるLeNet5の精度履歴を示します。その結果、提案手法を用いて低ランク勾配フローを積分すると、ネットワークのトレーニングプロセスにおいて、より高速でロバストな収束率が得られることがわかりました。
計算性能
タッカー分解における畳み込み層の推論とトレーニングにおける計算性能は、現在のテンソルのランクに依存します。VGGとAlexNetの推論時間とメモリフットプリントを、提案手法で使用されているタッカー因子分解で評価し、因子分解されていないベースラインモデルと比較しました。その結果、現実的な圧縮率の場合、提案手法の計算フットプリントは、対応するベースラインモデルよりも大幅に小さいことがわかりました。
LoRA風低ランクアダプターを用いたファインチューニング
提案手法のもう一つの応用として、低ランクアダプターを用いた事前トレーニング済みモデルのファインチューニングを紹介します。提案手法は、モデルの圧縮と適応を完全に区別しないため、アプローチは同じです。この2つの設定を示すために、提案手法をテストした2つの異なる設定を紹介します。1つ目は、GLUEベンチマークにおけるDeberta V3のファインチューニングです。2つ目は、Dreamboothを用いたStable Diffusionのファインチューニングです。
考察
本稿では、タッカーテンソル因子分解多様体の形状を活用して、圧縮されたタッカーフォーマットのニューラルネットワークのためのロバストで効率的なトレーニングアルゴリズムを構築しました。提案手法は、フルモデルへの近似誤差限界、損失の降下と期待値における定常点への収束の保証という理論的裏付けを持っています。様々なベンチマークとモデルに対する圧縮率と精度の比率で示されているように、この手法は、交互または同時勾配降下を用いた標準的な因子分解手法よりも優れています。この手法は、ハイパーパラメータをτという1つのパラメータに大幅に削減します。このパラメータは、ネットワーク全体のトレーニングダイナミクスにおける層の重要性に応じて、層ごとの圧縮率として明確に解釈することができます。
制限事項
提案手法は、参照ネットワーク内に性能の高い低ランクのタッカーサブネットが存在することを前提としています。経験的にこれが観察されていますが、完全結合線形層の場合と同様に、この仮説を裏付ける理論的証拠を提供するためには、さらなる調査が必要です。
통계
VGG16ベンチマークにおいて、提案手法は90%を超える圧縮率でベースラインの精度を維持することができ、τ = 0.03(すなわち95.3%の圧縮率)の場合、平均でベースラインを上回りました。
AlexNetはVGG16のパラメータ数の16.8%です。そのため、圧縮を達成することは当然ながらより困難になります。それにもかかわらず、提案手法はベースラインを上回り、フルネットワークの性能に匹敵します。
ResNet18でも同様の挙動が観察されました。
提案された評価では、提案手法は、スパース性に基づく枝刈り(公平な比較のためにウォームアップなしで実装)、テンソルトレイン(TT)とタッカー因子分解、タッカー分解のためのリーマン勾配降下(RGD)、行列値DLRTを含む、すべての代替案と同等以上の性能を発揮しました。これは、各テンソルモードに沿った圧縮を個別に行うことができるタッカーフォーマットの柔軟性が高いためです。
圧縮率(c.r.)は1−c/fとして計算されます。ここで、cはトレーニング後の圧縮モデルにおける畳み込みパラメータの数、fはフルモデルの畳み込みパラメータの数です。これはトレーニング後の圧縮率ですが、因子分解に基づく手法では、トレーニングプロセス全体を通して同様の圧縮率が得られることを強調しておきます。
CP分解は多様体構造を持たないため、動的低ランクトレーニングには適しておらず、CP分解のDLRTバージョンは示されていません。
tiny-imagenetベンチマークでも同様の結果が得られました。