적응형 베이지안 서브스페이스 최적화기로 구현한 견고하고 효율적인 제로오더 LLM 파인튜닝

초록

본 논문은 제로오더(Gradient‑free) 방식으로 대형 언어 모델을 파인튜닝할 때, 기존 1차원 업데이트의 한계를 극복하고자 다중 방향의 유한 차분 정보를 베이지안 서브스페이스에 통합한다. 칼만 필터 기반의 적응형 베이지안 추정기로 잡음이 큰 저정밀(fp16/bf16) 환경에서도 안정적인 수렴을 보이며, 이론적으로는 기존 ZO 방법 대비 k/γ 배 빠른 수렴률을 달성한다. 실험 결과 RoBERTa, Mistral, OPT‑13B 등 다양한 모델·태스크에서 평균 6.67%p 이상의 성능 향상을 기록하고, 메모리 사용량은 추론 전용 기준과 거의 동일하게 유지한다.

상세 분석

BSZO(Adaptive Bayesian Subspace Zeroth‑Order Optimizer)는 제로오더 파인튜닝의 핵심 문제인 “그라디언트 추정의 고잡음·저차원성”을 두 단계로 해결한다. 첫 번째는 파라미터 공간을 저차원 서브스페이스(차원 k)로 투사하고, 해당 서브스페이스 내에서 다중 방향으로 유한 차분(ε·u_i) 샘플을 수행한다. 각 방향 u_i는 정규화된 랜덤 벡터이며, ε는 스케일 파라미터다. 두 번째는 이들 차분값을 “노이즈가 섞인 관측값”으로 간주하고, 칼만 필터를 이용해 베이지안 사후분포를 업데이트한다. 구체적으로, 사전분포는 평균 μ_t와 공분산 Σ_t 로 표현되며, 관측 모델은 y_i = u_i^T g + η_i (η_i∼𝒩(0,σ²)) 형태이다. 칼만 이득 K_t = Σ_t U (U^T Σ_t U + σ²I)^{-1} 를 통해 사후 평균 μ_{t+1}=μ_t+K_t(y-U^T μ_t) 와 공분산 Σ_{t+1}=(I-K_t U^T) Σ_t 를 계산한다. 여기서 U는 현재 배치의 방향 행렬이다.

핵심 적응 메커니즘은 “잔차 기반 스케일 조정”이다. 각 스텝에서 관측 잔차 r_t = y - U^T μ_t 의 제곱 노름을 추정해 σ²를 동적으로 업데이트한다. 이는 잡음 수준이 변동하는 저정밀 연산(fp16, bf16) 상황에서도 필터가 과도하게 신뢰하거나 과소신뢰하지 않게 만든다. 또한, 서브스페이스 차원을 k 로 조절함으로써 메모리와 연산량을 선형적으로 제어한다.

이론적 분석에서는 기존 ZO‑SGD가 기대 수렴 속도 O(d/ε²) (d는 전체 차원)인 반면, BSZO는 서브스페이스 차원 k 와 잡음 감소 계수 γ(γ≤1) 를 고려해 O((k/γ)/ε²) 로 개선됨을 증명한다. 즉, 동일한 함수 평가 횟수에서 k 배 더 많은 유용한 정보가 활용되며, γ 은 칼만 필터가 잡음을 억제하는 정도를 나타낸다.

실험에서는 RoBERTa‑base, Mistral‑7B, OPT‑13B 를 대상으로 GLUE, SuperGLUE, SQuAD, 그리고 텍스트 생성 태스크를 수행했다. BSZO는 MeZO, ZO‑Adam, ZO‑Sign 등 기존 ZO 기반 파인튜닝 방법보다 평균 3.2%~6.7%p 높은 정확도/점수를 기록했으며, 특히 저정밀(fp16, bf16) 환경에서 성능 저하가 거의 없었다. 메모리 사용량은 파라미터 저장 외에 추가적인 버퍼(공분산 행렬)만을 필요로 하므로, 전체 메모리 오버헤드는 1.00×~1.08× 로 추론 전용 메모리와 거의 차이가 없었다.

BSZO의 장점은 (1) 다중 방향 정보를 효율적으로 결합해 고차원 파라미터 공간에서도 안정적인 그라디언트 추정, (2) 칼만 필터 기반 적응형 잡음 추정으로 저정밀 연산에 강인함, (3) 서브스페이스 차원 조절을 통한 메모리·연산 효율성, (4) 이론적 수렴 가속도 보장이다. 한계점으로는 서브스페이스 차원 k 와 초기 공분산 Σ_0 설정이 성능에 민감할 수 있으며, 매우 큰 모델에서는 공분산 행렬 저장 비용이 여전히 제한 요인이 될 수 있다. 향후 연구는 스파스 공분산 근사, 자동 k 선택, 그리고 비선형 서브스페이스(예: 커널 PCA) 확장을 탐색할 여지가 있다.