💻 computer science

KeyMemRT Compiler and Runtime: Unlocking Memory-Scalable FHE

KeyMemRT는 데이터 흐름 분석을 활용하여 완전 동형 암호(FHE) 회전 키의 수명을 자동으로 관리함으로써 기존의 최첨단 컴파일러들과 비교하여 메모리 소비를 크게 줄이고 실행 속도를 향상시키는 MLIR 기반의 컴파일러 및 런타임 프레임워크입니다.

원저자: Eymen Ünay, Björn Franke, Jackson Woodruff

게시일 2026-01-27

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Eymen Ünay, Björn Franke, Jackson Woodruff

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 퍼즐을 풀려고 노력 중이라고 상상해 보세요. 하지만 당신은 눈가리개를 쓴 채로만 퍼즐 조각들을 만질 수 있습니다. 당신은 외부의 눈으로부터 퍼즐을 숨겨주는 특별한 규칙들을 사용하며, 보이지 않는 상태에서 조각들을 조작해야 합니다. 이것이 바로 **완전 동형 암호(Fully Homomorphic Encryption, FHE)**가 하는 일입니다. 즉, 데이터가 완전히 암호화된 상태에서도 컴퓨터가 계산을 수행할 수 있게 해주는 것입니다.

하지만 큰 문제가 하나 있습니다. 이 "눈가리개를 쓴" 수학 계산을 수행하려면, 컴퓨터는 거대한 양의 특별한 "마법 열쇠"(회전 키/rotation keys라고 불림)를 가지고 있어야 합니다.

문제점: "열쇠 수집광"

이 마법 열쇠들을 수천 개의 객실이 있는 호텔의 거대한 물리적 열쇠 모음이라고 생각해 보세요.

기존 방식 (ANT-ACE): 호텔 매니저가 손님이 도착하기도 전에 건물에 있는 모든 열쇠를 몽땅 집어 들어 프런트 데스크에 쏟아붓는 상황을 상상해 보세요. 손님이 101호실만 방문하면 되는데도, 매니저는 102호부터 5,000호까지의 열쇠를 모두 데스크 위에 올려둡니다.
- 결과: 데스크(메모리)가 지저분해지고 넘쳐납니다. 호텔이 매우 크다면, 데스크가 불필요한 열쇠들로 너무 가득 차서 다른 일을 할 공간조차 남지 않게 됩니다. 시스템은 과도한 열쇠를 관리하려다 보니 충돌이 발생하거나 느려지게 됩니다.
"느린" 방식 (Fhelipe): 또 다른 매니저는 공간을 아끼기 위해 몇 개의 "마스터 키"만 보관하려고 합니다. 101호실을 열기 위해 직접적인 열쇠가 없으면, 마스터 키를 사용해 1번 방을 열고, 그다음 2번 방을 여는 식으로 101호에 도달할 때까지 열쇠들을 줄줄이 연결하여 사용합니다.
- 결과: 데스크는 깔끔해지지만, 손님은 매니저가 긴 열쇠 사슬을 뒤적거리는 동안 오랫동안 기다려야 합니다. 이 과정은 매우 느립니다.

해결책: KeyMemRT

이 논문의 저자들은 열쇠를 완벽하게 관리하는 초스마트 자동 컨시어지 서비스인 KeyMemRT라는 새로운 시스템을 구축했습니다.

당신에게 무엇이 필요한지 정확히 압니다: 한꺼번에 모든 열키를 가져오는 대신, KeyMemটি는 손님의 일정(프로그램 코드)을 분석하여 그들이 정확히 어떤 방을 어떤 순서로 방문할지 파악합니다.
적시 전달 (Just-in-Time Delivery): 데스크를 깨끗하게 유지합니다. 손님이 101호실을 필요로 하는 바로 직전에 101호의 특정 열쇠를 가져옵니다.
즉각적인 정리: 손님이 101호실을 떠나자마자, 컨시어지는 즉시 그 열쇠를 다시 가져가서 치움으로써 다음 열쇠를 위한 공간을 만듭니다.
"프리페치(Prefetch)" 기술: 손님이 기다리지 않도록, 컨시어지는 손님이 현재 방에 머무는 동안 배경에서 다음 열쇠를 가져오기 시작합니다. 이 작업은 백그라운드에서 일어나므로 과정이 빠르고 매끄럽게 느껴집니다.

