RL 기반 FHE 코드 자동 최적화 프레임워크 CHEHAB

초록

CHEHAB RL은 깊이 강화학습을 이용해 완전동형암호(FHE) 프로그램을 자동으로 벡터화하고, 연산 지연 및 노이즈 성장을 최소화하는 새로운 컴파일러 최적화 프레임워크이다. LLM으로 생성한 대규모 FHE 표현식 데이터셋을 학습에 활용하며, 기존 최첨단 컴파일러 Coyote 대비 실행 속도 5.3배, 노이즈 2.54배 감소, 컴파일 시간 27.9배 단축을 달성한다.

상세 분석

본 논문은 FHE 코드 최적화를 순차적 의사결정 문제로 정의하고, 이를 해결하기 위해 심층 강화학습(RL) 에이전트를 설계하였다. 행동 공간은 “어떤 변환 규칙을 적용할지”와 “그 규칙을 어느 위치에 적용할지” 두 단계로 구성된 다중 이산 계층형 정책으로 설계돼, 수백 개에 달하는 변환 규칙과 다중 적용 지점을 효율적으로 탐색한다. 보상 함수는 실제 FHE 실행 대신 연산 비용, 회로 깊이, 곱셈 깊이 등을 정량화한 분석적 모델을 사용해 빠르게 계산되며, 노이즈 성장과 실행 지연을 동시에 최소화하도록 가중치를 조정한다. 데이터 부족 문제를 해결하기 위해 대형 언어 모델(LLM)을 활용해 15,855개의 다양한 FHE 표현식을 자동 생성했으며, 이 합성 데이터가 무작위 생성 데이터보다 학습 효율과 최종 성능에서 현저히 우수함을 실험적으로 입증한다.

RL 에이전트는 정책 네트워크와 위치 네트워크로 구성되며, 정책 네트워크는 현재 IR(Intermediate Representation) 상태를 입력으로 받아 변환 규칙을 선택하고, 위치 네트워크는 선택된 규칙의 적용 후보 위치를 평가한다. 이 구조는 평면 행동 공간을 사용하는 기존 방법에 비해 학습 수렴 속도가 빠르고, 탐색 비용이 크게 감소한다. 또한, 변환 규칙 자체는 기존 FHE 컴파일러에서 사용되는 벡터화, 회전, 마스킹 등 핵심 연산을 포괄하며, 규칙 간 의존성을 고려한 순차 적용이 가능하도록 설계되었다.

실험에서는 CHEHAB RL을 CHEHAB 기존 컴파일러와 Coyote에 통합해 12개의 벤치마크(신경망, 행렬 연산, 통계 함수 등)를 평가하였다. 결과는 평균 실행 시간 5.3배 향상, 평균 노이즈 2.54배 감소, 컴파일 시간 27.9배 단축을 보여, RL 기반 최적화가 전통적인 휴리스틱·ILP·프로그램 합성 방식보다 전반적인 효율성을 크게 높임을 증명한다. 특히 구조화된 루프 기반 코드뿐 아니라 비구조화된 스칼라 코드에서도 높은 벡터화 성공률을 기록하였다.

한계점으로는 현재 정책이 특정 FHE 스킴(BFV)과 파라미터 설정에 최적화돼 있어, 다른 스킴(예: CKKS)이나 보안 수준 변화에 대한 일반화가 필요하다. 또한, 보상 함수가 근사 모델에 기반하므로 실제 하드웨어 실행 시 미세한 차이가 발생할 수 있다. 향후 연구에서는 멀티스킴 지원, 보상 함수의 하드웨어 기반 피드백 루프 도입, 그리고 메타러닝을 통한 정책 전이 학습 등을 제안한다.