[논문 리뷰] Towards Compact Neural Networks via End-to-End Training: A Bayesian Tensor Approach with Automatic Rank Determination

Cole Hawkins, Xing Liu외 1명|arXiv (Cornell University)|2020. 10. 17.

Tensor decomposition and applications인용 수 4

한 줄 요약

이 논문은 훈련 중에 최적의 텐서 랭크를 자동으로 결정하는 엔드 투 엔드 베이지안 텐서라이제이션 훈련 프레임워크를 제안한다. 이는 초초소형 신경망을 처음부터 구축할 수 있도록 한다. 변분 추론과 저랭크 텐서 분해(CP, Tucker, TT)를 조합하여 모델 파라미터를 최대 26,000배(예: 4.25B에서 1.6×10⁵로) 감소시키며 정확도 손실을 최소화하여 에너지 효율적인 현장 내 훈련을 가능하게 한다.

ABSTRACT

While post-training model compression can greatly reduce the inference cost of a deep neural network, uncompressed training still consumes a huge amount of hardware resources, run-time and energy. It is highly desirable to directly train a compact neural network from scratch with low memory and low computational cost. Low-rank tensor decomposition is one of the most effective approaches to reduce the memory and computing requirements of large-size neural networks. However, directly training a low-rank tensorized neural network is a very challenging task because it is hard to determine a proper tensor rank {\it a priori}, which controls the model complexity and compression ratio in the training process. This paper presents a novel end-to-end framework for low-rank tensorized training of neural networks. We first develop a flexible Bayesian model that can handle various low-rank tensor formats (e.g., CP, Tucker, tensor train and tensor-train matrix) that compress neural network parameters in training. This model can automatically determine the tensor ranks inside a nonlinear forward model, which is beyond the capability of existing Bayesian tensor methods. We further develop a scalable stochastic variational inference solver to estimate the posterior density of large-scale problems in training. Our work provides the first general-purpose rank-adaptive framework for end-to-end tensorized training. Our numerical results on various neural network architectures show orders-of-magnitude parameter reduction and little accuracy loss (or even better accuracy) in the training process. Specifically, on a very large deep learning recommendation system with over $4.2 imes 10^9$ model parameters, our method can reduce the variables to only $1.6 imes 10^5$ automatically in the training process (i.e., by $2.6 imes 10^4$ times) while achieving almost the same accuracy.

연구 동기 및 목표

• 대규모 과도하게 파ram터화된 딥 뉴럴 네트워크 훈련의 높은 계산 비용과 에너지 소비 문제를 해결한다.
• 저랭크 분해에서 최적의 텐서 랭크를 수동으로 선택하는 문제를 극복한다. 이는 모델 정확도와 효율성에 매우 중요하다.
• 전체 정밀도 사전 훈련 없이도 초소형 신경망을 직접 처음부터 훈련할 수 있도록 한다.
• 다양한 텐서 형식(CP, Tucker, 텐서 트레인)을 지원하고 자동 랭크 적응 기능을 갖춘 확장 가능한 엔드 투 엔드 프레임워크를 개발한다.
• 대규모 모델에서 압축 비율을 극대화하고 추론 비용을 최소화하면서도 정확도를 유지하거나 향상시킨다.

제안 방법

• CP, Tucker, 텐서트레인, TTM 등의 저랭크 텐서 분해에 대해 베이지안 확률 모델을 수립하여 훈련 중에 자동으로 랭크를 결정할 수 있도록 한다.
• 텐서 요소와 랭크에 대한 사후 분포를 효율적으로 최적화하기 위해 맞춤형 확률적 변분 추론 알고리즘을 개발한다.
• 스pars리티 유도 랭크 적응을 가능하게 하기 위해 자동 관련성 결정(ARD)을 통합한 계층적 사전 분포 구조를 도입한다.
• 고차원이고 구조화된 파ram터 공간에서의 사후 근사 향상을 위해 투영된 커널 기반 스틸 차이(discrepancy)를 사용한다.
• 완전한 크기의 임bedding 테이블을 명시적으로 구성하지 않는 미분 가능 훈련 파이프라인을 설계하여 대규모 모델에서 메모리 오버헤드를 줄인다.
• 딥 러닝 추천 모델(DLRM)의 엔드 투 엔드 훈련과 통합하여 빌리언 파라미터 모델에서의 확장성을 입증한다.

실험 결과

연구 질문

• RQ1 베이지안 텐서라이제이션 훈련 프레임워크는 사전 초모수 조정 없이 엔드 투 엔드 훈련 중에 최적의 텐서 랭크를 자동으로 결정할 수 있는가?
• RQ2 자동 랭크 결정은 대규모 딥 러닝 모델에서 예측 정확도를 유지하거나 향상시키면서 얼마나 많은 파라미터 감소를 이룰 수 있는가?
• RQ3 메모리, 에너지 소비, 추론 효율성 측면에서 엔드 투 엔드 텐서라이제이션 훈련은 사후 압축 기법보다 어떻게 비교되는가?
• RQ4 제안된 방법은 FPGA와 같은 자원 제약이 있는 플랫폼에서 실용적인 현장 내 초소형 모델 훈련을 가능하게 하는가?
• RQ5 실제 벤치마크에서 다양한 텐서 형식(CP, Tucker, TT)이 압축 비율, 정확도, 훈련 효율성에 미치는 영향은 어떠한가?

주요 결과

• 대규모 DLRM에서 제안된 방법은 모델 파라미터를 4.25B에서 1.6×10⁵로 감소시켜 26,000배의 압축 비율을 달성했으며, 정확도 손실은 최소한이었다.
• DLRM 벤치마크에서 ARD-LU 버전은 단지 227만 개의 훈련 파라미터로 78.67%의 정확도를 달성하여 파라미터 수가 더 적은 상수 랭크 기반 베이스라인보다 뛰어난 성능을 보였다.
• FPGA에서 비텐서라이제이션 훈련 대비 임베디드 CPU에서의 비텐서라이제이션 훈련보다 에너지 효율성이 최대 123배 높고, 속도는 59배 빨라졌다.
• ARD를 통한 자동 랭크 결정은 고정 랭크 방법보다 더 우수한 일반화 성능을 보였으며, ARD-HC는 TT 형식에서 끝내운 파라미터 수가 163,976개에 불과했지만 78.73%의 정확도를 달성했다.
• 완전한 크기의 임베딩 테이블을 명시적으로 구성하지 않아도 프레임워크가 초초소형 모델을 성공적으로 훈련시켰으며, 사후 압축의 계산 비용을 피할 수 있었다.
• 베이지안 솔버는 자동 랭크 적응 기능을 갖춘 저랭크 텐서라이제이션을 통해 빌리언 파라미터 모델의 안정적이고 확장 가능한 훈련을 가능하게 했다.