기술 요약: PowerSoftmax: 암호화된 데이터를 통한 안전한 LLM 추론을 위한 접근

문제 제기

현대적인 프라이버시 보호형 대규모 언어 모델 (LLM) 은 종종 암호화된 데이터에 대한 복호화 없이 연산을 가능하게 하는 동형 암호 (HE), 특히 CKKS 방식을 의존합니다. 그러나 CKKS 는 다항식 연산만 지원합니다. 표준 트랜스포머 아키텍처는 지수, 나눗셈, 로그 연산 (예: 수치적 안정성을 위한 log-sum-exp 트릭) 을 포함하는 자기 어텐션의 Softmax 함수와 같은 비다항식 구성 요소에 크게 의존합니다.

이 격차를 해소하기 위한 기존 접근법은 다음과 같은 중대한 한계에 직면해 있습니다:

1. 학습 후 근사 (PTA): 사전 학습된 모델에서 비다항식 함수를 고차 다항식으로 직접 근사하는 방식입니다. 이는 높은 곱셈 깊이와 누적된 노이즈로 인해 비효율적인 추론을 초래합니다.
2. 근사 인식 학습 (AAT): 학습 전에 점 단위 어텐션이나 스케일된 ReLU 와 같이 근사가 쉬운 대안으로 비다항식 원시 연산을 대체하는 방식입니다. 이는 소형 모델에는 효과적이지만, 안정성과 확장성 측면에서 어려움을 겪었으며, 수십억 개의 파라미터를 가진 LLM 이나 32 층과 같은 심층 아키텍처를 지원하지 못했습니다.

핵심적인 과제는 표준 Softmax 기반 트랜스포머의 표현력, 추론 능력, 그리고 컨텍스트 학습 (ICL) 능력을 수십억 파라미터 규모에서 유지하면서, 동형 암호 제약 조건 (저차 다항식, 안정적인 수치 범위) 과 본질적으로 호환되는 트랜스포머 아키텍처를 설계하는 것입니다.

방법론

저자들은 Softmax 의 핵심 원리를 재현하면서도 다항식 근사에 적합하도록 설계된 새로운 HE 친화적 자기 어텐션 변형인 PowerSoftmax를 제안합니다. 방법론은 세 가지 명확한 단계로 구성됩니다:

1. HE 친화적 어텐션 설계

저자들은 Softmax 의 필수 속성인 $[0, 1]$로의 정규화, 점수 차이 증폭, 그리고 순서 보존을 추출합니다. 그리고 Softmax 의 지수 함수 $\exp(x)$를 $p$가 양의 짝수인 정수인 멱함수 $x^p$로 대체합니다.
제안된 어텐션 메커니즘은 다음과 같이 정의됩니다:
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
이 공식은 Softmax 의 직접적인 다항식 근사에 내재된 지수 함수와 0 으로 나누기 문제를 피합니다.

2. 학습 안정성 및 수치적 견고성

입력이 1 보다 크면 폭발하고 1 보다 작으면 소멸하는 멱함수로 인한 학습 불안정성을 해결하기 위해, 저자들은 두 가지 주요 수정 사항을 도입합니다:

• 안정적 변형: log-sum-exp 트릭에서 영감을 받아, 멱함수를 적용하기 전에 입력을 상수 $c = \|x\|_\infty + \delta$로 정규화합니다. 이는 입력이 관리 가능한 범위 $(0, 1)$ 내에 유지되도록 하여 오버플로우/언더플로우를 방지하고 부동 소수점 정밀도를 보존합니다.
• $\epsilon$-바운드 나눗셈: 0 근처에서 민감한 비다항식 연산인 나눗셈을 처리하기 위해, 분모를 $\epsilon + \sum x_i^p$로 수정합니다. 이는 $1/\epsilon^2$-리프시츠 연속 함수를 생성하여 후속 나눗셈의 다항식 근사를 훨씬 더 다루기 쉽게 만들고 필요한 곱셈 깊이를 줄입니다.

3. 길이 무관 근사

표준 나눗셈 근사 (예: Goldschmidt 알고리즘) 는 시퀀스 길이 $L$에 선형적으로 비례하는 어텐션 점수의 무제한 합으로 인해 어려움을 겪습니다. 저자들은 길이 무관 PowerSoftmax를 제안합니다:
$$\text{Length-Agnostic}(x)_j = \frac{x_j^p}{L \cdot \text{Mean}(x^p)}$$
시퀀스 길이 $L$은 공개된 값이므로, $1/L$ 항은 근사 없이 직접 계산할 수 있습니다. 모델은 합이 아닌 평균의 나눗셈을 근사하여, 시퀀스 길이에 관계없이 근사 도메인이 유계로 유지되도록 합니다.

4. 구축 파이프라인 (알고리즘 1)

다항식 LLM 을 생성하기 위한 전체 파이프라인은 다음과 같습니다:

