[論文レビュー] Neural Networks Efficiently Learn Low-Dimensional Representations with SGD

Alireza Mousavi-Hosseini, Sejun Park+3 more|arXiv (Cornell University)|Sep 29, 2022

Stochastic Gradient Optimization Techniques被引用数 6

ひとこと要約

この論文は、入力次元 $d$ が大きい場合でも、重み減衰を用いた確率的勾配降下法（SGD）が、2層ReLUニューラルネットワークにおいて低次元表現を効率的に学習できることを示している。理論的に、1層目の重みが真のモデルの方向で張られる $k$-次元主部分空間に収束することを証明し、ネットワーク幅に依存しない一般化境界と、単一インデックスモデルでは $d$ に線形に比例する標本複雑度を達成する。これはカーネル法を上回る性能を示している。

ABSTRACT

We study the problem of training a two-layer neural network (NN) of arbitrary width using stochastic gradient descent (SGD) where the input $\boldsymbol{x}\in \mathbb{R}^d$ is Gaussian and the target $y \in \mathbb{R}$ follows a multiple-index model, i.e., $y=g(\langle\boldsymbol{u_1},\boldsymbol{x} angle,...,\langle\boldsymbol{u_k},\boldsymbol{x} angle)$ with a noisy link function $g$. We prove that the first-layer weights of the NN converge to the $k$-dimensional principal subspace spanned by the vectors $\boldsymbol{u_1},...,\boldsymbol{u_k}$ of the true model, when online SGD with weight decay is used for training. This phenomenon has several important consequences when $k \ll d$. First, by employing uniform convergence on this smaller subspace, we establish a generalization error bound of $O(\sqrt{{kd}/{T}})$ after $T$ iterations of SGD, which is independent of the width of the NN. We further demonstrate that, SGD-trained ReLU NNs can learn a single-index target of the form $y=f(\langle\boldsymbol{u},\boldsymbol{x} angle) + ε$ by recovering the principal direction, with a sample complexity linear in $d$ (up to log factors), where $f$ is a monotonic function with at most polynomial growth, and $ε$ is the noise. This is in contrast to the known $d^{Ω(p)}$ sample requirement to learn any degree $p$ polynomial in the kernel regime, and it shows that NNs trained with SGD can outperform the neural tangent kernel at initialization. Finally, we also provide compressibility guarantees for NNs using the approximate low-rank structure produced by SGD.

研究の動機と目的

• ターゲットが $k \ll d$ 個の方向に依存する場合、広い2層ニューラルネットワークにおけるSGDがどのように低次元構造を学習するかを理解すること。
• ネットワーク幅に依存しないSGDで訓練されたネットワークの一般化保証を確立すること。
• SGDで訓練されたReLUネットワークが、$d$ に線形に比例する標本複雑度で特定の多項式ターゲットを学習できることを示し、カーネル法に比べて優れていること。
• SGDが誘導する低ランク構造に基づいた圧縮可能性保証を提供すること。

提案手法

• 任意の幅 $m$ を持つ2層ReLUニューラルネットワークの学習ダイナミクスを、オンラインSGDと重み減衰を用いて分析する。
• 入力がガウス分布に従い、ノイズのあるリンク関数を持つ複数インデックス関数 $y = g(\langle \mathbf{u}_1, \mathbf{x} \rangle, \dots, \langle \mathbf{u}_k, \mathbf{x} \rangle)$ としてターゲットをモデル化する。
• SGDの1層目の重みが $\mathbf{u}_1, \dots, \mathbf{u}_k$ で張られる $k$-次元主部分空間に収束することを証明する。
• 低次元部分空間における一様収束を用いて、$T$ 回の反復後における一般化誤差境界 $\mathcal{O}(\sqrt{kd/T})$ を導出する。
• サブガウスおよびサブ指数型集中不等式を用いて、ランダム特徴量と勾配更新の挙動を制御する。
• ラデマッハ複雑度と対称化を用いて一般化ギャップをバウンドし、標本複雑度の結果を確立する。

実験結果

リサーチクエスチョン

• RQ1 $k \ll d$ の場合、広い2層ReLUネットワークにおけるSGDは、真のモデルの方向で張られる主部分空間に収束するか？
• RQ2 SGDで訓練されたネットワークの一般化誤差は、ネットワーク幅に依存せずにバウンド可能か？
• RQ3 SGDで訓練されたReLUネットワークは、$d$ に線形に比例する標本複雑度で、$p$ 次多項式ターゲットを学習できるか？（カーネル法では $d^{\Omega(p)}$ 必要となるのと対照的）
• RQ4 SGDで訓練されたネットワークの圧縮可能性は何か？一般化とはどのように関連するか？

主な発見

• SGDで訓練された2層ReLUニューラルネットワークの1層目の重みは、真のモデルの方向 $\mathbf{u}_1, \dots, \mathbf{u}_k$ で張られる $k$-次元主部分空間に収束する。
• T 回の反復後の一般化誤差は $\mathcal{O}(\sqrt{kd/T})$ でバウンドされ、ネットワーク幅 $m$ に依存しない。
• 単一インデックスモデルで $f$ が単調かつ多項式成長を示す場合、SGDは $d$ に線形に比例する（対数要因を除いて）標本複雑度を達成し、カーネル法が要求する $d^{\Omega(p)}$ より優れている。
• SGDが誘導する低ランク構造により、圧縮可能性保証が可能となり、モデルの複雑度と圧縮表現のサイズが関連づけられる。
• SGDで訓練されたReLUネットワークは、初期状態のニューラルトランジットカーネルよりも、特定の $p$ 次多項式ターゲットをより効率的に学習できる。
• 主部分空間への収束により、特徴量の学習が向上し、ランダム特徴量よりも優れた一般化性能が得られる。