기술 요약: ZO-Act

문제 정의

대규모 언어 모델(LLM)의 미세 조정(Fine-tuning)은 일반적으로 1차 최적화(역전파)에 의존하는데, 이는 모델이 수십억 개의 파라미터로 확장됨에 따라 활성화 함수(activations), 옵티마이저 상태(optimizer states), 그래디언트 계산을 위한 저장 요구량으로 인해 메모리 비용이 매우 높아지는 문제가 발생합니다. 제로 차수(Zeroth-order, ZO) 최적화는 손실 함수(loss) 평가만을 사용하여 업데이트 방향을 추정하는 순방향 전용(forward-only) 대안을 제공하며, 이를 통해 역전파를 피할 수 있습니다. 그러나 기존의 ZO 방식들은 다음과 같은 중대한 과제에 직면해 있습니다:

1. 높적인 분산(High Variance): 전체 파라미터 공간을 섭동(perturbation)시키는 것(MeZO와 같은 방식)은 높은 분산의 그래디언트 추정치를 생성하여 수렴 속도를 늦추고 불안정성을 초래합니다.
2. 차원이 낮은 부적절한 부분 공간(Suboptimal Subspaces): LOZO나 SubZero와 같이 일부 방식은 저차원 부분 공간으로 섭동을 제한하지만, 이들은 종종 무작위로 구성된 부분 공간이나 실제 그래디언트 흐름과 일치하지 않는 무작위 섭동에 의존합니다.
3. 옵티마이저 호환성: 많은 ZO 방식은 섭동 구조가 명시적이고 가벼운 학습 가능한 변수를 노출하지 않기 때문에, 표준 모멘텀 기반 옵티마이저(예: Adam)를 효과적으로 활용하는 데 어려움을 겪습니다.
4. 양자화 제약(Quantization Constraints): 전체 가중치 섭동을 실체화(materialize)해야 하는 방식은 양자화된 LLM(예: INT4)을 미세 조정하기 어렵습니다.

방법론: ZO-Act

저자들은 ZO-Act라는 하나의 샷(one-shot), 활성화 정보 기반(activation-informed) ZO 미세 조정 방법을 제안합니다. 핵심 통찰은 선형 레이어의 가중치 그래디언트의 지배적인 성분들이 입력 활성화 행렬의 상위 우측 특이 벡터(top right singular vectors)가 형성하는 부분 공간 내에 존재한다는 것입니다.

핵심 메커니즘

1. One-Shot 초기화: 각 선형 레이어에 대해, ZO-Act는 캘리브레이션 배치(calibration batch)에 대한 단 한 번의 순방향 패스(forward pass)를 수행하여 입력 활성화 행렬 $X$를 계산합니다. 그 후, 얇은 특이값 분해(thin SVD)를 통해 $X$의 상위-$r$ 우측 특이 벡터($V_r$)를 계산합니다.
2. 고정된 부분 공간 파라미터화: 가중치 업데이트는 고정된 활성화 벡터들에 의해 정의되는 저차원 부분 공간으로 제한됩니다. 유효한 가중치 업데이트는 다음과 같이 파라미터화됩니다:
   $$\Delta W = V_r B$$
   여기서 $V_r \in \mathbb{R}^{m \times r}$은 고정된 활성화 기저(basis)이며, $B \in \mathbb{R}^{r \times n}$은 가벼운 학습 가능한 계수 행렬입니다. 원래의 사전 학습된 가중치 $W$와 기저 $V_r$은 훈련 과정 내내 고정된 상태를 유지합니다.
3. 계수 공간에서의 최적화: 전체 가중치 행렬 $W$를 섭동시키는 대신, ZO-Act는 저차원 계수 공간에서 무작위 섭동 $Z$를 샘플링합니다 ($Z \sim \mathcal{N}(0, I)$). 섭동된 유효 가중치는 $W_{eff} = W + V_r(B + \mu Z)$가 됩니다.
4. 순방향 전용 추정: $B$에 대한 그래디언트는 순방향 차분(forward-difference) 손실 평가를 사용하여 추정됩니다:
   $$\hat{g}_B = \frac{L(B + \mu Z) - L(B)}{\mu} Z$$
   이 추정된 그래디언트는 Adam과 같은 표준 1차 옵티마이저를 사용하여 $B$를 업데이트하는 데 사용됩니다.

이론적 분석

본 논문은 ZO-Act를 전체 가중치 공간(차원 $d$)이 아닌 제한된 계수 공간 $\beta$(차원 $k$)에 대한 제로 차수 최적화로 분석합니다.

• 분산 감소: ZO 추정치의 분산은 섭동 차원에 선형적으로 비례합니다. 차원을 $d$(전체 가중치)에서 $k = \sum r n_\ell$(계수, 여기서 $r \ll m$)로 줄임으로써, ZO-Act는 분산에 따른 수렴 항과 유한 차분 오차를 크게 감소시킵 direction.
• 편향-분산 트레이드오프(Bias-Variance Trade-off): 활성화 정보 기반의 부분 공간으로 업데이트를 제한하는 것은 부분 공간 근사 편향(subspace approximation bias)을 도입합니다. 그러나 저자들은 LLM의 활성화와 그래디언트가 저차원 구조를 보이기 때문에, 지배적인 업데이트 방향이 입력 활성화의 상위 특이 벡터들에 의해 잘 포착된다고 주장합니다. 따라서 이 편향은 제어 가능하며 데이터 기반의 부분 공간 특성에 의해 완화됩니다.