1. 아키텍처 수정: Softmax 를 Stable PowerSoftmax 변형으로 대체하고 범위 최소화 목적 함수를 사용하여 처음부터 (또는 지속적 사전 학습을 통해) 학습합니다.
2. 범위 최소화: 저차 다항식 근사가 정확하도록 보장하기 위해 비다항식 레이어 (PowerSoftmax, GELU, LayerNorm) 의 입력 범위를 제한하는 전용 손실 함수를 적용합니다.
3. 다항식 대체: 학습 후, 나머지 비다항식 연산 (나눗셈, 역제곱근) 을 Goldschmidt 방법과 같은 저차 다항식 근사로 대체하여 완전히 다항식인 모델을 만듭니다.

주요 기여

1. PowerSoftmax 메커니즘: 지수 스케일링을 다항식 스케일링으로 대체하는 새로운 자기 어텐션 변형으로, 어텐션의 핵심 설계 원리를 유지하면서 HE 환경에 특화되었습니다.
2. 확장 가능한 다항식 LLM: 10 억 개를 초과하는 파라미터 (특히 32 층의 14 억 파라미터) 를 가진 최초의 다항식 LLM 을 개발했습니다. 이러한 모델은 동급의 표준 트랜스포머와 비교할 수 있는 추론 및 컨텍스트 학습 (ICL) 능력을 입증했습니다.
3. 안정적인 학습 및 추론 전략: 이전 다항식 트랜스포머 접근법의 불안정성 및 확장성 문제를 극복하는 수치적으로 안정적인 학습 변형과 길이 무관 추론 전략을 도입했습니다.
4. 종합적 평가: 암호화된 연산에 대한 상세한 지연 시간 분석과 $p$ 멱수 파라미터 및 $\epsilon$-바운드 나눗셈의 영향을 포함한 설계 선택을 검증하는 제거 실험을 수행했습니다.

실험 결과

저자들은 인과 (GPT/Pythia) 및 양방향 (RoBERTa) 아키텍처에서 방법을 평가했습니다:

• 성능: 14 억 파라미터 다항식 모델은 제로샷 및 퓨샷 벤치마크 (Lambada, PIQA, WinoGrande, ARC 등) 에서 원래 Softmax 기반 Pythia 모델과 비교할 수 있는 성능을 달성했습니다. 특히, 이전 다항식 LLM 에서 관찰되지 않았던 ARC-Challenge 와 같은 추론 벤치마크에서 경쟁력 있는 성능을 보였습니다.
• 양방향 모델: 다항식 RoBERTa-Base 변형은 GLUE 작업 (SST-2, QNLI, MNLI) 에서 원래 Softmax 기반 RoBERTa 대비 약 1% 의 성능 저하만 보였으며, Nexus 와 같은 이전 다항식 접근법보다 우수한 성능을 발휘했습니다.
• HE 추론: HElayers 프레임워크를 사용한 Pythia 70M 모델에 대한 실험은 암호화된 추론이 평문 추론과 거의 동일한 결과를 산출함을 확인했습니다 (최대 MSE 손실 0.005).
• 효율성: 제안된 방법은 최신 대안 (Nexus, Cho 등) 에 비해 곱셈 깊이를 크게 줄였습니다. 128 토큰 시퀀스의 경우, PowerSoftmax 는 Taylor/Goldschmidt 접근법 대비 약 9.7 배, square-strategy 접근법 대비 약 20 배의 지연 시간 감소를 보였습니다.

중요성 및 주장

이 논문은 완전 다항식 LLM 이 추론 및 컨텍스트 학습 능력을 희생하지 않고 수십억 파라미터 규모로 확장될 수 있음을 입증함으로써 프라이버시 보호 머신러닝 분야에서 획기적인 발전을 이루었다고 주장합니다.

• 확장성: 이전 연구가 소형 모델이나 단순 분류 작업으로 제한되었던 것과 달리, 이 접근법은 14 억 파라미터를 가진 32 층 심층 트랜스포머의 학습에 성공했습니다.
• 비상호적 보안: 완전 다항식 모델을 사용하여 완전 동형 암호 (FHE) 를 활용함으로써, 다자간 계산 (MPC) 이나 상호작용 프로토콜과 관련된 통신 오버헤드 및 잠재적 취약점을 피하는 비상호적 추론을 가능하게 합니다.
• 실용성: 상세한 지연 시간 분석과 14 억 파라미터 모델의 성공적인 배포는 의료 및 금융과 같은 데이터 민감 분야에서 안전한 프라이버시 보호 LLM 추론이 실제로 가능해지고 있음을 시사하며, 이를 위한 길을 열고 있습니다.

저자들은 모델이 잘 작동하지만, 암호화된 데이터에 대한 순차적 디코딩에서의 자동회귀 생성 능력에 대한 완전한 평가는 향후 연구 과제로 남아 있음을 인정합니다.