[논문 리뷰] TT-Rec: Tensor Train Compression for Deep Learning Recommendation Models

Chunxing Yin, Bilge Acun외 2명|arXiv (Cornell University)|2021. 01. 25.

Tensor decomposition and applications참고 문헌 32인용 수 13

한 줄 요약

TT-Rec는 딥러닝 추천 모델(DLRMs)의 임bedding 테이블을 텐서 트레이스(TT) 분해 기반 방법으로 압축하는 방법을 제안한다. 이는 정확도 손실 없이 모델 크기를 최대 117배까지 줄이며, 학습 시간 오버헤드는 오직 13.9%에 불과하다. 최적화된 TT-EmbeddingBag 커널, 적응형 캐싱, 그리고 새로운 가우시안 기반 초기화 전략을 통해 대규모 환경에서 성능과 정확도를 유지한다.

ABSTRACT

The memory capacity of embedding tables in deep learning recommendation models (DLRMs) is increasing dramatically from tens of GBs to TBs across the industry. Given the fast growth in DLRMs, novel solutions are urgently needed, in order to enable fast and efficient DLRM innovations. At the same time, this must be done without having to exponentially increase infrastructure capacity demands. In this paper, we demonstrate the promising potential of Tensor Train decomposition for DLRMs (TT-Rec), an important yet under-investigated context. We design and implement optimized kernels (TT-EmbeddingBag) to evaluate the proposed TT-Rec design. TT-EmbeddingBag is 3 times faster than the SOTA TT implementation. The performance of TT-Rec is further optimized with the batched matrix multiplication and caching strategies for embedding vector lookup operations. In addition, we present mathematically and empirically the effect of weight initialization distribution on DLRM accuracy and propose to initialize the tensor cores of TT-Rec following the sampled Gaussian distribution. We evaluate TT-Rec across three important design space dimensions -- memory capacity, accuracy, and timing performance -- by training MLPerf-DLRM with Criteo's Kaggle and Terabyte data sets. TT-Rec achieves 117 times and 112 times model size compression, for Kaggle and Terabyte, respectively. This impressive model size reduction can come with no accuracy nor training time overhead as compared to the uncompressed baseline.

연구 동기 및 목표

• 딥러닝 추천 모델(DLRMs)의 급격히 증가하는 메모리 요구량을 해결한다. 현재 이러한 모델들은 TB 규모의 임베딩 테이블을 필요로 한다.
• 정확도를 훼손하지 않거나 학습 시간 오버헤드가 과도하게 높아지지 않도록 모델 메모리 용량을 줄인다.
• 모델 품질을 유지하면서도 임베딩을 압축하기 위한 원칙적이고 파arameterized인 접근법으로 텐서화(tensorization)를 탐색한다.
• 산업 규모의 DLRMs에 적합한 계산 효율적이고 하드웨어에 민감한 압축 프레임워크를 설계한다.
• 최적화된 커널 설계와 캐싱 전략를 통해 대규모 추천 모델의 실용적이고 저비용 학습을 가능하게 한다.

제안 방법

• 딥러닝 추천 모델의 큰 임베딩 테이블을 텐서 트레이스(TT) 분해를 이용한 저랭크 행렬 곱의 시퀀스로 대체한다.
• TT 기반의 임베딩 룩업 성능을 향상시키기 위해 최적화된 커스텀 TT-EmbeddingBag 커널을 구현하여 기존 최고 수준의 TT 구현보다 3배 빠른 성능을 달성한다.
• 자주 액세스되는 임베딩 벡터를 압축 해제된 형태로 캐시하여 재계산을 줄이고 정확도를 복구하는 캐싱 메커니즘을 도입한다.
• 학습 안정성과 모델 정확도 향상을 위해 가우시안 분포에서 샘플링하는 새로운 가중치 초기화 전략을 제안한다.
• GPU 가속 시스템에서 추론 및 학습 성능을 추가로 최적화하기 위해 배치 행렬 곱셈을 활용한다.
• 메모리, 정확도, 속도 등의 다양한 설계 차원에서 TT-랭크와 압축 파라미터를 튜닝하여 파레토 최적의 구성 요건을 식별한다.

실험 결과

연구 질문

• RQ1 텐서 트레이스 분해가 정확도를 떨어뜨리지 않으면서도 DLRMs에서 의미 있는 모델 압축을 달성할 수 있는가?
• RQ2 가중치 초기화 전략의 선택이 TT-Rec의 추천 모델 성능과 수렴에 어떤 영향을 미치는가?
• RQ3 캐싱 전략이 학습 중 TT 구조의 임베딩을 복구하는 데 발생하는 성능 오버헤드를 어느 정도 완화할 수 있는가?
• RQ4 다양한 TT-랭크와 임베딩 차원 설정에서 메모리 감소, 학습 시간, 모델 정확도 간의 상호 교환 관계는 어떠한가?
• RQ5 TT-Rec는 산업 규모의 학습 파이프라인과 하드웨어 가속기와의 호환성을 유지하면서도 수십만 배 수준의 압축을 달성할 수 있는가?

주요 결과

• Criteo Kaggle 및 Terabyte 데이터셋에서 TT-Rec는 압축되지 않은 기준선과 비교해 각각 117배와 112배의 모델 크기 감소를 달성했으며, 정확도 손실 없이 성능을 유지했다.
• 제안된 TT-EmbeddingBag 커널은 최고 수준의 기존 TT 구현보다 3배 더 빠르며, 계산 효율성 향상에 기여했다.
• 90%의 캐시 히트율을 확보한 TT-Rec는 학습 시간 성능에서 표준 EmbeddingBag을 능가했으며, 캐싱 전략의 효과성을 입증했다.
• 극한의 압축 조건에서도 기준선과 동일한 정확도를 유지함으로써 제안된 초기화 및 압축 기법의 강건성을 입증했다.
• 기준선 대비 학습 시간 오버헤드가 오직 13.9%에 불과하여 TT-Rec가 대규모 온라인 학습에 실용적으로 적용 가능하다는 점을 입증했다.
• 메모리 용량, 정확도, 학습 속도 간의 파레토 최적의 균형을 달성했으며, 다양한 하드웨어 및 구현 조건에 맞게 조정 가능한 파라미터를 제공한다.