[논문 리뷰] Training (Overparametrized) Neural Networks in Near-Linear Time

Jan van den Brand, Binghui Peng외 2명|arXiv (Cornell University)|2020. 06. 20.

Stochastic Gradient Optimization Techniques참고 문헌 89인용 수 23

한 줄 요약

이 논문은 과다매개변수화된 ReLU 신경망을 훈련하기 위한 근사 선형 시간 알고리즘을 제시한다. 이는 가우스-뉴턴 방법을 ℓ2-회귀 문제로 재구성하고, 빠른 존슨-린든스트라우스(Fast-Johnson-Lindenstrauss, Fast-JL) 차원 감소 기법을 활용해 그램 행렬을 조건화하는 데 기반한다. 알고리즘은 반복당 eO(mnd + n³)의 계산 비용을 가지며, ǫ 이하의 훈련 손실을 달성하기 위한 총 실행 시간은 eO((mnd + n³) log(1/ǫ))이다. 이는 이전의 이차수 최적화 방법 대비 제곱 속도 향상을 달성하였으며, 심층 학습 최적화에서 고급 무작위 선형 대수 기법의 실현 가능성을 입증한다.

ABSTRACT

The slow convergence rate and pathological curvature issues of first-order gradient methods for training deep neural networks, initiated an ongoing effort for developing faster $\mathit{second}$-$\mathit{order}$ optimization algorithms beyond SGD, without compromising the generalization error. Despite their remarkable convergence rate ($\mathit{independent}$ of the training batch size $n$), second-order algorithms incur a daunting slowdown in the $\mathit{cost}$ $\mathit{per}$ $\mathit{iteration}$ (inverting the Hessian matrix of the loss function), which renders them impractical. Very recently, this computational overhead was mitigated by the works of [ZMG19,CGH+19}, yielding an $O(mn^2)$-time second-order algorithm for training two-layer overparametrized neural networks of polynomial width $m$. We show how to speed up the algorithm of [CGH+19], achieving an $ ilde{O}(mn)$-time backpropagation algorithm for training (mildly overparametrized) ReLU networks, which is near-linear in the dimension ($mn$) of the full gradient (Jacobian) matrix. The centerpiece of our algorithm is to reformulate the Gauss-Newton iteration as an $\ell_2$-regression problem, and then use a Fast-JL type dimension reduction to $\mathit{precondition}$ the underlying Gram matrix in time independent of $M$, allowing to find a sufficiently good approximate solution via $\mathit{first}$-$\mathit{order}$ conjugate gradient. Our result provides a proof-of-concept that advanced machinery from randomized linear algebra -- which led to recent breakthroughs in $\mathit{convex}$ $\mathit{optimization}$ (ERM, LPs, Regression) -- can be carried over to the realm of deep learning as well.

연구 동기 및 목표

• 심층 학습에서 이차수 최적화의 높은 계산 비용, 특히 과다매개변수화된 네트워크에 대해 이전의 가우스-뉴턴 방법이 가지는 O(mn²)의 반복당 비용을 해결하기 위해.
• 부드러운 활성화 함수를 가지는 네트워크보다 더 복잡하고 현실적인 ReLU 네트워크에 이차수 최적화를 확장하기 위해.
• 기존의 이차수 최적화 방법에서 헤시안 역행렬 계산의 성능 저하 문제를 해결하고, 전체 기울기 차원(mn)에 대해 근사 선형 시간 훈련을 달성하기 위해.
• 이전에 볼록 최적화에서 사용된 고급 무작위 선형 대수 기법이 비볼록 심층 학습 환경으로 효과적으로 이식될 수 있음을 입증하기 위해.

제안 방법

• 공액 기울기 방법을 통해 효율적으로 해를 구할 수 있도록, 자코비안 행렬에 대한 ℓ2-회귀 문제로 가우스-뉴턴 업데이트를 재구성한다.
• 그램 행렬 J_t J_t^T를 조건화하기 위해, 빠른 존슨-린든스트라우스(Fast-JL) 유형의 차원 감소 기법을 적용하여 크기를 줄이되, 해의 품질은 유지한다.
• 일반화된 첫 번째 순서 공액 기울기 방법을 사용하여, 원래 행렬 크기 M에 종속되지 않는 시간 내에 조건화된 회귀 문제를 해결한다.
• 신경 탄성 커널(NTK) 이론을 활용해 과다매개변수화된 네트워크의 국소 선형화를 정당화하고, 커널 회귀로의 축소를 가능하게 한다.
• 집중 부등식을 사용하여 자코비안 근사 오차와 회귀 하위문제의 해 오차를 제한함으로써 수렴성을 확보한다.
• 반복당 비용이 eO(mnd + n³)인 백프로파게이션 프레임워크에 회귀 해법기를 통합하며, 이는 주로 자코비안 평가와 회귀 해법에 의해 지배된다.

실험 결과

연구 질문

• RQ1 과다매개변수화된 ReLU 네트워크에 대한 이차수 최적화를 O(mn²) 이하의 반복당 비용으로 낮춤으로써 실용화가 가능한가?
• RQ2 Fast-JL 차원 감소 기법을 사용해 가우스-뉴턴 헤시안 근사에 조건화를 적용할 수 있으며, 수렴 보장을 유지할 수 있는가?
• RQ3 두 층의 ReLU 네트워크에 대해 전체 기울기 차원(mn) 기준 근사 선형 시간 훈련이 가능한가?
• RQ4 볼록 최적화에서 성공한 랜덤라이즈드 선형 대수 기법이 비볼록 심층 학습 훈련에 적응 가능한가?
• RQ5 제안된 알고리즘이 SGD보다 훨씬 더 빠른 수렴 속도를 달성하면서도 일반화 성능을 유지하는가?

주요 결과

• 이전의 [CGH+19] 논문에서 제시한 O(mn²) 방법에 비해 제곱 속도 향상을 달성하여, 반복당 비용을 eO(mnd + n³)로 감소시켰다.
• 이론적 과다매개변수화 조건(m = Ω(max{λ⁻⁴n⁴, λ⁻²n²d log(n/δ)})) 하에, 전체 기울기 차원(mn) 기준 근사 선형 시간 복잡도를 가지는 ReLU 네트워크에 대한 첫 번째 이차수 알고리즘이다.
• 반복당 ∥ft+1 − y∥² ≤ ½∥ft − y∥²를 보장하여, 목표 손실로의 선형 수렴을 보장한다.
• 훈련 손실을 ǫ 이하로 줄이기 위한 총 실행 시간은 eO((mnd + n³) log(1/ǫ))이며, 빠른 행렬 곱셈을 사용할 경우 eO((mnd + n^ω) log(1/ǫ))로 개선된다.
• 네트워크 가중치의 갱신이 유한하게 유지되어 초기화로부터의 분리가 크지 않게 하므로 일반화 성능을 지원한다.
• 이 방법은 볼록 최적화 문제에 성공적으로 적용되었으며, γ-강볼록이고 β-연속 미분 가능한 함수에 대해 뉴턴 방법의 실행 시간을 eO((nd log(κ) + d³) log(1/ǫ))로 향상시켰다.