[ICLR 2024] Compression of Neural Networks 연구 동향

Introduction

심층신경망(Deep Neural Networks, DNNs) 모델은 깊은 구조를 통해 데이터의 복잡한 비선형 관계(Non-linear Relationship)를 학습합니다. 성능이 뛰어나지만, 많은 연산 자원과 메모리가 필요하다는 단점도 있습니다. 이러한 단점을 해결하기 위해 모델 압축 기술에 대한 연구가 오랫동안 중요하게 다뤄져 왔습니다.

최근에는 거대 언어 모델(Large Language Model, LLM)의 발전으로 인해 고성능 DNNs의 경량화에 대한 관심이 더욱 높아지고 있으며, 모델 압축 연구는 모델의 크기와 연산 복잡성을 줄이면서도 성능을 유지하거나 향상시키는 것을 목표로 합니다. 특히 모바일 기기와 같이 자원이 제한된 환경에서 DNNs를 활용하기 위해서는 모델 경량화가 필수적입니다. 모델 경량화를 통해 DNNs를 보다 다양한 응용 분야에 효과적으로 적용할 수 있을 것으로 기대됩니다.

LG AI연구원 Data Intelligence(DI) Lab은 “ICLR 2024 The 5th Workshop on Practical ML for Low Resource Settings (PML4LRS)”에서 “LESS: LEARNING TO SELECT A STRUCTURED ARCHITECTURE OVER FILTER PRUNING AND LOW-RANK DECOMPOSITION^[1]” 연구를 발표했습니다. 또한, Qualcomm 등 다른 기업의 연구자들과 함께 Lightweight Neural Networks에 대한 연구 현황 및 앞으로 나아가야 할 방향에 대해 함께 논의하였습니다. LG AI연구원이 발표한 논문과, 학회에서 얻은 중요한 인사이트들을 2편의 논문과 함께 소개해 드리겠습니다.

1. LESS: LEARNING TO SELECT A STRUCTURED ARCHITECTURE OVER FILTER PRUNING AND LOW-RANK DECOMPOSITION

그림 1. Low-rank Compression, Pruning, Hybrid 방법들의 압축률에 따른 최고 성능 비교.
X축은 FLOPs를 나타내며, 파란색이 우리의 방법을 적용했을 때의 성능을 나타냄^[1]

Filter Pruning과 Low-rank Decomposition은 모두 신경망 모델을 압축하기 위한 방법이지만, 서로 다른 가정에 기반을 두고 있습니다. Filter Pruning은 일부 필터가 중복되어 있다고 가정하고 필터 단위로 제거하는 반면, Low-rank Decomposition은 가중치 행렬의 일부 차원이 덜 중요하다고 보고 텐서의 랭크를 낮추는 방식으로 압축을 수행합니다. 그러나 기존의 연구들은 그림 1에서 보이는 것처럼 이 두 가지 방법을 효과적으로 결합하지 못해 개별 방법을 사용했을 때보다 성능 향상이 제한적이었습니다.

이에 따라 본 논문에서는 Filter Pruning과 Low-rank Decomposition을 효율적으로 통합하는 방법을 연구하였습니다. 특히, 휴리스틱에 많이 의존하지 않고 적절한 Rank를 선택하기 위해서는 학습 기반의 방법이 필요하다고 판단했고, 이를 위해 랭크 선택을 위한 새로운 SGD 기반 학습 기법인 Threshold Learning을 제안했습니다. 또한 기존의 Mask Learning 기반 Filter Pruning 방법을 개선하여 사용하였습니다. 이를 통해 본 논문에서는 End-to-end로 학습 가능한 Differentiable한 모듈을 만들 수 있었습니다.

그림 2. LeSS 알고리즘의 Forward 과정 Overview^[1]

제안된 방법은(그림 2) 먼저 가중치 텐서를 재배열한 후 SVD를 수행하고, 특이값을 임계값(Threshold)으로 제어할 수 있는 스칼라 변수를 학습합니다. 이 스칼라 변수는 크로스 엔트로피 손실을 최소화하면서도 동시에 목표 Target 제약을 만족할 수 있도록 학습됩니다. Filter Pruning을 위해서는 각 채널에 게이트를 두어 마스크를 학습하는 방법을 사용하였으며, 여기서도 마스크 변수가 크로스 엔트로피 손실을 최소화하면서 Target 제약을 만족하도록 학습됩니다. 특히 본 논문에서는 Scheduled Sigmoid 함수를 사용하여 점진적으로 마스크가 0 또는 1에 가까운 극단적인 값을 가지도록 유도하였으며, 모든 게이트의 합을 이용해 희망하는 예산에 가깝도록 Sparsity를 조절하였습니다.

