ZorBA: Zeroth-order Federated Fine-tuning of LLMs with Heterogeneous Block Activation

이 논문은 대규모 언어 모델의 연방 미세 조정 시 VRAM 사용량과 통신 오버헤드를 동시에 줄이기 위해 0 차 최적화와 이종 블록 활성화 메커니즘을 결합한 ZorBA 프레임워크를 제안하고 그 이론적 분석 및 실험적 유효성을 입증합니다.

핵심

🏛️ 배경: 거대한 도서관과 작은 책상들

상상해 보세요. 전 세계에 흩어진 수백 명의 학생 (클라이언트) 이 **거대한 도서관 (대규모 AI 모델)**을 함께 수리하고 개선하려는 상황입니다.

• 문제 1 (메모리 부족): 도서관이 너무 커서, 각 학생이 가진 책상 (VRAM/메모리) 에는 책 전체를 올려둘 공간이 없습니다. 기존 방식은 책을 다 펼쳐놓고 수정해야 해서 책상이 꽉 차버립니다.
• 문제 2 (통신 비용): 학생들은 서로 "어디를 고쳤어?"라고 물어보며 데이터를 주고받아야 하는데, 책이 너무 두꺼워서 우편 비용 (통신량) 이 천문학적으로 듭니다.

🚀 ZorBA: "눈을 감고 만져서 배우는" 새로운 방식

ZorBA 는 이 문제를 해결하기 위해 세 가지 창의적인 전략을 사용합니다.

1. "눈을 감고 만져서" 배우기 (Zeroth-Order Optimization)

기존 방식은 책을 한 장씩 뒤적이며 (역전파) "어디가 잘못됐는지" 정확히 계산하는 방식이라 메모리를 많이 썼습니다.
하지만 ZorBA 는 책을 완전히 뒤적이지 않고, 살짝 만져서 (순방향만 실행) "이렇게 살짝 건드리면 결과가 좋아지나?"를 확인합니다.

• 비유: 어두운 방에서 벽을 더듬어 길을 찾는 것처럼, 정확한 지도 (기울기) 가 없어도 "왼쪽으로 살짝 가보면 벽이 멀어지네?"라고 추측하며 학습합니다. 이 방식은 메모리 사용량을 획기적으로 줄여줍니다.

2. "나만의 책장"만 고르기 (Heterogeneous Block Activation)

모든 학생이 도서관의 모든 책장을 고칠 필요는 없습니다.

• 전략: 중앙 관리자는 각 학생의 책상 크기에 맞춰 학습할 책장 (Transformer Block) 을 다르게 배정합니다.
  • 책상이 작은 학생은 책장 1 개만 고치고, 큰 학생은 책장 5 개를 고칩니다.
• 효과: 학생들은 자신에게 맞는 양의 책장만 학습하므로 메모리 과부하가 사라집니다.

3. "비밀 번호" 공유 (Shared Random Seeds)

기존 방식은 "내가 고친 부분"을 모두 서버에 보내야 했지만, ZorBA 는 **서로 같은 '비밀 번호 (랜덤 시드)'**를 공유합니다.

• 비유: 학생과 관리자가 "오늘은 3 번 책장을 고친다"는 약속 (비밀 번호) 만 주고받으면, 양쪽 모두 그 약속대로 같은 책장을 고칠 수 있습니다.
• 효과: 거대한 데이터 (고친 내용) 를 주고받을 필요가 없어 통신 비용이 거의 제로 (Zero) 에 수렴합니다.

---

⚖️ 핵심 기술: "최적의 배분" 찾기

이제 중요한 질문이 생깁니다. "누가 몇 개의 책장을 맡아야 가장 빨리 도서관이 완성될까?"

• 너무 적게 맡으면 학습이 느려집니다.
• 너무 많이 맡으면 책상이 부족해집니다.

저자들은 이 문제를 수학적으로 분석하여 가장 효율적인 배분 공식을 만들었습니다.

• 핵심 통찰: 단순히 책장 수를 늘리는 것보다, **누가 가장 적은 책장을 맡고 있는지 (최소 인기)**를 고려하여 균형을 맞추는 것이 중요합니다.
• 해결책: 컴퓨터가 자동으로 "누가 몇 개를 맡아야 메모리도 아끼고 학습 속도도 빠를까?"를 계산해내는 알고리즘을 개발했습니다.

---

🏆 결과: 얼마나 잘했을까?

실험 결과, ZorBA 는 기존 방법들보다 다음과 같은 성과를 냈습니다.

