[논문 리뷰] LEMON: Lossless model expansion

Yite Wang, Jiahao Su외 7명|arXiv (Cornell University)|2023. 10. 12.

Advanced Neural Network ApplicationsComputer Science인용 수 3

한 줄 요약

LEMON은 사전 훈련된 작은 모델을 사용하여 더 큰 트랜스포머 모델을 초기화하는 손실 없는 모델 확장 프레임워크이다. 이는 성능 저하 없이 더 빠른 훈련을 가능하게 하며, 대칭성 깨는 가중치 초기화 및 최적화된 교육률 스케줄러를 통해 사전 훈련된 모델에서부터 훈련하는 것에 비해 비전 트랜스포머의 경우 56.7%, BERT의 경우 33.2%의 계산 비용 절감을 달성한다.

ABSTRACT

Scaling of deep neural networks, especially Transformers, is pivotal for their surging performance and has further led to the emergence of sophisticated reasoning capabilities in foundation models. Such scaling generally requires training large models from scratch with random initialization, failing to leverage the knowledge acquired by their smaller counterparts, which are already resource-intensive to obtain. To tackle this inefficiency, we present $ extbf{L}$ossl$ extbf{E}$ss $ extbf{MO}$del Expansio$ extbf{N}$ (LEMON), a recipe to initialize scaled models using the weights of their smaller but pre-trained counterparts. This is followed by model training with an optimized learning rate scheduler tailored explicitly for the scaled models, substantially reducing the training time compared to training from scratch. Notably, LEMON is versatile, ensuring compatibility with various network structures, including models like Vision Transformers and BERT. Our empirical results demonstrate that LEMON reduces computational costs by 56.7% for Vision Transformers and 33.2% for BERT when compared to training from scratch.

연구 동기 및 목표

• 작은 사전 훈련된 모델의 지식을 활용하여 대규모 모델을 사전 훈련 없이 훈련하는 데 발생하는 비효율성을 해결하기 위해.
• 너비나 깊이 증가가 나누어떨어지지 않는 경우에도 가능한 다양한 트랜스포머 아키텍처와 호환되는 손실 없는 모델 확장 방법을 개발하기 위해.
• 모델 성능을 유지하면서 훈련 시간과 계산 비용을 줄이기 위해.
• 특히 맞춤형 학습률 스케줄러를 통해 확장된 모델의 훈련 레시피를 최적화하기 위해.

제안 방법

• LEMON은 확장된 레이어에서 복제된 뉴런에 대해 대칭성을 깨는 초기화 전략을 사용하여, 동일한 출력 가중치를 부여하지 않음으로써 부족한 중복을 방지하고 진정한 능력 향상을 가능하게 한다.
• 너비 확장 중에 LayerNorm 레이어를 처리하기 위해 평균 확장을 도입하여, 비가운데 너비 증가(예: 512 → 768)일 경우에도 호환성을 유지한다.
• 아키텍처적 차이에 맞게 확장 과정을 조정함으로써 표준 및 Pre-LN 트랜스포머를 포함한 다양한 트랜스포머 유형을 지원한다.
• 사전 훈련에서 사용된 최대 학습률을 유지하면서 더 빠른 감쇠를 적용하는 맞춤형 학습률 스케줄러를 도입하여, 확장된 모델의 수렴을 향상시킨다.
• 작은 사전 훈련된 모델에서 더 큰 대상 모델로 직접 가중치를 전이할 수 있도록 하여 기능적 동치성을 유지한다.
• 실증적 검증을 통해 확장된 모델가 사전 훈련에서부터 훈련한 모델와 동일한 성능을 달성하며, 훨씬 줄어든 훈련 시간을 기록함을 확인하였다.

실험 결과

연구 질문

• RQ1 비가운데 너비나 깊이 증가가 있을 경우에도, 성능 저하 없이 사전 훈련된 작은 트랜스포머 모델을 더 큰 모델로 확장할 수 있는가?
• RQ2 모델 확장 과정에서 대칭성을 깨는 방법은 무엇인가? 이를 통해 확장된 모델이 원본 모델을 초월한 진정한 능력을 획득할 수 있는가?
• RQ3 확장된 모델의 빠른 수렴을 위해 최적의 훈련 레시피—특히 학습률 스케줄링—는 무엇인가?
• RQ4 제안된 확장 방법은 ViT와 BERT를 포함한 다양한 트랜스포머 아키텍처로 일반화될 수 있는가?
• RQ5 사전 훈련에서부터 훈련하는 것에 비해 손실 없는 모델 확장이 계산 비용을 얼마나 절감하는가?

주요 결과

• LEMON은 사전 훈련에서부터 훈련하는 것에 비해 비전 트랜스포머의 경우 56.7%, BERT의 경우 33.2%의 계산 비용 절감을 달성한다.
• 확장된 ViT(12,768)는 사전 훈련에서부터 훈련한 모델와 동일한 성능을 85 에포크 만에 달성하여 훈련 시간이 28.3% 감소하였다.
• 이 방법은 원본 모델과 정확한 기능적 동치성을 유지하여 확장 후 성능 저하 없이 기능을 보존한다.
• 동일한 최대 학습률을 사용하지만 더 빠른 감쇠를 적용하는 최적화된 학습률 스케줄러는 확장된 모델의 훈련 효율성을 크게 향상시킨다.
• 비가운데 너비 증가가 있을 경우에도 ViT, BERT, Pre-LN 유형을 포함한 다양한 트랜스포머 아키텍처와 호환되며, 이는 확장 방법의 일반화 가능성을 뒷받침한다.
• 실증 결과는 확장된 모델가 원본 모델의 모든 지식을 손실 없이 이관함을 확인하여, 확장의 손실 없는 성격을 검증하였다.