라반: 연합 학습을 위한 다중 헤드 저랭크 적응

초록

라반은 연합 학습 환경에서 파라미터 효율성을 유지하면서 모델 표현력을 높이기 위해 다중 LoRA 헤드를 도입한 방법이다. 고정된 B·A 행렬과 학습 가능한 H·s 스케일링 파라미터만을 업데이트함으로써, 통신량은 기존 LoRA와 동일하지만 비동질적인 데이터와 연산 자원을 가진 클라이언트에서도 높은 정확도를 달성한다. 실험 결과, 비IID 설정에서도 기존 PEFT 기법 대비 2~8%의 정확도 향상을 보였다.

상세 분석

라반(Ravan)은 연합 학습(Federated Learning, FL)에서 대규모 언어 모델(LLM)을 파라미터 효율적으로 미세조정하기 위한 새로운 저랭크 적응(LoRA) 변형이다. 기존 LoRA는 가중치 업데이트 ΔW를 두 개의 저랭크 행렬 B와 A의 곱으로 근사해 파라미터 수를 크게 줄이지만, 데이터 이질성 및 연산 이질성이 심한 FL 환경에서는 업데이트의 실질적인 랭크가 증가함에 따라 저랭크 근사에 의해 중요한 정보가 손실된다. 라반은 이를 해결하기 위해 ΔW를 다중 헤드 형태인 Σ_i s_i B_i H_i A_i 로 재파라미터화한다. 여기서 B_i와 A_i는 초기화 시 고정하고 서로 직교하도록 설계되며, 학습 가능한 부분은 각 헤드마다 스케일링 계수 s_i와 핵심 행렬 H_i뿐이다.

핵심 아이디어는 두 가지이다. 첫째, 여러 헤드를 동시에 사용함으로써 전체 파라미터 예산 N 내에서 효과적인 랭크를 √h · √N 만큼 확대한다. 수학적으로 rank(Σ_i B_i H_i A_i) ≤ √N · h 이며, 이는 단일 LoRA(head) 대비 √h 배 높은 랭크를 제공한다. 둘째, B_i와 A_i를 고정하고 H_i·s_i만을 평균화하면 정확한 집계(Exact Aggregation)가 보장된다. 이는 기존 LoRA 기반 FL 방법이 B와 A를 별도로 평균할 때 발생하는 근사 오차를 완전히 제거한다.

라반은 또한 연산 이질성을 고려한다. 클라이언트마다 메모리·연산 능력이 다를 경우, 제한된 자원을 가진 디바이스는 일부 헤드만 활성화하고 나머지는 동결(freeze)한다. 이렇게 하면 각 클라이언트가 처리해야 할 파라미터 양이 줄어들면서도 전체 모델은 모든 헤드의 조합을 통해 높은 랭크 근사를 유지한다. 통신 측면에서도 클라이언트는 s_i H_i 만을 서버에 전송하므로, 기존 LoRA와 동일한 비트 수를 사용한다.

실험에서는 CIFAR‑100, SVHN 등 비전 벤치마크와 언어 모델 파인튜닝 작업에서 라반이 FedIT, FedEx‑LoRA, FF‑A‑LoRA 등 최신 PEFT 기반 연합 학습 기법을 능가했다. 특히 비IID(Dirichlet α=0.3) 환경에서 2~8% 포인트의 정확도 향상을 기록했으며, 헤드 수와 스케일링 파라미터를 조절함으로써 다양한 클라이언트 자원 제약에 유연하게 대응할 수 있음을 보였다. 초기화 방법으로는 무작위 정규분포와 Gram‑Schmidt 정규직교화가 가장 좋은 성능을 보였으며, 이는 헤드 간의 서브스페이스가 충분히 독립적일 때 모델 표현력이 극대화된다는 가설을 실증한다.

요약하면, 라반은 (1) 파라미터 효율성 유지, (2) 높은 효과적 랭크 확보, (3) 정확한 집계 보장, (4) 연산 이질성에 대한 적응성이라는 네 가지 핵심 요구를 동시에 만족시키는 연합 학습용 PEFT 프레임워크이다.