대형 언어 모델(LLMs)은 자연어 이해와 생성에 점점 더 중요해지고 있다. 그러나 이러한 모델들을 다운스트림 작업에 맞게 미세 조정하는 것은 그 규모로 인해 여전히 도전적이다. 표준 접근법인 Adam과 같은 1차 최적화기 사용은 큰 양의 GPU 메모리를 소비한다. 13B 매개변수 모델의 경우, 이는 약 100GB 이상의 GPU 메모리가 필요하다고 해석되며, 추론 단독 비용의 대략 10배다. 이러한 요구 사항은 학계나 산업계 모두에서 완전한 미세 조정을 어렵게 만든다.

메모리 부담을 줄이기 위한 여러 전략이 제안되었다. 매개변수 효율적인 미세 조정(PEFT) 방법, LoRA와 Adapters를 포함하여 기본 모델은 동결하고 추가 매개변수 집합만 업데이트하는 방식이다. 하지만 이러한 방법들은 여전히 역전파에 의존하며 복잡한 작업에서는 전체 미세 조정보다 성능이 떨어질 수 있다. 다른 접근법으로는 제로쓰(0차) 최적화가 있는데, 이는 앞으로 진행되는 패스만 사용하여 기울기를 추정한다. MeZO는 이 방법을 통해 추론과 동일한 메모리 발자국을 가지면서도 합당한 정확도를 달성할 수 있음을 보여주었다. 하지만 제로쓰 방법은 높은 변동성을 가진 유한차분 기울기 추정 때문에 첫째 차 최적화기 대비 더 많은 반복이 필요하며, 수렴 속도가 느리다는 단점이 있다.

이는 다음과 같은 질문을 던진다: 수렴 속도와 메모리 사용량 사이의 균형은 어떻게 개선할 수 있을까? 기존 제로쓰 방법에는 세 가지 주요 약점이 있다. 첫째, 대부분의 제로쓰 최적화기는 배치 내에서 단일 무작위 방향으로 업데이트를 수행한다. 더 많은 앞으로 진행되는 패스와 페르튜베이션 방향을 사용하더라도 각 페르튜베이션은 고립되어 처리되고, 평균이나 독립적으로 활용된다—즉, 이러한 측정값들이 서로 어떻게 연관되는지에 대한 정보를 잃는다. 둘째, 학습 중 제로쓰 추정치의 노이즈 수준은 크게 변동하지만 대부분의 방법은 이 영향을 고려하지 않는다. 이것은 업데이트가 지역 최소값 주변에서 격렬하게 흔들리거나 분산되어 결국 학습 붕괴를 초래하는 부적응성을 초래한다. 또한, fp16/bf16과 같은 저정밀도 학습은 이러한 방법들이 붕괴하거나 성능이 크게 저하될 수 있음을 보여준다 (그림 1 및 표 [tab:llm-results] 참조).

/>

우리는 이러한 한계를 극복하기 위해 베이지안 서브스페이스 제로쓰 최적화(BSZO) 를 제안한다. 주요 아이디어는 기울기 추정을 추론 문제로 취급하는 것이다. 각 단계에서 우리는 $`k`$ 개의 무작위 방향을 샘플링하여 저차원 서브스페이스를 형성하고, 투영된 기울기를 잠재 변수로 모델링한다. 유한차분 쿼리를 독립적인 추정치로 취급하는 대신 칼만 필터를 사용하여 관측값을 집계한다—기본적으로 다음과 같은 질문을 한다: 지금까지 측정한 내용을 고려할 때, 진짜 기울기에 대한 가장 나은 추정치는 무엇인가? 이 베이지안 형식은 측정 노이즈를 고려하고 동일 수의 앞으로 진행되는 패스로부터 더 정확한 추정치를 생성한다. 우리는 또한 예측 잔차를 추적하고 훈련 중에 변동성을 실시간으로 조정하는 적응 메커니즘을 도입하여 알고리즘이 변화하는 곡률 조건에 적응하도록 한다.

우리의 기여는 다음과 같이 요약할 수 있다:

1. 우리는 베이지안 추론을 사용하여 서브스페이스 내에서 여러 페르튜베이션 방향 간의 기울기 정보를 집계하는 제로쓰 최적화기인 BSZO를 제안한다. 우리가 아는 한, 이는 LLMs에 대한 제로쓰 최적화에 베이지안 추론과 칼만 필터링을 처음으로 적용한 것이다.

2. 우리는 잔차 기반의 적응 방식을 설계하여 BSZO가 수동 조정 없이 매개변수 업데이트 크기를 자동으로 조정할 수 있게 한다.