주요 기여

1. 활성화 정보 기반 부분 공간: 무작위 또는 정적인 저차원 행렬을 사용하는 기존의 부분 공간 방식과 달리, ZO-Act는 입력 활성화로부터 직접 섭동 부분 공간을 도출하여 해당 레이어의 지배적인 활성화 방향과 일치하도록 보장합니다.
2. 명시적 계수 최적화: 기저를 고정하고 계수 행렬 $B$만을 최적화함으로써, ZO-Act는 ZO 미세 조정을 저차원 공간에 대한 명시적 최적화 문제로 변환합니다. 이를 통해 모멘텀 기반 옵티마이저(예: Adam)를 직접 적용할 수 있게 하며, 전체 가중치 섭동을 실체화하는 것을 방지합니다.
3. 양자화 호환성: ZO-Act는 저비트 가중치는 고정된 채로 저비트 계수 행렬만을 업데이트하므로, 양자화된 LLM(예: INT4)을 자연스럽게 지원합니다.
4. 이론적 보장: 본 논문은 ZO-Act가 전체 가중치 ZO와 비교했을 때, 관리 가능한 수준의 부분 공간 편향을 대가로 그래디언트 추정치의 분산과 유한 차분 오차를 모두 줄인다는 수렴 분석을 제공합니다.

실험 결과

저자들은 Llama-3-8B, OPT-13B, 그리고 INT4 양자화된 Llama-3-8B를 대상으로 언어 이해, 질의응답, 상식 추론 작업에서 성능을 평가했습니다.

• 성능: ZO-Act는 강력한 ZO 베이스라인들(MeZO, LOZO, SubZero, AGJO, ZO-Muon 포함)보다 일관되게 우수한 성능을 보였습니다.
  • Llama-3-8B에서 SST-2, RTE, BoolQ, WiC 작업에 대해 최고의 full-precision ZO 결과를 달성했습니다. 특히, RTE 성능을 81.2(ZO-Muon)에서 87.0으로, CB를 69.6에서 89.3으로 향상시켰습니다.
  • OPT-13B에서도 SST-2, RTE, CB, BoolQ에서 최상위 결과를 달성했습니다.
  • INT4 양자화 모델에서도 ZO-Act는 full-precision ZO 베이스라인들과 경쟁 가능한 수준을 유지하며, 메모리 제약이 있는 미세 조정에서의 효용성을 입증했습니다.
• 효율성: ZO-Act는 비교된 방법들 중 가장 낮은 메모리 사용량(Llama-3-8B 기준 약 16.1 GB, INT4의 경우 5.8 GB까지 하락)을 기록했으며, 고정된 순방향 쿼리 예산(forward-query budget) 하에서 가장 빠른 실행 시간(타 방법론들의 44.5분 이상 대비 42.0분)을 달성했습니다.
• Rank 민감도: 실험 결과, full-precision 모델의 경우 $r=1$의 랭크만으로도 충분히 성능을 낼 수 있었으며, 이는 섭동 차원을 최소화하면서도 최적의 정확도를 얻는 데 도움이 되었습니다. 반면, 양자화된 모델의 경우 감소된 용량을 보완하기 위해 더 높은 랭크($r=32$)가 유익했습니다.
• 안정성: 진단 연구를 통해 활성화 정보 기반의 부분 공간이 미세 조정 중에도 안정적으로 유지되며, 전체 그래디언트 에너지의 상당 부분을 포착한다는 것을 확인하여 "one-shot" 초기화 설계의 타당성을 검증했습니다.

의의 및 주장

본 논문은 ZO-Act가 대규모 및 양자화된 LLM을 위한 실용적이고 고성능인 제로 차수 미세 조정의 중요한 진전임을 주장합니다. 그 의의는 다음과 같습니다:

• 격차 해소: 데이터 기반의 부분 공간을 통해 추정치 분산을 줄임으로써, ZO와 1차 미세 조정 사이의 성능 격차를 좁힙니다.
• 옵티마이저 활용 가능성: 명시적이고 저차원인 학습 가능 변수를 노출함으로써, ZO 방식에 Adam과 같은 표준 옵티마이저를 적용하기 어렵다는 문제를 해결합니다.
• 확장성: 자원이 제한된 하드웨어에서 배포 시 필수적인 양자화 모델 미세 조정을 위한 메모리 효율적인 솔루션을 제공합니다.

저자들은 ZO-Act가 여전히 ZO 추정의 본질적인 노이즈와 부분 공간 편향으로 인해 1차 방식과의 성능 격차를 완전히 좁히지는 못했다는 한계를 인정합니다. 또한, 업데이트당 여러 번의 순방향 패스가 필요하기 때문에 ZO-Act가 역전파 기반 방식보다 실제 실행 시간(wall-clock time) 측면에서는 느리다는 점도 언급했습니다. 그럼에도 불구하고, 이 방법은 역전파가 메모리 제약으로 인해 불가능하거나 사용할 수 없는 시나리오에서 우수한 선택지로 자리매김하고 있습니다.