기술 요약: 운동 보상 가중치 압축 (MCWC)

문제 제기

신경망 배포는 연산 능력이 아닌, 체크포인트를 가속기에 스테이징하는 데 필요한 저장 공간, 대역폭, 콜드 스타트 오버헤드에 의해 병목 현상이 발생하고 있습니다. 기존 압축 파이프라인은 훈련 후 양자화 (post-training quantization) 와 가지치기 (pruning) 가 주류를 이루며, 일반적으로 각 레이어를 독립적인 객체로 취급합니다. 이러한 접근법은 모델의 입력 - 출력 함수를 변경하지 않고 숨겨진 단위나 어텐션 헤드를 재색인할 수 있는 등 함수를 보존하는 대칭성으로 인해 발생하는 상당한 레이어 간 중복성을 간과합니다. 그 결과, 이러한 방법들은 트랜스포머와 같은 현대 아키텍처에 내재된 강력한 레이어 간 유사성과 예측 가능성을 활용하지 못합니다.

방법론

MCWC 는 모델 깊이를 비디오 프레임과 유사한 시퀀스 신호로 취급하는 가중치 전용 코덱을 제안합니다. 핵심 전제는 치환 대칭성으로 인해 발생하는 "운동"을 제거함으로써 레이어 간 가중치 진화가 매우 예측 가능해져 효율적인 잔차 코딩이 가능해진다는 것입니다. 파이프라인은 네 가지 주요 단계로 구성됩니다:

1. 기능 정렬을 통한 운동 보상

인코더는 가중치 블록 (예: FFN 숨겨진 단위, 어텐션 헤드) 에 대한 함수 보존 치환을 식별하여 레이어 간 대응을 최대화합니다.

• 메커니즘: 각 레이어 $\ell$에 대해, 이 방법은 이전 레이어 $\ell-1$의 대응 블록과 블록을 정렬하는 치환 $\Pi_\ell$을 선택합니다. 이는 후보 블록과 이전에 정렬된 참조 간의 유사성 (가중치 기반 코사인 유사성과 활성화 통계 유사성의 하이브리드) 을 최대화하는 선형 할당 문제를 해결함으로써 달성됩니다.
• 결과: 이는 가중치 시퀀스를 인접 레이어가 기하학적으로 비교 가능한 "정렬된 좌표계"로 변환하여, 레이어 간 차이 에너지 (energy) 를 크게 감소시킵니다.

2. 레이어 순차 예측 코딩

정렬된 후, 가중치는 시간적 시퀀스로 취급됩니다.

• 예측: 경량 레이어 순차 예측기 ($g_\theta$) 는 이전 디코드 컨텍스트 (레이어 $\ell-1$의 정렬된 블록) 에 기반하여 현재 레이어의 정렬된 블록을 추정하며, 레이어 및 블록 유형 임베딩에 조건부입니다.
• 잔차: 코덱은 양자화된 예측 잔차 ($R = \text{AlignedTarget} - \text{Prediction}$) 만을 인코딩합니다.
• 키프레임: 긴 시퀀스 동안 오차 누적 (드리프트) 을 방지하기 위해 주기적인 "키프레임"이 잔차가 아닌 절대값으로 인코딩됩니다. 이는 디코더가 컨텍스트를 재설정할 수 있도록 하는 동기화 지점 역할을 합니다.

3. 학습된 양자화 및 엔트로피 모델링

• 양자화: 잔차는 채널별 스텝 크기와 평균을 갖는 학습된 스칼라 양자화를 사용하여 이산 코드 인덱스로 매핑됩니다.
• 엔트로피 코딩: 학습된 엔트로피 모델은 양자화된 심볼에 확률을 할당하여 산술 코딩을 가능하게 합니다. 이 모델은 레이어 인덱스, 블록 유형, 그리고 예측기 출력에서 유도된 로컬 통계에 조건부입니다.
• 훈련: 전체 시스템 (예측기, 양자화기, 엔트로피 모델) 은 재구성된 가중치가 작업 성능을 유지하도록 보장하는 증류 손실 (distillation loss) 인 $L_{task}$가 포함된 레이트 - 왜곡 목적 함수 ($L = L_{task} + \lambda E[-\log p_\psi(c)]$) 하에서 엔드 - 투 - 엔드로 훈련됩니다.

4. 사이드 정보 및 디코딩

비트스트림에는 엔트로피 코딩된 잔차/키프레임, 치환 사이드 정보 (오버헤드를 줄이기 위해 델타 코딩된 레허머 코드로 인코딩됨), 그리고 양자화 파라미터가 포함됩니다. 디코더는 이 과정을 역으로 수행합니다: 엔트로피 디코딩, 역양자화, 예측을 통한 잔차 재구성, 그리고 원래 텐서 구조를 복원하여 배포할 수 있도록 역치환을 적용합니다.

