FHE를 이용한 암호화된 대규모 언어 모델을 향하여

TL;DR

Hugging Face와 Zama는 완전 동형 암호(Fully Homomorphic Encryption, FHE)를 사용하여 대규모 언어 모델(LLM)의 일부를 암호화된 데이터 위에서 실행하는 방법을 시연했습니다. 이 접근 방식은 서비스 제공자가 원본 데이터를 볼 수 없는 상태에서 LLM이 민감한 사용자 쿼리를 처리할 수 있도록 하며, 동시에 온프레미스 배포를 통해 모델 소유자의 지적 재산이 유출되는 것을 방지합니다.

완전 동형 암호를 통한 LLM 프라이버시 해결

완전 동형 암호(FHE)는 암호화된 데이터에서 직접 함수를 실행할 수 있게 하여, 처리 전에 정보를 복호화할 필요를 없애줍니다. Hugging Face의 transformers 라이브러리에 있는 GPT-2 구현을 조정하고 Concrete-Python을 사용함으로써, 모델 추론 프로세스의 특정 섹션을 FHE 대응 방식으로 변환할 수 있습니다.

TFHE (Torus Fully Homomorphic Encryption) 스킴에서 모델 가중치와 활성화 값은 정수로 표현됩니다. 비선형 함수는 Programmable Bootstrapping (PBS)을 통해 처리되는데, 이는 암호화된 데이터에 대해 테이블 조회(TLU)를 수행하고 암호문을 갱신하여 임의의 계산이 가능하도록 합니다. PBS는 선형 연산보다 계산 비용이 더 많이 들지만, LLM 계산의 모든 하위 부분 또는 전체를 FHE로 표현할 수 있게 해줍니다.

기술적 구현: 어텐션 헤드 암호화

암호화된 LLM 레이어를 구현하기 위해 시스템은 클라이언트가 초기 로컬 추론을 수행하고, 중간 연산을 암호화하여 서버로 전송하는 하이브리드 방식을 사용합니다. 서버는 암호화된 데이터에 어텐션 메커니즘의 일부를 적용한 다음, 클라이언트가 복호화하여 로컬 추론을 계속할 수 있도록 결과를 반환합니다.

양자화 요구사항

FHE는 정수 기반으로 작동하므로 모델 가중치와 활성화 값은 양자화되어야 합니다. 재학습이 필요 없는 사후 양자화(post-training quantization)를 사용하여, Zama는 4비트 양자화가 원래 모델 정확도의 96%를 유지한다는 것을 발견했습니다(약 80개 문장 데이터셋 기준).

GPT-2와의 통합

구현 과정에는 양자화된 연산자를 포함하도록 특정 모듈의 forward pass를 다시 작성하는 작업이 포함됩니다. 예를 들어, GPT2LMHeadModel은 첫 번째 multi-head attention 모듈을 QGPT2SingleHeadAttention 모듈로 교체하여 수정할 수 있습니다.

이 설정에서는 query, key, value 행렬에 대한 프로젝션을 포함한 multi-head attention 메커니즘의 첫 번째 헤드가 FHE 친화적인 연산자를 사용하여 수행됩니다. 다른 계산은 부동 소수점 상태로 유지되며 클라이언트 측에서 암호화되지 않은 상태로 실행됩니다.

계산 복잡도 및 성능

어텐션 메커니즘은 query, key, value의 곱셈으로 인해 트랜스포머 모델에서 가장 계산 집약적인 부분입니다. FHE에서는 암호화 도메인 곱셈의 복잡성과 시퀀스 길이가 길어짐에 따라 연산이 제곱으로 증가하기 때문에 이 비용이 더욱 증가합니다.

주요 성능 관찰 사항은 다음과 같습니다:

  • 계산 부하: 길이가 6인 시퀀스는 11,622회의 PBS 연산이 필요합니다.
  • 현재 상태: 현재 구현은 최적화되지 않은 첫 번째 실험 단계로, CPU에서 수 초 내에 실행되지만 상당한 컴퓨팅 파워를 필요로 합니다.
  • 미래 전망: Zama는 향후 ASIC 하드웨어가 지연 시간을 1,000배에서 10,000배까지 개선하여, CPU에서 몇 분 걸리던 처리 시간을 100ms 미만으로 줄일 수 있을 것으로 전망합니다.

결론

LLM에 FHE를 통합하는 것은 사용자 프라이버시가 완전히 존중되고 모델 지적 재산이 보호되는 클라우드 기반 AI 서비스로 가는 경로를 제공합니다. transformers 라이브러리와 Zama의 Concrete 및 Concrete-ML 라이브러리를 활용함으로써, 개발자는 ML 모델을 FHE 대응 방식으로 변환하여 암호화된 데이터 위에서 예측을 수행하는 작업을 시작할 수 있습니다.

Sources