순전파만으로 언어 모델 미세 조정하기

초록

언어 모델(LM)의 미세 조정(fine-tuning)은 다양한 다운스트림 작업에서 성공을 거두었지만, LM의 크기가 커짐에 따라 역전파(backpropagation)는 엄청난 양의 메모리를 요구하게 된다. 제로차(zeroth-order, ZO) 방법은 원칙적으로 두 번의 순전파(forward pass)만으로도 그래디언트를 추정할 수 있지만, 대규모 모델을 최적화하는 데 있어 치명적으로 느리다고 이론적으로 예측되어 왔다. 본 연구에서는 메모리 효율적인 제로차 최적화기(MeZO)를 제안하며, 이는 고전적인 ZO-SGD 방법을 인플레이스(in-place) 방식으로 동작하도록 조정하여 LM을 추론(inference)과 동일한 메모리 사용량으로 미세 조정할 수 있게 한다. 예를 들어, 단일 A100 80GB GPU를 사용할 경우, MeZO는 300억 개의 파라미터를 가진 모델을 학습할 수 있는 반면, 역전파를 통한 미세 조정은 동일한 예산으로 27억 개의 파라미터를 가진 LM만을 학습할 수 있다. 우리는 다양한 모델 유형(마스킹된 LM과 자기회귀적 LM), 모델 규모(최대 660억 개의 파라미터), 그리고 다운스트림 작업(분류, 다중 선택, 생성)에 걸쳐 포괄적인 실험을 수행하였다. 실험 결과는 (1) MeZO가 컨텍스트 내 학습(in-context learning)과 선형 탐사(linear probing)를 크게 능가함을 보여주며, (2) MeZO가 여러 작업에서 역전파를 통한 미세 조정과 비슷한 성능을 달성하면서 최대 12배의 메모리 절감을 가능하게 함을 보여준다. 또한 (3) MeZO는 전체 파라미터 튜닝과 LoRA 및 프리픽스 튜닝(prefix tuning)과 같은 파라미터 효율적 튜닝 기법과 모두 호환되며, (4) MeZO는 미분 불가능한 목표(예: 정확도 또는 F1 점수 최대화)를 효과적으로 최적화할 수 있다. 우리는 이러한 실험적 결과를 이론적 통찰로 뒷받침하며, 고전적인 ZO 분석이 예측한 것과 달리 충분한 사전 학습과 작업 프롬프트가 MeZO가 대규모 모델을 미세 조정할 수 있게 하는 이유를 강조한다.

English

Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a prohibitively large amount of memory. Zeroth-order (ZO) methods can in principle estimate gradients using only two forward passes but are theorized to be catastrophically slow for optimizing large models. In this work, we propose a memory-efficient zerothorder optimizer (MeZO), adapting the classical ZO-SGD method to operate in-place, thereby fine-tuning LMs with the same memory footprint as inference. For example, with a single A100 80GB GPU, MeZO can train a 30-billion parameter model, whereas fine-tuning with backpropagation can train only a 2.7B LM with the same budget. We conduct comprehensive experiments across model types (masked and autoregressive LMs), model scales (up to 66B), and downstream tasks (classification, multiple-choice, and generation). Our results demonstrate that (1) MeZO significantly outperforms in-context learning and linear probing; (2) MeZO achieves comparable performance to fine-tuning with backpropagation across multiple tasks, with up to 12x memory reduction; (3) MeZO is compatible with both full-parameter and parameter-efficient tuning techniques such as LoRA and prefix tuning; (4) MeZO can effectively optimize non-differentiable objectives (e.g., maximizing accuracy or F1). We support our empirical findings with theoretical insights, highlighting how adequate pre-training and task prompts enable MeZO to fine-tune huge models, despite classical ZO analyses suggesting otherwise.

순전파만으로 언어 모델 미세 조정하기

Fine-Tuning Language Models with Just Forward Passes

초록

Support