FLoRG: 그람 행렬과 프로크루스테스 정렬을 통한 효율적인 연합 파인튜닝

초록

대규모 언어 모델(LLM)의 효율적인 연합 파인튜닝을 위한 새로운 프레임워크인 FLoRG를 소개한다. 기존 LoRA 방식의 두 저랭크 행렬을 연합 학습에 적용할 때 발생하는 집계 오류와 분해 드리프트 문제를 해결하기 위해, 단일 저랭크 행렬을 사용하고 그람 행렬을 집계하며, 프로크루스테스 정렬을 도입하여 분해 불일치를 최소화한다. 이를 통해 더 높은 정확도와 최대 2041배의 통신 오버헤드 감소를 달성했다.

상세 분석

본 논문이 제안하는 FLoRG 프레임워크의 기술적 핵심은 기존 연합 LoRA의 근본적인 한계를 해결하는 세 가지 혁신적 접근법에 있다.

첫째, 단일 저랭크 행렬과 그람 행렬 집계를 통한 집계 오류 제거이다. 기존 FedLoRA 방식은 클라이언트가 로컬로 업데이트한 두 개의 저랭크 행렬 B와 A를 서버가 각각 평균 집계한다. 이때, 집계된 두 행렬의 곱(BA)은 각 클라이언트의 로컬 업데이트(B_n * A_n)를 평균한 것과 수학적으로 일치하지 않는 ‘집계 오류’가 발생한다. FLoRG는 이 문제를 근본적으로 해결하기 위해 파인튜닝 매개변수를 단일 저랭크 행렬 A로 재파라미터화한다. 서버는 클라이언트로부터 전달받은 A_n으로부터 그람 행렬 Q_n = (A_n)^T A_n을 계산하여 집계한다. 그람 행렬의 집계는 선형 연산이며, 집계된 그람 행렬 Q_global은 로컬 그람 행렬들의 정확한 평균이므로 집계 오류가 완전히 제거된다. 또한, 클라이언트가 업로드하는 행렬이 하나로 줄어 통신 부하도 감소한다.

둘째, 프로크루스테스 정렬을 활용한 분해 드리프트 최소화이다. 서버가 집계된 그람 행렬 Q_global을 다시 다음 라운드의 파라미터 A_{t+1}로 분해해야 하는데, 행렬 분해(예: 고유값 분해)는 일반적으로 유일하지 않다. 특히 랭크가 낮거나 고유값이 중복될 경우 무수히 많은 분해가 가능하며, 이로 인해 선택된 분해에 따라 A_{t+1}의 방향이 달라져 학습 경로가 불안정해지는 ‘분해 드리프트’가 발생한다. FLoRG는 이 문제를 해결하기 위해 분해 후 얻은 후보 행렬 Â_{t+1}에 프로크루스테스 변환 행렬 S를 적용한다. 구체적으로, S는 새로운 분해 Â_{t+1}을 이전 라운드의 파라미터 A_t와 가능한 한 정렬(align)시키도록, 즉 ||SÂ_{t+1} - A_t||F를 최소화하도록 최적화된다. 이 최적화 문제는 특이값 분해(SVD)를 통해 폐쇄형 해를 구할 수 있으며, 그 결과 A_t의 서브스페이스를 최대한 보존하면서도 집계된 그람 행렬 정보는 그대로 유지하는 새로운 파라미터 A{t+1} = SÂ_{t+1}을 생성한다. 이는 연속된 라운드 간의 업데이트 일관성을 높여 학습 안정성을 증대시킨다.

셋째, 임의의 차원에 대한 적응성을 보장하는 공유 기저 행렬 L, R의 도입이다. FLoRG의 파인튜닝 업데이트는 ΔW = L Q R = L (A^T A) R 형태로 표현된다. 여기서 L과 R은 사전에 정의된 준정규 직교 행렬로, 모든 클라이언트와 라운드에서 고정되어 있다. 이 설계는 파인튜닝 대상 가중치 행렬 W의 입력/출력 차원(d_in, d_out)에 관계없이, 내부 저랭크 연산을 항상 동일한 크기(k = min(d_in, d_out))의 그람 행렬 Q를 통해 수행할 수 있게 해준다. 따라서 FLoRG는 모델의 어떤 계층에도 동일한 방식으로 적용 가능한 일반적인 프레임워크가 된다.

이론적 분석에서는 비볼록 손실 함수 하에서 FLoRG의 수렴 속도를 분석하고, 프로크루스테스 정렬이 분해 드리프트를 줄여 더 타이트한 수렴 상한을 제공함을 증명한다. 실험적으로는 GLUE 벤치마크에서 FedIT, FeDeRA 등 5개의 최신 기법을 능가하는 정확도를 보여주며, 통신량 측면에서도 압도적인 효율성(최대 2041배 감소)을 입증한다.