ENCRYPTION-FRIENDLY LLM ARCHITECTURE

1. Introduction

•

user interactions을 기반으로 personalized response을 생성  → 개인정보 유출 우려

•

Privacy-preserving machine learning (PPML): data privacy를 보호하면서 머신러닝을 사용하는 방법

→ 암호화 가정을 기반으로 증명 가능한 보안을 제공한 것은 아래 2가지 방법론

◦

multi-party computation: parties간의 통신을 활용

◦

homomorphic encryption(HE): 통신 없이도 암호화된 상태에서 산술 연산을 지원

•

HE 기술을 Transformer에 적용하는게 matrix multiplications와 non-polynomial operations (softmax)때문에 어렵다고 함

⇒ inference에 초점을 둔 HE-friendly transformer architecture 제시 및 personalized fine-tuning 학습법 제시

⇒ LoRA와 softmax의 Gaussian Kernel 대체로 성능 향상을 이끌고자 함

1.1 Prior Works

•

Transformer-based language models and LoRA.

◦

LoRA를 통해 ~1% 미만의 파라미터를 업데이트 하고 LLM을 finetuning할 수 있게됨

•

Privacy-preserving transformer using HE.

◦

Interactions: MPC와 HE를 동시에 활용 (parties간의 communications 감소)

▪

large-scale model에서 적용 X, online learning이 필요함

◦

Non-Interactions: original transformer → HE friendly structure

▪

pre-trained weight은 건들이지 X, inference에만 적용되는 구조 제안

⇒ fine-tuning을 통해 LM의 성능을 개선할 수 있는 것은 자명하나 prior works는 주로 보안 추론에 초점을 맞추었기 때문에 계산 복잡성이 많은 fine-tuning 세팅보다는 inference setting에서 실험을 집중적으로 진행

2. SERVER-CLIENT COMPUTATION MODEL AND PRELIMINARIES

2.1 SERVER-CLIENT COMPUTATION MODEL

•

(그림에서 볼 수 있듯이) 

◦

client: private data

◦

server: fine-tuning & inference outsourcing

Client-side에서 ‘encrypted token-embedded data’를 server로 보냄

Server에서는 model을 fine-tuning하면서 ‘encrypted token-embedded data’를 얻음

•

LoRA weight가 Server-side에 있어 사용자와 공유되지 않지만, client private data가 직접적으로 유출이 되는 system은 아님

2.2 HOMOMORPHIC ENCRYPTION AND CKKS

•

동형암호(HE): 복호화 없이 암호화된 상태에서 연산을 수행할 수 있는 암호화 방식

◦

CKKS: 동형 암호화(homomorphic encryption) 체계로, 실수나 복소수 데이터를 메시지 다항식으로 암호화하고 암호화된 상태에서 근사 연산을 수행할 수 있게 해주는 암호화 방식

characteristics

덧셈 (Add)

•

두 암호문 간의 연산 또는 평문(plaintext)과 암호문(ciphertext) 간의 연산 가능

•

요소별 덧셈(component-wise addition) 수행

곱셈 (Mult)

•

두 암호문 간의 연산 또는 평문과 암호문 간의 연산 가능

•

요소별 곱셈(component-wise multiplication) 수행

•

평문-암호문 간의 요소별 곱셈은 특별히 pMult로 표현

회전 (Rot)

•

메시지 벡터 (m0, m1, ..., mN/2-1)를 암호화한 암호문 ct가 주어졌을 때

•

결과값으로 (mi, ..., mN/2-1, m0, ..., mi-1)와 거의 동일한 메시지 벡터를 암호화한 ctroti를 반환

•

i번째 왼쪽 회전은 Rot(·, i)로 표기

•

평문에서의 동일한 연산은 pRot으로 표현

⇒ 연산수의 제한이 있다고 함, 자세히는 모르겠지만 bootstrapping을 통해 연산수 제한을 극복하는데 비효율적이라고…

•

Transformer에 HE를 적용해볼 때 고민해볼 사항은 2가지 존재

◦

plaintext-ciphertext matrix multiplication (PCMM) 

◦

ciphertext-ciphertext matrix multiplication (CCMM)

▪

Homomprhic matrix multiplication에서 Add, Mulit, Rot이 다 필요한데 다 들어가기 때문에 CCMM > PCMM 의 속도라고.

▪

Homomprhic matrix multiplication을 효율화하기 위해 matrix element가 ciphertext인 상태에서도 병렬적으로 처리가 가능한 JKLS라는 알고리즘 활용

3. SPEEDUP WITH LORA: AVOIDING LARGE CCMM

⇒ LoRA를 통해 ciphertext이 일어나는 연산 자체를 줄이자!

•

Encryption을 뜯어보면

Personalized fine-tuning data가 server-side에 암호화되어 전달되고

암호화된 데이터가 pre-training data를 fine-tuning되는데 사용됨

inference시에는 encrypted inference data가 CCMM computations을 통해 암호화된 연산을 수행

X(W + \Delta W) = XW + X\Delta W, \\ W \in \mathbb{R}^{n\times n} \\ X \in \mathbb{R}^{L\times n}

→

W

: PCMM을 수행하기 위한 weight

→

\Delta W

: CCMM 을 수행하기 위한 weight

•

Accelerating homomorphic matrix-multiplication with LoRA.

◦

A∈Rn×rB∈Rr×nA \in \mathbb{R}^{n\times r}\\B \in \mathbb{R}^{r\times n}A∈Rn×rB∈Rr×n

◦

LoRA를 통해 CCMMs의 크기를 줄임으로 CCMMs의 속도를 향상시킬 수 있음 (fine-tuning & inference 모두)

