[논문 리뷰] Training with Multi-Layer Embeddings for Model Reduction

Benjamin Ghaemmaghami, Zihao Deng외 5명|arXiv (Cornell University)|2020. 06. 10.

Recommender Systems and Techniques참고 문헌 24인용 수 5

한 줄 요약

이 논문은 정확도 대 모델 크기의 균형을 향상시키기 위해 단일 레이어 임베딩 테이블을 두 레이어의 선형 아키텍처로 대체하는 다층 임베딩 훈련(MLET)을 제안한다. 더 넓은 은닉 레이어 크기 k로 훈련한 후 최종 임베딩 차원 d로 투영함으로써 MLET는 학습된 표현의 효과적 랭크를 낮추어 정확도 손실 없이 임베딩 테이블 크기를 4–8배 감소시킨다. 이는 평균 25%의 훈련 시간 증가만을 수반한다.

ABSTRACT

Modern recommendation systems rely on real-valued embeddings of categorical features. Increasing the dimension of embedding vectors improves model accuracy but comes at a high cost to model size. We introduce a multi-layer embedding training (MLET) architecture that trains embeddings via a sequence of linear layers to derive superior embedding accuracy vs. model size trade-off. Our approach is fundamentally based on the ability of factorized linear layers to produce superior embeddings to that of a single linear layer. We focus on the analysis and implementation of a two-layer scheme. Harnessing the recent results in dynamics of backpropagation in linear neural networks, we explain the ability to get superior multi-layer embeddings via their tendency to have lower effective rank. We show that substantial advantages are obtained in the regime where the width of the hidden layer is much larger than that of the final embedding (d). Crucially, at conclusion of training, we convert the two-layer solution into a single-layer one: as a result, the inference-time model size scales as d. We prototype the MLET scheme within Facebook's PyTorch-based open-source Deep Learning Recommendation Model. We show that it allows reducing d by 4-8X, with a corresponding improvement in memory footprint, at given model accuracy. The experiments are run on two publicly available click-through-rate prediction benchmarks (Criteo-Kaggle and Avazu). The runtime cost of MLET is 25%, on average.

연구 동기 및 목표

• 추천 시스템에서 고차원 임베딩 테이블의 증가하는 메모리 비용을 해결하기 위해.
• 추론 아키텍처를 수정하지 않고 딥러닝 기반 추천 모델의 정확도 대 모델 크기 균형을 향상시키기 위해.
• 다층 선형 변환을 통해 단일 레이어 테이블보다 나은 임베딩을 생성할 수 있는지 탐구하기 위해.
• 실세계 성능을 평가하기 위해 페이스북의 오픈소스 DLRM 프레임워크 내에서 MLET를 프로토타입화하고 평가하기 위해.

제안 방법

• MLET는 ℝⁿˣᵈ에 속하는 단일 임베딩 테이블 W를 두 레이어 선형 변환 qW₁W₂로 대체한다. 여기서 W₁ ∈ ℝⁿˣᵏ이고 W₂ ∈ ℝᵏˣᵈ이다.
• 모델은 백프로파게이션을 사용해 엔드 투 엔드로 훈련되며, 최종 추론 시에는 오직 W₂만을 사용하여 d차원 출력을 유지한다.
• 이 방법은 깊이 있는 선형 네트워크에서 효과적 랭크가 낮아지는 이론적 통찰을 활용한다. 이는 더 나은 일반화를 가능하게 한다.
• 이 방법은 Criteo-Kaggle 및 Avazu 데이터셋에서 페이스북의 파이토치 기반 DLRM 프레임워크 내에서 구현 및 평가되었다.
• 훈련 메모리는 단일 레이어 훈련 대비 k/d 배수 증가하지만, 추론 메모리는 여전히 d에 비례한다.
• 최종 모델은 추론 시 단일 레이어 형태로 변환되어 기존 시스템과의 호환성을 유지한다.

실험 결과

연구 질문

• RQ1 다층 선형 아키텍처가 정확도와 모델 크기 측면에서 단일 레이어 임베딩보다 뛰어난 성능을 내는가?
• RQ2 왜 다층 선형 네트워크가 단일 레이어보다 더 나은 임베딩을 생성하는가? 이에 대한 기반이 되는 동역학적 메커니즘은 무엇인가?
• RQ3 MLET를 사용해 얼마나 많은 임베딩 테이블 크기를 줄일 수 있으며, 정확도를 유지하거나 향상시킬 수 있는가?
• RQ4 추론 메모리 비용이 추가로 증가하지 않는 조건에서 MLET가 정확도와 모델 크기 양면에서 향상되는 구성이 존재하는가?

주요 결과

• Criteo-Kaggle 및 Avazu 데이터셋에서 MLET는 단일 레이어 임베딩과 동일한 정확도 수준에서 임베딩 테이블 크기를 4–8배 감소시켰다.
• MLET의 최대 LogLoss 개선 폭은 Criteo-Kaggle에서 0.0025, Avazu에서 0.006이며, k가 일정 수준을 초과하면 성능이 포화 상태에 도달한다.
• 일부 k/d 비율에서 MLET는 추가 훈련 메모리 비용 없이도 정확도 향상과 함께 모델 크기 감소를 동시에 달성한다 (즉, k ≤ d). 이는 Avazu에서는 관찰되었지만 Criteo-Kaggle에서는 그렇지 않았다.
• MLET 훈련의 런타임 비용은 단일 레이어 훈련 대비 평균 25% 높으며, 추론 시간은 표준 DLRM과 동일하다.
• 성능 향상은 주로 k/d 비율에 의해 결정되며, 높은 k/d 비율일수록 정확도 향상이 단조롭게 증가한다.
• 낮은 d에서도 이 방법은 효과적이며, 이는 상대적인 모델 크기 감소가 가장 두드러지기 때문이다. 이는 k ≤ d 조건을 활용해 비용 없이도 성과 향상을 달성할 수 있기 때문이다.