[논문 리뷰] Online Model Compression for Federated Learning with Large Models

Tien-Ju Yang, Yonghui Xiao외 4명|arXiv (Cornell University)|2022. 05. 06.

Speech Recognition and Synthesis인용 수 4

한 줄 요약

이 논문은 온라인 모델 압축(OMC)을 제안하며, 이는 대규모 신경망 파라미터를 압축된 정밀도 형식으로 저장하고 전송하면서도, 실시간 계산 중에만 복원하는 경량 프레임워크이다. OMC는 변수별 변환, 가중치 행렬 전용 압축, 부분 파라미터 압축을 통해 전체 정밀도 정확도와 학습 속도를 유지하면서 메모리 사용량과 통신 비용을 최대 59%까지 줄인다.

ABSTRACT

This paper addresses the challenges of training large neural network models under federated learning settings: high on-device memory usage and communication cost. The proposed Online Model Compression (OMC) provides a framework that stores model parameters in a compressed format and decompresses them only when needed. We use quantization as the compression method in this paper and propose three methods, (1) using per-variable transformation, (2) weight matrices only quantization, and (3) partial parameter quantization, to minimize the impact on model accuracy. According to our experiments on two recent neural networks for speech recognition and two different datasets, OMC can reduce memory usage and communication cost of model parameters by up to 59% while attaining comparable accuracy and training speed when compared with full-precision training.

연구 동기 및 목표

• 대규모 모델을 사용하는 피어드레이티드 학습에서 높은 디바이스 내 메모리 사용량과 통신 비용을 해결한다.
• 모델 정확도를 훼손하지 않으면서 엣지 디바이스에서 학습 중인 파라미터의 메모리 프로파일을 줄인다.
• 피어드레이티드 환경에서 빈번한 압축/해제 압축 작업으로 인한 학습 속도 저하를 최소화한다.
• 자원이 제한된 엣지 디바이스에 최신 기술 수준의 자동 음성 인식(ASR) 모델을 구현할 수 있도록 한다.
• 온라인 및 즉각적인 복원을 통해 피어드레이티드 학습에서 효율적인 파라미터 운반 및 저장을 실현한다.

제안 방법

• 메모리 및 통신 오버헤드를 줄이기 위해 전체 정밀도 FP32 대신 압축된 부동소수점 형식(예: S1E5M10)으로 모델 파라미터를 저장한다.
• 전방 또는 역방향 전파가 필요한 경우에만 실시간으로 파라미터를 복원하고, 사용 후 즉시 할당 해제한다.
• 압축 이전에 파라미터 분포를 정규화함으로써 학습 안정성을 높이기 위해 변수별 변환(PVT)을 적용한다.
• 메모리 소비가 가장 높은 구성 요소(즉, 선형 레이어)에만 압축을 제한하기 위해 가중치 행렬 전용 압축을 사용한다.
• 낮은 비트 폭(예: 11비트)으로 90%의 파라미터만 압축하고 나머지 10%는 높은 정밀도로 유지함으로써 부분 파라미터 압축(PPQ)을 구현한다.
• 압축 효율성과 오차 제어를 균형 잡기 위해 조정 가능한 지수 및 가수 비트 할당(예: S1E3M7, S1E5M7)을 사용한 부동소수점 압축을 적용한다.

Figure 1: The illustration of the framework of the proposed online model compression. The cubes with dashed borderlines are transient variables.

실험 결과

연구 질문

• RQ1 피어드레이티드 학습 중 모델 파라미터를 온라인으로 압축함으로써 전체 정밀도 학습 대비 디바이스 내 메모리 사용량을 크게 줄일 수 있을까?
• RQ2 변수별 변환은 압축된 피어드레이티드 학습의 안정성과 수렴에 어떤 영향을 미치는가?
• RQ3 가중치 행렬만 또는 일부 파라미터만 압축할 경우 비트 폭 감소와 모델 성능 간의 상호 상충 관계는 어떠한가?
• RQ4 낮은 비트 폭으로 부분 파라미터 압축이 더 높은 비트 폭의 전체 파라미터 압축보다 정확도와 효율성 측면에서 뛰어나게 성능을 발휘할 수 있는가?
• RQ5 OMC는 전체 정밀도 학습과 유사한 학습 속도를 유지하면서도 피어드레이티드 학습에서 통신 비용을 얼마나 줄일 수 있는가?

주요 결과

• OMC는 전체 정밀도 학습 대비 모델 파라미터의 메모리 사용량과 통신 비용을 최대 59%까지 줄였다.
• 변수별 변환을 적용한 OMC는 Non-Streaming Conformer 모델에서 안정적인 학습 수렴을 달성하였으며, 이를 적용하지 않은 경우와는 달리 정확도 저하가 발생하지 않았다.
• 11비트 형식(S1E3M7)으로 부분 압축한 경우, 13비트 형식(예: S1E5M7, S1E4M8)으로 전체 파라미터를 압축한 것보다 더 낮은 WER를 달성하였다.
• 11비트 부분 압축 설정은 테스트된 모든 13비트 전체 압축 설정보다 빠른 수렴 속도와 더 낮은 단어 오류율(WER)을 기록하였다.
• 빈번한 실시간 압축 및 복원 작업에도 불구하고 OMC는 전체 정밀도 학습과 유사한 학습 속도를 유지하였다.
• OMC는 메모리 및 통신 오버헤드를 효과적으로 관리함으로써 엣지 디바이스에서 대규모 ASR 모델(예: Non-Streaming Conformer)의 학습을 가능하게 하였다.

Figure 2: The illustration of per-variable transformation. The cubes with dashed borderlines are transient variables.