1. 메모리 사용량 62% 감소: 작은 책상에서도 거대한 도서관을 학습할 수 있게 되었습니다.
2. 통신 비용 대폭 절감: 데이터 주고받는 양이 거의 없어졌습니다.
3. 학습 속도 향상: 불필요한 작업을 줄이고 핵심에 집중해서 더 빨리 학습이 완료되었습니다.

📝 한 줄 요약

"ZorBA 는 거대한 AI 모델을 작은 컴퓨터들에서 학습시킬 때, '눈을 감고 만져서' 메모리를 아끼고, '비밀 번호'로 통신을 줄이며, 각자의 능력에 맞춰 '책장'을 나누어 가장 효율적으로 학습하는 똑똑한 방법입니다."

기술 요약

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 페더러티드 파인튜닝 (Federated Fine-tuning) 은 분산된 클라이언트 간 협업 학습을 가능하게 하지만, 다음과 같은 두 가지 주요 장애물에 직면해 있습니다.

• 높은 VRAM 사용량: 기존 페더러티드 러닝 (FL) 은 모델 업데이트를 위해 역전파 (Backpropagation, BP) 를 사용하여 1 차 미분 (기울기) 을 계산합니다. LLM 은 수억에서 수십억 개의 파라미터를 가지므로, 역전파 과정에서 필요한 그래디언트와 중간 활성화 (activations) 를 저장하는 데 막대한 비디오 랜덤 액세스 메모리 (VRAM) 가 필요합니다. 이는 리소스가 제한된 클라이언트 (예: 엣지 디바이스) 에서는 실행 불가능할 수 있습니다.
• 커뮤니케이션 오버헤드: LLM 의 고차원 파라미터를 클라이언트와 서버 간에 빈번하게 교환하면 통신 비용이 급증합니다. 또한, 1 차 미분이 불가능한 블랙박스 시스템이나 비미분 가능 연산자가 포함된 모델에서는 기존 FL 을 적용할 수 없습니다.

기존의 제로오더 최적화 (Zeroth-order Optimization, ZO) 기반 접근법은 역전파 없이 순전파 (Forward-pass) 만으로 기울기를 추정하여 VRAM 문제를 해결하려 시도했으나, 모든 블록을 활성화하는 방식은 여전히 높은 VRAM 사용량, 느린 수렴 속도 (고차원 공간에서의 분산 증가), 그리고 여전히 존재하는 통신 오버헤드를 유발했습니다.

2. 제안된 방법론: ZorBA (Methodology)

저자들은 ZorBA라는 새로운 페더러티드 파인튜닝 프레임워크를 제안합니다. 이는 제로오더 최적화와 **이종 블록 활성화 (Heterogeneous Block Activation)**를 결합한 방식입니다.

핵심 구성 요소:

1. 제로오더 최적화 (Zeroth-order Optimization):

   • 역전파를 사용하지 않고, 무작위 섭동 벡터 (perturbation vectors) 를 사용하여 손실 함수 값의 유한 차분 (finite differences) 으로 기울기를 추정합니다.
   • 이를 통해 클라이언트 측에서 그래디언트 저장 없이 순전파만 수행하여 VRAM 사용량을 획기적으로 줄입니다.

2. 이종 블록 활성화 (Heterogeneous Block Activation):

   • 모든 클라이언트가 모델의 모든 트랜스포머 블록을 업데이트하는 대신, 서버가 각 클라이언트의 VRAM 용량에 따라 서로 다른 블록의 부분집합을 활성화합니다.
   • 클라이언트 $n$은 활성화된 블록의 파라미터만 업데이트하고, 나머지 블록은 동결 (frozen) 시킵니다.
   • 이는 VRAM 사용량을 줄이면서도 전체 모델의 수렴을 돕는 균형을 맞춥니다.

3. 공유 무작위 시드 (Shared Random Seeds):

   • 클라이언트와 서버가 동일한 무작위 시드를 공유하여 동일한 섭동 벡터를 독립적으로 생성합니다.
   • 이로 인해 클라이언트와 서버 간에 고차원 파라미터나 섭동 벡터를 전송할 필요가 없으며, 기울기 유출 (gradient leakage) 을 방지하고 통신 오버헤드를 최소화합니다.
   • 클라이언트는 오직 유한 차분 값 (loss difference) 만을 서버로 전송합니다.