제안된 LeSS 방법은 기존의 Pruning, Low-rank Decomposition, 그리고 이를 결합한 하이브리드 방법들을 큰 차이로 앞섰습니다. ResNet56을 CIFAR10 데이터셋에 적용한 경우, LeSS는 모델의 FLOPs를 80%까지 줄이면서도 Baseline 모델보다 높은 성능을 달성하였습니다. ImageNet 데이터셋에서도 FLOPs를 65% 감소시키면서 Baseline과 유사한 성능을 보였습니다. 다양한 네트워크 구조에 대한 추가 실험에서도 LeSS는 항상 우수한 압축 성능을 보였습니다. 특히 경량 모델로 잘 알려진 MobileNetV2에 적용했을 때, 파라미터 수를 36%, FLOPs를 35% 감소시키면서도 Baseline보다 높은 성능을 달성할 수 있었습니다.

LeSS는 조합 최적화 문제(COP)를 효율적으로 해결하기 위해 휴리스틱을 피하고 미분 가능한 접근 방식을 사용하도록 설계되었습니다. 압축을 고려한 특별한 Regularizer 없이도, 필터와 Rank를 선택하는 효율적인 미분 학습을 통해 반복 없이 원하는 Target Resource를 만족하는 네트워크를 만들 수 있습니다. 본논문은 Vision Task에 한정하여 방법의 효과를 실험적으로 증명하였지만, 다른 Task에도 효과적으로 적용 가능할 것으로 생각됩니다.

2. SHEARED LLAMA: ACCELERATING LANGUAGE MODEL PRE-TRAINING VIA STRUCTURED PRUNING^[2]

구조화된 가지치기(Structured Pruning)는 모델의 파라미터 그룹을 제거하여 모델을 압축하고 추론 속도를 높이는 방법론입니다. 그러나 기존의 구조화된 Pruning 방식은 종종 널리 사용되는 아키텍처에서 벗어나는 비정상적인 모델 구성을 만들어냅니다. 예를 들어, 레이어마다 다른 수의 헤드를 가진 모델을 생성하면, 이는 일반적인 균일한 크기의 레이어 구성에 비해 추론 시 오버헤드를 발생시킵니다^[3]. 본 연구에서는 소스 모델을 앞서 지정한 모든 Target Structure를 만족하도록 Pruning 하는 것을 목표로 합니다.

본 논문의 Contribution은 세 가지로 요약될 수 있습니다. (1) LLM분야에서 처음으로 Pre-defined된 아키텍처를 Target 모델로 선정하고, 같은 구성을 갖도록 현재 아키텍처를 Pruning 하였습니다. (2) Pruning시에 Target Model의 모양을 랑그랑주를 사용하여 Pruning 모델의 모양에 직접적으로 제약 조건을 부과하여 학습하였습니다. (3) 마지막으로 동적 배치 로딩으로 사용하여 Pre-train Model에 사용하였던 학습데이터 분포를 그대로 Pruned Model에 사용할 경우, 성능저하가 일어날 수 있는 문제점을 보완하였습니다.

그림 3. Target Structure Pruning 방법의 작동 모식도. 밝은 색은 Pruned된 Substructure을 나타냄^[2]

더 자세하게는, 트랜스포머 아키텍처에 기반한 LLM 모델들은 모든 차원을, 다양한 Granularity를 갖는 구조를 축소해야 하기 때문에 어려운 문제입니다. 이를 해결하기 위해 본 논문에서는 기존 사전 학습 모델의 구성을 목표 아키텍처로 사용합니다. 이는 이러한 Pre-Defined된 구성이 이미 모델의 표현력과 추론 효율성 사이의 균형을 잘 최적화하고 있다는 직관에 기반합니다. 예를 들어, 2.7B 모델을 생성한다고 하면, INCITE-Base-3B^[4] 아키텍처를 목표 구조로 사용합니다.

본 논문은 모든 레이어에 걸쳐 지속되는 레이어 및 Hidden 차원과 같은 전역적인 것부터 어텐션 헤드와 중간 차원과 같은 지역적인 것까지 다양한 Granularity에서 모델 매개변수에 대한 Pruning 마스크를 학습합니다. 각 마스크 변수는 연관된 하위 구조를 Pruning할지 유지할지를 제어합니다. 예를 들어, 해당 레이어의 z_layer가 0이면 해당 레이어를 제거합니다. 그림 3은 Pruning 마스크가 어떻게 Pruning 구조를 제어하는지 예시를 보여줍니다.

