[논문 리뷰] Robust and Efficient Zeroth-Order LLM Fine-Tuning via Adaptive Bayesian Subspace Optimizer

Jian Feng, Zhihong Huang|arXiv (Cornell University)|2026. 01. 04.

Metaheuristic Optimization Algorithms Research인용 수 0

한 줄 요약

BSZO는 Kalman 필터링으로 다 방향 그래디언트 정보를 모으는 베이지안 서브스페이스 제로차 순서 최적화기로, LLM 미세조정의 수렴성과 견고성을 향상시키면서 메모리 사용량은 낮게 유지합니다.

ABSTRACT

Fine-tuning large language models (LLMs) with zeroth-order (ZO) optimization reduces memory by approximating gradients through function evaluations. However, existing methods essentially perform updates in a one-dimensional space, and suffer from collapse or substantial performance degradation under low-precision training. We introduce BSZO, an adaptive extbf{B}ayesian extbf{S}ubspace extbf{Z}eroth-Order extbf{O}ptimizer, which applies Kalman filtering to combine finite-difference information across multiple perturbation directions within a subspace. By treating each finite-difference measurement as a noisy observation, BSZO builds a posterior distribution over the subspace-projected gradient and updates it through Bayesian inference, with a residual-based adaptive mechanism to adapt to noise variations. Theoretical analysis shows that BSZO improves the convergence rate by a factor of $k/γ$ compared to standard ZO methods. Experiments on RoBERTa, Mistral, and OPT models show that BSZO outperforms the baselines across various tasks, achieving up to 6.67\% absolute average improvement on OPT-13B while remaining robust under fp16/bf16 precision and keeping memory usage close to inference-only baselines (1.00$ imes$--1.08$ imes$ of MeZO).

연구 동기 및 목표

• 대규모 언어 모델의 역전파 기반 그래디언트 없이 메모리 효율적인 미세조정을 동기화합니다.
• 저정밀 학습에서 기존 제로차(order) 방법의 불안정성과 성능 저하를 해결합니다.
• Bayesian 추론을 통해 여러 perturbation 방향에서 유한 차이 신호를 융합하도록 BSZO를 제안합니다.
• RoBERTa, Mistral, OPT 모델에 대한 이론적 수렴 보장 및 실험적 검증을 제공합니다.

제안 방법

• k개의 무작위 방향을 샘플링하여 k-차원 서브스페이스를 형성하고 프로젝션된 그래디언트를 잠재 변수로 모델링합니다.
• 각 유한 차이 측정치를 정규화된 서브스페이스 그래디언트의 노이즈가 섞인 선형 관측으로 간주하고 Kalman filtering을 통해 Gaussian 후방분포를 업데이트합니다.
• 학습 중 관측 잡음 분산을 동적으로 조정하는 잔차 기반 적응 기법을 사용합니다.
• 서브스페이스의 후방 평균을 내려가며 매 배치마다 k번의 업데이트를 수행하도록 파라미터를 업데이트합니다.
• 초기 perturbation 결과를 캐시하고 재사용하여 순방향 패스를 줄이며, 낮은 정밀도에서 탐색 공간을 더 잘 포착하기 위해 기본 버전(BSZO-B)에서 추가적인 순방향 패스를 선택적으로 수행합니다.

실험 결과

연구 질문

• RQ1 제로차 LLM 미세조정이 다중 perturbation 방향에서 정보를 집계함으로써 안정화되고 데이터 효율이 향상될 수 있는가?
• RQ2 서브스페이스로 투영된 그래디언트에 대한 Kalman 필터링 기반의 베이지안 추론이 fp16/bf16 정밀도에서 수렴 속도와 견고성을 개선하는가?
• RQ3 BSZO가 기존 제로차 및 일차 최적화 방법과 비교해 메모리 및 계산에서 어떤 트레이드오프를 가지는가?
• RQ4 적응적 잔차 기반 노이즈 추정이 모델 규모 및 작업에 따라 성능에 어떤 영향을 미치는가?

주요 결과

• BSZO는 RoBERTa, OPT, Mistral 모델에서 안정적이고 경쟁력 있는 정확도를 달성하며 여러 작업에서 종종 베이스라인을 능가합니다.
• 수렴 속도는 표준 ZO 방법과 비교하여 이론적으로 k/γ 배로 개선됩니다.
• BSZO의 메모리 사용은 추론 전용 베이스라인(1.00×–1.08× MeZO)과 가까우며 HiZOO 및 MeZO-Adam에 비해 훨씬 더 메모리 효율적입니다.
• 감소된 정밀도에서 BSZO 및 BSZO-B는 견고하게 작동하는 반면, 몇몇 베이스라인은 붕괴되거나 현저히 저하됩니다.
• 디코더 전용 모델에서 BSZO는 일관되게 최상위 또는 근사 최상위 평균 정확도를 달성하며, 모델 크기가 커질수록 더 큰 이점을 보입니다.