기술 요약: CHEHAB RL: 완전 동형 암호 연산 최적화를 위한 학습

1. 문제 정의

완전 동형 암호(Fully Homomorphic Encryption, FHE)는 데이터를 복호화하지 않고도 암호화된 상태에서 연산을 가능하게 하여 프라이버시를 보호한다. 그러나 FHE는 여전히 계산 비용이 매우 높다. 예를 들어, 암호문 곱셈은 평문 연산보다 $10^6$배 더 느릴 수 있다. 이러한 성능 격차는 고차 다항식 곱셈과 상당한 노이즈 증가를 포함하는 동형 연산의 본질적인 복잡성에서 기인한다.

효율적인 FHE 코드를 작성하는 것은 깊은 암호학적 전문 지식을 요구하는 복잡하고 오류가 발생하기 쉬운 작업이다. FHE 코드 생성을 자동화하기 위해 컴파일러들이 개발되어 왔으나, 기존 솔루션들은 다음과 같은 중대한 한계에 직면해 있다:

• 구조적 코드 vs 비구조적 코드: HECO 및 CHET와 같은 컴파일러는 구조적(루프 기반) 코드에 집중하는 반면, Coyote 및 Porcupine과 같은 다른 도구들은 구조적 코드와 비구조적(임의의, 루프 기반이 아닌) 코드를 모두 처리하려고 시도한다.
• 탐색 공간의 복잡성: 현재의 최첨단 컴파일러(예: Coyote, Porcupine)는 벡터화를 조합 탐색 문제로 프레임화한다. 이들은 방대한 설계 공간(예: 어떤 하위 표현식을 패킹할지, 데이터 레이아웃, 회전 삽입 결정 등)을 탐색하기 위해 수동으로 튜닝된 휴리스틱, 정수 선형 계획법(ILP) 또는 프로그램 합성(program synthesis)에 의존한다.
• 확장성 및 지역 최적해(Local Optima): 이 복잡한 공간에서 로컬 결정은 비로컬 효과(예: 패킹 결정이 향방의 회전 비용을 증가시키는 경우)를 초래할 수 있다. 단순한 휴리스틱은 종종 좋지 않은 지역 최적해에 갇히기 쉬우며, 프로그램 크기가 커짐에 따라 전수 조사(exhaustive search)는 계산적으로 불가능해진다.

2. 방법론

본 논문은 딥 러닝 강화 학습(Deep Reinforcement Learning, RL)을 활용하여 FHE 코드 최적화를 자동화하는 새로운 프레임워크인 CHEHAB RL을 제안한다. 조합 탐색 대신, 이 방식은 RL 에이전트가 일련의 재작성 규칙(rewriting rules)을 적용하는 정책을 학습하도록 하는 순차적 의사결정 과정으로 문제를 정형화한다.

2.1 프레임워크 개요

CHEHAB RL은 CHEHAB FHE 컴파일러에 통합된다. 워크플로우는 다음과 같다:

1. 입력: CHEHAB 도메인 특화 언어(DSL)로 작성되어 추상 구문 트리(AST) 및 데이터 흐름 그래프(DAG) 형태의 중간 표현(IR)으로 낮추어진(lowered) 프로그램.
2. RL 최적화: RL 에이전트가 용어 재작성 시스템(Term Rewriting System, TRS)에서 일련의 재작성 규칙을 선택하여 IR을 순회하며, 스칼라 코드를 벡터화하고 식을 단순화하며 노이즈/지연 시간을 최소화한다.
3. 출력: 최적화된 IR은 Microsoft SEAL 백엔드(BFV 스킴)를 타겟팅하는 C++ 코드를 생성하는 데 사용된다.

2.2 주요 기술 구성 요소

A. 상태 표현 (State Representation)

RL 에이전트를 위해 프로그램 IR을 표현하기 위해, 저자들은 식별자 및 상수 불변성(Identifier and Constant Invariant, ICI) 토큰화를 채택한다.

