기술 요약: Dominant-Layer ZO

문제 정의

제로 차수(Zeroth-order, ZO) 최적화는 대규모 언어 모델(LLM)을 미세 조정할 때 역전파의 메모리 오버헤드를 피하고 순방향 패스(forward pass)만을 사용하여 적응할 수 있게 하는 메모리 효율적인 대안으로 부상했습니다. MeZO와 같은 방법들은 LLM이 추론 수준의 메모리로도 미세 조정될 수 있음을 입증했지만, 유용한 적응(adaptation)이 모델 아키텍처의 어느 부분에 어떻게 분포되어 있는지에 대한 내부 메커니즘은 여전히 불분명합니다. 기존의 ZO 접근 방식은 대부분 미세 조정을 전체 모델 프로세스로 취급하며, 이로 인해 "순방향 전용 업데이트 상황에서 유용한 최적화가 LLM 아키텍처 내의 어디에서 실제로 발생하는가?"라는 근본적인 질문을 남겨두고 있습니다.

방법론

저자들은 ZO 미세 조정의 계층별(layer-wise) 동작에 대한 체계적인 조사를 제안하며, 이를 통해 "지배적 계층(Dominant Layer)"의 발견과 활용을 이끌어냈습니다.

1. 지배적 계층의 경험적 발견

LLaMA2-7B 및 Qwen3-8B를 대상으로 여러 다운스트림 태스크(분류, 다중 선택, 생성)에 걸쳐 계층별 격리 연구를 수행한 결과, 저자들은 다음을 관찰했습니다:

• 불균등한 적응: 대부분의 계층은 ZO 하에서 개별적으로 업데이트될 때, 미세 조정이 없는 베이스라인보다 개선이 거의 없거나 전혀 나타나지 않았습니다.
• 단일 계층의 지배력: 특정 단일 계층이 일관되게 전체 모델 ZO 미세 조정과 대등하거나 이를 상회하는 성능을 달립다.
• 태스크 불변성 및 모델 특이성: 주어진 모델에 대해 동일한 계층이 다양한 태스크에서 지배적인 역할을 합니다. 그러나 이 계층의 인덱스는 모델 제품군(예: LLaMA2-7B의 경우 레이어 1 vs. Qwen3-8B의 경우 레이어 6)에 따라 달라집니다.
• 1차 미세 조정(First-Order, FO)과의 대조: 매칭된 FO 미세 조정에서는 개선이 고르게 분포되며 단일 계층이 지배적이지 않습니다. 이는 이러한 현상이 ZO 최적화 고유의 특성임을 나타냅니다.

2. 활성화 이상치(Activation Outliers)를 통한 식별

비싼 계층별 미세 조정을 거치지 않고 지배적 계층을 식별하기 위해, 저자들은 활성화 이상치(activation outliers)—입력 전반에 걸쳐 일관되게 나타나는 매우 큰 크기를 가진 특정 차원—현상을 활용합니다.

• 관찰: 지배적 계층은 항상 활성화 이상치가 나타나는 첫 번째 잔차 스트림(residual stream) 계층과 일치합니다.
• 선택 방법: 간단한 추론 전용 분석이 제안됩니다. 작은 캘리브레이션 세트가 주어지면, 사전 훈련된 모델을 순방향으로 실행하여 계층별 활성화 통계를 측정합니다. 명확한 이상치 패턴(비정상적으로 큰 최대 활성화 크기)을 보이는 첫 번째 계층이 미세 조정을 위한 지배적 계층으로 선택됩니다.

3. 이론적 설명

논문은 왜 이 계층이 중요한지를 다음과 같이 설명합니다:

• 섭동 민감도(Perturbation Sensitivity): ZO는 무작위 파라미터 섭동에 의해 발생하는 최종 손실(loss)의 차이를 기반으로 업데이트를 추정합니다. 지배적 계층은 섭동 하에서 가장 큰 손실 변화를 유도합니다.
• 잔차 스트림 전파(Residual Stream Propagation): 지배적 계층은 모델 초기에 등장하며 활성화 이상치 계층과 일치하기 때문에, 여기에 주입된 섭동은 잔차 스트림에 진입하여 이후의 모든 디코딩 계층으로 전파됩니다.
• 신호 축적(Signal Accumulation): 이러한 섭동은 네트워크를 통과하면서 반복적으로 변형되고 축적되어, 최종 손실에서 불균형하게 크고 안정적인 변화를 만들어냅니다. 이는 후속 계층의 섭동이 손실 변화를 유도할 수는 있으나 안정적인 유용한 업데이트를 형성할 전파 깊이가 부족한 것과 달리, 순방향 전용 업데이트를 위한 강력한 최적화 신호를 제공합니다.

주요 기여