${\tilde{\mathcal{L}}}^{\text{head}}(\lambda ,\varphi ,z)={\lambda }^{\text{head}}\cdot (\sum z^{\text{head}}-H_{\mathcal{T}})+{\varphi }^{\text{head}}\cdot {(\sum z^{\text{head}}-H_{\mathcal{T}})}^2.$

위 식에서처럼 본 논문은 Target Model의 구조 정보를 바로 Loss Function에 구성하여 Explicit하게 모델 구조에 맞게 학습시킵니다.

학습할 때 사용되는 전체 Loss는 다음과 같습니다. 모든 Granularity에 해당하는 Loss가 포함되어 구성되어 있음을 확인할 수 있습니다(Head, Int, Layer, Hidden).

${\mathcal{L}}_{\text{prune}}(\theta ,z,\lambda ,\varphi )=\mathcal{L}(\theta ,z)+\sum _{j=1}^{L_{\mathcal{S}}}{\tilde{\mathcal{L}}}_j^{\text{head}}+\sum _{j=1}^{L_{\mathcal{S}}}{\tilde{\mathcal{L}}}_j^{\text{int}}+{\tilde{\mathcal{L}}}^{\text{layer}}+{\tilde{\mathcal{L}}}^{\text{hidden}}$

마지막으로 동적 배치 로딩을 통해 Pre-trained 모델과 Pruned 모델의 Loss 차이를 기반으로 Data의 Proportion을 동적으로 바꾸어 성능향상에 기여하였습니다. 특히 이 방법은 On-the-fly 방법으로 데이터에 제약이 없이 Open Source 데이터셋도 바로 활용 가능하다는 장점이 있습니다. 본 논문에서 제시하는 방법을 다양한 Dataset에 대해 평가해 보았고, 그 우수함을 확인할 수 있었습니다. (그림 4)

그림 4. Sheared-LLaMA의 성능 결과표^[2]

본 연구에서는 구조화된 가지치기(Structured Pruning) 기법을 활용하여 언어 모델의 압축과 추론 속도 향상을 이루고자 하였습니다. 기존의 구조화된 Pruning 방식은 종종 비정상적인 모델 구성을 만들어내는 한계점이 있었으나, 본 연구에서는 이를 극복하기 위해 소스 모델을 미리 정의된 Target 구조에 맞춰 Pruning하는 접근 방식을 제안하였습니다. 본 연구의 접근 방식은 기존 구조화된 Pruning 기법의 한계를 극복하고, 효과적으로 언어 모델을 압축하며 추론 속도를 높일 수 있는 새로운 방향을 제시하였다는 점에서 의의가 있습니다. 제안된 방법론은 실제 응용 분야에서 LLM의 효율적인 배포와 활용을 가능케 할 것으로 기대됩니다. 향후 연구에서는 본 연구에서 제시한 아이디어를 더욱 발전시켜, 다양한 아키텍처와 데이터셋에 적용해 봄으로써 그 효과를 검증하고 일반화할 수 있을 것으로 생각됩니다.

3. SPARSE MODEL SOUPS: A RECIPE FOR IMPROVED PRUNING VIA MODEL AVERAGING^[5]

이번 논문은 Pruning 방법과 Ensemble방법을 결합하여 새로운 Compression Paradigm을 제시했습니다. 앞서 말씀드렸던 것처럼 Pruning은 오래전부터 불필요한 파라미터들을 제거하는 방법이 많이 사용되어 왔습니다. 이 방법을 적용하면, 파라미터 저장 공간과 연산량을 모두 줄이면서도 압축 이전 모델과 유사한 성능을 유지할 수 있다는 큰 장점이 있습니다. 더불어 여러 모델을 결합한 Ensemble을 통해 예측 성능을 개선하는 방법들이 제안되어 왔습니다^[6]. 특히 최근에는 Ensemble 모델의 추론 비용 저감을 위해 pruned된 모델의 ensemble을 활용하는 방법이 연구되었습니다. 파라미터들을 모두 평균하여 단일 모델을 구성하는 Model Soups 방법^[7]도 제안되었으나, 이러한 방법은 모든 매개변수를 평균하면 전체 희소성이 감소할 수 있다는 문제점이 존재합니다. 이에 본 연구에서는 Pruning된 모델을 재 학습하여 평균화에 적합한 모델을 생성하는 방법을 제안합니다. 이는 전이 학습의 Finetuning 단계와 유사하게 다양한 하이퍼파라미터 구성을 탐색하여 이루어지게 됩니다. 이 논문의 Contribution은 다음과 같습니다.

