[논문 리뷰] Training Multi-Layer Over-Parametrized Neural Network in Subquadratic Time

Zhao Song, Lichen Zhang외 1명|arXiv (Cornell University)|2021. 12. 14.

Advanced Neural Network Applications인용 수 6

한 줄 요약

이 논문은 구조화된 가중치 행렬과 효율적인 데이터 구조를 활용하여, 깊이 있는 과다 매개변수화된 신경망을 반복당 서브제곱 시간 이내에 훈련하는 새로운 프레임워크를 제안한다. 초기화 단계의 사전처리와 적응형 그래디언트 계산을 통해 반복당 비용을 O(m²⁻Ω(1))으로 줄여, 표준 O(m²)보다 크게 낮추며, 대규모 언어 모델의 빠른 피지테이닝을 가능하게 한다.

ABSTRACT

We consider the problem of training a multi-layer over-parametrized neural network to minimize the empirical risk induced by a loss function. In the typical setting of over-parametrization, the network width $m$ is much larger than the data dimension $d$ and the number of training samples $n$ ($m=\mathrm{poly}(n,d)$), which induces a prohibitive large weight matrix $W\in \mathbb{R}^{m imes m}$ per layer. Naively, one has to pay $O(m^2)$ time to read the weight matrix and evaluate the neural network function in both forward and backward computation. In this work, we show how to reduce the training cost per iteration. Specifically, we propose a framework that uses $m^2$ cost only in the initialization phase and achieves \emph{a truly subquadratic cost per iteration} in terms of $m$, i.e., $m^{2-Ω(1)}$ per iteration. Our result has implications beyond standard over-parametrization theory, as it can be viewed as designing an efficient data structure on top of a pre-trained large model to further speed up the fine-tuning process, a core procedure to deploy large language models (LLM).

연구 동기 및 목표

• 큰 m×m 가중치 행렬로 인해 발생하는 훈련 시 반복당 금지적인 O(m²) 비용 문제를 해결한다.
• 훈련 반복당 비용을 진정으로 서브제곱 이하, 즉 O(m²⁻Ω(1))로 줄이는 방법을 개발한다.
• 과다 매개변수화와 구조화된 계산을 활용하여 대규모 언어 모델(Langauge Models)의 효율적 피지테이닝을 가능하게 한다.
• 이전 방법들이 O(nm²) 비용을 부담하거나 입력 차원 d에 대해 지수적 의존성을 보이는 등의 한계를 극복한다.
• 사전 훈련된 모델 위에 데이터 구조를 설계하여 피지테이닝을 가속화하면서도 수렴 보장을 유지한다.

제안 방법

• 신경망 레이어의 활성화에서 희소성을 유도하기 위해 이동된 ReLU 활성화 함수를 사용하여 각 레이어의 실제 계산량을 감소시킨다.
• 활성화된 뉴런 출력의 분포와 노름을 모델링하기 위해 잘라낸 가우시안 랜덤 변수를 활용한다.
• 잘라낸 카이제곱 및 서브가우시안 분포에 대한 농도 부등식을 적용하여 레이어 간 노름 변동을 제한한다.
• 초기화 단계에서 한 번만 O(m²) 비용을 부담하는 사전처리 단계를 설계하여, 이후 모든 훈련 반복이 서브제곱 시간 이내로 수행되도록 한다.
• 모든 데이터 포인트와 레이어에 대한 유니온 바운드를 적용하여 네트워크 전반에서 고확률로 노름 안정성을 확보한다.
• 사전에 계산된 통계적 성질에 기반해 활성화된 뉴런을 인덱싱하는 데이터 구조를 사용하여, 빠른 순방향 및 역방향 전파를 구현한다.

실험 결과

연구 질문

• RQ1 깊이 있는 과다 매개변수화된 신경망에서 m×m 가중치 행렬을 사용할 때 반복당 서브제곱 시간을 달성할 수 있는가?
• RQ2 과다 매개변수화와 활성화 희소성을 어떻게 활용하여 표준 O(m²) 장벽을 초월해 계산량을 줄일 수 있는가?
• RQ3 입력 차원 d에 대해 지수적 의존성이 없이 O(m²⁻Ω(1)) 반복 비용을 보장하는 사전처리 체계를 설계할 수 있는가?
• RQ4 이론적으로 도출된 서브제곱 비용이 대규모 언어 모델의 피지테이닝에 실제로 얼마나 실현 가능한가?
• RQ5 랜덤 초기화 조건 하에서 서브제곱 훈련 비용을 달성하면서도 레이어 간 노름 안정성을 유지할 수 있는가?

주요 결과

• 제안된 프레임워크는 깊이 있는 과다 매개변수화된 네트워크 훈련에 대해 반복당 O(m²⁻Ω(1)) 비용을 달성하여 표준 O(m²)보다 크게 낮춘다.
• 이 방법은 초기화 단계에서만 O(m²) 비용을 부담하며, 이후 모든 훈련 반복은 서브제곱 시간 이내로 수행된다.
• 고확률로 모든 레이어와 데이터 포인트에서 은닉 표현의 ℓ₂ 노름이 [1−ε, 1+ε] 범위 내에 유지되어 안정적인 훈련을 보장한다.
• 이 프레임워크는 깊이 L ≥ 2 이고 너비 m = poly(n, d)인 네트워크에 적용 가능하며, 일반적인 과다 매개변수화 조건을 만족한다.
• 과다 매개변수화된 구조를 효율적으로 활용함으로써 대규모 언어 모델의 더 빠른 피지테이닝을 가능하게 한다.
• 분석은 잘라낸 가우시안 및 카이제곱 분포에 대한 농도 경계에 기반하여, 노름 안정성에 대한 이론적 보장을 확립한다.