1. 발견: 단일 계층을 튜닝함으로써 전체 모델 ZO 미세 조정의 성능을 회복하거나 능가할 수 있는 "지배적 계층" 현상을 규명했습니다.
2. 효율적 식별: 지배적 계층은 모델마다 다르며, 첫 번째 활성화 이상치 계층에 대한 경량화된 추론 전용 분석을 통해 훈련 전에 식별될 수 있음을 입증했습니다.
3. 메커니즘 통찰: 지배적 계층의 지배력이 높은 섭동 민감도와 잔차 스트림 내 초기 배치(early placement)의 결 조합에서 기인하며, 이를 통해 섭동 효과가 네트워크 전체에 걸쳐 축적될 수 있음을 설명했습니다.
4. 검증: 9개의 벤치마크에 대해 LLaMA2-7B 및 Qwen3-8B를 사용하여 광범위한 실험을 수행했습니다.

실험 결과

LLaMA2-7B 및 Qwen3-8B를 대상으로 SST-2, RTE, CB, BoolQ, WSC, MultiRC, COPA, SQuAD, DROP 데이터셋을 사용하여 실험을 진행했습니다.

• 성능:

  • LLaMA2-7B: Dominant-layer ZO 미세 조정은 전체 모델 MeZO 대비 평균 0.57%, MeZO LoRA 대비 **0.09%**의 점수 향상을 보였습니다. 특히 RTE(+2.17%)와 SQuAD(+2.32%, LoRA 변형 기준)에서 주목할 만한 이득이 관찰되었습니다.
  • Qwen3-8B: Dominant-layer ZO 미세 조정은 전체 모델 MeZO 대비 평균 1.15%, MeZO LoRA 대비 **1.12%**의 점수 향상을 보였으며, CB(+3.57%)와 SST-2(+2.04%)에서 상당한 이득을 보였습니다.
  • 비교: 이 방법은 Sparse-MeZO(요소별 희소 마스킹 사용)보다 평균 0.13% 높은 성능을 보였으며, 이는 계층 수준의 선택이 파라미터 수준의 선택만큼 중요하다는 것을 시사합니다.

• 훈련 효율성:

  • 업데이트를 단일 계층(예: LLaMA2-7B의 32개 중 1개)으로 제한함으로써 파라미터 측 연산에서 상당한 속도 향상을 달성했습니다.
  • 섭동 시간은 27.34배~31.28배 감소했으며, 업데이트 시간은 31.56배~32.71배 감소했습니다.
  • 엔드 투 엔드(End-to-End) 속도 향상: 순방향 패스에 소요되는 시간과 파라미터 연산에 소요되는 시간의 비율에 따라, 긴 입력 태스크(DROP)에서는 1.12배에서 짧은 입력 태스크(COPA)에서는 4.52배까지 나타났습니다.

• 절제 연구(Ablation Studies):

  • 두 번째 계층(예: 더 늦은 단계의 민감한 계층)을 추가해도 성능이 일관되게 향상되지 않았으며, 이는 병목 현상이 업데이트되는 계층의 수가 아니라 지배적 계층의 특정 전파 특성에 있음을 시사합니다.
  • 지배적 계층 내에서 활성화 이상치와 관련된 MLP 채널만을 미세 조정했을 때 대부분의 성능 이득을 회복하였으며, 이는 해당 특정 채널들의 결정적인 역할을 확인시켜 줍니다.

의의 및 주장

본 논문은 전체 모델 ZO 미세 조정이 단순히 너무 많은 파라미터를 업데이트하는 것이 아니라, 유용한 적응을 단일 계층에 불균등하게 집중시키고 있다고 주장합니다.

• 실질적 함의: Dominant-layer ZO 미는 단일 지배적 계층에 집중함으로써 전체 모델의 성능을 유지하면서도 훈련 비용(메모리 및 시간)을 크게 줄일 수 있습니다.
• 개념적 함의: 본 연구 결과는 유용한 업데이트가 어디에서(계층 정체성) 그리고 계층 내부의 어디에서(이상치 채널) 발생하는지를 명시적으로 고려하는 것이 균등한 전체 모델 업데이트보다 더 효과적임을 시사하며, 향후 ZO 방법 설계에 대한 통찰을 제공합니다.
• 한계점: 저자들은 Dominant-layer ZO와 1차(First-order) 미세 조정 사이의 여전한 성능 격리를 인정했습니다. 또한, 본 연구가 특정 모델과 옵티마이저로 제한되었다는 점을 언급하며, 향후 연구에서 이러한 접근 방식을 다른 ZO 옵티마이저(예: ZO-AdaMM, FZOO) 또는 더 많은 모델 제품군과 결합하는 것을 탐구할 수 있다고 밝혔습니다.