• 정규화: 변수 이름을 변경하고(예: $v_0, v_1$), 0과 1을 제외한 상수는 일반적인 토큰($c_0, c_1$)으로 추상화한다. 이는 명명 규칙이나 특정 상수 값에 무관한 정준 형태(canonical form)를 생성하여, 의미론적으로 동일한 프로그램들을 동일한 표현으로 붕괴시킨다.
• 인코딩: 토큰화된 시퀀스는 트랜스포머 인코더(4개 레이어, 8개 어텐션 헤드)에 의해 처리되어 고정된 길이의 256차원 임베딩을 생성한다. 이 아키텍처는 펼쳐진(unrolled) FHE 프로그램의 장거리 의존성을 모델링하는 데 있어 RNN(LSTM/GRU)보다 우수한 능력을 갖추고 있어 선택되었다.

B. 행동 공간 (Action Space)

행동 공간은 많은 수의 규칙과 적용 위치를 처리할 수 있도록 설계되었다.

• 계층적 정책: 에이전트는 평탄한 행동 공간(규칙 + 위치) 대신 두 단계의 계층적 정책을 사용한다:
  1. 규칙 선택 네트워크: MLP가 84개의 규칙 세트 중에서 적용할 재작성 규칙을 선택한다.
  2. 위치 선택 네트워크: 두 번째 MLP가 IR 내에서 규칙이 적용될 특정 위치(매치 인덱스)를 선택한다.
• 종료: 특수 END 액션은 에이전트가 최적화 에피소드를 종료할 수 있게 한다.

C. 보상 모델 (Reward Model)

에이전트는 실행 시간과 노이즈 증가 사이의 균형을 맞추는 복합 비용 함수를 최소화하도록 훈련된다.

• 비용 함수: $Cost(e) = w_{ops} \cdot C_{ops}(e) + w_{depth} \cdot D_{circuit}(e) + w_{mult} \cdot D_{mult}(e)$.
  • $C_{ops}$: 연산 비용의 가중 합 (벡터 덧셈: 1, 벡터 곱셈: 100, 회전: 50, 스칼라 연산: 250).
  • $D_{circuit}$: 회로 깊이.
  • $D_{mult}$: 곱셈 깊이 (BFV에서 노이즈 성장에 결정적임).
• 보상 구조:
  • 단계 보상 ($R_{step}$): 각 액션 이후 비용 감소율에 기반한 즉각적인 보상.
  • 최종 보상 ($R_{final}$): 초기 상태에서 최종 상태까지의 총 비용 감소에 기반한 글로벌 보상. 이는 에이전트가 국소적으로 탐욕적이지 않도록 하고 전역적으로 최적화된 회로를 생성하도록 유도한다.

D. 훈련 데이터 합성 (Training Data Synthesis)

기존 FHE 최적화 데이터셋의 부재를 인식하여, 저자들은 대규모 언어 모델(LLM, Gemini 2.5 Flash)을 사용하여 15,855개의 고유한 식으로 구성된 훈련 데이터셋을 합성하였다.

• LLM 가이드: LLM은 구문 규칙, 재작성 규칙 예시, 그리고 실제 커널 모티프(예: 행렬 덧셈, 제곱 차이)를 프롬프트로 전달받아 다양하고 최적화 가능한 코드 패턴을 생성했다.
• 검증: 생성된 식들은 구문적 유효성, 고유성(ICI 토큰화를 통해 확인), 그리고 벤치마크 프로그램 제외 여부를 거쳐 필터링되어 일반화 성능을 보장했다.

E. 훈련 알고리즘

에이전트는 단일 멀티 코어 CPU에서 2백만 타임스텝 동안 근사 정책 최적화(Proximal Policy Optimization, PPO) 알고리즘을 사용하여 훈련된다.

3. 주요 기여