3. BSZO의 수렴성을 분석하고, 표준 제로쓰 방법에 비해 $`k/\gamma`$ 배 개선된 수렴 속도를 보여준다.

4. 다양한 LLMs과 벤치마크에서 실험은 BSZO가 다양한 작업에서 강력한 성능을 달성하고, 저정밀도 학습 하에서도 안정적이며 MeZO와 유사한 메모리 소비량을 유지함을 보여준다.

관련 연구

제로쓰 최적화를 위한 LLMs.

전통적인 도함수 없는 방법들은 대체 모델링을 통해 강력한 샘플 효율성을 달성하지만, 차원이 증가할수록 각 반복의 비용이 급격히 증가하여 LLM 규모에서는 사용하기 어려워진다. SPSA 추정기는 무작위 페르튜베이션을 통해 기울기를 근사화함으로써 확장 가능한 대안을 제공한다. 이를 바탕으로 MeZO는 메모리 효율적인 제로쓰 미세 조정을 LLMs에 도입하여, 임의의 시드에서 페르튜베이션을 재생성함으로써 추론 시간 메모리를 일치시킨다. 후속 방법들은 다른 병목 현상을 타겟팅한다: Sparse-MeZO는 영향력 있는 매개변수에 대한 업데이트를 제한하고, HiZOO는 적응적 사전 조건화를 위해 대각 행렬 추정치를 사용하며, LOZO는 낮은 랭크 기울기 구조를 활용하고, TeZO는 반복 간의 시간 상관성을 포착한다. 이러한 발전에도 불구하고 대부분의 방법들은 배치당 업데이트를 수행하는 “한 배치, 한 업데이트” 패러다임을 따르며, 배치 내에서 여러 함수 평가가 여러 매개변수 업데이트를 지원할 수 있다는 가능성은 간과한다. 또한 일부 방법들은 상당한 메모리 오버헤드를 초래하며, 이는 여전히 전체 미세 조정보다 낮지만 제로쓰 최적화의 원래 목표—메모리 소비량을 최소화하는 것—와 충돌할 수 있다. 저정밀도 학습은 메모리 제약 상황에서 필수적이므로 이러한 방법들의 안정성에 대한 추가 평가가 필요하다.

인구 기반 기울기 추정.

다른 접근법으로는 반복마다 여러 페르튜베이션을 평가하고 이를 단일 업데이트로 집계하는 방법이 있다. 진화 전략과 Augmented Random Search는 강화 학습에서 이 패러다임을 인기화시켰다. 하지만 이러한 방법들은 신뢰할 수 있는 기울기 추정치를 얻기 위해 배치당 많은 함수 평가를 필요로 한다. LLM의 각 앞으로 진행되는 패스가 이미 계산적으로 비싸므로, 그러한 샘플 집약적인 접근법은 언어 모델 미세 조정에 대해 실용적이지 않다. 이는 자연스럽게 다음과 같은 질문을 제기한다: 한정된 수의 함수 평가에서 더 많은 정보를 어떻게 추출할 수 있을까? 우리의 연구에서는 유한차분 측정값을 기저 기울기에 대한 임의적인 선형 관측치로 취급하고 베이지안 추론을 사용하여 방향 간에 정보를 융합하는 방법으로 이 질문에 답한다.

최적화를 위한 베이지안 추론.

베이지안 방법은 관찰값과 사전 지식을 통합하면서 불확실성을 정량화하는 원칙적인 방법을 제공한다. 칼만 필터링은 새로운 측정치가 도착할 때마다 숨겨진 상태에 대한 가우시안 믿음을 순차적으로 업데이트하는 표준 예이다. 가우시안 프로세스는 이 아이디어를 함수 공간 모델링으로 확장하고 베이지안 최적화의 기반을 제공한다. 우리의 연구에서는 칼만 관점을 제로쓰 기울기 추정에 적용한다: 투영된 기울기를 숨겨진 상태로 모델링하고 각 페르튜베이션 쿼리를 임의적인 선형 관측치로 해석하며, 단계 내에서 샘플된 모든 방향을 통합하는 사후 확률을 업데이트한다. 베이지안 프레임워크의 유연성을 활용하여 우리는 역사적 정보와 현재 배치 정보를 융합하는 적응 잔차 메커니즘을 설계한다. 이는 추가적인 메모리 오버헤드 없이 개선된 기울기 추정치를 제공한다.

방법