그들은 무엇을 달성했나요?

연구진은 다양한 복잡한 작업(의료 스캔이나 금융 데이터의 이미지를 인식하는 등)을 사용하여 이 새로운 시스템을 기존 방식들과 비교 테스트했습니다.

메모리 절감: "열쇠 수집광" 방식(ANT-ACE)과 비교했을 때, KeyMemRT는 필요한 메모리 양을 1.74배 줄였습니다. 이는 호텔 전체의 열쇠를 창고 대신 단 하나의 서류 가방 안에 담아 넣는 것과 같습니다.
속도: "느린 사슬" 방식(Fhelipe)과 비교했을 때, KeyMemRT는 1.73배 더 빨랐습니다. 단순히 공간만 아낀 것이 아니라, 열쇠를 처음부터 만드는 데 시간을 낭비하지 않았기 때문에 작업이 훨씬 더 빠르게 진행되었습니다.

이것이 왜 중요한가요?

현재 FHE는 이 열쇠들을 보관하기 위해 엄청난 양의 메모리(수백 기가바이트)를 필요로 하기 때문에 사용하기 어렵습니다. 이로 인해 비용이 많이 들고 실제 활용하기에 비현실적입니다.

Key-MemRT는 번역가이자 교통 통제사 역할을 합니다. 복잡하고 무질서한 코드를 최적화하여, 열쇠를 보관하기 위해 슈퍼컴퓨터를 필요로 하지 않고도 일반 컴퓨터에서 실행될 수 있도록 만듭니다. 이는 개인정보를 보호하는 컴퓨팅(원래의 숫자를 전혀 보지 않고도 민감한 의료 또는 금융 데이터를 분석하는 것과 같은)을 더 확장 가능하고 효율적으로 만들어 줍니다.

요약하자면: 그들은 FHE 프로그램이 쓸데없는 열쇠를 모으거나 열쇠를 새로 만드는 데 시간을 낭비하는 것을 막아주는 스마트한 시스템을 구축했으며, 이를 통해 프라이버시를 보호하는 컴퓨팅이 더 빠르고 더 작고 저렴한 기기에서도 실행될 수 있도록 만들었습니다.

기술 요약: KeyMemRT 컴파일러 및 런타임

문제 정의

완전 동형 암호(Fully Homomorphic Encryption, FHE)는 암호화된 데이터에 대한 개인정보 보호 연산을 가능하게 하지만, 과도한 계산 및 메모리 요구 사항으로 인해 실제 세계의 도입에 큰 장벽이 있습니다. 주요 병목 현상은 회전 키(rotation keys) 관리입니다. 배치형 FHE 스킴(예: CKKS, BGV, BFV)에서 회전 연산은 벡터 내 연산 및 부트스트래핑(bootstrapping)을 위해 필수적입니다. 그러나 이러한 스킴들은 특정 회전 인덱스마다 고유한 회전 키를 필요로 합니다.

FHE 프로그램이 확장됨에 따라 필요한 회전 인덱스의 수가 증가하며, 이는 모든 키를 동시에 저장하기 위해 수백 기가바이트 또는 심지어 테라바이트 단위의 메모리를 필요로 하게 만듭니다. 기존의 FHE 컴파일러들은 일반적으로 회전 키를 단일한 정적 엔티티로 취급하여, 프로그램 실행 기간 동안 모든 키를 메모리에 로드합니다. 이러한 방식은 메모리 확장성이 부족합니다. 반대로, 메모리를 줄이기 위해 시도된 기존 전략인 Fhelipe는 임의의 회전 인덱스를 2의 거듭제곱 체인으로 분해하는 방식에 의존합니다. 이 방식은 저장되는 키의 수는 줄여주지만, 체이닝 오버헤드로 인해 계산 지연 시간을 크게 증가시킵니다. 메모리 사용량과 실행 속도의 균형을 맞출 수 있는 세밀하고 자동화된 키 관리(fine-grained, automatic key management) 기능을 제공하는 컴파일러는 현재 부재한 상태입니다.