1. 정형화: FHE 코드 최적화를 강화 학습으로 해결 가능한 순차적 의사결정 문제로 정형화한 최초의 사례.
2. CHEHAB RL 프레임워크: 구조적 및 비구조적 FHE 코드를 자동으로 벡터화하면서 노이즈 증가와 회로 지연 시간을 최소화하는 새로운 시스템. 이는 기존의 조합 탐색 방식보다 확장성이 뛰어나다.
3. LLM 합성 데이터셋: LLM을 통해 생성된 대규모의 다양한 훈련 데이터셋을 구축하고 검증함으로써, 이러한 데이터로 훈련하는 것이 무작위 프로그램으로 훈련하는 것보다 훨씬 나은 에이전트를 생성함을 입증함.
4. 성능 향상: CHEHAB RL이 최첨단인 Coyote 컴파일러보다 실행 속도는 5.3배 빠르고, 노이즈 축적은 2.54배 적으며, 컴파일 시간은 27.9배 빠르다는 실증적 근거를 제시함 (기하 평균 기준).
5. 오픈 소스: CHEHAB RL 프레임워크를 커뮤니티에 공개함.

4. 결과 및 평가

저자들은 이미지 처리 필터, 선형 대수 커널, 무작위로 생성된 불규칙 다항식을 포함한 벤치마크 세트를 통해 CHEHAB RL을 최첨단 벡터화 FHE 컴파일러인 Coyote와 비교 평가하였다.

• 실행 시간: CHEHAB RL은 Coyote보다 평균적으로 5.3배 빠른 코드를 생성했다. 특정 사례(예: Linear Regression 32)에서는 속도가 114배에 달했다. 이는 에이전트가 발견한 더 적은 연산과 더 단순한 데이터 레이아웃 덕분이었다.
• 노이즈 예산: 생성된 회로는 평균적으로 2.54배 적은 노이즈 예산을 소비했다. Coyote가 노이즈 예산을 소진하여 실행에 실패한 벤치마크에서도 CHEHAB RL은 유효하고 실행 가능한 회로를 성공적으로 생성했다.
• 컴파일 시간: RL 기반 컴파일 프로세스는 Coyote의 탐색 기반 방식보다 27.9배 빨랐다. RL 에이전트는 작은 벤치마크에서 신경망 추론 단계로 인한 오버헤드가 발생했으나, 조합 탐색이 지나치게 느려지는 대규모 회로에서는 뛰어난 확장성을 보여주었다.
• 절제 연구 (Ablation Studies):
  • LLM vs 무작위 데이터: LLM 생성 데이터로 훈련된 에이전트는 무작위 데이터로 훈련된 에이전트보다 특정 벤치마크에서 13배 빠른 회로를 생성하여, 현실적인 코드 패턴의 중요성을 입증했다.
  • 보상 설계: 단계 보상과 최종 보상을 결합하는 것이 단계 보상만 사용하는 것보다 우수했으며, 이는 지역 최적해를 방지했다.
  • 아키텍처: 트랜스포머 인코더는 프로그램 표현 학습에서 GRU 베이스라인보다 우수했으며, 계층적 행동 공간은 학습 속도와 보상 측면에서 평탄한 행동 공간보다 우수했다.

5. 의의 및 주장

본 논문은 CHEHAB RL이 FHE 컴파일링의 패러다임 변화를 나타낸다고 주장한다. 휴리스틱 기반의 조합 탐색에서 학습 기반의 정책 최적화로 전환함으로써, 이 프레임워크는 기존 방식의 고질적인 문제인 확장성 및 지역 최적해 문제를 해결한다.

저자들은 이것이 FHE 최적화를 위해 RL을 사용한 첫 사례임을 강조하며, 벡터화, 암호학적 노이즈 축적, 연산 비용 사이의 트레이드오프를 성공적으로 조율했다고 밝힌다. 이 연구는 RL 에이전트가 수동으로 튜닝된 휴리스틱 및 ILP 솔버보다 특히 비구조적 코드에 대해 더 효과적인 전역 전략을 학습할 수 있음을 보여준다. 또한 LLM을 활용한 데이터셋 합성은 데이터가 부족한 도메인에서 고품질 훈련 데이터를 생성하는 경로를 제시한다.

결론적으로 CHEHAB RL은 암호화된 연산의 실행 시간과 컴파일 시간을 모두 크게 단축하면서도, 기초가 되는 암호 체계의 보안 보장을 유지함으로써 FHE의 실질적인 도입을 가능하게 한다.