IPU 메모리 절감을 위한 나비형 분해 기법

초록

본 논문은 메모리 용량이 제한적인 인텔리전스 프로세싱 유닛(IPU)에서 모델 압축을 위해 나비형(Factorization) 구조를 적용하는 방법을 제시한다. 완전 연결층과 합성곱층을 나비형 및 픽셀화 나비형으로 대체함으로써 98.5%의 압축률을 달성하고, GPU 대비 1.3~1.6배의 연산 성능 향상을 기록하였다. CIFAR‑10 실험에서는 학습 시간이 1.62배 빨라지는 효과를 확인하였다.

상세 분석

IPU는 수천 개의 작은 코어와 온칩 메모리를 고속 패브릭으로 연결한 대규모 병렬 처리 장치로, 기존 GPU와는 메모리 용량이 현저히 작다는 구조적 차이를 가진다. 이러한 특성은 대규모 딥러닝 모델을 그대로 포팅할 경우 메모리 부족 현상을 초래하므로, 모델 압축 기술이 필수적이다. 논문은 이 문제 해결을 위해 나비형 분해(butterfly factorization)를 선택한다. 나비형 구조는 O(N log N) 연산 복잡도와 O(N log N) 파라미터 수를 갖는 계층적 행렬 분해 방식으로, 완전 연결층을 여러 단계의 스위치와 대각 행렬 곱으로 대체한다. 이때 각 단계는 2×2 블록 교환(swap)과 스케일링을 수행하므로, 메모리 접근 패턴이 규칙적이며 IPU의 고속 로컬 메모리와 잘 맞는다.

구현 측면에서 저자들은 두 가지 변형을 제안한다. 첫 번째는 전통적인 나비형 구조이며, 두 번째는 픽셀화 나비형(pixelated butterfly)으로, 입력 특성을 공간적으로 재배열해 더 작은 스케일링 행렬을 사용한다. 픽셀화는 특히 합성곱층을 대체할 때 유용한데, 공간 차원을 유지하면서 파라미터 수를 추가로 감소시킨다. 두 구조 모두 파라미터를 1/64 수준으로 축소하면서도 정확도 손실을 최소화한다.

성능 평가에서는 CIFAR‑10 데이터셋을 이용해 ResNet‑18 기반 모델에 나비형 레이어를 삽입하였다. 압축 후 모델은 원본 대비 메모리 사용량이 98.5% 감소했으며, IPU에서의 실행 시간은 나비형 경우 1.3배, 픽셀화 나비형 경우 1.6배가량 단축되었다. GPU와 비교했을 때도 동일한 압축 비율에서 IPU가 더 높은 처리 효율을 보였다. 학습 시간 측면에서는 전체 파이프라인이 1.62배 빨라졌으며, 이는 메모리 대역폭 제한이 완화된 결과로 해석된다.

한계점으로는 나비형 구조가 특정 네트워크 아키텍처에만 자연스럽게 적용 가능하다는 점과, 파라미터 재배치 과정에서 발생하는 수치적 불안정성을 완전히 해소하지 못했다는 점을 들 수 있다. 또한, 현재 구현은 정밀도 16비트(FP16) 기준이며, 더 낮은 정밀도(INT8)로 확장할 경우 추가적인 양자화 기법이 필요하다. 향후 연구에서는 나비형 구조와 스파스 연결을 결합하거나, 자동화된 구조 탐색을 통해 다양한 모델에 대한 최적 압축 비율을 찾는 방향이 제시된다.