방법론

저자들은 개별적인 회전 키의 수명(lifetime)을 관리하도록 설계된 컴파일러 및 런타임 프레임워크인 KeyMemRT를 제시합니다. 이 시스템은 MLIR (Multi-Level Intermediate Representation) 인프라를 기반으로 구축되었으며, 새로운 디렉트(dialect)와 런타임 아키텍처를 도입합니다.

1. `kmrt` MLIR 디렉트

KeyMemRT는 키 관리를 IR 내에 직접 표현하기 위해 kmrt라는 새로운 디렉트를 도입합니다.

새로운 타입: 회전 인덱스에 의해 매개변수화된 rot_key 타입을 정의하여, 컴파일러가 서로 다른 인덱스에 대한 키를 구분할 수 있도록 합니다.
새로운 연산: 이 디렉트는 명시적인 키 관리를 위한 연산들을 포함합니다:
- kmrt.load_key: 특정 회전 키를 로드합니다.
- kmrt.clear_key: 메모리에서 키를 제거합니다.
- kmrt.prefetch_key: 비동기 로딩을 위한 힌트를 런타임에 제공합니다.
- kmrt.use_key / kmrt.assume_key: 키의 범위(scope)와 생존성(liveness)을 관리합니다.
통합: 컴파일러는 고수준 연산(예: HEIR를 통한 ckks 디렉트로부터의 연산)을 kmrt 연산을 삽입하면서 openfhe 디렉트로 낮춥니다(lowering).

2. 컴파일러 최적화 파이프라인

KeyMemRT 컴파일러는 성능 저하 없이 메모리 풋프린트를 줄이기 위해 일련의 로워링(lowering) 및 최적화 패스를 수행합니다:

CKKS to OpenFHE 로워링: 모든 회전 연산을 특정 rot_key 값과 연관시켜, 문제를 생존 범위(liveness ranges)를 생성하는 것에서 병합하는 것으로 역전시킵니다.
Baby-Step Giant-Step (BSGS) 분해: 연속적인 회전 인덱스에 대해 회전을 체이닝함으로써 필요한 키의 수를 $N$ 에서 $\approx 2\sqrt{N}$ 으로 줄입니다. 이는 특히 행렬-벡터 곱셈 패턴에서 효과적입니다.
회전 키 병합: 데이터 흐름 분석을 사용하여 키가 삭제된 직후 동일한 인덱스로 즉시 다시 로드되는 시퀀스를 감지하고, 불필요한 로드/클리어 쌍을 제거합니다.
회전 호이스팅(Rotation Hoisting): 공통 입력에 적용되는 회전 패턴을 식별하고 이를 OpenFHE 라이브러리에서 지원하는 더 빠른 대안으로 변환합니다.
부트스트랩 키 관리: 암호문 레벨을 갱신하는 데 사용되는 내부 회전 키를 IR로 구체화(materialize)합니다. 이를 통해 컴파일러가 명시적 회전 키와 동일한 세밀한 제어로 부트스트랩 키를 관리할 수 있게 하며, 이는 이전에 다뤄지지 않았던 기능입니다.
공통 부분식 제거 (Common Subexpression Elimination, CSE): 키 최적화로 인해 발생하는 데드 코드(dead code)를 제거합니다.
힌트 배치: 런타임이 I/O와 계산을 중첩(overlap)할 수 있도록 prefetch_key 힌트를 삽입합니다.

3. KeyMemRT 런타임

런타임은 컴파일된 C++ 코드와 FHE 라이브러리(OpenFHE) 사이의 중개자 역할을 합니다. 이는 컴파일러가 삽입한 힌트에 따라 디스크에서 RAM으로 키를 동적으로 로드하고 언로드(unload)하는 것을 관리합니다.

