LocoProp：レイヤー毎に最適化を行いバックプロパゲーションを効率化(1/2)

１．LocoProp：レイヤー毎に最適化を行いバックプロパゲーションを効率化(1/2)まとめ

・ディープラーニングではモデルや学習データほどパラメータ更新手法は研究されていない
・高度な更新手法は単純な手法より効率的だが計算量と必要メモリが多くなるという欠点がある
・LocoPropは各レイヤー毎に最適化する事で性能と計算量の削減を両立する新手法

２．LocoPropとは？

以下、ai.googleblog.comより「Enhancing Backpropagation via Local Loss Optimization」の意訳です。元記事は2022年7月29日、Ehsan AmidさんとRohan Anilさんによる投稿です。

first-order derivativeという言い回しが出てきます。日本語では一階微分や一次微分と訳される事が多いようで、いずれにしても一回だけ微分するの意です。

高性能な重み更新手法は高次な微分が必要になって計算量が多くなるけれども、レイヤー毎に最適化してやる事で、一次微分を使ってもほぼほぼ同程度の性能がより少ない計算量で実現できます！ってお話です。

アイキャッチ画像はlatent diffusionでプロンプトはLocoProp。Locoから機関車を連想してるんでしょうかね、左上が良い感じです。

ディープニューラルネットワーク(DNN:Deep Neural Network)の成功には、モデル設計と学習データが重要な要素となりますが、モデルパラメータ(重み)の更新に使用する特定の最適化方法については、あまり議論されません。

DNNの学習には、真実のラベルとモデルの予測値との間の不一致を測定する損失関数を最小化することが含まれます。学習はバックプロパゲーションによって行われ、勾配降下法によってモデルの重みが調整されます。勾配降下法では、重みに対する損失の勾配(すなわち微分)を用いて重みが更新されます。

最も単純な重みの更新は確率的勾配降下法に対応し、すべてのステップで、勾配に対して(適切なステップサイズ、別名、学習率で)重みを負の方向に移動させます。

より高度な最適化手法では、過去のステップからの情報、および／または現在の重みの周りの損失関数の局所特性(曲率情報など)を用いて、重みを更新する前に負の勾配の方向を修正します。

例えば、モメンタム・オプティマイザーは過去の更新の平均方向に沿って移動することを奨励し、AdaGradオプティマイザーは過去の勾配に基づいて各座標をスケーリングします。これらのオプティマイザは、一般的に一次微分(すなわち勾配)からの情報のみを用いて更新方向を修正することから、一般的に一次法(first-order methods)と呼ばれます。さらに重要なことは、重みパラメータの成分が互いに独立に扱われることです。

ShampooやK-FACなどのより高度な最適化では、パラメータの勾配間の相関を捉え、収束を改善し、反復回数を減らし、解の質を向上させることが示されています。これらの手法は、損失の導関数の局所的な変化、すなわち勾配の変化に関する情報を捕捉します。

この追加情報を用いて、高次オプティマイザは、異なるパラメータ群間の相関を考慮することで、学習モデルの更新方向をより効率的に発見することができます。一方、高次更新の方向性を計算するためには、一次更新よりも計算コストがかかるという欠点があります。また、統計量を格納するために多くのメモリを使用し、行列の反転を伴うため、高次オプティマイザの実際の適用には支障があります。

論文「LocoProp: Enhancing BackProp via Local Loss Optimization」ではDNNモデルを学習するための新しい枠組みを紹介します。

私達の新しい枠組みであるLocoPropは、ニューラルネットワークをレイヤーのモジュール構成として考えています。一般に、ニューラルネットワークの各レイヤーは、入力に対して線形変換を行い、その後、非線形の活性化関数を適用します。

新しい枠組みでは、各レイヤーに重み正則化器、出力目標、損失関数が割り当てられます。各レイヤーの損失関数は、そのレイヤーの活性化関数と一致するように設計されています。この定式化により、学習は与えられたミニバッチのサンプルに対する局所的な損失を、レイヤーを超えて反復的かつ並列的に最小化します。

本手法では、(RMSPropのような)一次最適化器を用いて、バッチごとに複数の局所更新を行い、高次最適化器に必要な行列反転のような計算量の多い操作を回避することができます。しかし、局所的な更新を組み合わせた結果は、むしろ高次の更新のように見えることを示します。経験的に、LocoPropはディープオートエンコーダベンチマークにおいて一次法を上回り、ShampooやK-FACのような高次最適化器に匹敵する性能を示し、高いメモリと計算量が必要ないことを示します。

手法

ニューラルネットワークは一般に、モデルの入力をレイヤーごとに出力特徴表現に変換する複合関数と見なされています。LocoPropはこの考え方を採用し、ネットワークをレイヤーに分解しています。特に、出力における損失関数を最小化するためにレイヤーの重みを更新するのではなく、LocoPropは各レイヤーに特有のあらかじめ定義された局所損失関数を適用します。

与えられたレイヤーに対して、損失関数は活性化関数に一致するように選択されます。例えば、tanh活性化を持つレイヤーに対してはtanh損失が選択されます。各レイヤーごとの損失は、(与えられたミニバッチのサンプルに対する)レイヤーの出力と、そのレイヤーの目標出力の概念との間の不一致を測定します。さらに、正則化項により、更新された重みが現在の値から大きくずれることがないようにします。レイヤー別損失関数(ローカルターゲット付き)と正則化の組み合わせが、各レイヤーの新しい目的関数として使用されます。

バックプロパゲーションと同様に、LocoProp は活性化を計算するためにフォワードパスを適用します。バックワードパスでは、LocoProp は各レイヤーのニューロン毎に「ターゲット」を設定します。最後に、LocoPropはモデルの学習をレイヤーごとに独立した問題に分割し、各レイヤーの重みにいくつかの局所更新を並行して適用することができます。