「深層ニューラルネットワークの高速化」の感想と残された謎
だいぶ遅れてしまいましたが実務でも使える「深層ニューラルネットワークの高速化」とそこで議論されるDNNの汎化性能の考察に関する感想です
本書は深層ニューラルネットワークの高速化、パラメーター圧縮に関する実用的な内容となっています。pytorch, tensorflow,Nvidiaの高速化コンパイラTensorRTや関連するpythonライブラリの使い方が付属コードとともに詳しく具体的に説明されています。pytorchなどの学習済みネットワークを使えば手軽に実験、検証をすることができます。GPUを用いた推論以外にも圧縮によって小さくなったニューラルネットをCPUで実行するためのライブラリや低ビット演算などの手法が紹介されています。
量子化、枝刈り、蒸留、低ランク近似、アーキテクチャと推論時のみの工夫で実施できるものから学習時のネットワーク設計の時点で考慮すべき方法まで順番に解説されていて、そしてパフォーマンスを見るためのツール、学習時の効率化についても解説されています。実際にはこれらの手法は重複したり同時に使ったりすることができ例えば量子化後の誤差を考慮した学習(Quantization Awareness Training)があったり、蒸留において量子化後もネットワークを生徒モデルにすることもあります。
実践的な説明やTransformer(flash attention)などの最先端のアーキテクチャにおける高速化の紹介とともになぜDNNはパラメーターが非常に多いにも関わらず過学習せず汎化性能が出るのか、そしてそれの圧縮、高速化の性質とどのように関係しているのかについて深い考察がなされています。
汎化と平坦解
ある意味奇妙な現象としてパラメーターを圧縮、量子化することで精度が向上する場合があると報告されています。古典的な機械学習モデルではパラメーター数が多すぎると過学習を起こし汎化性能が低下しますがそれと似たような現象が起こる場合があるということです。この上手く圧縮することで情報が整理されるような性質と性能との関係に関してはPACベイズ理論を用いて解説されますがこれがDNNで成立するのかどうかは考察の余地があるように思います。なぜ上手くいくDNNが平坦性を持つのか、パラメーター数、データ数との関係などについてもPACベイズ理論に基づいた研究がなされているはずですが関連研究をあまり調べきれていません。 loss関数の底が平坦であれば、多少のパラメーターの変化は精度に影響を及ぼしません。この範囲内で容量の小さいネットワークを探索する余地があります。この平坦性はDNNの他の性質をも説明できる概念です。画像生成AIで画風を合成するのによく使われるモデルマージはパラメーターの加算(平均)によってなされますが自然な画像を生成するようなマージ元のネットワークはともに平坦な領域に存在しているためその平均もまた平坦な領域に存在するだろうという仮説です。
有名な宝くじ仮説(ネットワークの一部のみが性能に関連し、ランダム初期化した場合はその一部のネットワークが含まれている)とその解の頑健性、初期値依存性が枝刈りのところで触れられています。汎化との関係でいえば一部のネットワーク以外は性能に関係せずパラメーターを自由に触れる余地があるのであればその自由度が平坦解の自由度に相当すると言えるかもしれません。「大規模言語モデルは新たな知能か」でも汎化と平坦性、宝くじ仮説が整合性のある性質であることが説明されていたと思います。
平坦解とその周辺のloss landscape
一方でパラメーター空間の様々な場所から近似最適解に到達できるという性質には別の説明を要します。多次元の複雑な形状のloss関数は一般に多数の局所最小解(鞍点)を持つとされますが一説には多次元のパラメーター空間におけるloss関数の鞍点は多次元ゆえにヘッシアンの正の固有値を持ち別の鞍点への移動を繰り返しているヘテロクリニック軌道を形成しているとされています。しかし鞍点近傍では学習軌道の速度は非常に遅くなります。「情報幾何の新展開」では初期のニューラルネットの研究でも学習が非常に遅れてしまう現象が報告されておりそれは素子の置換対称性が複数の局所最小解を作っておりそれが
Resnetではある程度解決されたとされています。LLMにおいては同様の現象が腑落ち(Grokking)として見えているのかもしれませんが小規模なネットワークで見える現象とLLMの現象が類似のものであるかのは考察、研究の余地がありそうです。
初期値から平坦解への学習経路やloss landscapeの形状を図示するのはその多次元性からも困難です。敢えて描くとすると下図のように高次元のパラメーター空間内の各小領域は初期値を出発して赤点の鞍点を経由して次元を低めながら比較的低次元(だが平坦解よりは高次元)の青色の領域、そして灰色の平坦解に到達するというものになるかもしれません。
実際には平坦解も高次元でありその体積は大きくなり、パラメーター空間全体に分布しているかもしれません。またこの図では後述のNTKの解析に置いて見られた初期値の無限小の近傍に最適値があるという性質は表現できていません。平坦解は下図のようにパラメーター空間の広範囲に分布していると言えるかもしれません。
Hessianの固有値
そこで解の平坦性をloss関数のHessian行列の固有値の分布から見てとろうという研究があります。
「情報幾何の新展開」でも紹介されている唐木田先生の研究では一層の要素数を無限大にした極限のニューラルタンジェントカーネル(NTK)におけるloss関数のヘッシアンであるフィッシャー情報行列の固有値の分布から平坦性について論じた研究
初期値の無限小の近傍に最適値がある(これが平坦解かもしれない)という研究、また勾配消失を起こさないような条件であるDynamical Isometryの場合の固有値の分布に関する研究もされています。
Hessianの固有値に着目する方法は学習においてはSharpness-Aware Minimization(SAM)
として知られ高い性能を発揮し、量子化においても層やパラメーターの重要度の指標として用い、Hessianの近似値計算法を使って高速にハードウェアに対応した混合精度量子化パラメーターを決定する方法が有ります(HAWQ)。
パラメーター空間の地形を探索するという挙動は拡散モデルにおいても存在し、パラメーター空間の探索は学習過程で成され、生成過程では学習されたパラメーターの地形を辿ってデータの分布を再現しましす。拡散モデルでどのようにして学習、汎化が成されているかの挙動をスコア関数のヤコビアンの固有値から考察し、データの分布する低次元多様体への漸近の様子と対応付けた研究もあります。 平坦な領域とそこへ至る低次元の経路を通る時の軌道周辺の小領域の次元の変化は数値的に観察すると面白そうで拡散モデルにおける方法を利用できないかとも考えられます。
解の特異性という観点からは”局所的な実対数しきい値(学習係数)λが小さい＝特異性が高い点(図のW1)のほうがが近傍が平らである=事後分布が大きいく選ばれやすい”という説明が成されています。これとHessianの固有値との関連は考察の余地が有ります。
まとめ
DNNのパラメーター数を大きくすればするほど高くなる汎化性能(double descent)はLLMの際限ない大規模化をもたらしエネルギー消費の点でも社会に影響を与えています。一方で圧縮、高速化、高性能化が同時に実現できる場合があるというのは新しい時代のムーアの法則のようでもあり、DNN研究開発だけでなく産業全体の中で少なくない位置を占めるかもしれません。この2つの傾向は独立したものなのか、平坦性のような何らかの共通点、依存関係があるのかはより詳しく知りたいです。
Discussion