1. Pruning된 모델을 다양한 하이퍼파라미터로 재 학습하여 평균하면, 개별모델보다 일반화 성능 및 OOD 성능이 뛰어난 모델을 만들 수 있음을 확인했습니다.

2. Sparse Model Soups (SMS)방법을 제안하여, 각 Pruning - 재 학습 단계를 평균 모델에서 시작하는 아이디어를 제안하였습니다.

3. SMS 방법을 다양한 최신 방식과 통합하여, 학습 중 pruning 성능을 개선하였습니다.

그림 5. (왼쪽) SMS 알고리즘의 모식도 / (오른쪽) SMS 수도코드^[5]

본 연구는 매우 Simple한 방법을 제안합니다. 그림 5(오른쪽)에서 보는 바와 같이 Prune - Retrain - Merge를 반복하며 파라미터 수를 줄이면서 성능을 향상시키는 방법입니다. 여기서 특이점은 같은 모델에서 시작하여 Seed, Weight Decay, Retraining Epochs, Learning Rate 등을 차별화하여 m개의 같은 구조이지만 서로 다른 특성을 갖도록 학습된 모델을 얻고, 이들의 파라미터를 평균하여 단일 모델로 병합한다는 점입니다. 이 방법의 경우, 모두 같은 구조로 Pruning된 모델에서 구조적 변경없이 재학습이 일어나기 때문에 단일모델로 병합된 후에도 같은 Pruning Ratio를 유지할 수 있다는 장점이 있습니다. 특히 m개의 모델을 병렬적으로 학습할 수 있기 때문에 학습속도 측면에서도 큰 이점을 갖습니다. 또한 이 방법은 Scratch부터 모델을 학습할 필요 없이, Pre-trained 모델에 적용 가능한 방법이라는 것도 큰 이점입니다. 파라미터를 평균할 때는 각 모델에 동일한 가중치 1/m을 부여하여 평균값(UniformSoup)으로 사용하거나, 검증 정확도에 따라 모델을 나열하고 이전 검증 정확도보다 높은 모델의 경우에만 평균하는 방법(GreedySoup)을 사용하였습니다.

그림 6. WideResNet-20 모델의 CIFAR-100 성능
(a) UniformSoup과 GreedySoup의 다양한 sparsity에 따른 성능비교
(b) 다양한 Retraining Epochs에 따른 성능비교^[5]

본 연구는 Vision Task에 방법론을 적용하여, 그 효과를 입증하였습니다. 우선적으로 Uniform하게 m개의 모델을 평균하는 방법과 모델의 성능에 따라 개선된 모델들에 한정하여 평균화한 방법을 비교해 보았을 때, 유의미한 차이는 없는 것으로 보입니다. 두 번째로 Retraining Epochs수가 40 이상이면 성능이 Saturation되는 것을 확인할 수 있습니다.

그림 7. ImageNet 데이터셋에 대해 ResNet-50모델에 SMS를 적용한 성능 결과표^[5]

ImageNet에 적용한 결과를 보면 이전 연구들에 비해 성능향상이 이루어졌고, 특히 기존 다른 방법들과 Orthogonal하게 병합하여 사용할 경우, Consistent하게 성능향상이 이루어지는 것을 확인 할 수 있었습니다^[5].

본 연구에서 제안된 방법은 기존의 Pruning 및 Ensemble 기법의 장점을 결합하여, 모델 압축과 성능 향상을 동시에 이룰 수 있는 새로운 접근법을 제시하였다는 점에서 의의가 있습니다. 특히 이 논문은 최근 LLM 에서도 LoRA모듈을 효과적으로 압축하기 위한 연구에서도 활발하게 연구되고 있는 Model Soups를 Efficient하게 활용했다는 점도 돋보이며, LLM 등의 다른 Domain에도 적용 시 이를 통해 거대 언어 모델의 효율적인 배포와 활용이 가능해질 것으로 기대됩니다. 향후 연구에서는 제안된 방법을 다양한 모델 아키텍처와 데이터셋에 적용하여 그 효과를 검증하고, 더욱 발전된 모델 압축 기법으로 발전시켜 나갈 수 있을 것입니다.