본 절에서는 베이지안 서브스페이스 제로쓰 최적화(BSZO) 알고리즘을 제시한다. 이 알고리즘은 베이지안 방법을 사용하여 서브스페이스의 단계 크기를 제어한다.

기본 개념

우리는 다음과 같은 확률적 최적화 문제를 고려한다:

\begin{equation}
  \min_{\theta \in \mathbb{R}^n} \mathcal{L}(\theta) := \mathbb{E}_{\xi \sim \mathcal{D}}[\mathcal{L}(\theta; \xi)],
\end{equation}

여기서 $`\theta \in \mathbb{R}^n`$는 모델 매개변수, $`\mathcal{D}`$는 훈련 데이터셋이고, $`\mathcal{L}(\theta; \xi)`$는 미니배치 $`\xi`$에 대한 손실이다. 최적값은 $`\mathcal{L}^* := \min_{\theta} \mathcal{L}(\theta)`$로 표기한다.

\begin{equation}
    \| \mathcal{L}(\theta) - \mathcal{L}(\theta') \| \leq L \| \theta - \theta' \|.
\end{equation}

\begin{equation}
    \mathcal{L}(\theta)\le\mathcal{L}(\theta') + \nabla \mathcal{L}(\theta')^\top (\theta - \theta') + \frac{L}{2} \| \theta - \theta' \|^2.
\end{equation}

\begin{equation}
    \mathbb{E}_\xi[\|\nabla \mathcal{L}(\theta; \xi) - \nabla \mathcal{L}(\theta)\|^2] \leq \sigma_g^2, \quad \forall \theta \in \mathbb{R}^n
\end{equation}

\begin{equation}
    \hat{y}(\theta; \xi, d) = \frac{\mathcal{L}(\theta + \varepsilon Bd; \xi) - \mathcal{L}(\theta; \xi)}{\varepsilon},
\end{equation}

베이지안 기울기 추정

$\mathcal{L}(\theta+\varepsilon Bd)$에 대해, 서브스페이스 기울기는 연쇄 법칙을 통해 얻어진다:

\begin{equation}
  g_d := \nabla_d \mathcal{L}(\theta+\varepsilon Bd\mid d=0) = \varepsilon B^T g,
\end{equation}

여기서 $`g:=\nabla \mathcal{L}`$는 $\mathcal{L}$의 실제 기울기이다. 수치 정확성을 유지하기 위해, 정규화된 서브스페이스 기울기를 도입한다: $`\tilde{g} := B^\top g = \frac{g_s}{\varepsilon}\in\mathbb{R}^k`$.

\begin{equation}
    \mathbb{E}[\hat{y}(d)] = d^\top B^\top g + O(\varepsilon L)\approx d^\top \tilde{g}.
\end{equation}

보조정리 5에 근거하여, 단면 차이 $\hat{y}(d)$는 정규화된 서브스페이스 기울기 $\tilde{g}$의 가우시안 노이즈가 있는 선형 관측치로 모델링할 수 있다:

\begin{equation}
  \hat{y}(d) = d^\top \tilde{g} + \nu, \quad \nu \sim \mathcal{N}(0, \sigma_{e}^2 \|d\|^2),
\end{equation}

여기서 $\nu$는 종합적인 노이즈 항을 나타낸다. 이 방식의 분산 정의에 대한 설명은 부록 8.2에서 제공한다. 그런 다음, 베이지안 접근법으로 $\tilde{g}$에 가우시안 사전 확률을 두어 $`\tilde{g} \sim \mathcal{N}(0, \sigma_p^2 I_k)`$를 만족하게 하여 사후 확률을 폐포형에서 계산할 수 있다.

서브스페이스 내의 사후 업데이트

표준 베이지안 선형 회귀 이론에 따르면, $`m`$ 개의 페르튜베이션과 관측치 $(d^{(1)}, \hat{y}^{(1)}), \ldots, (d^{(m)}, \hat{y}^{(m)})$ 이후 사후 확률 $\tilde{g} | Y \sim \mathcal{N}(\mu^{(m)}, \Sigma^{(m)})$도 가우시안 분포를 따른다. 여기서

\begin{equation}
\begin{aligned}
  \Sigma^{(m)} &= \left( \sigma_p^{-2} I_k + D^\top R^{-1} D \right)^{-1}, \\
  \mu^{(m)} &= \Sigma^{(m)} D^\top R^{-1} Y.
\end{aligned}
\end{equation}

여기서 $`D = [d^{(1)}, \ldots, d^{(m)}]^\top \in \mathbb{R}^{m \times k}`$는 설계 행렬이고, $`Y = [\hat