주요 기여

1. 가중치 공간의 운동 보상: 비디오 압축의 운동 추정과 유사하게, 표현된 함수를 유지하면서 레이어 간 대응을 증가시키기 위한 블록 단위 치환 선택의 도입.
2. 깊이를 위한 예측 코덱: 깊이를 시간 축으로 취급하는 새로운 아키텍처로, 명시적인 키프레임과 잔차 코딩을 활용하여 정렬된 레이어를 예측 가능한 "프레임"으로 변환.
3. 엔드 - 투 - 엔드 레이트 - 왜곡 훈련: 디코딩 후 정확하고 학습된 확률 모델 하에서 낮은 엔트로피를 갖도록 이산 잔차 스트림을 보장하기 위해 양자화 및 엔트로피 모델을 공동으로 훈련.
4. 실증적 검증: 트랜스포머 언어 모델 (Pythia) 과 비전 분류기 (ViT) 모두에서 개선된 레이트 - 정확도 파레토 프론티어 입증 및 정렬, 예측, 엔트로피 모델링이 전체 이득을 위해 모두 필요함을 확인하는 제거 실험 (ablation studies).

결과

MCWC 는 균일 INT8/INT4 양자화, SmoothQuant, GPTQ, AWQ, 그리고 학습된 스칼라 양자화를 포함한 강력한 베이스라인에 대해 Pythia-1.4B(언어 모델링) 와 ViT-B/16(비전 분류) 에서 평가되었습니다.

• 언어 모델링: 공격적인 압축 비율 (예: 약 2.8 비트/파라미터) 에서 MCWC 는 16.40 의 퍼플렉시티를 달성하여 Prune+Quant(18.10) 와 학습된 스칼라 양자화 (17.30) 를 능가합니다. 중간 비율 영역 (약 4.2 비트/파라미터) 에서 15.95 의 퍼플렉시티를 달성하여 균일 INT4(16.85) 를 상회하고 GPTQ(16.40) 및 AWQ(16.25) 와 같은 헤시안 인식 방법과 경쟁력을 유지합니다.
• 비전 분류: 2.8 비트/파라미터에서 MCWC 는 ImageNet-1k 에서 80.15% 의 Top-1 정확도를 유지하여 Prune+Quant(77.60%) 와 학습된 스칼라 양자화 (78.70%) 를 능가합니다. 4.2 비트/파라미터에서 80.92% 의 정확도를 달성하여 균일 INT4(79.35%) 를 초과합니다.
• 디코딩 효율성: 예측 재구성과 역치환은 선형 시간 연산이며, 주요 비용은 엔트로피 디코딩과 텐서 물질화이므로, 디코딩 시간은 표준 양자화 베이스라인 (예: CPU 에서 Pythia-1.4B 기준 약 2.5 초) 과 경쟁력 있게 유지됩니다.
• 진단: 제거 실험은 기능 정렬이 인접 레이어 간 코사인 유사성을 증가시킴 (Pythia-1.4B 의 경우 약 0.28 에서 약 0.67 로) 과 정규화된 잔차 에너지를 감소시켜, 정렬이 더 예측 가능한 좌표계를 노출함을 확인합니다.

중요성 및 주장

이 논문은 MCWC 를 저정밀도 추론 커널 (예: INT8/INT4 행렬 곱셈) 을 대체하는 것이 아니라, 특수화된 체크포인트 저장 및 전송 코덱으로 위치시킵니다. 그 주요 가치는 대규모 모델 가중치를 배포하고 반복적으로 물질화하는 것과 관련된 저장 공간, 네트워크 대역폭, 그리고 콜드 스타트 지연 시간을 줄이는 데 있습니다.

저자들은 운동 보상을 통해 레이어 간 중복성을 활용함으로써 MCWC 가 기존 양자화와 직교하는 병목 현상을 해결한다고 주장합니다. 이 방법은 체크포인트가 여러 번 이동, 저장 또는 복제되는 배포 환경 (예: 서버리스 콜드 스타트, 모델 레지스트리, 엣지 배포) 에서 가장 효과적이며, 일회성 오프라인 인코딩 비용을 반복적인 배포 절감으로 상쇄할 수 있게 합니다. 이 연구는 신경 가중치를 고립된 텐서가 아닌 깊이를 통한 구조화된 궤적으로 취급함으로써 효율적인 모델 압축을 위한 새로운 경로를 제시합니다.