X(W + AB) = XW + (XA)B,\\X \in \mathbb{R}^{L\times n}

→

\mathcal{O}(nr)

2번만 하면 연산이 끝남

•

Reducing optimizer states and inverse square root with LoRA.

◦

암호화된 상태에서 연산을 하면 오류 수정을 위해 bootstrapping (BTS)을 많이 수행하고 이전에 언급한것처럼 이게 bottleneck을 초래

▪

AdamW에 있는 x→1xx \rightarrow \frac{1}{\sqrt{x}}x→x​1​가 delicate 연산을 요구하는데, HE-operation에서의 부담을 줄이기 위해  x→1x+ϵx \rightarrow \frac{1}{\sqrt{x+\epsilon}}x→x+ϵ​1​으로 구함

▪

또한 LoRA를 사용함으로써 BTS를 사용하는 횟수자체를 줄임

•

2-layer BERT model → 368 ciphertexts containing weights

◦

\w LoRA → 15 ciphertexts containing weights

4. SPEEDUP WITH GAUSSIAN KERNEL: POLY-APX-FRIENDLY DESIGN

⇒ HE연산에서 부하를 가져오는 softmax를 Gaussian kernel로 대체

\text{Softmax}(x_1, x_2, \cdots, x_n)i = \frac{\exp(x_i - \alpha)}{\sum_j \exp(x_j - \alpha)},\\ where \ \alpha = \max_{1\leq j\leq n} {x_j}

•

softmax를 안정화시키기 위해 위와 같은 방법론들이 쓰이지만, HE를 안정적으로 수행하기 위해서는 exponential, division, max functions 모두 non-polynomial이어서 approximate시켜야한다고 함

→ 선행 연구에서 attention에서 division없이 transformer의 성능 유지

\text{GK-Attention}(Q, K, V) = S(Q, K)V \\ S(Q, K){ij} = \exp\left(-\frac{1}{2\sqrt{n}}|Q{i:} - K_{j:}|_2^2\right), \\i,j = 1,\ldots,L \\ \text{n is the hidden dimension} \\ \text{Q\{i:\} and K\{j:\} mean the ith and jth row of Q and K}

•

논문에서는 softmax의 연산을 scaled exponential kernel을 normalized한 것과 동일하다고 함

•

그리고 normalizing score가 필요없다는 주장과 함께 위의 수식을 softmax의 대체로 쓰겠다고 주장

•

12n\frac{1}{2\sqrt{n}}2n​1​는 scalar라 위의 수식에서 non-poly는 exp만 존재

exp(x) \approx p_k(x) := \left(1 + \frac{x}{2^k}\right)^{2^k} \\ for \ k \in \mathbb{N}

•

또한 위의 exp(x)연산은 범위를 [−2k,0][-2^k, 0][−2k,0]로 제한시키기 때문에 안정적인 연산이 가능하다고 함

→GAUSSIAN KERNEL의 실험적인 개선을 보여준 논문을 근거로 제시

5. EXPERIMENTAL RESULTS

•

Implementation of homomorphic encryption (HE).

◦

자체적인 HE implemention을 위해 C++ HEaaN library 사용

◦

더 넓은 입력 범위가 보장되는 ExtBTS를 활용했다고 함

•

Architecture.

◦

(PoC라 그런지) 2-layer BERT로 실험하고 q,k,v layer에 LoRA를 추가함

•

Non-polynomial approximation.

◦

아래와 같이 다양한 non-polynomial functions에 대한 polynomial approximation이 적용

▪

“worst prec”:∥f (x) − poly(x)||_∞ 의 최대값

▪

“Avg prec”: 평균(평균)

◦

optimzer의 1x\frac{1}{\sqrt{x}}x​1​의 대체로 newton’s method로 수렴이 가능해짐

⇒ 점화식 형태로 표현이 가능하고 수렴이 가능해짐

Time Measurements

→ CCMM이 PCMM보다 느림

→ softmaxfunction의 경우 여러 approxs를 쓴 기법을 활용했음에도 Gaussian kernerl을 쓴게 훨씬 빨랐음

→ LoRA가 가장 큰 training time efficient를 가져오고, GK가 이후 marginal한 efficiency를 가져옴

→ Inference에는 LoRA & GK 둘다 비슷한 크기의 efficiency를 가져옴

Downstream tasks.

→ Down Stream Task에서 Plaintext Setting에서 GK로 FT하면 성능이 오히려 조금오름

→ 모델 크기가 워낙 작은지라, LoRA를 적용하면 성능이 떨어짐

→ HE Setting에서 plainsetting으로 eval하거나 HE로 eval할 경우 둘다 비슷한 수준으로 Full+SM에 비해서 조금은 떨어지지만 성능을 보임.

→ 당연히 Full+SM에 비해서 inference 속도는 4배정도 빠른편

5. Conclusion

→ 개인정보다 api로 넘어가는걸 경계하는 이 시점에서 Homomorphic encryption (HE) friendly한 transformer architecture를 제시한 아주 좋은 논문

→ 너무 작은 모델 (2-Layer Bert), 고전적인 benchmark (GLUE), hardware specific setting (C++ HEaaN library), 그리고 이 분야 특유의 높은 진입장벽과 표현들 (e.g., CCMM, PCMM)이 limitation이자 아쉬움

→ BERT 자체를 뜯어야 Homomorphic encryption이 완성됨으로 사실상 api 공급업체들에게 유용한 논문이 아닐까 생각이 들었다. 사용자 입장에선 조금 아쉬운 논문..