Leveraging ASIC AI Chips for Homomorphic Encryption
이 논문은 동형 암호 연산의 에너지 효율성을 극대화하기 위해 TPU 의 저정밀도 매트릭스 엔진과 메모리 구조에 최적화된 새로운 컴파일러 프레임워크 'CROSS'를 제안하며, 기존 GPU 기반 라이브러리보다 월등한 성능을 입증합니다.
원저자:Jianming Tong, Tianhao Huang, Jingtian Dang, Leo de Castro, Anirudh Itagi, Anupam Golder, Asra Ali, Jeremy Kun, Jevin Jiang, Arvind, G. Edward Suh, Tushar Krishna
이 논문은 **"암호화된 데이터를 해독하지 않고도 계산할 수 있는 기술 (동형 암호)"**을 훨씬 더 빠르고 효율적으로 만들 수 있는 새로운 방법을 소개합니다.
비유를 들어 쉽게 설명해 드릴게요.
1. 문제 상황: "고급 레스토랑의 요리사" vs "대량 급식소"
동형 암호 (HE) 란?
마치 봉인된 상자에 들어간 재료를 요리사가 직접 손대지 않고도 요리할 수 있는 마법 같은 기술입니다. 데이터의 프라이버시를 해치지 않고 클라우드에서 계산을 할 수 있게 해주죠.
하지만 문제는 너무 느리고 비싸다는 것입니다. 일반 컴퓨터로 하면 1,000 배나 느려져서 실용성이 떨어집니다.
기존 해결책 (GPU): "고급 레스토랑의 셰프"
지금까지는 그래픽 카드 (GPU) 를 썼습니다. 이는 정교하고 섬세한 작업 (예: 한 알 한 알 쌀을 세는 일) 을 잘하는 셰프와 같습니다.
하지만 동형 암호는 엄청난 양의 데이터를 처리해야 하는데, 셰프가 하나하나 정성스럽게 하느라 시간이 너무 오래 걸립니다.
새로운 시도 (TPU): "대량 급식소"
구글의 TPU 같은 AI 가속기는 원래 수천 명의 학생에게 급식을 줄 때 쓰입니다. 한 사람 한 사람을 섬세하게 대하기보다, 한 번에 대량으로 효율적으로 처리하는 데 특화되어 있죠.
문제는 이 "대량 급식소"에 "셰프용 레시피"를 그대로 가져다 붙이면, 급식소 주방이 엉망이 된다는 것입니다. (예: 쌀 한 알을 세는 일을 시키면 급식소의 거대한 밥솥이 놀게 됩니다.)
2. 이 논문의 해결책: CROSS (크로스)
이 논문은 **"기존에 있는 AI 칩 (TPU) 을 동형 암호에 맞게 재해석하는 컴파일러 (CROSS)"**를 개발했습니다.
핵심 아이디어는 두 가지입니다:
① BAT (기반 정렬 변환): "빈 공간 제거하기"
상황: 기존 방식은 계산할 때 많은 '0'이 포함된 희박한 행렬을 사용합니다. 마치 빈 의자가 가득 찬 대형 극장에서 관객이 극히 드물게 앉아 있는 것과 같습니다. (계산 자원 낭비)
CROSS 의 해결: "빈 의자는 치우고, 관객만 모아서 앉게 하세요!"
미리 계산된 상수들을 활용하여, 빈 공간 (0) 을 없애고 데이터를 빽빽하게 채웁니다.
이렇게 하면 TPU 의 거대한 밥솥 (MXU) 이 비어있지 않고, 수천 개의 요리를 동시에 처리할 수 있게 되어 속도가 비약적으로 빨라집니다.
② MAT (메모리 정렬 변환): "배열 바꾸기 생략하기"
상황: 동형 암호 계산 중에는 데이터를 섞거나 뒤집는 작업이 자주 필요합니다. 기존 방식은 계산할 때마다 데이터를 한 번씩 꺼내서 재배열했습니다. 마치 요리할 때마다 재료를 한 번씩 꺼내서 다시 정리하는 것과 같습니다. (시간 낭비)
CROSS 의 해결: "요리 전에 재료를 미리 섞어두세요!"
계산이 시작되기 **전 (컴파일 단계)**에 필요한 데이터 순서를 미리 바꿔둡니다.
그래서 실제 계산할 때는 재료를 꺼내서 바로 요리만 하면 됩니다. 불필요한 이동 시간을 아껴서 속도를 높입니다.
3. 결과: "기적 같은 효율성"
이 방법을 적용한 결과, 구글의 최신 TPU 칩은 다음과 같은 성과를 냈습니다.
속도: 기존에 가장 빠르다고 알려진 GPU 기반 방식보다 최대 13 배 이상 빠른 동형 암호 계산 (NTT) 속도를 달성했습니다.
에너지 효율: 같은 전력을 썼을 때, 기존 방식들보다 최대 451 배 더 많은 계산을 처리했습니다. (에너지 효율이 압도적으로 좋음)
의미: 비싼 전용 칩 (ASIC) 을 새로 만들지 않고, 이미 존재하는 AI 칩을 활용해서 가장 효율적인 동형 암호 시스템을 만들 수 있음을 증명했습니다.
4. 요약: 왜 이것이 중요한가?
이 논문은 **"새로운 하드웨어를 사지 않아도, 기존 AI 칩을 똑똑하게 쓰면 데이터 보안과 속도 문제를 동시에 해결할 수 있다"**는 것을 보여줍니다.
과거: "보안을 지키려면 속도가 느려지고, 속도를 내려면 보안이 느슨해지거나 비싼 전용 장비를 사야 한다."
이제: "이미 있는 AI 칩 (TPU) 을 CROSS 라는 '요리법'으로 바꾸면, 보안은 지키면서 AI 가 처리하는 속도만큼이나 빠르고 에너지도 아끼는 암호화 계산이 가능하다."
결론적으로, 이 기술은 클라우드에서 우리의 개인 데이터를 해독하지 않고도 AI 가 분석할 수 있는 실용적인 길을 열어주었습니다.
1. 연구 배경 및 문제 정의 (Problem)
동형 암호화 (HE) 의 한계: 동형 암호화는 클라우드 환경에서 데이터 프라이버시를 보장하는 강력한 기술이지만, 평문 연산에 비해 엄청난 계산 오버헤드가 발생합니다. 기존 CPU 기반 구현은 1000 배 이상의 지연 시간을 보입니다.
기존 가속기 (GPU) 의 비효율성: 현재 HE 가속의 주류는 GPU 입니다. 그러나 GPU 는 HE 알고리즘에 최적화되지 않은 아키텍처를 가지고 있어, 전용 ASIC 에 비해 에너지 효율성 (Performance per Watt) 이 약 33 배 낮습니다.
AI 가속기 (TPU) 의 활용 가능성과 장애물: 구글의 TPU 와 같은 AI 전용 ASIC 은 대규모 행렬 연산 (MXU) 과 큰 온칩 메모리를 갖추고 있어 HE 에 유망해 보이지만, 기존 GPU 용 HE 라이브러리를 그대로 TPU 에 이식할 경우 심각한 성능 저하가 발생합니다. 그 이유는 다음과 같습니다:
산술 불일치 (Arithmetic Mismatch): 최신 GPU 기반 HE 알고리즘은 32 비트 정수 연산에 의존하거나 희소 행렬 (Sparse Matrix) 연산을 사용합니다. 이는 TPU 의 저정밀도 (INT8) 고처리량 행렬 엔진 (MXU) 을 활용하지 못하게 하거나, TPU 의 저처리량 벡터 유닛 (VPU) 에만 의존하게 만들어 MXU 를 유휴 상태로 만듭니다.
메모리 조작 불일치 (Memory Manipulation Mismatch): HE 의 핵심 연산인 NTT(수론적 변환) 는 세밀한 데이터 셔플링과 전치 (Transpose) 를 요구합니다. TPU 는 대규모 SIMD 레지스터 (Coarse-grained) 를 사용하므로, 이러한 세밀한 데이터 재배열이 비효율적이고 메모리 병목 현상을 초래합니다.
2. 제안 방법론: CROSS (Methodology)
저자들은 CROSS라는 컴파일러 프레임워크를 제안하여, 기존 AI 가속기 (TPU) 의 아키텍처 강점에 맞춰 HE 워크로드를 변환합니다. CROSS 는 두 가지 핵심 변환 기법을 사용합니다.
가. Basis-Aligned Transformation (BAT)
목적: 고평도 (High-precision, 예: 32 비트) 모듈러 산술을 TPU 의 MXU 가 효율적으로 처리할 수 있는 밀집된 저정밀도 (INT8) 행렬 곱셈으로 변환합니다.
작동 원리:
기존 GPU 방식은 곱셈을 희소 행렬 (Sparse Matrix) 형태로 표현하여 많은 0 을 포함하고 있어 연산 및 메모리 낭비가 큽니다.
BAT 는 미리 알려진 상수 (Twiddle factors, 평가 키 등) 를 오프라인에서 전처리하여, 곱셈을 밀집된 (Dense) 행렬 곱셈으로 재구성합니다.
이를 통해 불필요한 0 연산을 제거하고, TPU 의 MXU (INT8) 를 최대한 활용하여 연산 효율을 극대화합니다.
이론적으로 계산 및 메모리 사용량을 약 2 배 절감합니다.
나. Memory-Aligned Transformation (MAT)
목적: 런타임에서의 비용이 큰 데이터 재배열 (Transpose, Shuffling) 을 제거하고, 이를 컴파일 시 계산 과정에 내재화합니다.
작동 원리:
HE 알고리즘에서 필요한 데이터 순서 변경 (예: 비트 반전, 행/열 전치) 을 '순열 행렬 (Permutation Matrix)'로 표현합니다.
이 순열 연산을 미리 계산된 매개변수 (Twiddle factors 등) 에 적용하여 오프라인에서 처리합니다.
결과적으로 런타임에는 별도의 메모리 재배열 연산 없이, 계산만 수행하여 원하는 데이터 레이아웃을 얻는 '레이아웃 불변 (Layout-invariant)' 커널을 생성합니다.
이는 TPU 의 Coarse-grained 메모리 시스템에 최적화되어 메모리 지연을 획기적으로 줄입니다.
3. 주요 기여 (Key Contributions)
아키텍처 불일치 분석: GPU 최적화 HE 알고리즘을 AI 가속기 (TPU) 에 이식할 때 발생하는 산술 및 메모리 불일치를 체계적으로 분석하고 정량화했습니다.
BAT 및 MAT 기법 개발: 고평도 정수 연산을 저정밀도 행렬 연산으로 매핑하고 (BAT), 런타임 재배열을 제거하는 (MAT) 두 가지 아키텍처 보편적 최적화 기법을 제안했습니다.
실제 하드웨어 검증: 구글 TPU v6e(실제 단일 호스트) 에서 CROSS 를 구현하고, 기존 CPU, GPU, FPGA, 전용 ASIC 대비 성능을 평가했습니다.
새로운 패러다임 제시: 하드웨어 수정 없이 기존 AI ASIC 을 활용하여 동형 암호화 연산의 에너지 효율성을 SoTA(State-of-the-Art) 수준으로 끌어올리는 새로운 접근법을 제시했습니다.
4. 평가 결과 (Results)
실제 TPU v6e 환경에서 CROSS 는 다음과 같은 성과를 달성했습니다.
NTT 처리량: 기존 GPU 기반 솔루션 (WarpDrive on A100) 대비 1.43 배의 처리량 향상을 기록하며, 실질적으로 사용 가능한 장치 중 SoTA 기록을 경신했습니다.
에너지 효율성 (Throughput per Watt):
OpenFHE (CPU): 451 배 향상
WarpDrive (GPU): 7.81 배 향상
FIDESlib (GPU): 1.83 배 향상
FAB (FPGA): 1.31 배 향상
HEAP (FPGA): 1.86 배 향상
Cheddar (GPU): 1.15 배 향상
특히, 전용 HE ASIC (BASALISC, CraterLake) 과 비교했을 때에도 HE-Add 연산에서 더 나은 효율성을 보였으며, HE-Mult/Rotate 연산에서도 성능 격차를 크게 줄였습니다.
ML 워크로드: MNIST 분류 및 로지스틱 회귀 모델에서 기존 솔루션 대비 높은 처리량과 낮은 지연 시간을 달성했습니다.
5. 의의 및 결론 (Significance)
AI ASIC 의 HE 가속 가능성 입증: 전용 HE 칩을 개발하지 않고도, 기존에 널리 배포된 AI 가속기 (TPU) 를 활용하여 동형 암호화 연산을 SoTA 수준의 에너지 효율성으로 수행할 수 있음을 증명했습니다.
프라이버시 보호 AI 의 실용화: 클라우드 환경에서 민감한 데이터를 암호화된 상태로 처리하는 'Privacy-Preserving AI'의 실현 가능성을 높였습니다. 하드웨어 변경 없이 소프트웨어 (컴파일러) 레벨에서 최적화가 가능하므로 배포 장벽이 낮습니다.
향후 방향: 현재 TPU 와 전용 ASIC 간의 약 33 배의 성능 격차는 주로 전용 ASIC 의 '저비용 셔플링 엔진'과 '하드웨어 최적화된 모듈리' 부재에서 기인합니다. CROSS 는 이러한 격차를 줄이는 중요한 첫걸음이며, 향후 AI/HE 공동 가속 (Co-acceleration) 의 표준 플랫폼으로 자리 잡을 잠재력을 가집니다.
결론적으로, CROSS 는 동형 암호화의 계산 비용 문제를 해결하기 위해 AI 가속기의 아키텍처 강점을 극대화하는 혁신적인 컴파일러 프레임워크이며, 프라이버시 보호 컴퓨팅의 새로운 에너지 효율성 기준을 제시합니다.