Apple이 Core ML에 6-bit palettization 양자화 기법을 도입해 Stable Diffusion의 메모리 사용량을 대폭 감소시키고 추론 속도를 향상

Context

Stable Diffusion, Large Language Model 등 생성형 AI 모델을 온디바이스에서 실행하려는 개발자 수요가 증가했으나, 모델 크기와 메모리 요구량이 iPhone, iPad, Mac에서의 실행을 제약했다. 기존 Core ML은 모델 로드 시점에 가중치를 완전히 압축 해제했기 때문에 메모리 사용량이 압축 해제된 모델 크기에 달했다.

Technical Solution

• 6-bit palettization 양자화 도입: 16-bit 부동소수점 가중치를 6-bit 인덱스로 압축해 모델 크기 대폭 감소
• 레이어 단위 동적 압축 해제: 가중치를 6-bit로 유지했다가 추론 진행 중 필요한 레이어별로 실시간 변환하도록 변경
• coremltools.optimize 서브모듈 신규 추가: PyTorch/TensorFlow 모델을 Core ML 포맷으로 변환 시 --quantize-nbits 옵션으로 2-bit, 4-bit, 6-bit, 8-bit 양자화 적용 가능
• Attention 레이어 최적화(SPLIT_EINSUM_V2): Query 시퀀스를 512 청크 단위로 분할해 대규모 중간 텐서 생성 방지
• 학습 시점 압축(Training-time Compression): 사전 학습 모델 파인튜닝 중 미분 가능한 가중치 클러스터링으로 4-bit, 2-bit 압축 적용 시 품질 손실 최소화

Impact

iPhone 13에서 Stable Diffusion 실행 속도가 2023년 12월 대비 6-bit palettization 적용 후 측정 가능하게 향상됨. Attention 레이어 최적화(SPLIT_EINSUM_V2)로 Neural Engine 성능이 10~30% 개선. 모델 다운로드 시간 및 온디바이스 디스크 사용량이 대폭 감소(구체적 수치는 아티클에 미기재).

Key Takeaway

모델 압축 시 저수준 메모리 접근 특성(메모리 전송이 실행 경로의 병목)을 고려하면, 완전 압축 해제보다 동적 압축 해제가 더 효율적일 수 있다. 양자화는 사후 처리뿐 아니라 학습 중 적용하면 낮은 비트에서도 품질 손실을 최소화할 수 있다.