추가적으로, ICLR 2024에 참석하여 기업 연구원들과 심도 있는 논의를 통해 최근 언어모델(LLM) 압축 연구가 매우 활발하게 진행되고 있다는 것에 공감했습니다. 주목할 만한 점은, 기존 Vision과 Audio 분야에서 연구되었던 방법론에서 얻은 통찰을 바탕으로 LLM 모델의 고유한 특성에 맞게 이를 적용하려는 시도가 활발히 이루어지고 있다는 것입니다. 예를 들어, 지식 증류(Knowledge Distillation)와 같은 전통적인 모델 압축 기법이 여전히 유효하게 사용되고 있으며, 여기에 더해 교사 모델과 학생 모델 간의 상호작용을 최적화하는 새로운 기법들이 연구되고 있었습니다. 이러한 방법들은 모델의 성능을 유지하면서도 파라미터 수를 크게 줄일 수 있다는 점에서 많은 주목을 받고 있습니다.

또한, 양자화(Quantization)와 저정밀도 연산(Low-Precision Computation) 기술도 언어모델 압축에 있어 중요한 역할을 하고 있음을 확인할 수 있었습니다. 특히, 최신 연구들은 8비트 이하의 저정밀도 양자화가 모델의 성능 저하 없이 효과적으로 수행될 수 있음을 보여주고 있습니다. 이러한 기술들은 거대 언어모델의 배포 및 실시간 응답 속도 향상에 큰 기여를 할 것으로 기대됩니다.

이와 더불어, 분산 학습(Distributed Training) 및 모델 분할(Model Parallelism) 기법을 활용하여 거대 언어모델을 더욱 효율적으로 학습하고 운영할 수 있는 방안도 활발히 논의되고 있었습니다. 이러한 기법들은 특히 클라우드 환경에서 거대 언어모델을 운영하는 기업들에게 매우 유용한 솔루션을 제공할 수 있을 것으로 보입니다.

Conclusion

심층 신경망(DNNs)의 압축 기술은 이러한 혁신적인 언어모델의 실용화와 상용화를 위해 필수적인 연구 분야입니다. 모델 압축을 통해 메모리 사용량을 줄이고, 추론 속도를 높이며, 배포 비용을 절감하는 등의 중요한 이점을 얻을 수 있기 때문입니다. 이는 대규모 모델을 운영하는 데 있어 매우 핵심적인 요소로 작용합니다.

경량화된 모델은 특히 모바일 기기나 임베디드 시스템과 같이 자원이 제한된 환경에서도 우수한 성능을 발휘할 수 있게 해줍니다. 이는 실시간 음성 인식, 증강 현실, 자율 주행 등 다양한 분야에서의 활용 가능성을 크게 확장할 수 있습니다. 이러한 응용 분야들은 신속하고 효율적인 추론을 필요로 하는데, 경량화된 모델은 고성능 서버 기반 모델과 유사한 성능을 제공하면서도 자원 소모를 최소화하여 효율성을 극대화할 수 있습니다. 뿐만 아니라, 모델 압축은 배포 비용 절감이라는 경제적 이점도 있습니다. 대규모 모델을 클라우드 환경에서 운영할 때 저장 공간과 전송 대역폭이 주요 비용 요인으로 작용하는데, 압축 기술을 활용하면 이러한 비용을 효과적으로 절감할 수 있습니다. 이는 기업의 운영비 절감에 직접적으로 기여할 수 있습니다. 나아가, 경량화된 모델은 네트워크 환경이 불안정하거나 제한적인 상황에서도 원활히 동작할 수 있어, 보다 광범위한 사용자에게 서비스 접근성을 제공할 수 있게 됩니다.

종합해 보면, LLM을 비롯한 심층 신경망의 모델 압축 연구는 단순한 기술적 발전을 넘어, 실제 응용 분야에서의 혁신을 가능케 하는 매우 중요한 연구 영역입니다. 모델 압축 연구가 단순히 이론적 접근에 그치지 않고 실험적 검증과 적용 사례 중심으로 다뤄진다면, 실제 산업 현장에서의 적용 가능성도 높일 수 있을 것입니다. 이를 통해 언어모델 압축 기술도 앞으로 더욱 실용적이고 효과적으로 발전해 나갈 것입니다. LG AI연구원의 DI Lab에서는 앞으로 이 분야의 지속적인 연구 개발을 이어 나갈 예정이며, 더욱 다양하고 혁신적인 활용 사례를 만들어 나갈 것입니다.

 

▶[ICLR 2024]Foundation Model 분야 최신 연구 동향