활성화 기반 제로오더 최적화로 LLM 파인튜닝 메모리 절감

초록

본 논문은 대규모 언어 모델(LLM) 파인튜닝 시 메모리 사용을 최소화하기 위해, 순전파 과정에서 얻은 활성화 행렬이 정의하는 저차원 서브스페이스에 제한된 제로오더(Zero‑Order) perturbation을 적용하는 AGZO(Activation‑Guided Zeroth‑Order) 방법을 제안한다. 이 접근법은 기존의 등방성( isotropic) 혹은 무작위 저랭크 perturbation보다 실제 그래디언트와의 코사인 유사도가 높으며, 실험에서 MeZO·LOZO 등 최신 ZO 베이스라인을 지속적으로 능가하고 1차 미분 기반 파인튜닝에 근접한 성능을 보이면서 메모리 사용량은 거의 변하지 않는다.

상세 분석

AGZO의 핵심 아이디어는 “선형 레이어의 그래디언트는 해당 레이어에 입력된 활성화 행렬이 생성하는 열 공간에 완전히 포함된다”는 수학적 사실이다. 논문은 이를 식 (5)·(6)으로 명시하고, 실제 GPT‑2·SST‑2 실험에서 활성화 행렬의 상위 r (예: r = 10) 특이벡터만을 사용해도 그래디언트와의 코사인 유사도가 0.99에 육박함을 보여준다. 이는 고차원 파라미터 공간에서 대부분의 유용한 신호가 저차원 서브스페이스에 집중된다는 강력한 경험적 증거이며, 저랭크 구조가 자연스럽게 존재한다는 기존 연구(Aghajanyan et al., 2021 등)와도 일치한다.

이론적으로는 AGZO가 “서브스페이스 스무딩(subspace‑smoothed) 목표”를 최적화한다는 점을 증명한다. 즉, 무작위 전체 공간에 대한 가우시안 스무딩 대신, 활성화가 정의하는 서브스페이스에 대해 가우시안 스무딩을 수행함으로써, 기대 코사인 유사도가 기존 ZO 방법보다 크게 향상된다는 불평등을 도출한다. 이 불평등은 활성화 행렬의 스펙트럼이 급격히 감소한다는 가정 하에 성립하며, 실제 실험에서도 코사인 유사도가 평균 +0.15 ~ +0.25 정도 상승함을 확인한다.

알고리즘 측면에서 AGZO는 매 미니배치마다 파워 이터레이션을 몇 단계 수행해 활성화 행렬 Hℓ 의 상위 r 개의 좌측 특이벡터 Aℓ 를 추정한다. 이 과정은 Hℓ 을 즉시 폐기하고, Aℓ 만을 메모리에 유지함으로써 메모리 오버헤드를 최소화한다. 선형 레이어에 대해서는 Δℓ = Rℓ Aℓᵀ ( Rℓ 은 표준 정규분포를 따르는 d_out × r 행렬) 로 저랭크 perturbation을 생성하고, 비선형 레이어는 기존 가우시안 perturbation을 그대로 사용한다. 이렇게 구성된 Δ 는 전체 파라미터 공간이 아닌, 활성화가 정의하는 저차원 서브스페이스에 제한되므로, 동일한 쿼리 예산(함수 평가 횟수)에서도 더 유의미한 방향으로 파라미터를 업데이트한다.

실험에서는 Qwen‑3·7B·13B와 Pangu‑α·13B 모델을 대상으로 다양한 자연어 처리 벤치마크(예: GLUE, SuperGLUE, Alpaca‑Eval 등)에서 MeZO·LOZO 대비 평균 2‑4 % 포인트의 정확도 향상을 기록한다. 특히, 시퀀스 길이·배치 크기 증가에 따른 GPU 피크 메모리 사용량을 측정했을 때, AGZO는 MeZO와 거의 동일한 메모리 프로파일을 보이며, 1차 미분 기반 파인튜닝 대비 ≈ 30 % 정도의 메모리 절감을 유지한다. 또한, 직접 측정한 그래디언트 코사인 유사도에서도 AGZO가 0.85 ~ 0.92 수준을 달성해, 기존 ZO 방법(0.65 ~ 0.78)보다 현저히 높은 정렬도를 보여준다.

전체적으로 AGZO는 (1) 활성화‑구조 기반 저차원 서브스페이스를 실시간으로 추출, (2) 해당 서브스페이스에 제한된 저랭크 perturbation을 사용, (3) 메모리 효율성을 유지하면서도 그래디언트 근사 정확도를 크게 향상시키는 세 가지 핵심 설계 원칙을 성공적으로 구현한다. 이는 메모리 제한이 심한 현장 환경에서 LLM 파인튜닝을 실현할 수 있는 새로운 패러다임을 제시한다.