4. 최적화 알고리즘 ($\epsilon$-constraint Lexicographic Algorithm):

   • 수렴 속도와 VRAM 사용량을 동시에 최적화하기 위해 블록 활성화 결정 행렬을 최적화하는 문제를 다목적 최적화 문제로 formulat 합니다.
   • $\Lambda(A)$ 최소화: 수렴 속도에 영향을 미치는 편향 (bias) 항을 최소화하기 위해 각 클라이언트가 활성화하는 블록의 '최소 인기도 (least popularity)'를 극대화하는 방향으로 설계됩니다.
   • 알고리즘 단계:
     1. VRAM 제약 하에서 모든 클라이언트의 최소 인기도를 최대화하는 $\gamma^\star$를 계산 (Dinic's algorithm 사용).
     2. 초기 할당 후, 최소 인기도가 $\gamma^\star$인 클라이언트 수를 줄이기 위해 추가 블록을 할당하는 그리디 (Greedy) 알고리즘 수행.
     3. 파레토 프론트 (Pareto front) 상에서 수렴 속도와 VRAM 사용량 사이의 최적 균형을 선택.

3. 주요 기여 (Key Contributions)

• ZorBA 프레임워크 제안: LLM 의 페더러티드 파인튜닝을 위해 제로오더 최적화와 이종 블록 활성화 메커니즘을 통합한 최초의 프레임워크 중 하나입니다.
• 이론적 분석: 비볼록 (non-convex) 환경에서 ZorBA 의 수렴 한계를 증명했습니다. 특히, 블록 활성화 결정이 수렴 속도와 VRAM 사용량에 미치는 영향을 정량화하고, 블록 활성화의 분포 (인기도) 가 수렴에 미치는 영향을 규명했습니다.
• 효율적인 최적화 알고리즘: NP-난해 (NP-hard) 인 블록 활성화 문제를 해결하기 위해 $\epsilon$-constraint 기법과 레xicographic 알고리즘을 결합하여 근사 최적 해를 구하는 알고리즘을 개발했습니다.
• 실험적 검증: AG-News, SST-2, SNLI 데이터셋과 OPT-125M, OPT-1.3B 모델을 사용하여 기존 방법 (FedIT, FedZO, DeComFL) 과 비교 평가했습니다.

4. 실험 결과 (Results)

• VRAM 사용량 감소: ZorBA 는 기존 방법들 (FedIT, FedZO, DeComFL) 에 비해 **최대 62.41%**까지 총 VRAM 사용량을 줄였습니다. 이는 리소스 제한이 있는 클라이언트에서 LLM 파인튜닝을 가능하게 합니다.
• 수렴 속도: 제안된 방법은 제로오더 기반 베이스라인 (FedZO, DeComFL) 보다 더 빠른 수렴 속도를 보였습니다. 이는 최적화된 이종 블록 활성화가 수렴을 가속화함을 의미합니다.
• 통신 오버헤드: 공유 시드와 유한 차분 전송 방식을 통해 통신 오버헤드를 기존 1 차 미분 기반 방법 (FedIT) 에 비해 매우 낮게 유지했습니다. (예: OPT-1.3B 기준, FedIT 대비 $10^{12}$배 이상 감소).
• 트레이드오프 분석: 활성화된 블록의 수와 수렴 속도 사이의 트레이드오프 관계를 분석하여, VRAM 사용량을 줄이면서도 수렴 속도를 크게 저하시키지 않는 최적의 활성화 전략을 찾을 수 있음을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 대규모 언어 모델의 페더러티드 학습이 직면한 메모리 제약과 통신 비용이라는 두 가지 핵심 문제를 동시에 해결하는 실용적인 솔루션을 제시합니다.

• 접근성 향상: 고사양 GPU 가 없는 엣지 디바이스에서도 LLM 파인튜닝이 가능해져, 데이터 프라이버시를 유지하면서 분산된 환경에서 LLM 을 적응시킬 수 있는 길을 엽니다.
• 이론적 통찰: 단순히 블록을 무작위로 줄이는 것이 아니라, 클라이언트 간 블록 활성화의 '분포'와 '인기도'를 최적화함으로써 수렴 속도를 개선할 수 있다는 이론적 통찰을 제공합니다.
• 실용성: 공유 시드 메커니즘을 통해 통신 효율성을 극대화하고, 역전파 없이도 고품질의 모델 업데이트를 가능하게 하여 실제 배포 환경에 적용 가능한 강력한 프레임워크를 제시합니다.

요약하자면, ZorBA 는 제로오더 최적화의 메모리 효율성과 지능적인 블록 활성화 전략을 결합하여, 리소스 제약이 있는 환경에서도 빠르고 효율적으로 LLM 을 파인튜닝할 수 있는 새로운 패러다임을 제시합니다.