[논문 리뷰] Efficient Pruning for Machine Learning Under Homomorphic Encryption

Ehud Aharoni, Moran Baruch외 9명|arXiv (Cornell University)|2022. 07. 07.

Cryptography and Data Security참고 문헌 44인용 수 4

한 줄 요약

HE-PEx는 동형 암호화(HE) 기반 기계학습에서 암호문 크기와 추론 지연을 크게 줄이기 위해 가중치 프루닝, 순열, 확장 기법을 결합한 새로운 프루닝 프레임워크를 제안한다. 새로운 순열 알고리즘을 통해 가중치를 재정렬하고 희소 구조를 확장함으로써, 평균 제곱 재구성 오차가 1.5×10⁻⁵에 불과한 MNIST 오토인코더에서 메모리와 지연 시간을 최대 60% 감소시켰으며, HE 환경에서의 표준 프루닝 기법보다 뚜렷이 뛰어난 성능을 발휘한다.

ABSTRACT

Privacy-preserving machine learning (PPML) solutions are gaining widespread popularity. Among these, many rely on homomorphic encryption (HE) that offers confidentiality of the model and the data, but at the cost of large latency and memory requirements. Pruning neural network (NN) parameters improves latency and memory in plaintext ML but has little impact if directly applied to HE-based PPML. We introduce a framework called HE-PEx that comprises new pruning methods, on top of a packing technique called tile tensors, for reducing the latency and memory of PPML inference. HE-PEx uses permutations to prune additional ciphertexts, and expansion to recover inference loss. We demonstrate the effectiveness of our methods for pruning fully-connected and convolutional layers in NNs on PPML tasks, namely, image compression, denoising, and classification, with autoencoders, multilayer perceptrons (MLPs) and convolutional neural networks (CNNs). We implement and deploy our networks atop a framework called HElayers, which shows a 10-35% improvement in inference speed and a 17-35% decrease in memory requirement over the unpruned network, corresponding to 33-65% fewer ciphertexts, within a 2.5% degradation in inference accuracy over the unpruned network. Compared to the state-of-the-art pruning technique for PPML, our techniques generate networks with 70% fewer ciphertexts, on average, for the same degradation limit.

연구 동기 및 목표

• 패킹 제약으로 인해 희소 가중치가 암호문 수를 줄이지 못하는 동형 암호화(HE) 환경에서 표준 프루닝 기법의 비효율성을 해결한다.
• 다차원 텐서의 SIMD 패킹으로 인해 대부분의 요소가 0이 되어도 암호문을 제거할 수 없는 제약을 극복한다.
• 모델 정확도를 훼손하지 않으면서 HE 암호화된 신경망에서 높은 프루닝 비율을 달성할 수 있는 프레임워크를 개발한다.
• 프루닝을 HE 우아한 패킹(타일 텐서)과 비상호작용 추론 모델에 통합하여 실제 기밀성 보장 기반 기계학습 워크로드를 지원한다.
• 순열과 확장을 통해 정확도 손실을 복구하고, HE 보호 추론 환경에서 프루닝의 잠재력을 극대화할 수 있음을 입증한다.

제안 방법

• 비제로 값들을 더 적은 수의 더욱 조밀한 타일에 모으기 위해 네트워크 가중치를 재정렬하는 새로운 순열 알고리즘을 제안한다. 이는 프루닝 중에 전체 암호문을 제거할 가능성을 높인다.
• 희소 가중치를 확장된 텐서 차원에 따라 재분배하는 가중치 확장 기법을 도입하여, 프루닝 이후에도 모델 정확도를 유지한다.
• 각 타일이 단일 암호문에 패킹되는 타일 텐서 패킹(HeLayers에서 유래)과 프루닝을 통합함으로써 효율적인 SIMD 연산을 가능하게 한다.
• 지연 시간, 메모리, 재구성 오차를 균형 잡는 목적 함수를 사용하여 프루닝 비율, 타일 크기, 정확도 손실을 동시에 최적화하는 프레임워크를 설계한다.
• 최신 프루닝 기법인 Hunter(Cai 등, 2022)을 비상호작용 HE 환경에 적응시켜, HE 호환 모델에 대한 엔드 투 엔드 프루닝 및 재학습을 가능하게 한다.
• 비상호작용 추론과의 호환성을 유지하고 계산 오버헤드를 줄이기 위해 ReLU의 HE 우아한 대체로 제곱 활성화 함수를 사용한다.

실험 결과

연구 질문

• RQ1 암호문이 0인 요소가 대부분임에도 불구하고 암호문을 제거할 수 없는 동형 암호화 환경에서 표준 신경망 프루닝 기법을 효과적으로 적용할 수 있는가?
• RQ2 모델 가중치의 순열이 암호화된 신경망에서 프루닝의 효과를 높여 암호문 제거 가능 수를 늘리는 데 어떻게 기여하는가?
• RQ3 어떤 정도로 가중치 확장 기법이 HE 보호 모델에서 극단적인 프루닝 이후 정확도 손실을 완화할 수 있는가?
• RQ4 HE 암호화된 오토인코더에서 모델 압축(암호문 감소), 지연 시간, 메모리, 재구성 정확도 사이의 타협 가능한 균형은 어느 정도인가?
• RQ5 HE 환경 하에서 암호문 제거 및 추론 효율성 측면에서 제안된 HE-PEx 프레임워크는 기존 프루닝 기법보다 어떻게 우월한가?

주요 결과

• HE-PEx는 MNIST 오토인코더에서 2×2 타일 크기일 경우 최대 85% 감소, 16×16 타일 크기일 경우 60% 감소를 기록했으며, 평균 제곱 재구성 오차는 1.5×10⁻⁵에 불과했다.
• 가장 큰 오토인코더 모델에서 HE-PEx는 메모리와 추론 지연 시간을 각각 60% 감소시켜, 종단 간 HE 추론에서 뚜렷한 성능 향상을 보였다.
• 표준 프루닝 기법은 85%의 가중치 프루닝 시 암호문의 17%만 제거하지만, HE-PEx는 지능적인 가중치 순열과 패킹 덕분에 이 비율을 85% 이상으로 끌어올렸다.
• 제안된 순열 알고리즘 덕분에 95% 프루닝 이후에도 지연 시간이 30% 감소했으며, 이는 암호문 제거가 성능 향상의 유일한 요인은 아님을 시사한다.
• 프루닝, 순열, 확장을 조합함으로써 HE-PEx는 활성 암호문 수와 HE 연산 수를 극적으로 줄이면서도 높은 정확도를 유지할 수 있었다.
• 이 방법은 활성화 함수 선택과 수직적 관계를 가지며, 컨볼루션 레이어를 포함한 다른 레이어로도 확장 가능하지만, 타일 기반 패킹을 사용하는 완전히 연결된 레이어에서 가장 효과적이다.