[논문 리뷰] High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation

Jimmy Ba, Murat A. Erdogdu외 4명|arXiv (Cornell University)|2022. 05. 03.

Stochastic Gradient Optimization Techniques인용 수 11

한 줄 요약

이 논문은 비례적 점근적 조건(n, d, N → ∞, 비율 n/d, N/d → 일정)에서 이중층 신경망의 첫 번째 레이어 가중치에 대한 첫 번째 기울기 단계를 분석하며, 심지어 한 단계의 갱신만으로도 가중치 행렬에 질서가 있는 랭크-1 스파이크가 발생하고, 이는 교사 모델의 선형 성분과 정렬됨을 보여준다. 이러한 정렬은 갱신된 특징에 대해 리지 회귀를 수행할 경우 랜덤 특징을 초월하고, 충분히 큰 학습률을 사용할 경우 입력에 대한 최적의 선형 모델조차도 능가함을 의미하며, 이는 비례적 점근적 조건에서의 초기 특징 학습의 힘을 입증한다.

ABSTRACT

We study the first gradient descent step on the first-layer parameters $\boldsymbol{W}$ in a two-layer neural network: $f(\boldsymbol{x}) = \frac{1}{\sqrt{N}}\boldsymbol{a}^ opσ(\boldsymbol{W}^ op\boldsymbol{x})$, where $\boldsymbol{W}\in\mathbb{R}^{d imes N}, \boldsymbol{a}\in\mathbb{R}^{N}$ are randomly initialized, and the training objective is the empirical MSE loss: $\frac{1}{n}\sum_{i=1}^n (f(\boldsymbol{x}_i)-y_i)^2$. In the proportional asymptotic limit where $n,d,N o\infty$ at the same rate, and an idealized student-teacher setting, we show that the first gradient update contains a rank-1 "spike", which results in an alignment between the first-layer weights and the linear component of the teacher model $f^*$. To characterize the impact of this alignment, we compute the prediction risk of ridge regression on the conjugate kernel after one gradient step on $\boldsymbol{W}$ with learning rate $η$, when $f^*$ is a single-index model. We consider two scalings of the first step learning rate $η$. For small $η$, we establish a Gaussian equivalence property for the trained feature map, and prove that the learned kernel improves upon the initial random features model, but cannot defeat the best linear model on the input. Whereas for sufficiently large $η$, we prove that for certain $f^*$, the same ridge estimator on trained features can go beyond this "linear regime" and outperform a wide range of random features and rotationally invariant kernels. Our results demonstrate that even one gradient step can lead to a considerable advantage over random features, and highlight the role of learning rate scaling in the initial phase of training.

연구 동기 및 목표

• 이중층 신경망에서 첫 번째 레이어 가중치에 대한 첫 번째 기울기 갱신이 표현 학습을 어떻게 향상시키는지 이해하는 것.
• 비례적 점근적 영역(n, d, N → ∞, 비율 n/d, N/d → 일정)에서 이 갱신이 커널 리지 회귀 성능에 미치는 영향을 규명하는 것.
• 갱신된 특징이 랜덤 특징과 선형 모델을 초월할 수 있는 학습률 스케일링 조건을 규명하는 것.
• 초기 학습에서의 비커널 행동에 대한 이론적 기반을 마련하고, 이를 특징 적응과 연결하는 것.

제안 방법

• 랜덤 초기화를 가진 이중층 ReLU 또는 erf-활성화 신경망에서 첫 번째 기울기 하강 단계를 분석한다.
• 특징 학습의 영향을 분리하기 위해 단일 색인 교사 모델 f*를 사용하는 학생-교사 프레임워크를 도입한다.
• 예측 위험의 결정론적 근사값을 도출하기 위해 비례적 점근적 조건(n, d, N → ∞, 비율 n/d, N/d → 일정)에서 고차원 점근적 분석을 수행한다.
• 갱신된 가중치 행렬의 구조를 랭크-1 스파이크 + 노이즈로 유도하며, 교사 모델 f*의 선형 성분과의 정렬을 보여준다.
• 스펙트럼 분석과 랜덤 행렬 이론을 활용해 한 기울기 단계 이후 공액 커널에 대한 리지 회귀의 예측 위험을 계산한다.
• 학습률의 두 가지 스케일링을 고려한다: 작은 η (가우시안 등가성)과 큰 η (선형 영역을 초월함), 각 경우에 대해 위험 한계를 도출한다.

실험 결과

연구 질문

• RQ1 첫 번째 레이어 가중치에 대한 첫 번째 기울기 갱신이 랜덤 특징을 초월해 구조적인 변화를 유도하는가?
• RQ2 이 단일 갱신이 고차원 영역에서 랜덤 특징과 선형 모델을 초월하는 성능 향상을 이끌 수 있는가?
• RQ3 학습률 스케일링이 갱신된 특징이 선형 영역을 초월해 일반화할 능력에 어떻게 영향을 미치는가?
• RQ4 한 단계 갱신 후 커널의 정확한 통계적 행동은 무엇이며, 교사 모델의 구조와 어떻게 관련되는가?

주요 결과

• W에 대한 첫 번째 기울기 갱신은 가중치 행렬에 랭크-1 스파이크를 유도하며, 이는 교사 모델 f*의 선형 성분과 정렬된다.
• 작은 학습률에서는 갱신된 특징이 가우시안 등가성을 달성하고, 랜덤 특징보다 예측 위험이 향상되지만 선형 모델의 위험에 의해 제한된다.
• 충분히 큰 학습률을 사용할 경우, 갱신된 특징에 대한 리지 추정기는 입력에 대한 어떤 선형 모델보다도 뛰어난 성능을 보이며, 선형 영역을 벗어나는 것을 보여준다.
• 한 단계 후의 예측 위험은 10τ* + C′(√τ*·√(d/n) + d/n)로 상한선이 존재하며, 이때 τ*는 f*의 비선형 성분의 노름을 나타내어 교사 모델의 비선형성에 의존함을 보여준다.
• 이 개선은 커널 행렬의 최상위 특이벡터가 훈련 레이블과 정렬되기 때문에 발생하며, 이는 랭크-1 갱신에 의해 유도된다.
• 결과적으로, 한 기울기 단계만으로도 통계적으로 의미 있는 특징 학습이 가능하며, 고정 커널 방법을 능가할 수 있음을 입증한다.