저메모리 모드 (Low-Memory Mode): 지시사항을 순차적으로 실행하여, 피크 메모리 사용량을 최소화하기 위해 엄격히 필요한 경우에만 키를 로드합니다.
균형 모드 (Balanced Mode): 계산과 병렬로 키를 로드하기 위해 프리페칭(prefetching) 메커니즘을 활용합니다. 내부 큐를 유지하며, 사용자 정의 제한까지 키를 미리 로드하고 공간이 필요할 때 이를 클리어함으로써 지연 시간 오버헤드를 줄입니다.

주요 기여

새로운 MLIR 디렉트 (kmrt): 회전 키를 특정 타입과 메모리 관리 연산을 가진 IR 요소로 노출하여 세밀한 분석을 가능하게 한 최초의 디렉트입니다.
자동화된 세밀한 키 관리: 키의 생존성을 자동으로 추적하고, 부트스트랩 키를 최적화하며, 키를 글로벌 세트가 아닌 개별적으로 관리하는 최초의 FHE 컴플라이러입니다.
런타임 프레임워크: 컴파일러 힌트를 활용하여 I/O와 계산을 중첩시키는 온디맨드 키 로딩 및 클리어링을 지원하는 런타임입니다.
통합: 기존 FHE 프런트엔드(Orion 및 HEIR로 입증됨)가 타겟팅할 수 있는 포스트 최적화 컴파일러 툴체인입니다.

결과

저자들은 KeyMemRT를 두 가지 최첨단 접근 방식인 ANT-ACE(모든 키를 메모리에 유지) 및 Fhelipe(키를 줄이기 위해 분해를 사용하지만 지연 시간이 증가함)와 비교 평가했습니다. 벤치마크에는 링 차원(ring dimension) $2^{16}$ 및 곱셈 깊이(multiplicative depth) 30의 FHE 파라미터가 설정된 다양한 ML 모델(MLP, LoLa, LeNet, ResNet 변형, AlexNet)이 포함되었습니다.

메모리 감소:
- ANT-ACE와 비교했을 때, KeyMemRT는 기하 평균적으로 1.74배의 메모리 감소를 달과 성취했습니다. LeNet과 같은 특정 사례에서는 메모리 사용량이 2.59배 감소했습니다.
- Fhelipe와 비교했을 때, KeyMemRT는 1.16배의 메모리 감소를 달성했습니다.
지연 시간/속도 향상:
- ANT-ACE와 비교했을 때, KeyMemRT (Balanced 모드)는 1.20배의 속도 향상을 달성했습니다.
- Fhelipe와 비교했을 때, KeyMemRT는 기하 평균 1.73배의 유의미한 속도 향상을 달성했으며, AlexNet에서는 최대 2.5배의 개선을 보였습니다. 이는 KeyMemRT가 Fhelipe에서 요구되는 회전 체이닝의 계산 오버헤드를 피하기 때문입니다.
확장성: KeyMemRT는 ANT-ACE 하에서 100 GB 이상의 메모리 요구 사항을 초과하는 대규모 프로그램(예: ResNet-50)을 실행할 수 있음을 입증했습니다.

의의 및 주장

본 논문은 KeyMemRT가 자동화된 세밀한 키 관리 및 부트스트랩 키 관리를 제공하는 최초의 FHE 컴파일러라고 주장합니다. 저자들은 이 작업이 FHE의 메모리 확장성을 열어주어, 모든 회전 키를 동시에 저장해야 하는 과도한 메모리 부담 없이도 임의의 규모가 큰 프로그램을 실행할 수 있게 한다고 단언합니다.

메모리 사용량을 전체 키의 수로부터 분리하고 대신 키의 생존성에 기반하여 관리함으로써, KeyMemRT는 FHE 프로그램을 더 효율적으로 확장할 수 있게 합니다. 저자들은 이 방식이 성능을 희생하지 않는다는 점을 강조합니다. 메모리를 위해 속도를 희착하는 분해 기반 방식과 달리, KeyMemRT는 필요한 정확한 시점에 필요한 키만 로드함으로써 메모리 감소와 지연 시간 개선을 동시에 달성합니다. 이는 FHE 도입의 핵심적인 장벽을 해결합니다.